From 4a16b997f8b20ade899d9b271a6135c72d625753 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:00:18 -0700 Subject: [PATCH 1/8] [None][feat] Add MiniMax-M3 Rubin decode targets Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/CMakeLists.txt | 32 +- .../kernels/minimaxM3/routerGemv.cu | 67 +- .../kernels/minimaxM3/routerGemv.h | 4 + .../kernels/minimaxM3/specTail.cu | 25 +- .../minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu | 267 + .../minimaxM3Sm107Tp2/denseQkNormRopeFp8.h | 64 + .../kernels/minimaxM3Sm107Tp2/draftStep.cu | 620 ++ .../kernels/minimaxM3Sm107Tp2/draftStep.h | 111 + .../minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu | 189 + .../minimaxM3Sm107Tp2/drafterRopeFp8Kv.h | 79 + .../kernels/minimaxM3Sm107Tp2/foldTail.cu | 86 + .../kernels/minimaxM3Sm107Tp2/foldTail.h | 72 + .../minimaxM3Sm107Tp2/fullCtxReduce.cu | 223 + .../kernels/minimaxM3Sm107Tp2/fullCtxReduce.h | 71 + .../minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu | 1092 +++ .../minimaxM3Sm107Tp2/mcAllreduceNormTp2.h | 83 + .../kernels/minimaxM3Sm107Tp2/producerPw2.cu | 339 + .../kernels/minimaxM3Sm107Tp2/producerPw2.h | 76 + .../minimaxM3Sm107Tp2/sampledAccept.cu | 596 ++ .../kernels/minimaxM3Sm107Tp2/sampledAccept.h | 82 + .../kernels/minimaxM3Sm107Tp2/selectLong2.cu | 246 + .../kernels/minimaxM3Sm107Tp2/selectLong2.h | 50 + .../kernels/minimaxM3Sm107Tp2/siluAndMul.cu | 89 + .../kernels/minimaxM3Sm107Tp2/siluAndMul.h | 40 + .../minimaxM3Sm107Tp2/sparseDecodeLc2.cu | 640 ++ .../minimaxM3Sm107Tp2/sparseDecodeLc2.h | 77 + .../kernels/minimaxM3Sm107Tp2/specOffsets.cu | 58 + .../kernels/minimaxM3Sm107Tp2/specOffsets.h | 41 + .../kernels/minimaxM3Sm107Tp2/stepPrologue.cu | 111 + .../kernels/minimaxM3Sm107Tp2/stepPrologue.h | 47 + .../kernels/minimaxM3Sm107Tp2/swigluOai.cu | 100 + .../kernels/minimaxM3Sm107Tp2/swigluOai.h | 43 + .../minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu | 296 + .../minimaxM3Sm107Tp4/denseQkNormRopeFp8.h | 68 + .../kernels/minimaxM3Sm107Tp4/draftStep.cu | 667 ++ .../kernels/minimaxM3Sm107Tp4/draftStep.h | 107 + .../minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu | 189 + .../minimaxM3Sm107Tp4/drafterRopeFp8Kv.h | 79 + .../kernels/minimaxM3Sm107Tp4/foldTail.cu | 64 + .../kernels/minimaxM3Sm107Tp4/foldTail.h | 129 + .../minimaxM3Sm107Tp4/fullCtxReduce.cu | 268 + .../kernels/minimaxM3Sm107Tp4/fullCtxReduce.h | 77 + .../minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu | 1264 ++++ .../minimaxM3Sm107Tp4/mcAllreduceNormTp4.h | 83 + .../kernels/minimaxM3Sm107Tp4/producerPw4.cu | 340 + .../kernels/minimaxM3Sm107Tp4/producerPw4.h | 76 + .../minimaxM3Sm107Tp4/sampledAccept.cu | 944 +++ .../kernels/minimaxM3Sm107Tp4/sampledAccept.h | 87 + .../kernels/minimaxM3Sm107Tp4/selectLong2.cu | 312 + .../kernels/minimaxM3Sm107Tp4/selectLong2.h | 52 + .../kernels/minimaxM3Sm107Tp4/siluAndMul.cu | 89 + .../kernels/minimaxM3Sm107Tp4/siluAndMul.h | 40 + .../minimaxM3Sm107Tp4/sparseDecodeLc4.cu | 743 ++ .../minimaxM3Sm107Tp4/sparseDecodeLc4.h | 104 + .../kernels/minimaxM3Sm107Tp4/specOffsets.cu | 58 + .../kernels/minimaxM3Sm107Tp4/specOffsets.h | 41 + .../kernels/minimaxM3Sm107Tp4/stepPrologue.cu | 111 + .../kernels/minimaxM3Sm107Tp4/stepPrologue.h | 47 + .../kernels/minimaxM3Sm107Tp4/swigluOai.cu | 100 + .../kernels/minimaxM3Sm107Tp4/swigluOai.h | 43 + .../kernels/minimaxM3Sm107Tp8/draftStep.cu | 711 ++ .../kernels/minimaxM3Sm107Tp8/draftStep.h | 119 + .../minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu | 189 + .../minimaxM3Sm107Tp8/drafterRopeFp8Kv.h | 79 + .../kernels/minimaxM3Sm107Tp8/foldTail.cu | 64 + .../kernels/minimaxM3Sm107Tp8/foldTail.h | 129 + .../minimaxM3Sm107Tp8/fullCtxReduce.cu | 268 + .../kernels/minimaxM3Sm107Tp8/fullCtxReduce.h | 77 + .../minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu | 1096 +++ .../minimaxM3Sm107Tp8/mcAllreduceNormTp8.h | 83 + .../kernels/minimaxM3Sm107Tp8/producerPw8.cu | 340 + .../kernels/minimaxM3Sm107Tp8/producerPw8.h | 76 + .../minimaxM3Sm107Tp8/sampledAccept.cu | 946 +++ .../kernels/minimaxM3Sm107Tp8/sampledAccept.h | 87 + .../kernels/minimaxM3Sm107Tp8/siluAndMul.cu | 89 + .../kernels/minimaxM3Sm107Tp8/siluAndMul.h | 40 + .../minimaxM3Sm107Tp8/sparseDecodeLc8m.cu | 746 ++ .../minimaxM3Sm107Tp8/sparseDecodeLc8m.h | 105 + .../kernels/minimaxM3Sm107Tp8/specOffsets.cu | 58 + .../kernels/minimaxM3Sm107Tp8/specOffsets.h | 41 + .../kernels/minimaxM3Sm107Tp8/stepPrologue.cu | 111 + .../kernels/minimaxM3Sm107Tp8/stepPrologue.h | 47 + .../kernels/minimaxM3Sm107Tp8/swigluOai.cu | 100 + .../kernels/minimaxM3Sm107Tp8/swigluOai.h | 43 + cpp/tensorrt_llm/kernels/quantization.cu | 30 + cpp/tensorrt_llm/kernels/quantization.cuh | 58 + .../trtllmGenKernels/fmha/fmhaKernels.h | 135 +- .../thop/minimaxM3/minimaxM3RouterGemvOp.cpp | 23 +- .../minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp | 125 + .../minimaxM3Sm107Tp2DraftStepOp.cpp | 220 + .../minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp | 103 + .../minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp | 130 + .../minimaxM3Sm107Tp2FullCtxReduceOp.cpp | 74 + .../minimaxM3Sm107Tp2McAllreduceNormOp.cpp | 181 + .../minimaxM3Sm107Tp2ProducerPwOp.cpp | 185 + .../minimaxM3Sm107Tp2SampledAcceptOp.cpp | 110 + .../minimaxM3Sm107Tp2SelectLongOp.cpp | 85 + .../minimaxM3Sm107Tp2SiluAndMulOp.cpp | 66 + .../minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp | 169 + .../minimaxM3Sm107Tp2SpecOffsetsOp.cpp | 74 + .../minimaxM3Sm107Tp2StepPrologueOp.cpp | 68 + .../minimaxM3Sm107Tp2SwigluOaiOp.cpp | 65 + .../minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp | 128 + .../minimaxM3Sm107Tp4DraftStepOp.cpp | 215 + .../minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp | 103 + .../minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp | 89 + .../minimaxM3Sm107Tp4FoldTailParams.h | 111 + .../minimaxM3Sm107Tp4FullCtxReduceOp.cpp | 92 + .../minimaxM3Sm107Tp4McAllreduceNormOp.cpp | 181 + .../minimaxM3Sm107Tp4ProducerPwOp.cpp | 185 + .../minimaxM3Sm107Tp4SampledAcceptOp.cpp | 113 + .../minimaxM3Sm107Tp4SelectLongOp.cpp | 87 + .../minimaxM3Sm107Tp4SiluAndMulOp.cpp | 66 + .../minimaxM3Sm107Tp4SparseDecodeLcOp.cpp | 191 + .../minimaxM3Sm107Tp4SpecOffsetsOp.cpp | 74 + .../minimaxM3Sm107Tp4StepPrologueOp.cpp | 68 + .../minimaxM3Sm107Tp4SwigluOaiOp.cpp | 65 + .../minimaxM3Sm107Tp8DraftStepOp.cpp | 236 + .../minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp | 103 + .../minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp | 89 + .../minimaxM3Sm107Tp8FoldTailParams.h | 111 + .../minimaxM3Sm107Tp8FullCtxReduceOp.cpp | 92 + .../minimaxM3Sm107Tp8McAllreduceNormOp.cpp | 181 + .../minimaxM3Sm107Tp8ProducerPwOp.cpp | 185 + .../minimaxM3Sm107Tp8SampledAcceptOp.cpp | 113 + .../minimaxM3Sm107Tp8SiluAndMulOp.cpp | 66 + .../minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp | 189 + .../minimaxM3Sm107Tp8SpecOffsetsOp.cpp | 74 + .../minimaxM3Sm107Tp8StepPrologueOp.cpp | 68 + .../minimaxM3Sm107Tp8SwigluOaiOp.cpp | 65 + tensorrt_llm/_device_copy.py | 8 + .../models/minimax_m3/kernels/fake_ops.py | 82 + .../minimax_m3_nvfp4__sm_107__tp2/__init__.py | 5 + .../minimax_m3_nvfp4__sm_107__tp2/modeling.py | 6453 +++++++++++++++++ .../ops/__init__.py | 3 + .../ops/activation/__init__.py | 3 + .../ops/activation/m3_silu_and_mul.py | 83 + .../ops/attention/__init__.py | 3 + .../ops/attention/m3_dense_decode.py | 190 + .../ops/attention/m3_draft_attention.py | 131 + .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 + .../ops/attention/m3_fullctx_decode.py | 142 + .../ops/attention/m3_index_score.py | 239 + .../ops/attention/m3_msa_attention.py | 249 + .../ops/attention/m3_producer.py | 216 + .../ops/attention/m3_qk_norm_rope_fp8.py | 137 + .../ops/attention/m3_qkv_insert.py | 134 + .../ops/attention/m3_select_blocks.py | 82 + .../ops/attention/m3_select_long.py | 192 + .../ops/attention/m3_select_long_dsl.py | 142 + .../ops/attention/m3_sparse_decode.py | 207 + .../ops/attention/m3_sparse_decode_lc.py | 474 ++ .../ops/attention/msa_kv_update.py | 200 + .../ops/comm/__init__.py | 3 + .../ops/comm/m3_allreduce.py | 61 + .../ops/comm/m3_mc_allreduce_norm.py | 319 + .../ops/comm/m3_mc_finalize_allreduce_norm.py | 207 + .../ops/embedding/__init__.py | 3 + .../ops/embedding/m3_embedding.py | 48 + .../ops/embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../ops/embedding_firstnorm/op.py | 143 + .../ops/gemm/__init__.py | 3 + .../ops/gemm/m3_bf16_linear.py | 41 + .../ops/gemm/m3_dense_swiglu.py | 72 + .../ops/gemm/m3_draft_lowm.py | 169 + .../ops/gemm/m3_fp32_linear.py | 46 + .../ops/gemm/m3_mx_gemv.py | 170 + .../ops/gemm/m3_mx_gemv_bf16.py | 100 + .../ops/gemm/m3_mxfp8_gemm.py | 58 + .../ops/gemm/m3_mxfp8_linear.py | 137 + .../ops/gemm/m3_target_head.py | 346 + .../ops/gemm/m3_ugpu_head.py | 168 + .../ops/long_select.py | 70 + .../ops/m3_attn_lc.py | 164 + .../ops/m3_draft_step.py | 341 + .../ops/m3_drafter_attn.py | 413 ++ .../ops/m3_fold_tail.py | 172 + .../ops/m3_force_vec.py | 63 + .../ops/m3_fullctx_decode.py | 297 + .../ops/m3_index_score.py | 670 ++ .../ops/m3_locality.py | 448 ++ .../ops/m3_lowm_direct.py | 497 ++ .../ops/m3_lowm_splitk.py | 1249 ++++ .../ops/m3_mc_allreduce.py | 392 + .../ops/m3_moe/__init__.py | 16 + .../ops/m3_moe/m3_moe_kernel.py | 3936 ++++++++++ .../ops/m3_moe/op.py | 808 +++ .../ops/m3_moe_finalize.py | 95 + .../ops/m3_prequant.py | 111 + .../ops/m3_report.py | 34 + .../ops/m3_router_gemv.py | 107 + .../ops/m3_sampled_accept.py | 203 + .../ops/m3_select_long_dsl.py | 366 + .../ops/m3_shard_accept.py | 223 + .../ops/m3_spec_tail.py | 113 + .../ops/m3_target_head.py | 113 + .../ops/m3_tokmap_guard.py | 185 + .../ops/moe/__init__.py | 3 + .../ops/moe/m3_moe_engine.py | 342 + .../ops/moe/m3_router_gemv.py | 101 + .../ops/moe/m3_router_gemv_fp4.py | 63 + .../ops/moe/m3_router_tc.py | 114 + .../ops/msa_kv_update.py | 89 + .../ops/mx_gemv/__init__.py | 4 + .../ops/mx_gemv/m3_mx_gemv_kernel.py | 1237 ++++ .../ops/mx_gemv/op.py | 496 ++ .../ops/norm/__init__.py | 3 + .../ops/norm/m3_embedding_firstnorm.py | 98 + .../ops/spec/__init__.py | 4 + .../ops/spec/m3_capture_norm.py | 87 + .../ops/spec/m3_chunk_norm.py | 81 + .../ops/spec/m3_draft_input.py | 140 + .../ops/spec/m3_draft_silu.py | 58 + .../ops/spec/m3_draft_step.py | 432 ++ .../ops/spec/m3_force_vec.py | 136 + .../ops/spec/m3_sampled_accept.py | 308 + .../ops/spec/m3_sampled_accept_ws_words.py | 43 + .../ops/spec/m3_shard_accept.py | 342 + .../ops/spec/m3_slot_scatter.py | 72 + .../ops/spec/m3_spec_offsets.py | 60 + .../ops/spec/m3_step_gather.py | 160 + .../ops/spec/m3_step_prologue.py | 76 + .../ops/triton_sparse_decode.py | 1747 +++++ .../ops/triton_sparse_decode_mq.py | 237 + .../minimax_m3_nvfp4__sm_107__tp2/weights.py | 865 +++ .../minimax_m3_nvfp4__sm_107__tp4/__init__.py | 5 + .../minimax_m3_nvfp4__sm_107__tp4/modeling.py | 5936 +++++++++++++++ .../ops/__init__.py | 3 + .../ops/activation/__init__.py | 3 + .../ops/activation/m3_silu_and_mul.py | 83 + .../ops/attention/__init__.py | 3 + .../ops/attention/m3_dense_decode.py | 190 + .../attention/m3_dense_qk_norm_rope_fp8.py | 134 + .../ops/attention/m3_draft_attention.py | 111 + .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 + .../ops/attention/m3_fullctx_decode.py | 142 + .../ops/attention/m3_index_score.py | 159 + .../ops/attention/m3_index_score_t2sc.py | 168 + .../ops/attention/m3_msa_attention.py | 224 + .../ops/attention/m3_producer.py | 223 + .../ops/attention/m3_qk_norm_rope_fp8.py | 137 + .../ops/attention/m3_qkv_insert.py | 134 + .../ops/attention/m3_select_blocks.py | 82 + .../ops/attention/m3_select_long.py | 186 + .../ops/attention/m3_select_long_dsl.py | 140 + .../ops/attention/m3_sparse_decode.py | 215 + .../ops/attention/m3_sparse_decode_lc.py | 436 ++ .../ops/attention/msa_kv_update.py | 200 + .../ops/comm/__init__.py | 3 + .../ops/comm/m3_allreduce.py | 61 + .../ops/comm/m3_mc_allreduce_norm.py | 329 + .../ops/comm/m3_mc_finalize_allreduce_norm.py | 197 + .../ops/embedding/__init__.py | 3 + .../ops/embedding/m3_embedding.py | 48 + .../ops/embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../ops/embedding_firstnorm/op.py | 142 + .../ops/gemm/__init__.py | 3 + .../ops/gemm/m3_bf16_linear.py | 41 + .../ops/gemm/m3_dense_swiglu.py | 72 + .../ops/gemm/m3_draft_lowm.py | 154 + .../ops/gemm/m3_fp32_linear.py | 46 + .../ops/gemm/m3_mx_gemv.py | 166 + .../ops/gemm/m3_mx_gemv_bf16.py | 106 + .../ops/gemm/m3_mxfp8_gemm.py | 58 + .../ops/gemm/m3_mxfp8_linear.py | 137 + .../ops/gemm/m3_mxfp8_pinned_gemm.py | 69 + .../ops/gemm/m3_mxfp8_pinned_quantize.py | 47 + .../ops/gemm/m3_target_head.py | 346 + .../ops/long_select.py | 72 + .../ops/m3_attn_lc.py | 148 + .../ops/m3_draft_lowm.py | 143 + .../ops/m3_draft_step.py | 222 + .../ops/m3_drafter_attn.py | 405 ++ .../ops/m3_fold_tail.py | 172 + .../ops/m3_force_vec.py | 64 + .../ops/m3_fullctx_decode.py | 314 + .../ops/m3_index_score.py | 638 ++ .../ops/m3_index_score_t2sc.py | 682 ++ .../ops/m3_lowm_splitk.py | 1249 ++++ .../ops/m3_mc_allreduce.py | 422 ++ .../ops/m3_moe/__init__.py | 16 + .../ops/m3_moe/m3_moe16_kernel.py | 3871 ++++++++++ .../ops/m3_moe/m3_moe_kernel.py | 3914 ++++++++++ .../ops/m3_moe/op.py | 779 ++ .../ops/m3_moe/op16.py | 288 + .../ops/m3_prequant.py | 111 + .../ops/m3_report.py | 34 + .../ops/m3_sampled_accept.py | 174 + .../ops/m3_select_long_dsl.py | 367 + .../ops/m3_shard_accept.py | 176 + .../ops/m3_spec_tail.py | 116 + .../ops/m3_target_head.py | 113 + .../ops/m3_tokmap_guard.py | 185 + .../ops/moe/__init__.py | 3 + .../ops/moe/m3_moe16_engine.py | 67 + .../ops/moe/m3_moe_engine.py | 326 + .../ops/moe/m3_router_tc.py | 119 + .../ops/msa_kv_update.py | 89 + .../ops/mx_gemv/__init__.py | 3 + .../ops/mx_gemv/m3_mx_gemv_kernel.py | 1299 ++++ .../ops/mx_gemv/op.py | 521 ++ .../ops/norm/__init__.py | 3 + .../ops/norm/m3_embedding_firstnorm.py | 87 + .../ops/spec/__init__.py | 4 + .../ops/spec/m3_capture_norm.py | 80 + .../ops/spec/m3_chunk_norm.py | 70 + .../ops/spec/m3_draft_input.py | 140 + .../ops/spec/m3_draft_silu.py | 58 + .../ops/spec/m3_draft_step.py | 417 ++ .../ops/spec/m3_force_vec.py | 86 + .../ops/spec/m3_force_vec_req.py | 62 + .../ops/spec/m3_sampled_accept.py | 289 + .../ops/spec/m3_sampled_accept_ws_words.py | 59 + .../ops/spec/m3_shard_accept.py | 284 + .../ops/spec/m3_slot_scatter.py | 72 + .../ops/spec/m3_spec_offsets.py | 60 + .../ops/spec/m3_step_gather.py | 152 + .../ops/spec/m3_step_prologue.py | 76 + .../ops/triton_sparse_decode.py | 1730 +++++ .../ops/triton_sparse_decode_mq.py | 237 + .../minimax_m3_nvfp4__sm_107__tp4/weights.py | 865 +++ .../minimax_m3_nvfp4__sm_107__tp8/__init__.py | 5 + .../minimax_m3_nvfp4__sm_107__tp8/modeling.py | 6140 ++++++++++++++++ .../ops/__init__.py | 3 + .../ops/activation/__init__.py | 3 + .../ops/activation/m3_silu_and_mul.py | 83 + .../ops/attention/__init__.py | 3 + .../ops/attention/m3_dense_decode.py | 190 + .../ops/attention/m3_draft_attention.py | 131 + .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 + .../ops/attention/m3_fullctx_decode.py | 142 + .../ops/attention/m3_index_score.py | 167 + .../ops/attention/m3_msa_attention.py | 249 + .../ops/attention/m3_producer.py | 223 + .../ops/attention/m3_qk_norm_rope_fp8.py | 137 + .../ops/attention/m3_qkv_insert.py | 134 + .../ops/attention/m3_select_blocks.py | 82 + .../ops/attention/m3_select_long.py | 193 + .../ops/attention/m3_select_long_dsl.py | 142 + .../ops/attention/m3_sparse_decode.py | 215 + .../ops/attention/m3_sparse_decode_lc8.py | 487 ++ .../ops/attention/msa_kv_update.py | 200 + .../ops/comm/__init__.py | 3 + .../ops/comm/m3_allreduce.py | 60 + .../ops/comm/m3_mc_allreduce_norm.py | 300 + .../ops/comm/m3_mc_finalize_allreduce_norm.py | 215 + .../ops/embedding/__init__.py | 3 + .../ops/embedding/m3_embedding.py | 48 + .../ops/embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../ops/embedding_firstnorm/op.py | 142 + .../ops/gemm/__init__.py | 3 + .../ops/gemm/m3_bf16_linear.py | 41 + .../ops/gemm/m3_dense_swiglu.py | 72 + .../ops/gemm/m3_draft_lowm.py | 156 + .../ops/gemm/m3_fp32_linear.py | 46 + .../ops/gemm/m3_mx_gemv.py | 160 + .../ops/gemm/m3_mx_gemv_bf16.py | 100 + .../ops/gemm/m3_mxfp8_gemm.py | 60 + .../ops/gemm/m3_mxfp8_linear.py | 137 + .../ops/gemm/m3_target_head.py | 346 + .../ops/long_select.py | 70 + .../ops/m3_attn_lc8.py | 152 + .../ops/m3_draft_step.py | 355 + .../ops/m3_drafter_attn.py | 404 ++ .../ops/m3_fc_wait.py | 40 + .../ops/m3_fold_tail.py | 180 + .../ops/m3_force_vec.py | 63 + .../ops/m3_fullctx_decode.py | 328 + .../ops/m3_index_score.py | 692 ++ .../ops/m3_layer/__init__.py | 14 + .../ops/m3_layer/phase_producer.py | 487 ++ .../ops/m3_locality.py | 448 ++ .../ops/m3_lowm_splitk.py | 1592 ++++ .../ops/m3_mc_allreduce.py | 431 ++ .../ops/m3_moe/__init__.py | 16 + .../ops/m3_moe/m3_moe_kernel.py | 3841 ++++++++++ .../ops/m3_moe/op.py | 693 ++ .../ops/m3_prequant.py | 111 + .../ops/m3_report.py | 34 + .../ops/m3_router_gemv.py | 107 + .../ops/m3_sampled_accept.py | 230 + .../ops/m3_select_long_dsl.py | 367 + .../ops/m3_shard_accept.py | 176 + .../ops/m3_spec_tail.py | 113 + .../ops/m3_target_head.py | 113 + .../ops/m3_tokmap_guard.py | 185 + .../ops/moe/__init__.py | 3 + .../ops/moe/m3_moe_engine.py | 341 + .../ops/moe/m3_router_gemv.py | 101 + .../ops/moe/m3_router_gemv_fp4.py | 63 + .../ops/moe/m3_router_tc.py | 106 + .../ops/msa_kv_update.py | 89 + .../ops/mx_gemv/__init__.py | 4 + .../ops/mx_gemv/m3_mx_gemv_kernel.py | 1472 ++++ .../ops/mx_gemv/op.py | 645 ++ .../ops/norm/__init__.py | 3 + .../ops/norm/m3_embedding_firstnorm.py | 90 + .../ops/spec/__init__.py | 4 + .../ops/spec/m3_capture_norm.py | 80 + .../ops/spec/m3_chunk_norm.py | 74 + .../ops/spec/m3_draft_input.py | 140 + .../ops/spec/m3_draft_silu.py | 58 + .../ops/spec/m3_draft_step.py | 440 ++ .../ops/spec/m3_force_vec.py | 136 + .../ops/spec/m3_sampled_accept.py | 282 + .../ops/spec/m3_sampled_accept_ws_words.py | 59 + .../ops/spec/m3_shard_accept.py | 266 + .../ops/spec/m3_slot_scatter.py | 72 + .../ops/spec/m3_spec_offsets.py | 60 + .../ops/spec/m3_step_gather.py | 160 + .../ops/spec/m3_step_prologue.py | 76 + .../ops/triton_sparse_decode.py | 1754 +++++ .../ops/triton_sparse_decode_mq.py | 237 + .../minimax_m3_nvfp4__sm_107__tp8/weights.py | 884 +++ .../modeling_v2/models/minimax_m3/routing.py | 75 +- .../sparse/minimax_m3/cache_manager.py | 128 +- .../sparse/minimax_m3/kernels/paged_cache.py | 42 +- .../kernels/trtllm_gen_dense_decode.py | 5 +- .../backends/sparse/minimax_m3/msa_backend.py | 103 +- .../_torch/attention/backends/trtllm.py | 9 +- .../custom_ops/mxfp8_cute_pinned_tactics.py | 160 + .../_torch/custom_ops/torch_custom_ops.py | 44 +- .../cute_dsl_kernels/minimax_m3/__init__.py | 3 + .../embedding_firstnorm/__init__.py | 3 + .../m3_embedding_firstnorm_kernel.py | 128 + .../minimax_m3/embedding_firstnorm/op.py | 146 + .../minimax_m3/moe/__init__.py | 16 + .../minimax_m3/moe/m3_moe16_kernel.py | 3871 ++++++++++ .../minimax_m3/moe/m3_moe_kernel.py | 3714 ++++++++++ .../cute_dsl_kernels/minimax_m3/moe/op.py | 602 ++ .../cute_dsl_kernels/minimax_m3/moe/op16.py | 288 + .../minimax_m3/mx_gemv/__init__.py | 3 + .../minimax_m3/mx_gemv/m3_mx_gemv_kernel.py | 1297 ++++ .../cute_dsl_kernels/minimax_m3/mx_gemv/op.py | 521 ++ tensorrt_llm/_torch/late_compile.py | 90 + tensorrt_llm/_torch/pyexecutor/_util.py | 13 +- .../_torch/pyexecutor/cuda_graph_runner.py | 18 +- .../engine/runners/decoder/inputs.py | 48 +- .../engine/runners/decoder/metadata.py | 8 + .../engine/runners/decoder/runner.py | 42 +- tensorrt_llm/_torch/pyexecutor/h2d_cache.py | 118 +- tensorrt_llm/_torch/pyexecutor/profiling.py | 31 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 73 +- .../_torch/pyexecutor/runtime_hooks.py | 62 +- tensorrt_llm/_torch/pyexecutor/step_inputs.py | 304 +- tensorrt_llm/_torch/speculative/interface.py | 17 +- .../_torch/speculative/spec_sampler_base.py | 129 +- .../accuracy/modeling_v2/test_minimax_m3.py | 64 +- .../defs/accuracy/test_llm_api_pytorch.py | 2 + .../test_lists/test-db/l0_b300.yml | 2 +- .../integration/test_lists/test-db/l0_cpu.yml | 8 + .../test_lists/test-db/l0_gr100.yml | 34 + .../msa/test_minimax_m3_kv_size_estimate.py | 84 + .../attention/sparse/msa/test_msa_backend.py | 47 +- .../sparse/msa/test_msa_runtime_policy.py | 225 +- .../sparse/msa/test_msa_skip_cache_routing.py | 89 + .../test_context_fmha_kernel_presence.py | 31 + .../test_fallback_fmha_extra_page_sizes.py | 58 + .../test_allreduce_tune_nccl_symmetric.py | 34 + .../test_mxfp8_cute_pinned_tactics.py | 100 + .../kernels/parallel/test_mxfp8_mxfp8_gemm.py | 17 +- .../test_modeling_v2_hist_topk_op_matrix.py | 1 + .../test_modeling_v2_draft_ehproj_table.py | 1 + .../test_modeling_v2_draft_sample_glue.py | 1 + .../test_modeling_v2_m3_silu_and_mul.py | 49 + .../attention/_m3_attn_runtime_tp2.py | 275 + .../attention/_m3_nvfp4_cache.py | 108 + ...est_modeling_v2_m3_block_table_contract.py | 106 + ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 143 + .../test_modeling_v2_m3_draft_attention.py | 148 + ...ing_v2_m3_draft_attention_paged_context.py | 163 + .../test_modeling_v2_m3_drafter_fullctx.py | 271 + ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 406 ++ ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 49 + .../test_modeling_v2_m3_fp8_saturate.py | 110 + .../test_modeling_v2_m3_fullctx_decode.py | 215 + .../test_modeling_v2_m3_identity_selection.py | 102 + .../test_modeling_v2_m3_index_score.py | 107 + .../test_modeling_v2_m3_msa_attention.py | 224 + .../test_modeling_v2_m3_n_valid_contract.py | 138 + .../attention/test_modeling_v2_m3_producer.py | 218 + .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 81 + .../test_modeling_v2_m3_qkv_insert.py | 80 + .../test_modeling_v2_m3_select_blocks.py | 45 + .../test_modeling_v2_m3_select_long.py | 179 + .../test_modeling_v2_m3_select_long_dsl.py | 191 + .../test_modeling_v2_m3_sparse_decode.py | 451 ++ .../test_modeling_v2_m3_sparse_decode_lc.py | 490 ++ .../test_modeling_v2_m3_tokmap_guard.py | 226 + .../test_modeling_v2_msa_kv_update.py | 128 + .../comm/_m3_allreduce_op_matrix.py | 440 ++ .../comm/_m3_mc_allreduce_norm_op_matrix.py | 755 ++ .../comm/_m3_moe_push_op_matrix.py | 411 ++ .../comm/_rank_job.py | 87 + .../comm/test_modeling_v2_m3_allreduce.py | 39 + ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 21 + ...deling_v2_m3_mc_finalize_allreduce_norm.py | 33 + .../test_modeling_v2_m3_moe_push_op_matrix.py | 25 + .../test_modeling_v2_m3_embedding.py | 37 + .../gemm/_m3_mxfp8_operands.py | 30 + .../gemm/test_modeling_v2_m3_bf16_linear.py | 31 + .../gemm/test_modeling_v2_m3_dense_seam.py | 142 + .../gemm/test_modeling_v2_m3_dense_swiglu.py | 142 + .../gemm/test_modeling_v2_m3_draft_lowm.py | 183 + .../gemm/test_modeling_v2_m3_fp32_linear.py | 32 + .../gemm/test_modeling_v2_m3_mx_gemv.py | 87 + .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 63 + .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 38 + .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 46 + .../gemm/test_modeling_v2_m3_pinned_gemm.py | 49 + .../gemm/test_modeling_v2_m3_ugpu_head.py | 116 + .../gemm/test_r1b1_gemv_order_gate_a.py | 82 + .../moe/_m3_moe_operands.py | 209 + .../moe/_m3_router.py | 86 + .../test_modeling_v2_m3_fp32_router_rows.py | 36 + .../moe/test_modeling_v2_m3_moe_engine.py | 170 + .../moe/test_modeling_v2_m3_moe_epoch_wrap.py | 59 + .../test_modeling_v2_m3_moe_finalize_rows.py | 63 + .../moe/test_modeling_v2_m3_moe_flat_gate.py | 97 + .../moe/test_modeling_v2_m3_moe_sh_take.py | 73 + ...test_modeling_v2_m3_moe_triton_finalize.py | 74 + .../moe/test_modeling_v2_m3_router_tc.py | 58 + ...test_modeling_v2_m3_embedding_firstnorm.py | 103 + .../spec/_m3_spec_inputs.py | 106 + .../spec/test_modeling_v2_m3_capture_norm.py | 57 + .../test_modeling_v2_m3_capture_workspaces.py | 56 + .../spec/test_modeling_v2_m3_chunk_norm.py | 50 + ...test_modeling_v2_m3_draft_ar_push_words.py | 107 + .../spec/test_modeling_v2_m3_draft_input.py | 98 + .../test_modeling_v2_m3_draft_push_take.py | 43 + .../spec/test_modeling_v2_m3_draft_silu.py | 88 + .../spec/test_modeling_v2_m3_draft_step.py | 814 +++ .../spec/test_modeling_v2_m3_fold_tail.py | 295 + .../spec/test_modeling_v2_m3_force_al.py | 63 + .../spec/test_modeling_v2_m3_force_vec.py | 47 + .../spec/test_modeling_v2_m3_force_vec_req.py | 48 + .../spec/test_modeling_v2_m3_lowm_splitk.py | 65 + .../test_modeling_v2_m3_sampled_accept.py | 276 + .../spec/test_modeling_v2_m3_shard_accept.py | 127 + .../spec/test_modeling_v2_m3_slot_scatter.py | 57 + .../spec/test_modeling_v2_m3_spec_offsets.py | 138 + .../spec/test_modeling_v2_m3_step_gather.py | 80 + .../spec/test_modeling_v2_m3_step_prologue.py | 192 + .../test_modeling_v2_m3_wide_draft_logits.py | 132 + .../test_modeling_v2_m3_capture_boundary.py | 163 + .../test_modeling_v2_m3_decode_dense_route.py | 85 + .../test_modeling_v2_m3_dispatch.py | 771 ++ ...odeling_v2_m3_draft_attention_bcg_route.py | 69 + ...st_modeling_v2_m3_drafter_fullctx_cells.py | 34 + .../test_modeling_v2_m3_rank_agreement.py | 179 + .../test_modeling_v2_m3_wide_lc_cap.py | 46 + .../test_modeling_v2_m3_silu_and_mul.py | 49 + .../attention/_m3_attn_runtime_tp4.py | 274 + .../attention/_m3_nvfp4_cache.py | 104 + .../test_modeling_v2_m3_dense_decode.py | 153 + ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 180 + .../test_modeling_v2_m3_draft_attention.py | 138 + .../test_modeling_v2_m3_drafter_fullctx.py | 289 + ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 371 + ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 49 + .../test_modeling_v2_m3_fp8_saturate.py | 102 + .../test_modeling_v2_m3_fullctx_decode.py | 312 + .../test_modeling_v2_m3_index_score.py | 133 + .../test_modeling_v2_m3_index_score_t2sc.py | 112 + .../test_modeling_v2_m3_msa_attention.py | 214 + .../test_modeling_v2_m3_n_valid_contract.py | 138 + .../attention/test_modeling_v2_m3_producer.py | 209 + .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 81 + .../test_modeling_v2_m3_qkv_insert.py | 80 + .../test_modeling_v2_m3_select_blocks.py | 45 + .../test_modeling_v2_m3_select_long.py | 183 + .../test_modeling_v2_m3_select_long_dsl.py | 191 + .../test_modeling_v2_m3_sparse_decode.py | 442 ++ .../test_modeling_v2_m3_sparse_decode_lc.py | 526 ++ .../test_modeling_v2_m3_tokmap_guard.py | 226 + .../test_modeling_v2_msa_kv_update.py | 133 + .../comm/_m3_allreduce_op_matrix.py | 440 ++ .../comm/_m3_capture_in_boundary_op_matrix.py | 307 + .../comm/_m3_draft_ar_norm_op_matrix.py | 533 ++ .../comm/_m3_mc_allreduce_norm_op_matrix.py | 665 ++ .../comm/_rank_job.py | 87 + .../comm/test_modeling_v2_m3_allreduce.py | 39 + ...test_modeling_v2_m3_capture_in_boundary.py | 22 + .../comm/test_modeling_v2_m3_draft_ar_norm.py | 32 + ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 20 + ...deling_v2_m3_mc_finalize_allreduce_norm.py | 48 + .../test_modeling_v2_m3_embedding.py | 37 + .../gemm/_m3_mxfp8_operands.py | 30 + .../gemm/test_modeling_v2_m3_bf16_linear.py | 31 + .../gemm/test_modeling_v2_m3_dense_swiglu.py | 142 + .../gemm/test_modeling_v2_m3_draft_lowm.py | 182 + .../gemm/test_modeling_v2_m3_fp32_linear.py | 32 + .../gemm/test_modeling_v2_m3_mx_gemv.py | 176 + .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 68 + .../gemm/test_modeling_v2_m3_mx_gemv_order.py | 78 + .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 38 + .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 46 + .../test_modeling_v2_m3_mxfp8_pinned_gemm.py | 82 + ...st_modeling_v2_m3_mxfp8_pinned_quantize.py | 48 + .../gemm/test_modeling_v2_m3_target_head.py | 253 + .../moe/_m3_moe_operands.py | 170 + .../moe/_m3_router.py | 39 + .../moe/test_modeling_v2_m3_moe16_engine.py | 63 + .../moe/test_modeling_v2_m3_moe_engine.py | 139 + .../test_modeling_v2_m3_moe_flat_gate_tp4.py | 74 + .../moe/test_modeling_v2_m3_router_tc.py | 60 + ...test_modeling_v2_m3_embedding_firstnorm.py | 76 + .../spec/_m3_spec_inputs.py | 106 + .../spec/test_modeling_v2_m3_capture_norm.py | 62 + .../test_modeling_v2_m3_capture_normed.py | 70 + .../spec/test_modeling_v2_m3_chunk_norm.py | 55 + .../spec/test_modeling_v2_m3_draft_input.py | 103 + .../test_modeling_v2_m3_draft_push_take.py | 42 + .../spec/test_modeling_v2_m3_draft_silu.py | 98 + .../spec/test_modeling_v2_m3_draft_step.py | 694 ++ ...t_modeling_v2_m3_draft_step_fixed_order.py | 148 + .../test_modeling_v2_m3_exact_numerics.py | 289 + .../spec/test_modeling_v2_m3_fold_tail.py | 302 + .../spec/test_modeling_v2_m3_force_vec.py | 47 + .../spec/test_modeling_v2_m3_force_vec_req.py | 93 + .../spec/test_modeling_v2_m3_lowm_splitk.py | 63 + .../test_modeling_v2_m3_sampled_accept.py | 297 + .../test_modeling_v2_m3_sampled_accept_1r.py | 186 + .../spec/test_modeling_v2_m3_shard_accept.py | 125 + .../spec/test_modeling_v2_m3_slot_scatter.py | 62 + .../spec/test_modeling_v2_m3_spec_offsets.py | 138 + .../spec/test_modeling_v2_m3_step_gather.py | 85 + .../spec/test_modeling_v2_m3_step_prologue.py | 192 + .../test_modeling_v2_m3_capture_boundary.py | 154 + .../test_modeling_v2_m3_dispatch.py | 286 + .../test_modeling_v2_m3_prebuild.py | 314 + .../test_modeling_v2_m3_rank_agreement.py | 181 + .../test_modeling_v2_m3_silu_and_mul.py | 49 + .../attention/_m3_attn_runtime_tp8.py | 274 + .../attention/_m3_nvfp4_cache.py | 105 + ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 144 + .../test_modeling_v2_m3_draft_attention.py | 148 + .../test_modeling_v2_m3_drafter_fullctx.py | 271 + ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 371 + ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 49 + .../test_modeling_v2_m3_fp8_saturate.py | 102 + .../test_modeling_v2_m3_fullctx_decode.py | 310 + .../test_modeling_v2_m3_identity_selection.py | 103 + .../test_modeling_v2_m3_index_score.py | 111 + .../test_modeling_v2_m3_msa_attention.py | 224 + .../test_modeling_v2_m3_n_valid_contract.py | 138 + .../attention/test_modeling_v2_m3_producer.py | 209 + .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 81 + .../test_modeling_v2_m3_qkv_insert.py | 80 + .../test_modeling_v2_m3_select_blocks.py | 45 + .../test_modeling_v2_m3_select_long.py | 179 + .../test_modeling_v2_m3_select_long_dsl.py | 191 + .../test_modeling_v2_m3_sparse_decode.py | 626 ++ .../test_modeling_v2_m3_sparse_decode_lc8.py | 360 + .../test_modeling_v2_m3_tokmap_guard.py | 226 + .../test_modeling_v2_msa_kv_update.py | 128 + .../comm/_m3_allreduce_op_matrix.py | 454 ++ .../comm/_m3_mc_allreduce_norm_op_matrix.py | 788 ++ .../comm/_m3_moe_push_op_matrix.py | 403 + .../comm/_rank_job.py | 111 + .../comm/test_modeling_v2_m3_allreduce.py | 42 + ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 62 + .../test_modeling_v2_m3_embedding.py | 37 + .../gemm/_m3_mxfp8_operands.py | 30 + .../gemm/test_modeling_v2_m3_bf16_linear.py | 31 + .../gemm/test_modeling_v2_m3_dense_swiglu.py | 142 + .../gemm/test_modeling_v2_m3_draft_lowm.py | 177 + .../gemm/test_modeling_v2_m3_fp32_linear.py | 32 + .../gemm/test_modeling_v2_m3_mx_gemv.py | 87 + .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 63 + .../test_modeling_v2_m3_mx_gemv_producer.py | 166 + .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 38 + .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 46 + .../gemm/test_modeling_v2_m3_pinned_gemm.py | 52 + .../test_modeling_v2_m3_swiglu_oai_bf16x2.py | 149 + .../gemm/test_r1b1_gemv_order_gate_a.py | 85 + .../moe/_m3_moe_operands.py | 172 + .../moe/_m3_router.py | 86 + .../moe/test_modeling_v2_m3_moe_engine.py | 139 + .../moe/test_modeling_v2_m3_router_tc.py | 58 + ...test_modeling_v2_m3_embedding_firstnorm.py | 102 + .../spec/_m3_spec_inputs.py | 106 + .../spec/test_modeling_v2_m3_capture_norm.py | 57 + .../test_modeling_v2_m3_capture_normed.py | 71 + .../test_modeling_v2_m3_capture_target.py | 82 + .../test_modeling_v2_m3_capture_workspaces.py | 56 + .../spec/test_modeling_v2_m3_chunk_norm.py | 50 + .../spec/test_modeling_v2_m3_di_fc_flag.py | 91 + .../spec/test_modeling_v2_m3_draft_input.py | 98 + .../test_modeling_v2_m3_draft_push_take.py | 43 + .../spec/test_modeling_v2_m3_draft_silu.py | 88 + .../spec/test_modeling_v2_m3_draft_step.py | 747 ++ ...t_modeling_v2_m3_draft_step_fixed_order.py | 150 + .../spec/test_modeling_v2_m3_fold_tail.py | 376 + .../spec/test_modeling_v2_m3_force_al.py | 63 + .../spec/test_modeling_v2_m3_force_vec.py | 47 + .../spec/test_modeling_v2_m3_force_vec_req.py | 48 + .../spec/test_modeling_v2_m3_lowm_splitk.py | 98 + .../test_modeling_v2_m3_sampled_accept.py | 267 + .../test_modeling_v2_m3_sampled_accept_1r.py | 186 + .../spec/test_modeling_v2_m3_shard_accept.py | 120 + .../spec/test_modeling_v2_m3_slot_scatter.py | 57 + .../spec/test_modeling_v2_m3_spec_offsets.py | 138 + .../spec/test_modeling_v2_m3_step_gather.py | 80 + .../spec/test_modeling_v2_m3_step_prologue.py | 228 + .../test_modeling_v2_m3_capture_boundary.py | 103 + .../test_modeling_v2_m3_dispatch.py | 479 ++ .../test_modeling_v2_m3_rank_agreement.py | 190 + .../modeling_v2/test_modeling_v2_claims.py | 9 + .../modeling_v2/test_modeling_v2_routing.py | 98 +- .../engine/test_decoder_runtime_hooks.py | 18 +- .../pyexecutor/test_cuda_graph_runner.py | 44 + .../pyexecutor/test_first_content_emission.py | 113 + .../_torch/pyexecutor/test_h2d_cache.py | 102 + .../pyexecutor/test_h2d_changed_span.py | 157 + .../test_kv_connector_executor_lifetime.py | 1 + .../pyexecutor/test_kv_lens_update_once.py | 92 + .../_torch/pyexecutor/test_runtime_opt_ins.py | 58 + .../_torch/pyexecutor/test_step_inputs.py | 340 + .../test_one_model_sampling_rejection.py | 29 + .../hw_agnostic/test_ragged_spec_sampler.py | 2 +- .../speculative/test_rng_upload_pack.py | 93 + .../test_spec_sampler_step_fold.py | 116 + 726 files changed, 168457 insertions(+), 216 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DraftStepOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SelectLongOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DraftStepOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4McAllreduceNormOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SelectLongOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp create mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/long_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_attn_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_direct.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe_finalize.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_report.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_spec_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept_ws_words.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_qk_norm_rope_fp8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/long_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_attn_lc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score_t2sc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe16_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_report.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept_ws_words.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/weights.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/long_select.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_attn_lc8.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fc_wait.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_report.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_spec_tail.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept_ws_words.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/weights.py create mode 100644 tensorrt_llm/_torch/custom_ops/mxfp8_cute_pinned_tactics.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py create mode 100644 tensorrt_llm/_torch/late_compile.py create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_msa_skip_cache_routing.py create mode 100644 tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py create mode 100644 tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py create mode 100644 tests/unittest/_torch/kernels/parallel/test_mxfp8_cute_pinned_tactics.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_attn_runtime_tp2.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_nvfp4_cache.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_identity_selection.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_allreduce.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_moe_push_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_moe_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_push_take.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_attn_runtime_tp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_nvfp4_cache.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_saturate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_allreduce.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_order.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_moe_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_normed.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept_1r.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_attn_runtime_tp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_nvfp4_cache.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_saturate.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_identity_selection.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_rank_job.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_allreduce.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_moe_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_push_take.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py create mode 100644 tests/unittest/_torch/pyexecutor/test_first_content_emission.py create mode 100644 tests/unittest/_torch/pyexecutor/test_h2d_changed_span.py create mode 100644 tests/unittest/_torch/pyexecutor/test_kv_lens_update_once.py create mode 100644 tests/unittest/_torch/pyexecutor/test_runtime_opt_ins.py create mode 100644 tests/unittest/_torch/speculative/test_rng_upload_pack.py diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index 6a64c469411d..b79d6abee1bc 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -91,6 +91,9 @@ list(FILTER SRC_CU EXCLUDE REGEX "marlin/.*") # architecture-scoped OBJECT library below. list(FILTER SRC_CU EXCLUDE REGEX "minimaxM3/(fusedDecodeAttn|attnMega|attnC1)\\.cu$") +# The MiniMax-M3 sm_107 decode kernels are built as their own sm_107-only OBJECT +# library below. +list(FILTER SRC_CU EXCLUDE REGEX "minimaxM3Sm107Tp[248]/.*") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") @@ -141,11 +144,38 @@ set_property(TARGET minimax_m3_attn_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET minimax_m3_attn_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) set_cuda_architectures(minimax_m3_attn_src 100 103) +# MiniMax-M3 sm_107 decode kernels (minimaxM3Sm107Tp{2,4,8}/): written for +# sm_107 only (PDL griddepcontrol, sm_10x-only PTX), so they are compiled for +# the sm_10x entries of this build's architectures only (the same targets, e.g. +# the sm_100f family, that kernels_src compiles its sm_10x code for), not for +# the older ones. A build without an sm_10x architecture compiles them for +# sm_107a: their torch ops link against the host launchers, and only the sm_107 +# targets launch them. +file(GLOB MINIMAX_M3_SM107_SRC minimaxM3Sm107Tp2/*.cu minimaxM3Sm107Tp4/*.cu + minimaxM3Sm107Tp8/*.cu) +add_library(minimax_m3_sm107_src OBJECT ${MINIMAX_M3_SM107_SRC}) +set_property(TARGET minimax_m3_sm107_src PROPERTY POSITION_INDEPENDENT_CODE ON) +set_property(TARGET minimax_m3_sm107_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS + ON) +target_include_directories( + minimax_m3_sm107_src + PRIVATE + $ +) +target_link_libraries(minimax_m3_sm107_src PRIVATE trtllm_gen_fmha_interface) +set(MINIMAX_M3_SM107_ARCHS ${CMAKE_CUDA_ARCHITECTURES}) +list(FILTER MINIMAX_M3_SM107_ARCHS INCLUDE REGEX "^10[0-9]") +if(NOT MINIMAX_M3_SM107_ARCHS) + set(MINIMAX_M3_SM107_ARCHS 107a-real) +endif() +set_property(TARGET minimax_m3_sm107_src PROPERTY CUDA_ARCHITECTURES + ${MINIMAX_M3_SM107_ARCHS}) + add_library( kernels_src STATIC ${SRC_CPP} ${SRC_CU} ${MOE_KERNELS_SRC} $<$:$> - $) + $ $) set_property(TARGET kernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET kernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) target_include_directories( diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu index d12b765e5ecf..9ef178897f9f 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.cu @@ -520,36 +520,13 @@ void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, i launchGemv<1>(x, ldx, m, k, w, e, out, qOut, sfOut, sfScale, pdl, stream); } -void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out, - uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream) +int minimaxM3RouterTcCluster(int32_t k) { - int64_t const M = m, K = k; - TLLM_CHECK_WITH_INFO( - M >= 1 && M <= kTcMaxM && K % (kTcSlices * 64) == 0, "unsupported router shape M=%d K=%d", m, k); - int const ks = K / kTcSlices; - auto launch_tc = [&](auto kernel, int cl) -> cudaError_t - { - int const kc = (kTcSlices / cl) * ks; - size_t const smem = static_cast(kTcExperts) * (kc + 4) * 4 + static_cast(M) * (kc + 8) * 2; - cudaLaunchConfig_t cfg = {}; - cfg.gridDim = dim3(8 * cl); - cfg.blockDim = dim3(256); - cfg.dynamicSmemBytes = smem; - cfg.stream = stream; - cudaLaunchAttribute attr[2]; - attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - attr[1].id = cudaLaunchAttributeClusterDimension; - attr[1].val.clusterDim.x = cl; - attr[1].val.clusterDim.y = 1; - attr[1].val.clusterDim.z = 1; - cfg.attrs = attr; - cfg.numAttrs = 2; - return cudaLaunchKernelEx( - &cfg, kernel, x, ldx, w, out, static_cast(M), static_cast(K), qOut, sfOut, sfScale); - }; - // 16-CTA clusters (128 CTAs, one K slice each) when the device co-schedules all 8, else 8-CTA clusters. - static int const cluster = [&]() + // 16-CTA clusters (128 CTAs, one K slice each) when the device co-schedules all 8, else 8-CTA clusters; decided + // once per process (device-static, for the first K seen: MiniMax-M3 has one hidden size) and configures both + // kernels. + int const ks = k / kTcSlices; + static int const cluster = [ks]() { int device = 0, max_optin = 0; cudaGetDevice(&device); @@ -592,6 +569,38 @@ void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int } return clusters >= 8 ? 16 : 8; }(); + return cluster; +} + +void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out, + uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream) +{ + int64_t const M = m, K = k; + TLLM_CHECK_WITH_INFO( + M >= 1 && M <= kTcMaxM && K % (kTcSlices * 64) == 0, "unsupported router shape M=%d K=%d", m, k); + int const ks = K / kTcSlices; + auto launch_tc = [&](auto kernel, int cl) -> cudaError_t + { + int const kc = (kTcSlices / cl) * ks; + size_t const smem = static_cast(kTcExperts) * (kc + 4) * 4 + static_cast(M) * (kc + 8) * 2; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(8 * cl); + cfg.blockDim = dim3(256); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = cl; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + return cudaLaunchKernelEx( + &cfg, kernel, x, ldx, w, out, static_cast(M), static_cast(K), qOut, sfOut, sfScale); + }; + int const cluster = minimaxM3RouterTcCluster(k); cudaError_t const err = cluster == 16 ? launch_tc(m3_router_tc_kernel<16>, 16) : launch_tc(m3_router_tc_kernel<8>, 8); if (err != cudaSuccess) diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h index 6f4467883b24..6548cb8ca83d 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3/routerGemv.h @@ -41,6 +41,10 @@ void invokeMinimaxM3RouterGemv(__nv_bfloat16 const* x, int32_t ldx, int32_t m, i void invokeMinimaxM3RouterTc(__nv_bfloat16 const* x, int32_t ldx, int32_t m, int32_t k, float const* w, float* out, uint8_t* qOut, uint8_t* sfOut, float const* sfScale, bool pdl, cudaStream_t stream); +//! The cluster size `invokeMinimaxM3RouterTc` launches with (16 or 8): decided once per process from the device's +//! co-scheduling (device-static) and configures both kernels. SM107 callers use the router only when this is 16. +int minimaxM3RouterTcCluster(int32_t k); + } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu index 4069da14faf0..45526249947f 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu @@ -43,6 +43,9 @@ namespace specTail namespace { constexpr int kVec = 8; // bf16 per 16-byte access +// draft_input loads its constant norm weight row before the grid dependency wait (the same values, bit for bit). +// Kill: false (load the row after the wait, with the normalization). +constexpr bool kDraftInputPrewaitWeight = true; __device__ __forceinline__ float block_sum(float v, float* red) { @@ -72,6 +75,20 @@ __global__ void __launch_bounds__(kThreads) __shared__ float red[kThreads / 32]; int const row = blockIdx.x; bool const is_hidden = blockIdx.y == 1; + // The norm weight row is constant (no grid of the step writes it): load it before the grid dependency wait, so + // after the wait only the dependent loads (the id / hidden row) remain ahead of the reduction. The same values and + // arithmetic. + __nv_bfloat16 const* const w = is_hidden ? w_hidden : w_embed; + uint4 wv[kIters]; + if constexpr (kDraftInputPrewaitWeight) + { +#pragma unroll + for (int it = 0; it < kIters; ++it) + { + int const c = (it * kThreads + threadIdx.x) * kVec; + wv[it] = c < H ? *reinterpret_cast(w + c) : make_uint4(0u, 0u, 0u, 0u); + } + } #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("griddepcontrol.wait;" ::: "memory"); #endif @@ -112,7 +129,6 @@ __global__ void __launch_bounds__(kThreads) asm volatile("griddepcontrol.launch_dependents;"); #endif float const rs = rsqrtf(block_sum(ss, red) / H + (is_hidden ? eps_hidden : eps_embed)); - __nv_bfloat16 const* w = is_hidden ? w_hidden : w_embed; __nv_bfloat16* dst = out + static_cast(row) * 2 * H + (is_hidden ? H : 0); #pragma unroll for (int it = 0; it < kIters; ++it) @@ -120,8 +136,11 @@ __global__ void __launch_bounds__(kThreads) int const c = (it * kThreads + threadIdx.x) * kVec; if (c < H) { - uint4 const wr = *reinterpret_cast(w + c); - __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + if constexpr (!kDraftInputPrewaitWeight) + { + wv[it] = *reinterpret_cast(w + c); + } + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wv[it]); uint4 o; __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); #pragma unroll diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu new file mode 100644 index 000000000000..a66c20a0c867 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.cu @@ -0,0 +1,267 @@ +/* + * 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. + */ + +// The TP2 dense layers' QK norm + RoPE with FP8 output. The norm and RoPE are the generic fusedQKNormRopeKernel's at +// head_dim 128 without YaRN or mRoPE, expression for expression; the store rounds Q/K to BF16 and then to E4M3 +// (saturating), the two roundings of the stock BF16 kernel followed by its FP8 cast. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace denseQkNormRopeFp8 +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kVecSize = kElemsPerThread * sizeof(__nv_bfloat16) / 4; +constexpr int kBlockSize = 256; +static_assert(kVecSize == 2, "a lane loads 4 BF16 values as one uint2"); + +// tensorrt_llm::common::warpReduceSum: the same butterfly, so the same sum. +__device__ __forceinline__ float warpReduceSum(float val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + val = val + __shfl_xor_sync(0xffffffff, val, mask, 32); + } + return val; +} + +__device__ __forceinline__ void storeFp8ViaBf16(__nv_fp8_e4m3* out, long long offsetThread, float const (&elements)[4]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + float2 const rounded = __bfloat1622float2(__float22bfloat162_rn(make_float2(elements[i], elements[i + 1]))); + __nv_fp8x2_e4m3 const packed(rounded); + reinterpret_cast<__nv_fp8x2_storage_t*>(&out[offsetThread])[i / 2] = packed.__x; + } +} + +// The K (kv 0) or V (kv 1) row of one head of token tokenIdx in the HND paged pool, at the token's cache slot. +__device__ __forceinline__ __nv_fp8_e4m3* cacheRow(KvCacheView const& cache, int kv, int headIdx, int tokenIdx) +{ + long long const slot = static_cast(cache.outCacheLoc[tokenIdx]); + long long const page = slot / cache.tokensPerBlock; + long long const within = slot % cache.tokensPerBlock; + return cache.base + page * cache.pageStride + kv * cache.kvStride + headIdx * cache.headStride + + within * cache.tokenStride; +} + +template +__global__ void denseQkNormRopeFp8Kernel(__nv_bfloat16 const* qkvIn, __nv_fp8_e4m3* qkvOut, int numHeadsQ, + int numHeadsK, int numHeadsV, int rotaryDim, float eps, __nv_bfloat16 const* qWeight, __nv_bfloat16 const* kWeight, + float base, int const* positionIds, int numTokens, bool useGemma, KvCacheView cache) +{ + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId; + int const totalQkHeads = numHeadsQ + numHeadsK; + int const totalHeads = totalQkHeads + numHeadsV; + int const tokenIdx = globalWarpIdx / totalHeads; + int const localHeadIdx = globalWarpIdx % totalHeads; + if (tokenIdx >= numTokens) + { + return; + } + bool const isQ = localHeadIdx < numHeadsQ; + bool const isV = localHeadIdx >= totalQkHeads; + int headIdx; + int segStart; + if (isQ) + { + headIdx = localHeadIdx; + segStart = 0; + } + else if (!isV) + { + headIdx = localHeadIdx - numHeadsQ; + segStart = numHeadsQ * kHeadDim; + } + else + { + headIdx = localHeadIdx - totalQkHeads; + segStart = totalQkHeads * kHeadDim; + } + + using vec_T = uint2; + int const offsetThread + = tokenIdx * totalHeads * kHeadDim + segStart + headIdx * kHeadDim + laneId * kElemsPerThread; + float elements[kElemsPerThread]; + float sumOfSquares = 0.0f; + { + vec_T vec = *reinterpret_cast(&qkvIn[offsetThread]); +#pragma unroll + for (int i = 0; i < kVecSize; i++) + { + float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&vec) + i)); + sumOfSquares += vals.x * vals.x; + sumOfSquares += vals.y * vals.y; + elements[2 * i] = vals.x; + elements[2 * i + 1] = vals.y; + } + } + if (isV) + { + // V is BF16 already: its single E4M3 rounding is the stock cast's. + storeFp8ViaBf16(qkvOut, offsetThread, elements); + if (kWriteCache && tokenIdx < cache.numLiveTokens && cache.outCacheLoc[tokenIdx] >= 0) + { + storeFp8ViaBf16(cacheRow(cache, 1, headIdx, tokenIdx), laneId * kElemsPerThread, elements); + } + return; + } + + sumOfSquares = warpReduceSum(sumOfSquares); + float rms_rcp = rsqrtf(sumOfSquares / static_cast(kHeadDim) + eps); + for (int i = 0; i < kElemsPerThread; i++) + { + int dim = laneId * kElemsPerThread + i; + float weight = isQ ? __bfloat162float(qWeight[dim]) : __bfloat162float(kWeight[dim]); + elements[i] *= rms_rcp * (useGemma ? (1.0f + weight) : weight); + } + + float elements2[kElemsPerThread]; + float cos_vals[kElemsPerThread]; + float sin_vals[kElemsPerThread]; + float const neg2_log2base_over_rd = -2.0f * __log2f(base) / static_cast(rotaryDim); + int const rd_mask = rotaryDim - 1; + bool const rd_is_pow2 = ((rotaryDim & rd_mask) == 0); + float const pos_id = static_cast(positionIds[tokenIdx]); + if constexpr (kInterleave) + { + for (int i = 0; i < kElemsPerThread; i++) + { + if (i % 2 == 0) + { + elements2[i] = -elements[i + 1]; + } + else + { + elements2[i] = elements[i - 1]; + } + int dim_idx = laneId * kElemsPerThread + i; + int half_dim = dim_idx / 2; + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); + } + } + else + { + __syncwarp(); + int pairOffset = (rotaryDim / 2) / kElemsPerThread; + for (int i = 0; i < kElemsPerThread; i++) + { + elements2[i] = __shfl_xor_sync(0xffffffff, elements[i], pairOffset); + if (laneId < pairOffset) + { + elements2[i] = -elements2[i]; + } + int dim_idx = laneId * kElemsPerThread + i; + dim_idx = rd_is_pow2 ? ((dim_idx * 2) & rd_mask) : ((dim_idx * 2) % rotaryDim); + int half_dim = dim_idx / 2; + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); + } + __syncwarp(); + } + + // attention_factor is 1 without YaRN: (x * cos + x2 * sin) * 1 rounds as the generic kernel's product does. + float const attention_factor = 1.0f; + if (rotaryDim == kHeadDim) + { + for (int i = 0; i < kElemsPerThread; i++) + { + elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor; + } + } + else + { + for (int i = 0; i < kElemsPerThread; i++) + { + int dim_idx = laneId * kElemsPerThread + i; + if (dim_idx < rotaryDim) + { + elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor; + } + } + } + storeFp8ViaBf16(qkvOut, offsetThread, elements); + if constexpr (kWriteCache) + { + // Rows past the live count (padded graph rows) own no slot, nor does a negative (sentinel) slot. + if (!isQ && tokenIdx < cache.numLiveTokens && cache.outCacheLoc[tokenIdx] >= 0) + { + storeFp8ViaBf16(cacheRow(cache, 0, headIdx, tokenIdx), laneId * kElemsPerThread, elements); + } + // The compact query [tokens, q heads, 128] the decode reads, beside the packed row. + if (isQ && cache.qCompact != nullptr) + { + storeFp8ViaBf16(cache.qCompact, + (static_cast(tokenIdx) * numHeadsQ + headIdx) * kHeadDim + laneId * kElemsPerThread, + elements); + } + } +} + +} // namespace + +void launch(void const* qkvIn, void* qkvOut, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, int rotaryDim, + float eps, void const* qWeight, void const* kWeight, float base, bool interleave, int const* positionIds, + bool useGemma, KvCacheView const* cache, cudaStream_t stream) +{ + int const totalWarps = numTokens * (numHeadsQ + numHeadsK + numHeadsV); + int const warpsPerBlock = kBlockSize / 32; + dim3 const grid((totalWarps + warpsPerBlock - 1) / warpsPerBlock); + auto const* in = static_cast<__nv_bfloat16 const*>(qkvIn); + auto* out = static_cast<__nv_fp8_e4m3*>(qkvOut); + auto const* qw = static_cast<__nv_bfloat16 const*>(qWeight); + auto const* kw = static_cast<__nv_bfloat16 const*>(kWeight); + KvCacheView const view = cache != nullptr ? *cache : KvCacheView{}; + auto const run = [&](auto kernel) + { + kernel<<>>(in, out, numHeadsQ, numHeadsK, numHeadsV, rotaryDim, eps, qw, kw, base, + positionIds, numTokens, useGemma, view); + }; + if (cache != nullptr) + { + interleave ? run(denseQkNormRopeFp8Kernel) : run(denseQkNormRopeFp8Kernel); + } + else + { + interleave ? run(denseQkNormRopeFp8Kernel) : run(denseQkNormRopeFp8Kernel); + } +} + +} // namespace denseQkNormRopeFp8 +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h new file mode 100644 index 000000000000..ec94c1dbd550 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h @@ -0,0 +1,64 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ +namespace denseQkNormRopeFp8 +{ + +constexpr int kHeadDim = 128; + +//! A layer's FP8 K/V pool in the HND paged layout [pages, 2 (K, V), kv heads, tokens per block, 128] (element +//! strides), with the step's per-token cache slots (page * tokensPerBlock + within) for its first numLiveTokens rows. +struct KvCacheView +{ + __nv_fp8_e4m3* base = nullptr; + int const* outCacheLoc = nullptr; + long long pageStride = 0; + long long kvStride = 0; + long long headStride = 0; + long long tokenStride = 0; + int tokensPerBlock = 1; + int numLiveTokens = 0; + //! Optional: every row's Q also stored compact, [tokens, q heads, 128]. + __nv_fp8_e4m3* qCompact = nullptr; +}; + +//! The MiniMax-M3 TP2 dense layers' per-head QK RMSNorm + partial RoPE of a packed BF16 [tokens, (q + k + v) * 128] +//! row, written out of place as FP8 E4M3 with V copy-cast. Q and K are rounded to BF16 before E4M3, as the stock path +//! (the BF16 fused_qk_norm_rope output cast to FP8) rounds them. With cache, the K and V rows of the live tokens are +//! also stored into the paged pool at their slots (the bytes of the qkvOut rows). +void launch(void const* qkvIn, void* qkvOut, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, int rotaryDim, + float eps, void const* qWeight, void const* kWeight, float base, bool interleave, int const* positionIds, + bool useGemma, KvCacheView const* cache, cudaStream_t stream); + +} // namespace denseQkNormRopeFp8 +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.cu new file mode 100644 index 000000000000..67e054e74712 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.cu @@ -0,0 +1,620 @@ +/* + * 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/minimaxM3Sm107Tp2/draftStep.h" + +#include +#include +#include + +// One pass of the fused greedy Eagle3 draft loop of MiniMax-M3 (GB300, TP4, <= 8 requests per pass). +// +// One kernel replaces the draft lm_head GEMM, the logits all-gather / float cast / argmax of the in-tree greedy draft +// sampling and the draft loop's per-pass bookkeeping, at the in-tree precision: +// +// * The vocab-sharded BF16 lm_head weight is stored in mma.m16n8k16 A-fragment order (the same values): per 16-row +// tile and k16 step, 512 bytes; lane (g, j) holds 16 bytes = {row g k 2j..2j+1, row g+8 k 2j.., row g k 2j+8.., +// row g+8 k 2j+8..}. +// * Or, by default, in a lossless 12-bit coding of those fragments (75% of the bytes; the target head's coding, see +// targetHead.cu): per 16-row tile and unit of 4 k16 steps a lane reads three 16-byte words, the sign/mantissa bytes +// of steps 0-1, of steps 2-3, then 4 code words (exponent - base, 4 bits; 15 = an exception). The decoded values +// are the checkpoint's BF16 bits; an exception decodes to +0 and the segment's flush adds the exact product x * w +// from the tile's list (products of BF16 values are exact in FP32). +// * A persistent grid (one 16-warp CTA per SM) streams whole tiles over the full K with BF16 MMAs and FP32 +// accumulation; the BF16 head input rows sit in shared memory in B-fragment order, unchanged. +// * Each logit is rounded to BF16 (the in-tree logits are the BF16 GEMM output) and becomes a 64-bit key +// ordered(value) << 32 | ~global_index, with -0 taken as +0: the larger value wins and ties pick the lowest index, as +// torch.max / torch.argmax over the gathered logits do. Keys are max-reduced per row: warp shuffles, then shared +// memory, then atomicMax. +// * The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping +// (Lamport: 0 = empty, a key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the +// slots, and runs the pass's bookkeeping. Call n uses slot half n % 2. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace draftStep +{ +namespace +{ +constexpr int kUnroll = 8; // k16 chunks in flight per warp +constexpr int kChunkBytes = 512; // one tile's k16 step: 32 lanes x 16 bytes +constexpr long long kMaxSpins = 1LL << 31; // cross-rank key poll bound + +// Coded units (4 k16 steps, 3 x 16 bytes per lane) in flight per warp: 3 up to 4 rows, 2 at 8 rows. +template +constexpr int codedUnroll() +{ + return MP >= 8 ? 2 : 3; +} + +__device__ __forceinline__ void mma_bf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + 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};\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)); +} + +// BF16 pair p (values 2p, 2p + 1) of a lane's 8 coded values: sign/mantissa bytes sm (4 per word), codes word cw. +__device__ __forceinline__ uint32_t coded_decode_pair(uint32_t sm_word, int byte0, uint32_t cw, int p, uint32_t base) +{ + uint32_t const b = __byte_perm(sm_word, 0u, byte0 == 0 ? 0x4140 : 0x4342); // {b0, 0, b1, 0} + uint32_t const q = (cw >> (4 * p)) & 0x000f000fu; + uint32_t const e = q * 128u + base; // exponent field of both halves + // Code 15 (an exception, sign/mantissa byte 0) decodes to +0. + return ((b & 0x007f007fu) | ((b << 8) & 0x80008000u) | e) & __vcmpne2(q, 0x000f000fu); +} + +// One coded unit (4 k16 steps) against the BF16 B fragments: steps alternate between the two accumulators, as the +// k16 steps of the plain BF16 head do. +template +__device__ __forceinline__ void chunk_mma_coded(uint4 const& sm01, uint4 const& sm23, uint4 const& cw, uint32_t base, + uint2 const* xs, int unit, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint32_t const smw[8] = {sm01.x, sm01.y, sm01.z, sm01.w, sm23.x, sm23.y, sm23.z, sm23.w}; + uint32_t const cws[4] = {cw.x, cw.y, cw.z, cw.w}; +#pragma unroll + for (int ci = 0; ci < 4; ++ci) + { + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[((unit * 4 + ci) * 4 + j) * MP + g]; + uint32_t const w0 = coded_decode_pair(smw[2 * ci], 0, cws[ci], 0, base); + uint32_t const w1 = coded_decode_pair(smw[2 * ci], 2, cws[ci], 1, base); + uint32_t const w2 = coded_decode_pair(smw[2 * ci + 1], 0, cws[ci], 2, base); + uint32_t const w3 = coded_decode_pair(smw[2 * ci + 1], 2, cws[ci], 3, base); + if (ci & 1) + mma_bf16(db, w0, w1, w2, w3, b.x, b.y); + else + mma_bf16(da, w0, w1, w2, w3, b.x, b.y); + } +} + +// Exceptions of `tile` with k in [kb, ke): v += x * w (exact) for the owner lanes (row g or g + 8; tokens 2j, +// 2j + 1 < M), x from the staged B fragments. Warp-uniform (shuffles). +template +__device__ __forceinline__ void coded_exceptions( + CodedArgs const& cd, int tile, int kb, int ke, uint2 const* xs, int M, int lane, float (&v)[4]) +{ + int const g = lane >> 2, j = lane & 3; + int const e0 = __ldg(cd.exc_off + tile), e1 = __ldg(cd.exc_off + tile + 1); + for (int b0 = e0; b0 < e1; b0 += 32) + { + uint32_t pos = 0xffffffffu; + float2 wv = make_float2(0.f, 0.f); + if (b0 + lane < e1) + { + pos = __ldg(cd.exc_pos + b0 + lane); + wv = __ldg(cd.exc_w + b0 + lane); + } + int const cnt = min(32, e1 - b0); + for (int q = 0; q < cnt; ++q) + { + uint32_t const pq = __shfl_sync(0xffffffffu, pos, q); + float const w0 = __shfl_sync(0xffffffffu, wv.x, q); + int const r = static_cast(pq >> 16), k = static_cast(pq & 0xffffu); + if (k >= kb && k < ke && (r & 7) == g) + { + int const st = k >> 4, rr = k & 15; + bool const hi = r >= 8; +#pragma unroll + for (int q2 = 0; q2 < 2; ++q2) + { + int const n = 2 * j + q2; + if (n < M) + { + uint2 const bw = xs[(st * 4 + ((rr & 7) >> 1)) * MP + n]; + uint32_t const word = rr >= 8 ? bw.y : bw.x; + float const xv = __uint_as_float((rr & 1) ? (word & 0xffff0000u) : (word << 16)); + float const d = __fmul_rn(xv, w0); + if (hi) + v[2 + q2] += d; + else + v[q2] += d; + } + } + } + __syncwarp(); + } + } +} + +__device__ __forceinline__ uint4 ldg_stream(uint4 const* p) +{ + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.u32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +// The argmax value of a logit: its BF16 rounding (the in-tree logits' dtype), -0 as +0. +__device__ __forceinline__ float key_value(float v) +{ + float const r = __bfloat162float(__float2bfloat16(v)); + return r == 0.f ? 0.f : r; +} + +__device__ __forceinline__ unsigned long long make_key(float v, uint32_t index) +{ + uint32_t u = __float_as_uint(v); + u = (u & 0x80000000u) ? ~u : (u | 0x80000000u); + return (static_cast(u) << 32) | static_cast(0xFFFFFFFFu - index); +} + +__device__ __forceinline__ unsigned long long shfl_max(unsigned long long k, int o) +{ + unsigned long long const ok = __shfl_xor_sync(0xffffffffu, k, o); + return ok > k ? ok : k; +} + +// One k16 step of a tile: the lane's A fragment w against the activation B fragment of step `chunk` (rows g < MP). +// Even and odd steps accumulate separately (two independent MMA chains); the tile's sum is da + db. +template +__device__ __forceinline__ void chunk_mma( + uint4 const& w, uint2 const* xs, int chunk, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[(chunk * 4 + j) * MP + g]; + if (chunk & 1) + mma_bf16(db, w.x, w.y, w.z, w.w, b.x, b.y); + else + mma_bf16(da, w.x, w.y, w.z, w.w, b.x, b.y); +} + +// Add a warp's segment of a tile into the CTA's accumulator (with a coded head, the segment's exceptions first). +template +__device__ __forceinline__ void flush_tile( + float4* acc, float (&da)[4], float (&db)[4], Params const& p, int tile, int kb, int ke, uint2 const* xs, int lane) +{ + float v[4] = {da[0] + db[0], da[1] + db[1], da[2] + db[2], da[3] + db[3]}; + if constexpr (CODED) + coded_exceptions(p.cd, tile, kb, ke, xs, p.M, lane, v); + float* a = reinterpret_cast(acc); + atomicAdd(a + 0, v[0]); + atomicAdd(a + 1, v[1]); + atomicAdd(a + 2, v[2]); + atomicAdd(a + 3, v[3]); +#pragma unroll + for (int e = 0; e < 4; ++e) + { + da[e] = 0.f; + db[e] = 0.f; + } +} + +// Balanced streaming GEMV over the fragment-ordered weight, persistent grid of 16-warp CTAs. CTA c owns tiles +// [t_begin, t_end); its (tile, chunk) space is split into 16 equal contiguous ranges, one per warp, so every warp +// streams the same number of bytes. A tile straddles at most two warps; their partial sums are added into shared +// memory (two addends: the order does not change the sum). MP = rows padded to 1 / 2 / 4 / 8. A chunk is one k16 +// step of the BF16 head, or one coded unit (4 k16 steps) of the coded head (CODED: no shared-memory prologue). +template +__global__ void __launch_bounds__(kWarps * 32, 1) + draft_step_kernel(Params const p, int kchunks, int pf_chunks, int pf_off) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(16) uint2 xs[]; // [K/16 steps][4 (j)][MP (g)], then the weight prologue + __shared__ float4 acc_s[kMaxTilesPerCta + 1][32]; // + 1: an equal-chunk span can touch one more tile + __shared__ unsigned long long red[kWarps][kMaxTokens]; + __shared__ __align__(8) unsigned long long pf_bar[kWarps]; + __shared__ bool s_last; + __shared__ int s_tok[kMaxTokens]; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const g = lane >> 2; + int const j = lane & 3; + constexpr int kChunkK = CODED ? 64 : 16; // k per chunk + constexpr int kWords = CODED ? 3 : 1; // 16-byte words per lane and chunk + constexpr int KU = CODED ? codedUnroll() : kUnroll; + int const K = kchunks * kChunkK; + int const M = p.M; + int const r_lo = p.tile_lo; + int const r_n = (p.tile_hi > 0 ? p.tile_hi : p.n_tiles) - r_lo; + bool const eq = p.bnd != nullptr; + // Whole tiles: CTA c owns tiles [t_begin, t_end). Equal chunks: CTA c streams chunks [c0, c1) of the range, which + // start off0 chunks into tile t_begin (its first and last tiles may be shared with the neighbouring CTAs). + long long const r_chunks = static_cast(r_n) * kchunks; + long long const c0 = eq ? (static_cast(blockIdx.x) * r_chunks) / gridDim.x : 0; + long long const c1 = eq ? (static_cast(blockIdx.x + 1) * r_chunks) / gridDim.x : 0; + int const t_begin = r_lo + + (eq ? static_cast(c0 / kchunks) + : static_cast((static_cast(blockIdx.x) * r_n) / gridDim.x)); + int const t_end = r_lo + + (eq ? static_cast((c1 + kchunks - 1) / kchunks) + : static_cast((static_cast(blockIdx.x + 1) * r_n) / gridDim.x)); + int const n_local = t_end - t_begin; + int const off0 = eq ? static_cast(c0 - static_cast(t_begin - r_lo) * kchunks) : 0; + int const span = eq ? static_cast(c1 - c0) : n_local * kchunks; + int const q_begin = off0 + warp * span / kWarps; + int const q_end = off0 + (warp + 1) * span / kWarps; + // Equal chunks: whether the first / last local tile is shared with the previous / next CTA of the launch. + bool const shared_first = eq && off0 != 0; + bool const shared_last = eq && ((c1 % kchunks) != 0); + long long const base = static_cast(t_begin) * kchunks; // first chunk of this CTA + uint4 const* const wf = p.wf; + + // Weights are static: start this warp's stream before waiting on the producer of x. The warp's first pf_chunks + // chunks come in by one bulk copy into shared memory (dynamic offset pf_off), which fills the grid-dependency + // wait and the staging; the register ring starts after them. + int const n_pf = CODED ? 0 : min(pf_chunks, q_end - q_begin); + uint4* const wpf = reinterpret_cast(reinterpret_cast(xs) + pf_off) + warp * pf_chunks * 32; + if (n_pf > 0 && lane == 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + uint32_t const bytes = static_cast(n_pf) * kChunkBytes; + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + static_cast(__cvta_generic_to_shared(wpf))), + "l"(wf + (base + q_begin) * 32), "r"(bytes), "r"(bar) + : "memory"); + } + __syncwarp(); + int const q_reg = q_begin + n_pf; + uint4 buf[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q_reg + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = ldg_stream(wf + ((base + q_reg + u) * kWords + l) * 32 + lane); + } + } + for (int i = threadIdx.x; i < n_local * 32; i += blockDim.x) + acc_s[i / 32][i % 32] = make_float4(0.f, 0.f, 0.f, 0.f); + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // Stage the head input rows in B-fragment order: item (k16 step S, row gg) = two 16-byte loads of x, four + // fragment words {x[2jj], x[2jj + 1]}, {x[2jj + 8], x[2jj + 9]} (BF16 pairs, moved unchanged). Rows >= M are 0. + constexpr int kSU = MP >= 4 ? 4 : 2; + int const n_items = (K / 16) * MP; + for (int i0 = threadIdx.x; i0 < n_items; i0 += kSU * blockDim.x) + { + uint4 xv[kSU][2]; +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + int const gg = i % MP; + int const S = i / MP; + xv[u][0] = make_uint4(0u, 0u, 0u, 0u); + xv[u][1] = make_uint4(0u, 0u, 0u, 0u); + if (i < n_items && gg < M) + { + long long const xr = p.x_rows != nullptr ? p.x_rows[gg] : gg; + uint4 const* xp = reinterpret_cast(p.x + xr * K + 16 * S); + xv[u][0] = xp[0]; + xv[u][1] = xp[1]; + } + } +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + if (i >= n_items) + continue; + int const gg = i % MP; + int const S = i / MP; + uint32_t const lo[4] = {xv[u][0].x, xv[u][0].y, xv[u][0].z, xv[u][0].w}; + uint32_t const hi[4] = {xv[u][1].x, xv[u][1].y, xv[u][1].z, xv[u][1].w}; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) + xs[(S * 4 + jj) * MP + gg] = make_uint2(lo[jj], hi[jj]); + } + } + __syncthreads(); + + bool const b_active = g < MP; + float da[4] = {0.f, 0.f, 0.f, 0.f}; + float db[4] = {0.f, 0.f, 0.f, 0.f}; + int tl = q_begin / kchunks; // local tile of the current accumulator + int c = q_begin - tl * kchunks; // chunk within that tile + int cs = c; // first chunk of the warp's segment of that tile + if (n_pf > 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + asm volatile( + "{\n .reg .pred p;\n PF_WAIT%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n" + " @!p bra PF_WAIT%=;\n}\n" ::"r"(bar) + : "memory"); + for (int q = q_begin; q < q_reg; ++q) + { + chunk_mma(wpf[(q - q_begin) * 32 + lane], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane); + c = 0; + cs = 0; + ++tl; + } + } + } + for (int q = q_reg; q < q_end; q += KU) + { + uint4 nxt[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + KU + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + nxt[u][l] = ldg_stream(wf + ((base + q + KU + u) * kWords + l) * 32 + lane); + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + u < q_end) + { + if constexpr (CODED) + chunk_mma_coded(buf[u][0], buf[u][kWords > 1 ? 1 : 0], buf[u][kWords > 2 ? 2 : 0], p.cd.base, + xs, c, j, g, b_active, da, db); + else + chunk_mma(buf[u][0], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane); + c = 0; + cs = 0; + ++tl; + } + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = nxt[u][l]; + } + if (c != 0) + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, c * kChunkK, xs, lane); + __syncthreads(); + + // Accumulator lane (g, j) of a tile: rows g / g + 8 (x, y / z, w) of tokens 2j / 2j + 1 (x, z / y, w). + unsigned long long best0 = 0ull, best1 = 0ull; // tokens 2j, 2j + 1 + for (int t = warp; t < n_local; t += kWarps) + { + float4 v = acc_s[t][lane]; + if ((t == 0 && shared_first) || (t == n_local - 1 && shared_last)) + { + // A tile shared with a neighbouring CTA: add this CTA's partial into bnd; the second CTA to arrive reads + // the sum (two addends), keys the tile and restores bnd and its counter to zero for the next call. + int const tile = t_begin + t; + float* const bp = reinterpret_cast(p.bnd + static_cast(tile) * 32 + lane); + atomicAdd(bp + 0, v.x); + atomicAdd(bp + 1, v.y); + atomicAdd(bp + 2, v.z); + atomicAdd(bp + 3, v.w); + __threadfence(); + __syncwarp(); + unsigned int prev = 0; + if (lane == 0) + prev = atomicAdd(p.bnd_cnt + tile, 1u); + prev = __shfl_sync(0xffffffffu, prev, 0); + if (prev == 0) + continue; // the neighbour keys it + __threadfence(); + v = __ldcg(p.bnd + static_cast(tile) * 32 + lane); + p.bnd[static_cast(tile) * 32 + lane] = make_float4(0.f, 0.f, 0.f, 0.f); + if (lane == 0) + p.bnd_cnt[tile] = 0u; + } + int const r0 = (t_begin + t) * 16 + g; + uint32_t const i0 = static_cast(p.vocab_offset + r0); + unsigned long long k0 = make_key(key_value(v.x), i0); + unsigned long long const k0b = make_key(key_value(v.z), i0 + 8); + unsigned long long k1 = make_key(key_value(v.y), i0); + unsigned long long const k1b = make_key(key_value(v.w), i0 + 8); + k0 = k0b > k0 ? k0b : k0; + k1 = k1b > k1 ? k1b : k1; + best0 = k0 > best0 ? k0 : best0; + best1 = k1 > best1 ? k1 : best1; + } +#pragma unroll + for (int o = 4; o < 32; o <<= 1) + { + best0 = shfl_max(best0, o); + best1 = shfl_max(best1, o); + } + asm volatile("griddepcontrol.launch_dependents;"); + if (lane < 4) + { + red[warp][2 * lane] = best0; + red[warp][2 * lane + 1] = best1; + } + __syncthreads(); + if (threadIdx.x < M) + { + unsigned long long b = 0ull; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + b = red[w][threadIdx.x] > b ? red[w][threadIdx.x] : b; + if (b != 0ull) + atomicMax(p.keys + threadIdx.x, b); + } + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned int const prev = atomicAdd(p.counters, 1u); + s_last = prev == (p.total_ctas != 0u ? p.total_ctas : gridDim.x) - 1u; + } + __syncthreads(); + if (!s_last) + return; + + __threadfence(); + unsigned int const call = *reinterpret_cast(p.counters + 1); + int const half = static_cast(call & 1u); + if (threadIdx.x < M) + { + int const t = threadIdx.x; + unsigned long long const mine = atomicExch(p.keys + t, 0ull); + unsigned long long* slot_mc = p.ws_mc + (half * kMaxTokens + t) * p.world; + unsigned long long* slot_uc = p.ws_uc + (half * kMaxTokens + t) * p.world; + slot_mc[p.rank] = mine; + unsigned long long best = 0ull; + for (int r = 0; r < p.world; ++r) + { + unsigned long long v; + long long spins = 0; + do + { + v = *reinterpret_cast(slot_uc + r); + if (++spins > kMaxSpins) + __trap(); // a missing peer key becomes a CUDA error, never a hang (the mc all-reduce's bound) + } while (v == 0ull); + best = v > best ? v : best; + slot_uc[r] = 0ull; + } + s_tok[t] = static_cast(0xFFFFFFFFu - static_cast(best & 0xFFFFFFFFull)); + if (p.out_keys != nullptr) + p.out_keys[t] = best; + } + __syncthreads(); + if (threadIdx.x == 0) + { + p.counters[0] = 0u; + p.counters[1] = call + 1u; + } + if (threadIdx.x < M) + { + StepArgs const& st = p.st; + int const b = threadIdx.x; + int const tok = s_tok[b]; + st.draft_out[b * st.draft_stride + st.step] = tok; + st.ids_next[b] = tok; + long long const r = p.x_rows != nullptr ? p.x_rows[b] : b; + st.pos_next[b] = st.pos_src[r] + 1; + if (st.step == 0 && st.num_acc != nullptr) + { + st.kv_lens[b] -= st.draft_len - st.num_acc[b]; + st.seq_lens[b] = 1; + } + else if (st.step > 0) + { + st.kv_lens[b] += 1; + } + } +#endif +} + +template +cudaError_t launch(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + auto kern = draft_step_kernel; + size_t const xs_bytes = static_cast(K / 16) * 4 * MP * sizeof(uint2); + // Shared-memory weight prologue: chunks per warp that fit next to the kernel's static + dynamic shared memory. + static int const sStaticSmem = [] + { + cudaFuncAttributes fa{}; + return cudaFuncGetAttributes(&fa, draft_step_kernel) == cudaSuccess + ? static_cast(fa.sharedSizeBytes) + : -1; + }(); + static int const sMaxSmemPerBlock = [] + { + int dev = 0, maxSmem = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + return maxSmem; + }(); + int pf = 0; + if (!CODED && sStaticSmem >= 0) + { + long long const room + = static_cast(sMaxSmemPerBlock) - sStaticSmem - static_cast(xs_bytes) - 1024; + long long const want = static_cast(kPrefetchKb) * 1024; + pf = static_cast(std::max(0ll, std::min(room, want)) / (kWarps * kChunkBytes)); + } + size_t const smem = xs_bytes + static_cast(pf) * kWarps * kChunkBytes; + cudaError_t err = cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + return err; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(grid)); + cfg.blockDim = dim3(kWarps * 32); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, kern, p, K / (CODED ? 64 : 16), pf, static_cast(xs_bytes)); +} + +template +cudaError_t dispatchRows(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + if (p.M <= 1) + return launch<1, CODED>(p, K, grid, pdl, stream); + if (p.M <= 2) + return launch<2, CODED>(p, K, grid, pdl, stream); + if (p.M <= 4) + return launch<4, CODED>(p, K, grid, pdl, stream); + return launch<8, CODED>(p, K, grid, pdl, stream); +} + +} // namespace + +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + bool const coded = p.cd.exc_off != nullptr; + if (p.M < 1 || p.M > kMaxTokens || K % (coded ? 64 : 16) != 0 || grid < 1) + return cudaErrorInvalidValue; + return coded ? dispatchRows(p, K, grid, pdl, stream) : dispatchRows(p, K, grid, pdl, stream); +} + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h new file mode 100644 index 000000000000..eb38b73488a4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h @@ -0,0 +1,111 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ +//! One pass of the fused greedy Eagle3 draft loop of the MiniMax-M3 NVFP4 TP4 decode path (sm_103a): the BF16 +//! vocab-sharded draft head, the cross-rank argmax and the draft loop's per-pass bookkeeping in one kernel. See +//! draftStep.cu. +namespace draftStep +{ + +constexpr int kMaxTokens = 8; // head rows (requests) of one launch +constexpr int kWarps = 16; // warps per CTA +constexpr int kMaxTilesPerCta = 32; // 16-row weight tiles one CTA can own +//! Shared-memory weight prologue per CTA, in KB (capped by the free shared memory). +constexpr int kPrefetchKb = 112; + +// Draft-pass epilogue (last CTA, one thread per request row b): stores the token as draft token `step`, stages the +// next pass's input id and position, and advances the draft KV lengths the way the in-tree draft loop does between +// passes (pass 0 with num_acc: drop the rejected verify rows and switch to one token per request; later passes: +1). +struct StepArgs +{ + int* draft_out; // [nreq, draft_stride]: column `step` receives the token + int draft_stride; + int step; + int* ids_next; // [nreq] + int const* pos_src; // positions of the rows the head read (indexed like x) + int* pos_next; // [nreq] = pos_src[x_rows[b] or b] + 1 + int* kv_lens; // [nreq] draft KV lengths, updated in place + int const* num_acc; // pass 0: accepted counts (nullptr: no rollback) + int* seq_lens; // pass 0: set to 1 + int draft_len; +}; + +// The 12-bit coded head (see draftStep.cu): the per-tile exception lists and the exponent base. exc_off == nullptr: +// the head is the plain BF16 fragment-ordered copy. +struct CodedArgs +{ + int const* exc_off; // [n_tiles + 1] exception list offsets + uint32_t const* exc_pos; // row in tile << 16 | k, per exception + float2 const* exc_w; // {w, 0} per exception: the exact weight + uint32_t base; // exponent base << 7 in both BF16 halves (base << 7 | base << 23) +}; + +// Everything a draft-step launch takes (see invokeDraftStep). +struct Params +{ + __nv_bfloat16 const* x; // [*, K] head input rows + long long const* x_rows; // optional [M]: head input row of request b (pass 0: the last accepted verify row) + int M; // requests, <= kMaxTokens + uint4 const* wf; // [n_tiles][K / 16][32 lanes] 16-byte BF16 A fragments, or the 12-bit coded stream + // [n_tiles][K / 64][3][32 lanes] when cd.exc_off is set (see draftStep.cu) + int n_tiles; + int vocab_offset; // global id of this rank's first head row + unsigned long long* keys; + unsigned int* counters; + unsigned long long* ws_uc; // this rank's exchange slots [2][kMaxTokens][world] + unsigned long long* ws_mc; // their multicast address + int rank; + int world; + unsigned long long* out_keys; // optional [M]: the winning global key of each row + StepArgs st; + CodedArgs cd; + // A pass split over several launches (one per locality domain, each on that domain's SMs): this launch streams + // tiles [tile_lo, tile_hi) (tile_hi 0 = n_tiles), and the last of total_ctas CTAs over all of the pass's launches + // runs the exchange and the bookkeeping (0 = this launch's grid: one launch per pass). + int tile_lo = 0; + int tile_hi = 0; + unsigned int total_ctas = 0; + // Equal-chunk split (set both, zero-filled [n_tiles][32] float4 / [n_tiles], restored to zero by the kernel): each + // CTA of a launch streams an equal share of its range's (tile, chunk) space instead of whole tiles, so a tile at a + // CTA boundary gets one partial sum from each of its two CTAs, added in bnd (two addends: the order does not change + // the sum) and keyed by the second to arrive (bnd_cnt). nullptr: whole tiles per CTA. + float4* bnd = nullptr; + unsigned int* bnd_cnt = nullptr; +}; + +//! Launch the draft step over `grid` 512-thread CTAs (K = the hidden size, a multiple of 16; of 64 for a coded head). +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream); + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu new file mode 100644 index 000000000000..179770d14a0e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu @@ -0,0 +1,189 @@ +/* + * 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. + */ + +// The TP2 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's GPT-NeoX +// path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin negated, +// dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale E4M3 +// store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query +// saturates the same way (448, not NaN). + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace drafterRopeFp8Kv +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kBlockSize = 256; + +template +__device__ __forceinline__ float rotate(float c, float x, float s, float pair) +{ + if constexpr (kRopeMode == 0) + { + return __fadd_rn(__fmul_rn(c, x), __fmul_rn(s, pair)); + } + else if constexpr (kRopeMode == 1) + { + return __fmaf_rn(c, x, __fmul_rn(s, pair)); + } + else + { + return __fmaf_rn(s, pair, __fmul_rn(c, x)); + } +} + +__device__ __forceinline__ void storeE4m3Pairs(__nv_fp8_e4m3* dst, __nv_bfloat16 const (&v)[kElemsPerThread]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + __nv_bfloat162 pair; + pair.x = v[i]; + pair.y = v[i + 1]; + __nv_fp8x2_e4m3 const packed(pair); + reinterpret_cast<__nv_fp8x2_storage_t*>(dst)[i / 2] = packed.__x; + } +} + +template +__global__ void __launch_bounds__(kBlockSize) drafterRopeFp8KvKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const lane = threadIdx.x % 32; + int const warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; + int const slots = p.numHeadsQ + p.numHeadsKv; + int const token = warp / slots; + if (token >= p.numTokens) + { + return; + } + int const slot = warp % slots; + bool const isQ = slot < p.numHeadsQ; + int const head = isQ ? slot : slot - p.numHeadsQ; + int const request = token / p.qLen; + int const row = token % p.qLen; + int const cacheLen = p.seqLens[request]; + int const kvPos = cacheLen - p.qLen + row; + int const ropePos + = p.specPositionOffsets != nullptr ? p.specPositionOffsets[request * p.qLen + row] + cacheLen - p.qLen : kvPos; + int const rowWidth = (p.numHeadsQ + 2 * p.numHeadsKv) * kHeadDim; + __nv_bfloat16 const* src + = p.qkv + static_cast(token) * rowWidth + (isQ ? head * kHeadDim : (p.numHeadsQ + head) * kHeadDim); + int const half = p.rotaryDim / 2; + float2 const* coef = p.cosSin + static_cast(ropePos) * half; + + __nv_bfloat16 rotated[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + int const d = lane * kElemsPerThread + i; + bool const firstHalf = d < half; + float const x = __bfloat162float(src[d]); + float const pair = __bfloat162float(src[firstHalf ? d + half : d - half]); + float2 const cs = d < p.rotaryDim ? coef[firstHalf ? d : d - half] : make_float2(1.0f, 0.0f); + float const s = firstHalf ? -cs.y : cs.y; + rotated[i] = __float2bfloat16(rotate(cs.x, x, s, pair)); + } + + if (isQ) + { + storeE4m3Pairs( + p.qOut + (static_cast(token) * p.numHeadsQ + head) * kHeadDim + lane * kElemsPerThread, rotated); + return; + } + + int const block = kvPos / p.tokensPerBlock; + long long const inPage = (static_cast(head) * p.tokensPerBlock + kvPos % p.tokensPerBlock) * kHeadDim + + lane * kElemsPerThread; + long long const pageElems = static_cast(p.numHeadsKv) * p.tokensPerBlock * kHeadDim; + int const* table = p.blockTable + request * p.blockTableRequestStride; + long long const kPage = table[block]; + long long const vPage = table[p.blockTableKvStride + block]; + storeE4m3Pairs(p.pool + kPage * pageElems + inPage, rotated); + + __nv_bfloat16 const* vSrc = p.qkv + static_cast(token) * rowWidth + + (p.numHeadsQ + p.numHeadsKv + head) * kHeadDim + lane * kElemsPerThread; + __nv_bfloat16 v[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + v[i] = vSrc[i]; + } + storeE4m3Pairs(p.pool + vPage * pageElems + inPage, v); +} + +//! With PDL and maxSharedCarveout, the maximum shared-memory carveout lets the full-context decode's large-smem CTAs +//! (the dependent this kernel releases right after its wait) co-reside on the SMs this small kernel occupies. Set once +//! per instantiation. +template +cudaError_t launchWithMaxSharedCarveout( + cudaLaunchConfig_t const& config, Params const& params, bool pdl, bool maxSharedCarveout) +{ + auto const kernel = drafterRopeFp8KvKernel; + if (pdl && maxSharedCarveout) + { + static cudaError_t const carveout = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (carveout != cudaSuccess) + { + return carveout; + } + } + return cudaLaunchKernelEx(&config, kernel, params); +} + +} // namespace + +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream) +{ + int const warps = params.numTokens * (params.numHeadsQ + params.numHeadsKv); + int const warpsPerBlock = kBlockSize / 32; + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((warps + warpsPerBlock - 1) / warpsPerBlock); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + switch (params.ropeMode) + { + case 0: return launchWithMaxSharedCarveout<0>(config, params, pdl, maxSharedCarveout); + case 1: return launchWithMaxSharedCarveout<1>(config, params, pdl, maxSharedCarveout); + case 2: return launchWithMaxSharedCarveout<2>(config, params, pdl, maxSharedCarveout); + default: return cudaErrorInvalidValue; + } +} + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h new file mode 100644 index 000000000000..3302c94390f3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h @@ -0,0 +1,79 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ +namespace drafterRopeFp8Kv +{ + +constexpr int kHeadDim = 128; + +struct Params +{ + //! [tokens, (numHeadsQ + 2 * numHeadsKv) * 128] BF16: the drafter layer's packed Q, K, V rows. + __nv_bfloat16 const* qkv; + //! [positions, rotaryDim / 2] (cos, sin) pairs: the layer's rotary_cos_sin table. + float2 const* cosSin; + //! [requests] cache lengths including this step's rows. + int const* seqLens; + //! [requests, qLen] spec-decoding position offsets of a multi-token pass, or null. + int const* specPositionOffsets; + //! [tokens, numHeadsQ, 128] E4M3 query out. + __nv_fp8_e4m3* qOut; + //! Base of the layer's K / V sub-pages, [*, numHeadsKv, tokensPerBlock, 128] E4M3. + __nv_fp8_e4m3* pool; + //! [requests, 2, maxBlocks] sub-page indices: K in row 0, V in row 1. + int const* blockTable; + long long blockTableRequestStride; + long long blockTableKvStride; + int numTokens; + int qLen; + int numHeadsQ; + int numHeadsKv; + int rotaryDim; + int tokensPerBlock; + //! How the rotation's two products are summed: 0 two products then the add, 1 fma(cos, x, sin * pair), + //! 2 fma(sin, pair, cos * x). + int ropeMode; +}; + +//! The MiniMax-M3 TP2 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, +//! rounded to BF16 (the in-tree preprocess's q output and cache input), then Q, K and V to saturating E4M3: K / V +//! as the in-tree preprocess stores them at unit KV scale, Q as torch's CUDA cast of the BF16 query. Rows map to +//! (request token / qLen, row token % qLen) at cache position seqLen - qLen + row. +//! +//! PDL (pdl): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction and the kernel +//! releases its dependent right after it. The dependent may read only step-invariant data before its own wait. +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream); + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu new file mode 100644 index 000000000000..dd0ad05437a6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu @@ -0,0 +1,86 @@ +/* + * 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/minimaxM3Sm107Tp2/foldTail.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +// Row i of the step (lane c = column), as the Triton fold tail's program i: the store's padded rows, then the next +// step's inputs read from them (zero past each source width). +__global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) +{ + int const i = blockIdx.x; + int const c = threadIdx.x; + if (i >= *p.ctrl) + { + return; + } + int64_t const s = p.slots[i]; + int32_t const nw = c < p.snW ? p.srcNew[i * p.sn0 + c * p.sn1] : 0; + if (c < p.dnW) + { + p.dstNew[s * p.dn0 + c * p.dn1] = nw; + } + int32_t const nxt = c < p.sxW ? p.srcNext[i * p.sx0 + c * p.sx1] : 0; + if (c < p.dxW) + { + p.dstNext[s * p.dx0 + c * p.dx1] = nxt; + } + int32_t const drf = c < p.sdW ? p.srcDraft[i * p.sd0 + c * p.sd1] : 0; + if (c < p.ddW) + { + p.dstDraft[s * p.dd0 + c * p.dd1] = drf; + } + int32_t const ln = p.srcLens[i]; + int32_t const t = p.tokensPerStep; + int32_t const d = t - 1; + if (c == 0) + { + p.dstLens[s] = ln; + p.kvOff[i] = ln - t; + } + if (c < t) + { + p.ids[i * t + c] = c < p.dxW ? nxt : 0; + p.posOff[i * t + c] = ln; + } + if (c < d) + { + p.draftIds[i * d + c] = c < p.ddW ? drf : 0; + } +} + +} // namespace + +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream) +{ + foldTailKernel<<>>(p); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h new file mode 100644 index 000000000000..929ac2fed9ed --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h @@ -0,0 +1,72 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +//! Operands of the speculative step's forward-tail store and next-step staging (kernels/m3_fold_tail.py). Every tensor +//! is int32 except `slots` (int64); strides are in elements. The store views are [slots, width], the sources +//! [rows, width]. +struct FoldTailParams +{ + int32_t const* ctrl; // ctrl[0]: the rows the step armed (0: write nothing) + int64_t const* slots; + int32_t const* srcNew; + int64_t sn0, sn1; + int32_t snW; + int32_t* dstNew; + int64_t dn0, dn1; + int32_t dnW; + int32_t const* srcNext; + int64_t sx0, sx1; + int32_t sxW; + int32_t* dstNext; + int64_t dx0, dx1; + int32_t dxW; + int32_t const* srcDraft; + int64_t sd0, sd1; + int32_t sdW; + int32_t* dstDraft; + int64_t dd0, dd1; + int32_t ddW; + int32_t const* srcLens; + int32_t* dstLens; + int32_t* ids; // [rows * T] + int32_t* draftIds; // [rows * (T - 1)] + int32_t* posOff; // [rows * T] + int32_t* kvOff; // [rows] + int32_t tokensPerStep; +}; + +//! One CTA of 32 lanes per row, `rows` CTAs; the first ctrl[0] rows write. Widths and tokensPerStep must be <= 32. +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu new file mode 100644 index 000000000000..236a13e66d1f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu @@ -0,0 +1,223 @@ +/* + * 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. + */ + +// The split-KV merge for the MiniMax-M3 NVFP4 TP2 (Rubin) full-context decode. The producer is the PrimTS split-KV +// decode with its separate-reduction scratch: per (request, KV head, split, row) one normalized 16-bit partial O and +// one fp32 log2-LSE. Active splits per row follow the producer's split ranges (the causal union of the row's grouped +// Q CTA, instruction-group aligned): splits past them hold nothing and are not read. +// +// One CTA per (request, KV head, kRowsPerCta rows); per row kGroups split groups x kThreadsPerRow threads x 8 dims. +// A thread issues the loads of every split it owns (strided over the groups) before using any, folds them under its +// group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to +// BF16 once. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace fullCtxReduce +{ +namespace +{ + +constexpr int kGroups = 8; +constexpr int kThreadsPerRow = kHeadDim / 8; +constexpr int kRowsPerCta = 2; +constexpr int kSplitsPerGroup = kMaxSplits / kGroups; +constexpr int kThreads = kGroups * kThreadsPerRow * kRowsPerCta; + +__device__ int activeSplits(Params const& p, int b, int row) +{ + int const seqKv = p.seqLensKv[b]; + int visible = seqKv; + if (p.causal) + { + int const token = row / p.headsPerKv; + int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; + int const end = min(base + p.qTokensPerCta, p.seqLenQ); + visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); + } + int const tiles = (visible + p.tileKv - 1) / p.tileKv; + int const tilesPerGroup = p.splits * p.instsKv; + int const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; + int const local = max(groups * p.instsKv, p.instsKv); + return (tiles + local - 1) / local; +} + +template +__device__ float2 unpack(uint32_t word) +{ + if constexpr (kBf16Partial) + { + return __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&word)); + } + else + { + return __half22float2(*reinterpret_cast<__half2 const*>(&word)); + } +} + +template +__global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const b = blockIdx.z; + int const h = blockIdx.y; + int const localRow = threadIdx.x / (kGroups * kThreadsPerRow); + int const row = blockIdx.x * kRowsPerCta + localRow; + int const g = (threadIdx.x / kThreadsPerRow) % kGroups; + int const t = threadIdx.x % kThreadsPerRow; + bool const validRow = row < p.rows; + int const active = validRow ? activeSplits(p, b, row) : 0; + + __shared__ float sMax[kRowsPerCta][kGroups]; + __shared__ float sSum[kRowsPerCta][kGroups]; + __shared__ float sAcc[kRowsPerCta][kGroups][kHeadDim]; + + uint4 const* partialO = static_cast(p.partialO); + size_t const kvIdx = static_cast(b) * p.numKvHeads + h; + float lse[kSplitsPerGroup]; + uint4 vec[kSplitsPerGroup]; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + int const s = g + i * kGroups; + if (s < active) + { + size_t const r = (kvIdx * p.splits + s) * p.rows + row; + lse[i] = __ldcg(p.partialStats + r); + vec[i] = __ldcg(partialO + r * kThreadsPerRow + t); + } + } + float groupMax = -CUDART_INF_F; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active) + { + groupMax = fmaxf(groupMax, lse[i]); + } + } + // A split of the CTA's causal union can hold no key of this row (log2-LSE -inf). Such a split, and a group of + // only such splits, is skipped rather than weighted 0: its partial O carries no defined value. + float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + float sum = 0.f; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active && lse[i] != -CUDART_INF_F) + { + float const w = exp2f(lse[i] - groupMax); + sum += w; + uint32_t const* words = reinterpret_cast(&vec[i]); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = unpack(words[j]); + acc[2 * j] += w * f.x; + acc[2 * j + 1] += w * f.y; + } + } + } + if (t == 0) + { + sMax[localRow][g] = groupMax; + sSum[localRow][g] = sum; + } +#pragma unroll + for (int j = 0; j < 8; ++j) + { + sAcc[localRow][g][t * 8 + j] = acc[j]; + } + __syncthreads(); + if (g != 0 || !validRow) + { + return; + } + float rowMax = -CUDART_INF_F; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + rowMax = fmaxf(rowMax, sMax[localRow][k]); + } + float total = 0.f; + float o[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + if (sMax[localRow][k] == -CUDART_INF_F) + { + continue; + } + float const scale = exp2f(sMax[localRow][k] - rowMax); + total += scale * sSum[localRow][k]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + o[j] += scale * sAcc[localRow][k][t * 8 + j]; + } + } + float const inv = total > 0.f ? 1.f / total : 0.f; + int const token = row / p.headsPerKv; + int const head = row % p.headsPerKv; + __nv_bfloat162 packed[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + packed[j] = __floats2bfloat162_rn(o[2 * j] * inv, o[2 * j + 1] * inv); + } + size_t const outRow + = (static_cast(b) * p.seqLenQ + token) * (p.numKvHeads * p.headsPerKv) + h * p.headsPerKv + head; + reinterpret_cast(p.out + outRow * kHeadDim)[t] = *reinterpret_cast(packed); +} + +} // namespace + +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) +{ + if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((params.rows + kRowsPerCta - 1) / kRowsPerCta, params.numKvHeads, params.batch); + config.blockDim = dim3(kThreads); + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + return params.bf16Partial ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) + : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); +} + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h new file mode 100644 index 000000000000..eca13e5d7dab --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h @@ -0,0 +1,71 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ +//! The split-KV merge of the MiniMax-M3 TP2 full-context decode (PrimTS split-KV decode, separate-reduction scratch) +//! straight to BF16. See fullCtxReduce.cu. +namespace fullCtxReduce +{ + +constexpr int kHeadDim = 128; +constexpr int kMaxSplits = 128; + +struct Params +{ + void const* partialO; // [B, Hkv, S, R, 128] normalized 16-bit partial O per split + float const* partialStats; // [B, Hkv, S, R] log2-LSE per split + int const* seqLensKv; // [B] + __nv_bfloat16* out; // [B, SQ, Hkv * hR, 128] + int batch; + int numKvHeads; + int splits; // S, the planned split capacity (<= kMaxSplits) + int rows; // R = hR * SQ, rows [token, local head] + int seqLenQ; // SQ + int headsPerKv; // hR + int qTokensPerCta; // query tokens one producer CTA owns (grouped Q) + int tileKv; // producer K/V tile + int instsKv; // producer K/V instructions per loop step + bool causal; + bool bf16Partial; +}; + +//! Merges the active splits of every row in fp32 with one max rescale and writes BF16. With pdl the launch is +//! programmatic: the grid waits on the producer (griddepcontrol.wait) before reading its scratch. +//! +//! PDL contract: griddepcontrol.wait is the kernel's first instruction, so nothing (scratch, lengths, out) is read or +//! written before every prerequisite grid has completed and flushed; a producer's early trigger +//! (griddepcontrol.launch_dependents) only lets this grid launch and wait sooner. The kernel never triggers its own +//! dependents early: the next kernel's post-wait code sees all of out, and its pre-wait prologue must not read out. +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream); + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu new file mode 100644 index 000000000000..cad2b93558ab --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu @@ -0,0 +1,1092 @@ +/* + * 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/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// SM107 (Rubin) copy of minimaxM3/mcAllreduceNorm.cu for the sm_107 TP2 target. Two changes in the arithmetic: the +// rank sum rounds to BF16 before the residual add, the rounding points of the stock all-reduce + fused residual +// RMSNorm on this target (bf16 all-reduce output, then bf16(residual + sum)); and every call but the MoE finalize runs +// the row kernel below, whose RMS statistic is the stock fused norm's reduction tree (bitwise the stock outputs). +// +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions): they are compiled for every sm_10x target of the build, the sm_100f family target +// included (which runs on SM107). A device pass for any other architecture compiles them to a trap, never to a silent +// empty body, and the host entry point checks that the kernel binary it launches holds sm_10x code (see +// checkBinary). +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace cg = cooperative_groups; + +namespace +{ + +constexpr int kEmpty = static_cast(0x80000000u); +constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k + +// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], +// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c +// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp +// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its +// arithmetic is the same source. +constexpr int kCapThreads = 256; +constexpr int kCapIters = 3; +constexpr int kCapWarps = kCapThreads / 32; +constexpr int kCapHidden = kCapThreads * kCapIters * 8; + +// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). +__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) +{ + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int q = 0; q < 8; ++q) + { + float const v = __bfloat162float(e[q]); + ss += v * v; + } + return ss; +} + +__device__ __forceinline__ float cap_warp_sum(float v) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. +__device__ __forceinline__ void cap_store( + float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) +{ + uint4 const wr = *reinterpret_cast(w); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int q = 0; q < 8; ++q) + oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); + *reinterpret_cast(cap) = o; +} + +__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) +{ + if ((w & 0xFFFFu) == 0x8000u) + w &= 0xFFFF0000u; + if ((w >> 16) == 0x8000u) + w &= 0x0000FFFFu; + return w; +} + +__device__ __forceinline__ float bf16_lo(uint32_t w) +{ + return __uint_as_float(w << 16); +} + +// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's +// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals +// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the +// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. +__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) +{ + float value; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); + uint32_t result; + asm("{\n" + " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" + " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" + " setp.le.f32 p_zero, %1, 0f00000000;\n" + " mov.b32 bits, %1;\n" + " shr.b32 exp_biased, bits, 23;\n" + " and.b32 exp_biased, exp_biased, 255;\n" + " and.b32 mantissa, bits, 0x7FFFFF;\n" + " setp.ne.u32 p_has_mant, mantissa, 0;\n" + " selp.u32 bump, 1, 0, p_has_mant;\n" + " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" + " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" + " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" + " @p_tiny_sub mov.u32 bump, 0;\n" + " add.u32 r, exp_biased, bump;\n" + " setp.gt.u32 p_ovf, r, 254;\n" + " selp.u32 r, 254, r, p_ovf;\n" + " selp.u32 %0, 0, r, p_zero;\n" + "}\n" + : "=r"(result) + : "f"(value)); + return result; +} + +__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) +{ + float inv; + asm("{\n" + " .reg .s32 new_exp;\n" + " .reg .b32 float_bits;\n" + " .reg .pred p_zero;\n" + " setp.eq.u32 p_zero, %1, 0;\n" + " sub.s32 new_exp, 254, %1;\n" + " max.s32 new_exp, new_exp, 0;\n" + " shl.b32 float_bits, new_exp, 23;\n" + " mov.b32 %0, float_bits;\n" + " @p_zero mov.b32 %0, 0;\n" + "}\n" + : "=f"(inv) + : "r"(ue8m0)); + return inv; +} + +// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. +__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) +{ + uint16_t pair; + asm("{\n" + " .reg .f32 f0, f1;\n" + " mul.f32 f0, %1, %3;\n" + " mul.f32 f1, %2, %3;\n" + " min.f32 f0, f0, 0f43E00000;\n" + " max.f32 f0, f0, 0fC3E00000;\n" + " min.f32 f1, f1, 0f43E00000;\n" + " max.f32 f1, f1, 0fC3E00000;\n" + " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" + "}\n" + : "=h"(pair) + : "f"(lo), "f"(hi), "f"(inv)); + return pair; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ float bf16_hi(uint32_t w) +{ + return __uint_as_float(w & 0xFFFF0000u); +} + +__device__ __forceinline__ int4 ld_volatile_v4(int const* p) +{ + int4 v; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +struct FinalizeArgs +{ + __nv_bfloat16 const* fc2; // [permuted rows, H] routed-expert FC2 output (null: plain input x) + int const* perm; // [M, topk] expanded index -> permuted row, -1 = dropped + void const* scales; // [M, topk] routing weights (fp32 or bf16) + int scale_fp32; + int topk; +}; + +// grid = min(M, MAXT) * CLUSTER CTAs, cluster (CLUSTER, 1, 1), THREADS threads. +// OPTS (each bit-identical): 1 residual + norm weight loaded at entry (before the poll); 2 cluster RMS exchange by +// st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated right after the poll; 8 max-smem +// carveout so the next kernel's CTAs co-reside early. +// OPTS 16 (NOT bit-identical with the others; a different boundary): the Eagle3 drafter's boundary, whose stock is a +// plain BF16 all-reduce, then FlashInfer fused_add_rmsnorm: the RMS statistic and the normalization take the +// unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded to BF16). +// CAP: the Eagle3 capture epilogue (cap_out) is compiled in only where a call passes cap_out, so every other call +// runs the kernel without the capture branches, shared arrays or parameters in its body. +template +__global__ void __launch_bounds__(THREADS) mc_allreduce_norm_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, FinalizeArgs fin, int M, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int world, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ alignas(16) float cluster_sq[CLUSTER]; + extern __shared__ __align__(16) uint4 cap_slice[]; // [cols / 8] this CTA's updated chunks (cap_out only) + __shared__ alignas(8) unsigned long long xbar; + cg::cluster_group cluster = cg::this_cluster(); + if constexpr ((OPTS & 2) != 0) + { + if (threadIdx.x == 0) + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(b) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + } + int const crank = static_cast(cluster.block_rank()); + int const t = blockIdx.x / CLUSTER; + int const cols = H / CLUSTER; + int const c0 = crank * cols + threadIdx.x * 8; + bool const active = t < M && threadIdx.x * 8 < cols; + int64_t const half_words = static_cast(MAXT) * world * (H / 2); + + // prepushed 1: the producer (m3_mx_gemv) already pushed this rank's rows into the exchange slots; 2: and it + // launched this grid only after its pushes (its grid wait passed), so the slots are polled without a grid wait; + // 3: as 2, with x in set A and x2 in set B. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const parity = *reinterpret_cast(flags + blockIdx.x) & 1; + uint4 rraw = make_uint4(0u, 0u, 0u, 0u); + uint4 wraw = make_uint4(0u, 0u, 0u, 0u); + if constexpr ((OPTS & 1) != 0) + { + // the residual row was written grids earlier (prepushed >= 2: the producer launched this grid only after its + // own grid wait; otherwise the grid wait above) + if (active) + { + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + wraw = *reinterpret_cast(norm_w + c0); + } + } + // Push this rank's slice into every rank's buffer. + if (active && !prepushed) + { + uint4 raw; + if (fin.fc2 != nullptr) + { + // MoE finalize: fp32 sum of the routing-weighted top-k expert rows plus the shared expert, rounded once. + float f[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + if (fin.topk <= kFinMaxTopk) + { + // All top-k indices and scales, then every expert-row load in flight at once, then the same in-order + // FP32 accumulation as the loop below (one memory round trip instead of topk; bit-identical). + int prow[kFinMaxTopk]; + float sc[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + prow[k] = -1; + sc[k] = 0.f; + if (k < fin.topk) + { + int const e = t * fin.topk + k; + prow[k] = fin.perm[e]; + sc[k] = fin.scale_fp32 ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + } + } + uint4 rv[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] >= 0) + rv[k] = *reinterpret_cast(fin.fc2 + static_cast(prow[k]) * H + c0); + } +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] < 0) + continue; + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv[k]); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc[k]; + f[2 * q + 1] += v.y * sc[k]; + } + } + } + else + for (int k = 0; k < fin.topk; ++k) + { + int const e = t * fin.topk + k; + int const prow = fin.perm[e]; + if (prow < 0) + continue; + float const sc = fin.scale_fp32 + ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + uint4 const rv = *reinterpret_cast(fin.fc2 + static_cast(prow) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc; + f[2 * q + 1] += v.y * sc; + } + } + if (x2 != nullptr) + { + // As stock: the routed output is rounded to BF16 first, then the shared expert is added (one more + // BF16 rounding of the FP32 sum). + uint4 const rv = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const routed = __bfloat1622float2(__floats2bfloat162_rn(f[2 * q], f[2 * q + 1])); + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] = routed.x + v.x; + f[2 * q + 1] = routed.y + v.y; + } + } + __nv_bfloat162* o = reinterpret_cast<__nv_bfloat162*>(&raw); +#pragma unroll + for (int q = 0; q < 4; ++q) + o[q] = __floats2bfloat162_rn(f[2 * q], f[2 * q + 1]); + } + else + { + raw = *reinterpret_cast(x + static_cast(t) * H + c0); + } + if (fin.fc2 == nullptr && x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * world + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float s[8] = {}; + float sq = 0.f; + if (active) + { + // The residual row (written grids earlier) and every rank's slot are loaded together; slots still empty are + // re-polled. Ranks are then summed in rank order, as before. + if constexpr ((OPTS & 1) == 0) + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + constexpr int kMaxWorld = 8; + int* const p0 = ws_uc + parity * half_words + static_cast(t) * world * (H / 2) + c0 / 2; + int4 v[kMaxWorld]; + int4 v2[kMaxWorld]; + bool const two = prepushed == 3; + int* const p1 = p0 + 2 * half_words; // set B + uint32_t pending = (1u << world) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (r < world) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr ((OPTS & 16) != 0) + { + float const x0 = s[2 * q] + rf.x, x1 = s[2 * q + 1] + rf.y; + s[2 * q] = x0; + s[2 * q + 1] = x1; + sq += x0 * x0 + x1 * x1; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + sq += uf.x * uf.x + uf.y * uf.y; + } + } + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr ((OPTS & 16) == 0) + { + if constexpr (CAP) + cap_slice[threadIdx.x] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + if constexpr ((OPTS & 4) != 0) + { + // every CTA (active or not) rotates its call parity once its own poll is done (read only by later grids) + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffffu, sq, off); + if ((threadIdx.x & 31) == 0) + warp_sq[threadIdx.x / 32] = sq; + __syncthreads(); + if constexpr ((OPTS & 2) != 0) + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); // peers' xbar initialized (long since) + if (threadIdx.x == 0) + { + float total = 0.f; +#pragma unroll + for (int w = 0; w < THREADS / 32; ++w) + total += warp_sq[w]; + if constexpr ((OPTS & 2) == 0) + for (int c = 0; c < CLUSTER; ++c) + *cluster.map_shared_rank(&cluster_sq[crank], c) = total; + else + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t const d = static_cast(__cvta_generic_to_shared(&cluster_sq[crank])); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * 4) + : "memory"); + for (int c = 0; c < CLUSTER; ++c) + { + uint32_t rd, rb; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), + "r"(__float_as_uint(total)), "r"(rb) + : "memory"); + } + } + } + if constexpr ((OPTS & 2) == 0) + cluster.sync(); + else + { + // every peer's sum landed (st.async complete_tx on this CTA's barrier); no cluster-wide release-arrive + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t done = 0; + while (!done) + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(b) + : "memory"); + } + if constexpr ((OPTS & 16) == 0) + { + if constexpr (CAP) + { + // The row's chunks are spread over the cluster (CTA c holds chunks [c * per, (c + 1) * per)); every CTA + // forms capture_norm's statistic itself, its 128 threads standing in for capture threads ct and ct + 128. + __shared__ float cap_red[kCapWarps]; + int const per = cols / 8; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + float ca = 0.f, cb = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + { + int const ka = it * kCapThreads + static_cast(threadIdx.x); + int const kb = ka + kCapThreads / 2; + ca = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[ka % per], ka / per), ca); + cb = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[kb % per], kb / per), cb); + } + // Every peer read is done before any CTA moves on (and later exits). + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + ca = cap_warp_sum(ca); + cb = cap_warp_sum(cb); + if ((threadIdx.x & 31) == 0) + { + cap_red[threadIdx.x / 32] = ca; + cap_red[threadIdx.x / 32 + kCapWarps / 2] = cb; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + if (active) + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + if (active) + { + float total = 0.f; +#pragma unroll + for (int c = 0; c < CLUSTER; ++c) + total += cluster_sq[c]; + float const inv = rsqrtf(total / static_cast(H) + eps); + if constexpr ((OPTS & 1) == 0) + wraw = *reinterpret_cast(norm_w + c0); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + } + if (out_q != nullptr) + { + // MXFP8 copy of the normed row for the next GEMM, as the CuTe DSL mxfp8 quantize of `normed` (the tree's + // decode quantize, MX-TAC; see mxfp8_ue8m0): four threads hold one 32-column scale block. + float amax = 0.f; + if (active) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + if (active) + { + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((threadIdx.x & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } + } + if constexpr ((OPTS & 4) == 0) + { + __syncthreads(); + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +// Refuses (raises) to launch a kernel whose loaded binary holds no sm_10x code (PTX and SASS version >= 100): the +// device passes for other architectures compile the kernels to a trap, and a build that left out every sm_10x target +// must not reach the GPU at all. +template +bool checkBinary(Kernel kernel, char const* name) +{ + cudaFuncAttributes attr{}; + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, reinterpret_cast(kernel))); + TLLM_CHECK_WITH_INFO(attr.ptxVersion >= 100 && attr.binaryVersion >= 100, + "%s: the loaded binary holds no sm_10x code (PTX %d, SASS %d): rebuild with an sm_100f / sm_107a target", name, + attr.ptxVersion, attr.binaryVersion); + return true; +} + +template +void launch(MinimaxM3Sm107Tp2McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) +{ + int const M = p.numTokens, H = p.hidden; + TLLM_CHECK_WITH_INFO(H % (CLUSTER * 8) == 0 && H / CLUSTER / 8 <= THREADS, "unsupported hidden size %d", H); + TLLM_CHECK_WITH_INFO(M >= 1 && M <= MAXT, "%d tokens do not fit the %d-token bucket", M, MAXT); + auto kernel = p.capOut != nullptr ? mc_allreduce_norm_kernel + : mc_allreduce_norm_kernel; + static bool const binaryOk + = checkBinary(mc_allreduce_norm_kernel, "mc_allreduce_norm_kernel") + && checkBinary( + mc_allreduce_norm_kernel, "mc_allreduce_norm_kernel (capture)"); + (void) binaryOk; + // Both instantiations (with and without the capture epilogue) get the same attributes, whichever runs first. + static bool const configured = []() + { + auto setup = [](auto k) + { + bool ok = CLUSTER <= 8 + || cudaFuncSetAttribute(k, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess; + if ((OPTS & 8) != 0) + ok = ok + && cudaFuncSetAttribute( + k, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared) + == cudaSuccess; + return ok; + }; + bool const ok = setup(mc_allreduce_norm_kernel) + && setup(mc_allreduce_norm_kernel); + if (!ok) + cudaGetLastError(); + return ok; + }(); + TLLM_CHECK_WITH_INFO(configured, "cluster size not supported"); + FinalizeArgs const fin{p.fc2, p.perm, p.scales, p.scaleFp32 ? 1 : 0, p.topk}; + cudaLaunchConfig_t cfg = {}; + // The grid is sized to the call's rows: min(M, MAXT) clusters. Each (row, slice) slot keeps its own parity flag, + // read by the producers before they push and rotated only by the slot's CTA after it polled: a slot that a call + // does not use keeps its parity, and every rank runs the same calls, so the parities stay consistent when M + // changes between calls. + cfg.gridDim = dim3((M < MAXT ? M : MAXT) * CLUSTER); + cfg.blockDim = dim3(THREADS); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO((OPTS & 16) == 0 && H == kCapHidden && 2 * THREADS == kCapThreads, + "the capture epilogue takes H %d on the 128-thread cluster kernel, not the drafter boundary", kCapHidden); + cfg.dynamicSmemBytes = static_cast(H / CLUSTER / 8) * sizeof(uint4); + } + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = CLUSTER; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, fin, M, H, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, p.world, p.outNormed, p.outResidual, p.outQuant, p.outSf, (H / 32 + 3) / 4, + p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +// ws_uc / ws_mc are this rank's unicast and multicast views of ONE exchange buffer (they alias): no __restrict__. +// prepushed >= 2 contract: the producer (m3_mx_gemv's o_proj push, the m3_moe engine's finalize push) passes its own +// grid dependency, pushes its rows into the exchange, and only then releases this grid (griddepcontrol +// .launch_dependents); this kernel therefore skips griddepcontrol.wait and reads only the pushed (Lamport-polled) +// words, its parity flags and the residual / norm weight, which were written two or more grids earlier. +// +// The TP2 row kernel: one CTA of H / 8 = 768 threads per token row, used for every call but the MoE finalize. Its +// residual add + RMSNorm is TRT-LLM's rms_norm_kernel (customAllReduceKernels.cu), the stock fused norm after the +// all-reduce on this target, operation for operation at H 6144: thread j owns columns [8j, 8j + 8), its sum of squares +// is acc += v * v over them in order, then the warp butterfly 16..1, the 24 warp sums through shared memory into a +// second 32-lane butterfly (lanes >= 24 zero), rsqrtf(acc / H + eps) and (v * denom) * w. With the BF16 rank sum +// before the residual add, the normed rows, the updated residual and the MXFP8 copy are bitwise the stock ones. The +// exchange slots and the per-(row, slice) parity flags keep the cluster kernel's layout (slice = 96 threads), so +// producers that push (the o_proj GEMV) are unchanged. +constexpr int kSlices = 8; +constexpr int kSliceThreads = 96; +constexpr int kRowThreads = 768; + +// grid = M CTAs (one per row), THREADS = H / 8 threads. WORLD ranks (compile time). DRAFTER (opts bit 16): the Eagle3 +// drafter boundary, whose stock is a plain BF16 all-reduce then FlashInfer fused_add_rmsnorm: the RMS statistic and +// the normalization take the unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded). +// ws_uc / ws_mc alias (one exchange buffer): no __restrict__. +template +__global__ void __launch_bounds__(THREADS, 1) mc_allreduce_norm_row_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int maxt, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ float total_sh; + extern __shared__ __align__(16) uint4 cap_row[]; // [THREADS] the updated row's chunks (cap_out only) + int const t = blockIdx.x; + int const j = threadIdx.x; + int const c0 = j * 8; + int const slice = j / kSliceThreads; + int64_t const half_words = static_cast(maxt) * WORLD * (H / 2); + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); + int const parity = *reinterpret_cast(flags + t * kSlices + slice) & 1; + uint4 const rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + uint4 const wraw = *reinterpret_cast(norm_w + c0); + if (!prepushed) + { + uint4 raw = *reinterpret_cast(x + static_cast(t) * H + c0); + if (x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * WORLD + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } + asm volatile("griddepcontrol.launch_dependents;"); + float s[8] = {}; + float sq = 0.f; + { + int* const p0 = ws_uc + parity * half_words + static_cast(t) * WORLD * (H / 2) + c0 / 2; + int* const p1 = p0 + 2 * half_words; // set B + bool const two = prepushed == 3; + int4 v[WORLD]; + int4 v2[WORLD]; + uint32_t pending = (1u << WORLD) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr (DRAFTER) + { + s[2 * q] += rf.x; + s[2 * q + 1] += rf.y; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + } + } +#pragma unroll + for (int e = 0; e < 8; ++e) + sq += s[e] * s[e]; + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr (!DRAFTER) + { + if constexpr (CAP) + cap_row[j] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + // TRT-LLM reduce_fusion::block_reduce_sum. + sq += __shfl_xor_sync(~0, sq, 16); + sq += __shfl_xor_sync(~0, sq, 8); + sq += __shfl_xor_sync(~0, sq, 4); + sq += __shfl_xor_sync(~0, sq, 2); + sq += __shfl_xor_sync(~0, sq, 1); + if ((j & 31) == 0) + warp_sq[j / 32] = sq; + __syncthreads(); + { + float val = (j & 31) < THREADS / 32 ? warp_sq[j & 31] : 0.f; + val += __shfl_xor_sync(~0, val, 16); + val += __shfl_xor_sync(~0, val, 8); + val += __shfl_xor_sync(~0, val, 4); + val += __shfl_xor_sync(~0, val, 2); + val += __shfl_xor_sync(~0, val, 1); + if (j == 0) + total_sh = val; + } + // The parity of this row's slices rotates once every thread has read the exchange (later grids only). + if (j % kSliceThreads == 0) + flags[t * kSlices + slice] = parity ^ 1; + __syncthreads(); + if constexpr (!DRAFTER) + { + if constexpr (CAP) + { + // The whole row is in this CTA: threads 0..255 run capture_norm's chains over the staged chunks. + __shared__ float cap_red[kCapWarps]; + if (j < kCapThreads) + { + float css = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + css = cap_chunk_sq(cap_row[it * kCapThreads + j], css); + css = cap_warp_sum(css); + if ((j & 31) == 0) + cap_red[j / 32] = css; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + float const inv = rsqrtf(total_sh / H + eps); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + if (out_q != nullptr) + { + float amax = 0.f; +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((j & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +void launchRow( + MinimaxM3Sm107Tp2McAllreduceNormParams const& p, int32_t maxTokens, bool drafter, bool pdl, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(p.hidden == kRowThreads * 8 && p.world == 2, "the row kernel takes H 6144 at 2 ranks"); + static bool const binaryOk + = checkBinary(mc_allreduce_norm_row_kernel<2, kRowThreads, false, false>, "mc_allreduce_norm_row_kernel") + && checkBinary( + mc_allreduce_norm_row_kernel<2, kRowThreads, true, false>, "mc_allreduce_norm_row_kernel (drafter)") + && checkBinary( + mc_allreduce_norm_row_kernel<2, kRowThreads, false, true>, "mc_allreduce_norm_row_kernel (capture)"); + (void) binaryOk; + TLLM_CHECK_WITH_INFO(p.numTokens >= 1 && p.numTokens <= maxTokens, "%d tokens do not fit the %d-token bucket", + p.numTokens, maxTokens); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(p.numTokens); + cfg.blockDim = dim3(kRowThreads); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO( + !drafter && p.hidden == kCapHidden, "the capture epilogue takes H %d, not the drafter", kCapHidden); + cfg.dynamicSmemBytes = static_cast(kRowThreads) * sizeof(uint4); + } + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + // The capture epilogue never runs on the drafter boundary (checked above), so there are three instantiations. + auto kernel = drafter ? mc_allreduce_norm_row_kernel<2, kRowThreads, true, false> + : p.capOut != nullptr ? mc_allreduce_norm_row_kernel<2, kRowThreads, false, true> + : mc_allreduce_norm_row_kernel<2, kRowThreads, false, false>; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, p.hidden, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, maxTokens, p.outNormed, p.outResidual, p.outQuant, p.outSf, + (p.hidden / 32 + 3) / 4, p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +constexpr int kCluster = 8; +constexpr int kThreads = 128; + +} // namespace + +bool isMinimaxM3Sm107Tp2McAllreduceNormSupported(int32_t maxTokens, int32_t opts) +{ + return (maxTokens == 8 && (opts == 7 || opts == 15 || opts == 23 || opts == 31)) || (maxTokens == 16 && opts == 15) + || (maxTokens == 48 && opts == 15); +} + +void invokeMinimaxM3Sm107Tp2McAllreduceNorm(MinimaxM3Sm107Tp2McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream) +{ + if (params.fc2 == nullptr && params.world == 2 && params.hidden == kRowThreads * 8) + { + launchRow(params, maxTokens, (opts & 16) != 0, pdl, stream); + return; + } + if (maxTokens == 8 && opts == 7) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 23) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 31) + { + launch(params, pdl, stream); + } + else if (maxTokens == 16 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 48 && opts == 15) + { + launch(params, pdl, stream); + } + else + { + TLLM_THROW("unsupported (bucket %d, opts %d)", maxTokens, opts); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h new file mode 100644 index 000000000000..e8ade9c62c3e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h @@ -0,0 +1,83 @@ +/* + * 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 +{ + +//! Operands of the MiniMax-M3 multicast all-reduce + residual + RMSNorm kernel (decode, at most 48 tokens). +//! +//! Per rank the kernel forms the partial `bf16(x + x2)` (or, in finalize mode, +//! `bf16(bf16(sum_k scale * fc2[perm]) + x2)`, the routed sum rounded before the shared expert is added, as stock), +//! pushes it through the multicast mapping `wsMc` into every rank's exchange, polls its own unicast view `wsUc`, +//! sums the ranks in FP32 in rank order, adds the residual (BF16 round) and applies RMSNorm. With `outQuant` it +//! also writes the MXFP8 copy of the normed rows (E4M3 [numTokens, hidden] and swizzled UE8M0 scales). +struct MinimaxM3Sm107Tp2McAllreduceNormParams +{ + __nv_bfloat16 const* x{nullptr}; //!< [numTokens, hidden] partial; in finalize mode the FC2 rows [rows, hidden]. + __nv_bfloat16 const* x2{nullptr}; //!< Optional second addend [numTokens, hidden]. + __nv_bfloat16 const* fc2{nullptr}; //!< Finalize mode: the permuted FC2 output (same buffer as x); null otherwise. + int32_t const* perm{nullptr}; //!< Finalize mode: [numTokens, topk] expanded index -> permuted row, -1 dropped. + void const* scales{nullptr}; //!< Finalize mode: [numTokens, topk] routing weights (fp32 or bf16). + bool scaleFp32{false}; + int32_t topk{0}; + int32_t numTokens{0}; + int32_t hidden{0}; + __nv_bfloat16 const* residual{nullptr}; //!< [numTokens, hidden] + __nv_bfloat16 const* normWeight{nullptr}; //!< [hidden] + float eps{0.F}; + int32_t* wsUc{nullptr}; //!< This rank's unicast view of the exchange (int32 words, re-armed after reading). + int32_t* wsMc{nullptr}; //!< The multicast mapping of the exchange. + int32_t* flags{nullptr}; //!< Per-CTA call-parity counters, [maxTokens * 8]. + int32_t rank{0}; + int32_t world{0}; + __nv_bfloat16* outNormed{nullptr}; + __nv_bfloat16* outResidual{nullptr}; + uint8_t* outQuant{nullptr}; //!< Optional E4M3 [numTokens, hidden]. + uint8_t* outSf{nullptr}; //!< Optional UE8M0 scales, 128x4-swizzled. + //! 0: push x (+ x2) here; 2: a producer already pushed x and launched this grid after its pushes; 3: as 2, with x + //! in exchange set A and x2 in set B. + int32_t prepushed{0}; + //! Optional Eagle3 capture of a capture layer's boundary: capOut[t * capStride + c] = the m3_spec_tail + //! capture_norm of the updated residual row with weight capWeight [hidden] and capEps, bit for bit (hidden 6144, + //! not the drafter boundary). + __nv_bfloat16* capOut{nullptr}; + int64_t capStride{0}; + __nv_bfloat16 const* capWeight{nullptr}; + float capEps{0.F}; +}; + +//! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). +bool isMinimaxM3Sm107Tp2McAllreduceNormSupported(int32_t maxTokens, int32_t opts); + +//! Launch with 8-CTA clusters of 128 threads, one cluster per token. `maxTokens` selects the exchange bucket; +//! `opts` the epilogue variant (every variant is bit-identical). +void invokeMinimaxM3Sm107Tp2McAllreduceNorm(MinimaxM3Sm107Tp2McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu new file mode 100644 index 000000000000..da30d840cc2e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu @@ -0,0 +1,339 @@ +/* + * 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/minimaxM3Sm107Tp2/producerPw2.h" + +#include +#include +#include + +// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP2 heads. +// +// The SM107 TP2 target's copy of minimaxM3/producerPw.cu (the sm_103 TP4 / TP8 targets keep theirs). Same math and +// outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with --use_fast_math like +// the rest of the library), specialized to 32 query, 2 KV and 2 index heads per rank (one index-K head), head_dim 128, +// rotary_dim 64, page 128, NVFP4 main KV and FP8 index-K. One warp per (token, packed head); K and V rows go to their +// head's block of the HND pools (the native kernel's head strides). +// +// The decode producer is latency bound: after griddepcontrol.wait the native kernel runs a chain of dependent global +// loads (packed row -> norm weight -> position id -> RoPE row -> cache slot -> quant scale). This kernel reads the +// constant norm weights and KV quant scale before griddepcontrol.wait and issues the packed row, position id and cache +// slot loads together after it, followed by the RoPE row; with `prewaitStepInputs` the position id, the slot and the +// RoPE row are read before the wait too, leaving only the packed row after it. With `early` it also calls +// griddepcontrol.launch_dependents on entry: consumers read its outputs only after their own griddepcontrol.wait, which +// waits for this grid to complete, so the early trigger only lets the next kernel launch and run its pre-wait prologue +// while this one runs. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace producerPw2 +{ +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotaryDim = 64; +constexpr int kPageSize = 128; +constexpr int kElems = kHeadDim / 32; +constexpr int kBlock = 256; + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ uint32_t sf_group_mask4() +{ + uint32_t laneId; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return 0xFu << (laneId & ~3u); +} + +__device__ __forceinline__ uint32_t fp32x8_to_e2m1(float (&a)[8]) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(a[0]), "f"(a[1]), "f"(a[2]), "f"(a[3]), "f"(a[4]), "f"(a[5]), "f"(a[6]), "f"(a[7])); + return val; +#else + return 0; +#endif +} + +__device__ __forceinline__ void store_fp8x4(__nv_fp8_e4m3* out, int64_t off, float const (&e)[kElems]) +{ + __nv_fp8x2_e4m3 const low(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const high(make_float2(e[2], e[3])); + *reinterpret_cast(out + off) = static_cast(low.__x) | (static_cast(high.__x) << 16); +} + +__device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(e[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = sf_group_mask4(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + *scaleOut = encodedScale.__x; + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float q[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + q[i] = __bfloat162float(rounded[i]) * outputScale; + return static_cast(fp32x8_to_e2m1(q)); +} + +template +__global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (kEarly) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * (kBlock / 32) + warpId; + int const tokenIdx = globalWarp / kTotalHeads; + int const localHead = globalWarp % kTotalHeads; + if (tokenIdx >= p.numTokens) + return; + + constexpr int kBegin = kHeadsQ; + constexpr int vBegin = kBegin + kHeadsKV; + constexpr int indexQBegin = vBegin + kHeadsKV; + constexpr int indexKHead = indexQBegin + kHeadsIndex; + bool const isQ = localHead < kBegin; + bool const isK = localHead >= kBegin && localHead < vBegin; + bool const isV = localHead >= vBegin && localHead < indexQBegin; + bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; + bool const isIndexK = localHead == indexKHead; + + // Constant data read before the wait: norm weights and the KV quant scale. + float w[kElems] = {}; + float quantScale = 0.0F; + __nv_bfloat16 const* normWeight + = isQ ? p.qWeight : (isK ? p.kWeight : (isIndexQ ? p.indexQWeight : p.indexKWeight)); + if (!isV) + { + uint2 const wv = *reinterpret_cast(normWeight + laneId * kElems); + float2 const w01 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.x)); + float2 const w23 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.y)); + w[0] = w01.x; + w[1] = w01.y; + w[2] = w23.x; + w[3] = w23.y; + } + if (isK || isV) + { + quantScale = p.kvQuantScale[isV ? 2 : 1]; + } + + int64_t const inputOffset = (static_cast(tokenIdx) * kTotalHeads + localHead) * kHeadDim + laneId * kElems; + bool const rotates = !isV && laneId * kElems < kRotaryDim; + bool const slotted = isK || isV || isIndexK; + int const coefficient = (laneId * kElems) % (kRotaryDim / 2); + // kPrewait: the position id, the cache slot and the RoPE row are read before the wait. They are per-step metadata + // written before the replay (host-to-device) or by plain-launched kernels earlier in the step (position_ids by the + // step's input preprocessing, the slots by msa_kv_update or its torch fallback), so they are final before this + // grid can start whatever its predecessors trigger; the RoPE table is constant. + int positionId = 0; + int slot = 0; + float4 c = make_float4(0.0F, 0.0F, 0.0F, 0.0F); + float4 s = c; + auto const loadStepInputs = [&]() + { + positionId = rotates ? p.positionIds[tokenIdx] : 0; + slot = slotted ? p.outCacheLoc[tokenIdx] : 0; + if (rotates) + { + int64_t const ropeRow = static_cast(positionId) * kRotaryDim; + c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); + s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); + } + }; + if constexpr (kPrewait) + { + loadStepInputs(); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + // The packed row comes from the predecessor grid; without kPrewait the step inputs are issued with it, so the + // native kernel's serial chain of dependent loads becomes two round trips. + uint2 const packed = *reinterpret_cast(p.packed + inputOffset); + if constexpr (!kPrewait) + { + loadStepInputs(); + } + float cosine[kElems] = {c.x, c.y, c.z, c.w}; + float sine[kElems] = {s.x, s.y, s.z, s.w}; + + float e[kElems]; + float sumSquares = 0.0F; +#pragma unroll + for (int pair = 0; pair < 2; ++pair) + { + float2 const v = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); + e[2 * pair] = v.x; + e[2 * pair + 1] = v.y; + if (!isV) + sumSquares += v.x * v.x + v.y * v.y; + } + + if (!isV) + { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + sumSquares += __shfl_xor_sync(0xffffffff, sumSquares, mask, 32); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kHeadDim) + p.eps); +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + constexpr int kPairOffset = (kRotaryDim / 2) / kElems; +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + int const dim = laneId * kElems + i; + float paired = __shfl_xor_sync(0xffffffff, e[i], kPairOffset); + if (dim < kRotaryDim) + { + if (dim < kRotaryDim / 2) + paired = -paired; + // The native kernel's contraction of e * cos + paired * sin; spelled out so this build matches it. + e[i] = __fmaf_rn(paired, sine[i], __fmul_rn(e[i], cosine[i])); + } + } + } + + if (isQ) + { + store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); + return; + } + if (isIndexQ) + { + store_fp8x4(p.indexQOut, + (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); + return; + } + + if (slot < 0) + return; + int const page = slot / kPageSize; + if (page >= p.numPages) + return; + int const withinPage = slot & (kPageSize - 1); + if (isIndexK) + { + store_fp8x4(p.indexK, + static_cast(page) * p.indexPageStride + static_cast(withinPage) * p.indexTokenStride + + laneId * kElems, + e); + return; + } + + int const head = isK ? localHead - kBegin : localHead - vBegin; + int const plane = isV ? 1 : 0; + int const scaleColumn = laneId / 4; + int64_t const scaleWithinPage = isV + ? static_cast(withinPage / 4) * 32 + static_cast(scaleColumn) * 4 + withinPage % 4 + : static_cast(withinPage) * 8 + scaleColumn; + int64_t const scaleOffset = static_cast(page) * p.scalePageStride + + static_cast(plane) * p.scalePlaneStride + static_cast(head) * p.scaleHeadStride + + scaleWithinPage; + uint8_t* scaleOut = laneId % 4 == 0 ? p.kvScale + scaleOffset : nullptr; + uint16_t const packedFp4 = quantize_fp4x4(e, quantScale, scaleOut); + int64_t const dataOffset = static_cast(page) * p.dataPageStride + + static_cast(plane) * p.dataPlaneStride + static_cast(head) * p.dataHeadStride + + static_cast(withinPage) * p.dataTokenStride + laneId * 2; + *reinterpret_cast(p.kvData + dataOffset) = packedFp4; +#endif +} + +template +cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) +{ + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; + int const warps = numTokens * kTotalHeads; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); + cfg.blockDim = dim3(kBlock); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); +} + +} // namespace + +cudaError_t invokeProducerPw2( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream) +{ + if (numTokens <= 0) + return cudaSuccess; + if (prewaitStepInputs) + { + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); + } + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); +} + +} // namespace producerPw2 +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h new file mode 100644 index 000000000000..9e94868af843 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h @@ -0,0 +1,76 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP2 heads (32 query, 2 KV and +//! 2 index heads per rank) with its constant loads before the PDL grid-dependency wait and an optional early dependents +//! trigger. Same outputs, bit for bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw2.cu. +namespace producerPw2 +{ + +struct Params +{ + __nv_bfloat16 const* packed; + __nv_fp8_e4m3* qOut; + __nv_fp8_e4m3* indexQOut; + uint8_t* kvData; + uint8_t* kvScale; + __nv_fp8_e4m3* indexK; + int const* outCacheLoc; + float const* kvQuantScale; + int64_t dataPageStride, dataPlaneStride, dataHeadStride, dataTokenStride; + int64_t scalePageStride, scalePlaneStride, scaleHeadStride; + int64_t indexPageStride, indexTokenStride, numPages; + int numTokens; + float eps; + __nv_bfloat16 const* qWeight; + __nv_bfloat16 const* kWeight; + __nv_bfloat16 const* indexQWeight; + __nv_bfloat16 const* indexKWeight; + float const* rotaryCosSin; + int const* positionIds; +}; + +//! The heads per rank this producer is built for. +constexpr int kHeadsQ = 32; +constexpr int kHeadsKV = 2; +constexpr int kHeadsIndex = 2; + +//! Launch over `numTokens` tokens, one warp per (token, packed head). `prewaitStepInputs` reads the position ids, the +//! cache slots and the RoPE rows before the grid-dependency wait (they must not be written by a PDL predecessor). +cudaError_t invokeProducerPw2( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream); + +} // namespace producerPw2 +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu new file mode 100644 index 000000000000..878282b18839 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu @@ -0,0 +1,596 @@ +/* + * 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/minimaxM3Sm107Tp2/sampledAccept.h" + +#include "tensorrt_llm/common/assert.h" + +#include +#include + +// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). +// +// sm_107 (Rubin) TP2 copy of kernels/minimaxM3/sampledAccept.cu: identical algorithm and arithmetic; the vocab shard +// limit is raised to the TP2 shard (200064 / 2 = 100032 -> kMaxVs 100352), so a CTA stages up to kMaxVs / kSplits +// BF16 logits (25 KB at 8 splits, 50 KB at 4) and the launch opts in to that dynamic shared memory. +// +// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +// active must stay on the generic path; the caller decides. +// +// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens +// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the +// nucleus with probability >= top_p, so most rows decide on their second exchange. +// +// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based +// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one +// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and +// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the +// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing +// and two call halves keep a rank one call ahead from overwriting a slot still being read. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ +namespace sampledAccept +{ +namespace +{ +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kMaxRec = kMaxWorld * kMaxSplits; +constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) +constexpr int kEx = 4; // words per record: key, logit of key, f2, f3 +constexpr int kTagBits = 64 - 32 - kIdBits; +constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard +// Record polls trap after this many spins (the multicast all-reduce's bound): a missing writer becomes a CUDA error, +// never a hang. +constexpr long long kMaxSpins = 1LL << 31; + +// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. +__host__ __device__ inline long long half_words(int world) +{ + return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * kEx; +} + +__device__ __forceinline__ uint32_t ordered(float v) +{ + uint32_t const b = __float_as_uint(v); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +// Philox4x32-10. +__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) +{ +#pragma unroll + for (int i = 0; i < 10; ++i) + { + uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; + uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; + c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); + k.x += 0x9E3779B9u; + k.y += 0xBB67AE85u; + } + return c; +} + +// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. +__device__ __forceinline__ float uniform(uint32_t bits) +{ + return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); +} + +// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. +__device__ __forceinline__ unsigned long long gkey(float lg, int id) +{ + return (static_cast(ordered(lg)) << kIdBits) + | static_cast((1u << kIdBits) - 1u - static_cast(id)); +} + +__device__ __forceinline__ int key_token(unsigned long long k) +{ + return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); +} + +__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) +{ + asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); +} + +__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) +{ + unsigned long long v; + asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); + return v; +} + +template +struct Shared +{ + __nv_bfloat16 row[kMaxVs / kSplits]; + float redf[kWarps]; + unsigned long long redk[kWarps]; + unsigned long long rk[kMaxRec]; + float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; + float wsum[kWarps]; + int owner; + int pick; + float pickx; + bool last; +}; + +template +__device__ __forceinline__ float block_max(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = s.redf[0]; + for (int i = 1; i < kWarps; ++i) + r = fmaxf(r, s.redf[i]); + __syncthreads(); + return r; +} + +// Fixed-order block sum: identical result for identical inputs. +template +__device__ __forceinline__ float block_sum(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = 0.f; + for (int i = 0; i < kWarps; ++i) + r += s.redf[i]; + __syncthreads(); + return r; +} + +template +__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + if ((threadIdx.x & 31) == 0) + s.redk[threadIdx.x >> 5] = v; + __syncthreads(); + unsigned long long r = s.redk[0]; + for (int i = 1; i < kWarps; ++i) + r = s.redk[i] > r ? s.redk[i] : r; + __syncthreads(); + return r; +} + +// Exclusive prefix of v over the threads in thread order, and the total. +template +__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + float inc = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + float const t = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) + inc += t; + } + if (lane == 31) + s.wsum[warp] = inc; + __syncthreads(); + float base = 0.f; + total = 0.f; + for (int i = 0; i < kWarps; ++i) + { + base += i < warp ? s.wsum[i] : 0.f; + total += s.wsum[i]; + } + __syncthreads(); + return base + inc - v; +} + +// A row's kSplits CTAs wait for each other's records, so they launch as one thread-block cluster: the hardware +// schedules a cluster's CTAs together, whatever other streams hold. Rows never wait on other rows. +template +__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) +{ + extern __shared__ __align__(16) unsigned char smem[]; + Shared& s = *reinterpret_cast*>(smem); + float constexpr kNegInf = -__builtin_huge_valf(); + int const tid = threadIdx.x; + int const row = blockIdx.x / kSplits; + int const split = blockIdx.x % kSplits; + int const world = a.world; + int const nrec = world * kSplits; + // The row's sampling parameters are step inputs (the spec metadata, written before the target forward), not the + // predecessor head GEMM's output, so with prewait_params they load before the dependency wait and overlap the + // head's tail; the same values either way. + float temperature = 0.f, topP = 0.f; + long long seed_in = 0, offset_in = 0; + auto load_params = [&]() + { + temperature = a.temperatures[row]; + topP = a.top_ps[row]; + seed_in = a.seeds[row]; + offset_in = a.offsets[row]; + }; + if (a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (!a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // early_trigger: the dependent grid (the first draft pass's input) may launch now and park at its own + // griddepcontrol.wait, which still returns only once this grid has completed and flushed; it skips the launch + // latency after this kernel ends. Otherwise dependents are released at the acceptance epilogue (the CTAs of + // splits 1.. exit without releasing). + if (a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const tempInv = 1.0f / (temperature + kTempEpsilon); + unsigned long long const seed = static_cast(seed_in); + unsigned long long const offset = static_cast(offset_in); + unsigned int const call = *reinterpret_cast(a.counters + 1); + // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of + // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round + // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten + // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. + auto ktag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; + }; + auto ftag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; + }; + long long const rec_stride = static_cast(world) * kMaxSplits * kEx; + unsigned long long* const xu + = a.ws_uc + static_cast(call & 1u) * half_words(world) + static_cast(row) * rec_stride; + unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * kEx; + long long const round_stride = static_cast(kMaxRows) * rec_stride; + + // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- + __nv_bfloat16 const* src = a.logits + row * a.ld; + int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; + int const s0 = min(split * sl_len, a.Vs); + int const n = min(s0 + sl_len, a.Vs) - s0; + int const id0 = a.vocab_offset + s0; + float tmax = kNegInf; + bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; + int const nvec = vec ? n / 8 : 0; + for (int v = tid; v < nvec; v += kThreads) + { + uint4 const w = *reinterpret_cast(src + s0 + 8 * v); + *reinterpret_cast(&s.row[8 * v]) = w; + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); +#pragma unroll + for (int j = 0; j < 8; ++j) + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); + } + for (int i = nvec * 8 + tid; i < n; i += kThreads) + { + s.row[i] = src[s0 + i]; + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); + } + float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier + float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z + unsigned long long ckey = 0ULL; + int const gid = a.rank * kMaxSplits + split; + for (int r = 0;; ++r) + { + // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) + float const ref = r == 0 ? smax : M; + float const wmin = r == 0 ? 0.f : wc; + float ts = 0.f; + for (int i = tid; i < n; i += kThreads) + { + float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); + ts += w > wmin ? w : 0.f; + } + float total; + float const excl = block_excl_scan(ts, total, s); + uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), + static_cast(offset), static_cast(offset >> 32)), + make_uint2(static_cast(seed), static_cast(seed >> 32))); + float const tgt = uniform(rnd.x) * total; + if (tid == 0) + s.owner = -1; + __syncthreads(); + if (ts > 0.f && excl <= tgt) + atomicMax(&s.owner, tid); + __syncthreads(); + if (tid == s.owner) + { + float acc = excl; + int pick = -1; + float px = 0.f; + if constexpr (kIlpWalk) + { + // SA-PW: the owner's strided tokens 8 at a time. The 8 x and w (the same expressions) are + // independent, so their MUFU work overlaps; the serial sum then adds them in the same order with + // the same early exit: the same acc sequence, pick and px as the one-token walk, bit for bit. + constexpr int kW = 8; + bool done = false; + for (int i0 = tid; i0 < n && !done; i0 += kW * kThreads) + { + float xs[kW], wv[kW]; +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + xs[k] = i < n ? __fmul_rn(__bfloat162float(s.row[i]), tempInv) : 0.f; + wv[k] = __expf(xs[k] - ref); + } +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + if (!done && i < n && wv[k] > wmin) + { + acc += wv[k]; + pick = i; + px = xs[k]; + if (acc > tgt) + done = true; + } + } + } + } + else + { + for (int i = tid; i < n; i += kThreads) + { + float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); + float const w = __expf(x - ref); + if (w > wmin) + { + acc += w; + pick = i; + px = x; + if (acc > tgt) + break; + } + } + } + s.pick = pick; + s.pickx = px; + } + __syncthreads(); + unsigned long long sb = 0ULL; + if (total > 0.f && s.owner >= 0) + sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); + float const smab = total; + unsigned long long const ktag = ktag_of(r); + unsigned long long const ftag = ftag_of(r); + unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; + if (tid == 0) + { + unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; + st_sys(xw + 0, (sb << kTagBits) | ktag); + st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); + st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); + st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); + } + // ---- every record of the row, every rank --------------------------------------------------------------- + if (tid < nrec * kEx) + { + int const q = tid / kEx, f = tid % kEx; + unsigned long long v; + if (f == 0) + { + long long spins = 0; + do + { + v = ld_sys(xr + (q / kSplits * kMaxSplits + q % kSplits) * kEx); + if (++spins > kMaxSpins) + __trap(); + } while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); + s.rk[q] = v >> kTagBits; + } + else + { + long long spins = 0; + do + { + v = ld_sys(xr + (q / kSplits * kMaxSplits + q % kSplits) * kEx + f); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFFFFF00000000ULL) != ftag); + float const fv = __uint_as_float(static_cast(v)); + (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; + } + } + __syncthreads(); + // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- + unsigned long long gk = 0ULL; + float gv = 0.f; + for (int q = 0; q < nrec; ++q) + { + if (s.rk[q] > gk) + { + gk = s.rk[q]; + gv = s.rv[q]; + } + } + bool accept; + if (r == 0) + { + M = kNegInf; + for (int q = 0; q < nrec; ++q) + M = fmaxf(M, s.r2[q]); + float Z = 0.f; + for (int q = 0; q < nrec; ++q) + Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); + target = topP * Z; + accept = !(topP < 1.0f); // no nucleus cut: the first draw stands + ckey = gk; + cv = gv; + } + else + { + float S = 0.f; // mass strictly above the current candidate + for (int q = 0; q < nrec; ++q) + S += s.r2[q]; + accept = S <= target || gk == 0ULL; + if (!accept) + { + ckey = gk; + cv = gv; + } + } + wc = __expf(cv - M); + __syncthreads(); // records consumed before the next round's writes to s + if (accept) + { + if (split != 0) + return; + if (tid == 0) + { + a.tokens[row] = key_token(ckey); + if (a.extra_attempts != nullptr && r > 1) + atomicAdd(a.extra_attempts, r - 1); + } + break; + } + } + + // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (!a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (tid == 0) + s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; + __syncthreads(); + if (!s.last) + return; + __threadfence(); + if (tid == 0) + { + a.counters[0] = 0u; + a.counters[1] = call + 1u; + } + if (tid < a.nreq) + { + int const* tok = a.tokens; + int const b = tid; + int const k1 = a.draft_len + 1; + int nacc = 1; + bool run = true; + for (int k = 0; k < a.draft_len; ++k) + { + int const tt = *reinterpret_cast(tok + b * k1 + k); + a.acc_tokens[b * k1 + k] = tt; + run = run && a.draft[b * a.draft_len + k] == tt; + nacc += run ? 1 : 0; + } + a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); + if (a.force_vec != nullptr) + nacc = a.force_vec[b]; + else if (a.force_total > 0) + nacc = a.force_total; + a.num_acc[b] = nacc; + a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; + if (a.next_new != nullptr) + a.next_new[b * a.next_stride] + = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); + } +} + +template +cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) +{ + static int max_clusters = 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(a.rows * kSplits); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Shared); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + if (max_clusters == 0) + { + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } + } + TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); + return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); +} + +} // namespace + +long long sampledAcceptWsWords(int world) +{ + return 2 * half_words(world); +} + +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) +{ + // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); + // more rows: 4 CTAs per row. + if (a.rows <= 32) + { + return a.ilp_walk ? launch<8, true>(a, usePdl, stream) : launch<8, false>(a, usePdl, stream); + } + return a.ilp_walk ? launch<4, true>(a, usePdl, stream) : launch<4, false>(a, usePdl, stream); +} + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h new file mode 100644 index 000000000000..cd1e8ec7fefa --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.h @@ -0,0 +1,82 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ +//! Fused sampled (temperature + top-p) verification of the MiniMax-M3 Eagle3 target head over the TP-sharded +//! vocabulary, decode rows <= kMaxRows; see sampledAccept.cu. +namespace sampledAccept +{ + +constexpr int kMaxSplits = 8; // CTAs per verify row (at most) +constexpr int kMaxRows = 48; // verify rows per call +constexpr int kMaxWorld = 8; // TP ranks +constexpr int kMaxVs = 100352; // vocab shard width (TP2: 200064 / 2 = 100032, padded to a multiple of 512) +constexpr int kIdBits = 18; // token ids < 2^18 + +struct Args +{ + __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard + long long ld; + int rows, Vs, vocab_offset; + float const* temperatures; + float const* top_ps; + long long const* seeds; + long long const* offsets; + unsigned long long* ws_uc; + unsigned long long* ws_mc; + unsigned int* counters; // [0] finished rows, [1] call + int* tokens; // [rows] sampled token per row + int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows + int rank, world; + int const* draft; + int const* batch_idx; + int* acc_tokens; + int* num_acc; + long long* gather_ids; + int draft_len, nreq, force_total; + int* next_new; + int next_stride; + int const* force_vec; + int prewait_params; // nonzero: the row's sampling parameters load before the PDL dependency wait + int early_trigger; // nonzero: dependents are released right after the dependency wait, not at the epilogue + int ilp_walk; // nonzero: the owner walk takes its tokens 8 at a time (same weights, same summation order) +}; + +//! Words (uint64) of the multicast exchange workspace for `world` ranks: two call halves. +long long sampledAcceptWsWords(int world); + +//! Launch the sampled verification over a.rows * (8 or 4) co-resident CTAs (8 per row up to 32 rows). +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream); + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.cu new file mode 100644 index 000000000000..e708083ef9ff --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.cu @@ -0,0 +1,246 @@ +/* + * 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. + */ + +// Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. +// +// One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +// thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups `block % 16` is a lower +// bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared +// memory and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every +// candidate against the whole row (re-read). + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr float kInitScore = 1.0e30f; +constexpr float kLocalScore = 1.0e29f; +constexpr uint32_t kNegInfKey = 0xff800000u; +constexpr int kItems = 9; // keys per thread: rows up to 9216 blocks (1M tokens + draft tokens) + +// Descending order key: a smaller key is a larger score (NaN above +inf), as in the reference op. +__device__ __forceinline__ uint32_t orderedKey(float x) +{ + uint32_t const b = __float_as_uint(x); + return (b & 0x80000000u) ? b : (~b & 0x7fffffffu); +} + +__device__ __forceinline__ float applyForcing(float raw, int block, int initBlocks, long long localStart) +{ + if (block >= localStart) + return kLocalScore; + if (block < initBlocks) + return kInitScore; + return raw; +} + +struct SelectSmem +{ + uint32_t groupMin[16]; + unsigned long long cand[1024]; + int numCand; + unsigned long long picked[16]; +}; + +// One row with up to kItems * 1024 valid blocks; the caller picks kItems from the row's valid count. +template +__device__ __forceinline__ void selectRow(SelectSmem& sm, float const* rowBase, long long blockStride, int valid, + int initBlocks, long long localStart, bool loadedTrigger) +{ + int const tid = threadIdx.x, lane = tid & 31; + uint32_t key[kItems]; + uint32_t localMin = 0xffffffffu; +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + key[s] = block < valid ? orderedKey(applyForcing(rowBase[block * blockStride], block, initBlocks, localStart)) + : 0xffffffffu; + localMin = min(localMin, key[s]); + } + if (valid <= 16) + { + if (tid < valid) + sm.picked[tid] = (static_cast(key[0]) << 32) | static_cast(tid); + if (loadedTrigger) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + return; + } + // Group g = block % 16 = tid % 16: lanes l and l ^ 16 of every warp. + localMin = min(localMin, __shfl_xor_sync(0xffffffffu, localMin, 16)); + if (lane < 16) + atomicMin(&sm.groupMin[lane], localMin); + __syncthreads(); + // Every score of the row is loaded: release the sparse decode, so its launch and its pre-wait prologue overlap + // the rank below. This comes after this grid's own dependency wait, which the sparse decode relies on: it reads + // q and the seq_lens / block-table row before its own wait (sparseDecodeLc2.cu), safe only while every kernel + // between the producer and it releases its dependents after its own wait. + if (loadedTrigger) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + uint32_t bound = 0; +#pragma unroll + for (int g = 0; g < 16; ++g) + bound = max(bound, sm.groupMin[g]); +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + bool const isCand = block < valid && key[s] <= bound; + unsigned const ballot = __ballot_sync(0xffffffffu, isCand); + int base = 0; + if (lane == 0 && ballot) + base = atomicAdd(&sm.numCand, __popc(ballot)); + base = __shfl_sync(0xffffffffu, base, 0); + if (isCand) + { + int const i = base + __popc(ballot & ((1u << lane) - 1)); + if (i < 1024) + sm.cand[i] = (static_cast(key[s]) << 32) | static_cast(block); + } + } + __syncthreads(); + int const n = sm.numCand; + if (n <= 1024) + { + for (int i = tid; i < n; i += 1024) + { + unsigned long long const mine = sm.cand[i]; + int rank = 0; + for (int j = 0; j < n; ++j) + rank += sm.cand[j] < mine; + if (rank < 16) + sm.picked[rank] = mine; + } + return; + } + // More than 1024 keys tied at the bound (degenerate rows): rank every candidate against the whole row. +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + if (block < valid && key[s] <= bound) + { + int rank = 0; + for (int j = 0; j < valid; ++j) + { + uint32_t const other = orderedKey(applyForcing(rowBase[j * blockStride], j, initBlocks, localStart)); + rank += other < key[s] || (other == key[s] && j < block); + } + if (rank < 16) + sm.picked[rank] = (static_cast(key[s]) << 32) | static_cast(block); + } + } +} + +__global__ void __launch_bounds__(1024) m3_select_long2_kernel(float const* __restrict__ scores, long long headStride, + long long blockStride, long long queryStride, int const* __restrict__ nValid, int* __restrict__ out, + long long outQueryStride, long long outHeadStride, int numKvHeads, int numBlocks, int initBlocks, int localBlocks, + bool earlyTrigger, bool loadedTrigger) +{ + __shared__ SelectSmem sm; + int const row = blockIdx.x, query = row / numKvHeads, head = row % numKvHeads; + int const tid = threadIdx.x, lane = tid & 31; + if (tid < 16) + { + sm.groupMin[tid] = 0xffffffffu; + sm.picked[tid] = ~0ull; + } + if (tid == 0) + sm.numCand = 0; +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); + // The sparse decode that follows stages its table-independent inputs before its own grid wait; past this grid's + // wait those inputs are final, so it may launch now rather than at this grid's exit. + if (earlyTrigger) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + int const raw = nValid[query]; + int const valid = max(0, min(raw, numBlocks)); + long long const localStart = max(static_cast(raw) - localBlocks, 0LL); + float const* rowBase = scores + head * headStride + query * queryStride; + __syncthreads(); + if (valid <= 2048) + selectRow<2>(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + else if (valid <= 4096) + selectRow<4>(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + else + selectRow(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + __syncthreads(); + if (tid >= 32) + return; + uint32_t blockKey = ~0u; + if (lane < 16) + { + unsigned long long const c = sm.picked[lane]; + if (c != ~0ull && static_cast(c >> 32) != kNegInfKey) + blockKey = static_cast(c & 0xffffffffu); + } +#pragma unroll + for (int size = 2; size <= 32; size <<= 1) + { +#pragma unroll + for (int stride = size >> 1; stride > 0; stride >>= 1) + { + uint32_t const other = __shfl_xor_sync(0xffffffffu, blockKey, stride); + bool const ascending = (lane & size) == 0, lower = (lane & stride) == 0; + blockKey = (lower == ascending) ? min(blockKey, other) : max(blockKey, other); + } + } + if (lane < 16) + out[query * outQueryStride + head * outHeadStride + lane] = blockKey == ~0u ? -1 : static_cast(blockKey); +} + +} // namespace + +cudaError_t invokeSelectLong2(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* out, int64_t outQueryStride, int64_t outHeadStride, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQ, int32_t initBlocks, int32_t localBlocks, bool usePdl, bool earlyTrigger, + int32_t triggerPhase, cudaStream_t stream) +{ + if (numBlocks > kSelectLong2MaxBlocks || triggerPhase < 0 || triggerPhase > 1) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(totalQ * numKvHeads); + cfg.blockDim = dim3(1024); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = usePdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_select_long2_kernel, scores, static_cast(headStride), + static_cast(blockStride), static_cast(queryStride), nValidBlocks, out, + static_cast(outQueryStride), static_cast(outHeadStride), numKvHeads, numBlocks, + initBlocks, localBlocks, earlyTrigger, !earlyTrigger && triggerPhase == 1); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h new file mode 100644 index 000000000000..9fba25f03c77 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/selectLong2.h @@ -0,0 +1,50 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ + +//! Largest row (blocks per query) the long-row top-16 select covers: 9 keys per thread of one 1024-thread CTA. +constexpr int32_t kSelectLong2MaxBlocks = 9216; + +//! MiniMax-M3 top-16 block select over per-KV-head scores [numKvHeads, numBlocks, totalQ] (element strides given), +//! one 1024-thread CTA per (query, KV head). Same contract as minimax_m3_select_blocks: init / local forcing, the +//! per-query valid-block count, ties to the smaller block, forced -inf picks emitted as -1, ascending ids with -1 +//! tail padding. Writes out[query * outQueryStride + head * outHeadStride + 0..15]. The SM107 TP2 copy of +//! minimaxM3/selectLong with a choice of where the dependent grid (the sparse decode) is released: at this grid's exit +//! (triggerPhase 0), right after every score of the row is loaded (1), or, with earlyTrigger, right after this grid's +//! dependency wait. Every release point follows this grid's own dependency wait. +cudaError_t invokeSelectLong2(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* out, int64_t outQueryStride, int64_t outHeadStride, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQ, int32_t initBlocks, int32_t localBlocks, bool usePdl, bool earlyTrigger, + int32_t triggerPhase, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu new file mode 100644 index 000000000000..cb7d83b89bb4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu @@ -0,0 +1,89 @@ +/* + * 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/minimaxM3Sm107Tp2/siluAndMul.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr int kThreads = 128; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves). The +// instruction sequence of the Triton kernel's PTX at its defaults (host/silu/silu_and_mul_plain.ptx): g * -log2(e) as +// one product, ex2.approx, + 1, div.full 1 / that, g * s, (u + 0.0) as an IEEE add, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t siluMulBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 gl, gh, ul, uh, t0, t1, e0, e1, s0, s1, p0, p1, q0, q1, r0, r1, one, zero;\n\t" + ".reg .b16 hl, hh;\n\t" + "mov.b32 {hl, hh}, %1;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mul.rn.f32 t0, gl, 0fBFB8AA3B;\n\tmul.rn.f32 t1, gh, 0fBFB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 p0, gl, s0;\n\tmul.rn.f32 p1, gh, s1;\n\t" + "mov.b32 zero, 0f00000000;\n\tadd.rn.f32 q0, ul, zero;\n\tadd.rn.f32 q1, uh, zero;\n\t" + "mul.rn.f32 r0, q0, p0;\n\tmul.rn.f32 r1, q1, p1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +// Eight outputs per thread: 16-byte loads of the gate and the up half, one 16-byte store. +__global__ void __launch_bounds__(kThreads) + siluAndMulKernel(uint4 const* __restrict__ x, uint4* __restrict__ out, int64_t rows, int64_t vecs) +{ + int64_t const row = blockIdx.y; + int64_t const v = static_cast(blockIdx.x) * kThreads + threadIdx.x; + if (row >= rows || v >= vecs) + { + return; + } + uint4 const g = x[row * 2 * vecs + v]; + uint4 const u = x[row * 2 * vecs + vecs + v]; + uint4 o; + o.x = siluMulBf16x2(g.x, u.x); + o.y = siluMulBf16x2(g.y, u.y); + o.z = siluMulBf16x2(g.z, u.z); + o.w = siluMulBf16x2(g.w, u.w); + out[row * vecs + v] = o; +} + +} // namespace + +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream) +{ + int64_t const vecs = d / 8; + dim3 const grid(static_cast((vecs + kThreads - 1) / kThreads), static_cast(rows)); + siluAndMulKernel<<>>( + reinterpret_cast(x), reinterpret_cast(out), rows, vecs); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h new file mode 100644 index 000000000000..01b32137ac4a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h @@ -0,0 +1,40 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ + +//! out[m, j] = bf16(silu(g) * (u + 0)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both +//! contiguous, d a multiple of 8, 16-byte aligned rows): bit for bit the Triton silu_and_mul kernel at its defaults +//! (no limit, alpha 1, beta 0), the plain SwiGLU of the Eagle3 drafter MLP. +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu new file mode 100644 index 000000000000..d96f51d5d150 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu @@ -0,0 +1,640 @@ +/* + * 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. + */ + +// MiniMax-M3 long-context sparse decode attention in one launch, SM107 TP2 geometry (32 query heads on 2 KV heads +// per rank; staircase copy of minimaxM3/sparseDecodeLc.cu with a KV-head grid dimension). +// +// Grid: (tokens * 16, kv heads), one cluster of 16 CTAs per (query token, KV head); CTA `r` attends the token's `r`-th +// selected block (the Triton producer's chunk `r` at 16 chunks), keeps the partial in shared memory, and after a +// cluster barrier merges head `r` from all 16 partials over DSMEM, writing the BF16 row and, optionally, its MXFP8 copy +// and swizzled UE8M0 scales exactly as the Triton merge's `_store_mxfp8_head` lays them out. +// +// Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +// PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +// log-domain recurrence. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace +{ + +constexpr int kHeads = 16; +constexpr int kDim = 128; +constexpr int kPage = 128; +constexpr int kTopk = 16; +constexpr int kThreads = 128; +constexpr int kPitch = kDim + 8; // halves per smem row: 272 B, conflict-free ldmatrix +constexpr float kNegInf = -__builtin_huge_valf(); +constexpr int kBtCap = 4096; // block-table entries staged in smem by PREWAIT (512K tokens) + +using Params = MinimaxM3Sm107Tp2SparseDecodeLcParams; + +struct Smem +{ + __half qs[kHeads * kPitch]; + __half vs[kPage * kPitch]; + __half ps[kHeads * kPitch]; + float recv[kTopk][kDim]; // head `rank` partials pushed by the 16 CTAs, chunk-major + float recv_lse[kTopk]; + float wgt[kTopk]; // the 16 merge weights, computed once per CTA + float red[4][kHeads]; + float rowm[kHeads]; + float rowl[kHeads]; + alignas(8) unsigned long long rbar; // the split-K partials of head `rank` landed (complete_tx) +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t e2m1x2_to_f16x2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 s;\n cvt.u8.u32 s, %1;\n cvt.rn.f16x2.e2m1x2 %0, s;\n}" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t e4m3x2_to_f16x2(uint16_t pair) +{ + uint32_t out; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(pair)); + return out; +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ void ldsm_x4(uint32_t (&r)[4], uint32_t addr) +{ + 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"(addr)); +} + +__device__ __forceinline__ void ldsm_x4_t(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma(float (&c)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.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"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint32_t map_cluster(void const* local, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smem_u32(local)), "r"(rank)); + return remote; +} + +// Posted DSMEM stores that complete_tx on the receiving CTA's mbarrier `bar`. +__device__ __forceinline__ void st_async_v2(uint32_t addr, float a, float b, uint32_t bar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(addr), + "f"(a), "f"(b), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ void st_async_f32(uint32_t addr, float a, uint32_t bar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(addr), + "r"(__float_as_uint(a)), "r"(bar) + : "memory"); +} + +// One packed NVFP4 row (64 bytes, 128 values) with its 8 E4M3 scales -> 128 fp16 values in smem. +__device__ __forceinline__ void dequant_row(__half* dst, uint8_t const* packed, uint8_t const (&sc)[8], bool valid) +{ + uint4 const* src = reinterpret_cast(packed); + uint32_t* d = reinterpret_cast(dst); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + uint32_t words[4] = {w.x, w.y, w.z, w.w}; + // 16 bytes = 32 values = scale columns 2c and 2c + 1. +#pragma unroll + for (int h = 0; h < 2; ++h) + { + uint8_t const s = valid ? sc[2 * c + h] : 0; + uint32_t const sh = e4m3x2_to_f16x2(static_cast(s | (s << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); +#pragma unroll + for (int j = 0; j < 2; ++j) + { + uint32_t const word = words[2 * h + j]; +#pragma unroll + for (int b = 0; b < 4; ++b) + { + uint32_t const v2 = e2m1x2_to_f16x2((word >> (8 * b)) & 0xFF); + __half2 const x = __hmul2(*reinterpret_cast<__half2 const*>(&v2), scale); + d[c * 16 + h * 8 + j * 4 + b] = *reinterpret_cast(&x); + } + } + } + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +template +constexpr size_t smemBytes() +{ + return PREWAIT ? sizeof(Smem) + kBtCap * sizeof(int) : sizeof(Smem); +} + +template +__global__ void __launch_bounds__(kThreads) m3_lc_attn_kernel(__grid_constant__ Params const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(128) uint8_t smem_raw[]; + Smem& s = *reinterpret_cast(smem_raw); + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const tok = blockIdx.x / kTopk; + int const kvh = blockIdx.y; + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + + if (tid == 0) + { + // One arrival (here) plus every cluster CTA's partial row (kDim floats) and LSE for head `rank`. + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smem_u32(&s.rbar)) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(&s.rbar)), + "r"(kTopk * (kDim + 1) * 4) + : "memory"); + } + // Every peer must be running before its shared memory takes a DSMEM store: arrive now, wait before pushing. + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + + // Pre-wait read contract: everything that does not depend on the top-16 table is read before the grid dependency. + // - seq_lens: written inside the step's graph by msa_kv_update (the KV-length correction). That kernel is a plain + // launch, not a programmatic dependent, so it has completed before any later kernel of the step starts. + // - the query rows: written by the producer, two grids back. This holds only while every kernel between the + // producer and this grid (the long-row select) triggers its dependents after its own griddepcontrol.wait, so + // this grid cannot start before the producer completed. A kernel inserted into that chain, or a select that + // triggers before its wait, breaks it: these reads then move after the griddepcontrol.wait below. + int const req = tok / p.dql; + int const qoff = tok - req * p.dql; + int const seq = p.seq_lens[req]; + int const kv_len = max(seq - p.dql + qoff + 1, 0); + int const nblocks = (kv_len + kPage - 1) / kPage; + int const real_topk = min(kTopk, nblocks); + // q: 16 heads x 128 e4m3 -> fp16; thread = (head, 16-value chunk). + { + int const h = tid >> 3; + int const c = (tid & 7) * 16; + uint4 const w = *reinterpret_cast(p.q + tok * p.q_tok + (kvh * kHeads + h) * p.q_head + c); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; + uint32_t* d = reinterpret_cast(s.qs + h * kPitch + c); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + d[2 * i] = e4m3x2_to_f16x2(static_cast(words[i] & 0xFFFF)); + d[2 * i + 1] = e4m3x2_to_f16x2(static_cast(words[i] >> 16)); + } + } + + // PREWAIT: the request's block-table row (written before the graph runs) is staged before the grid dependency. + int* const s_bt = reinterpret_cast(smem_raw + sizeof(Smem)); + bool bt_smem = false; + if constexpr (PREWAIT) + { + int const nbt = (max(seq, 0) + kPage - 1) / kPage; + bt_smem = nbt <= kBtCap; + if (bt_smem) + { + for (int i = tid; i < nbt; i += kThreads) + s_bt[i] = __ldg(p.block_table + req * p.bt_stride + i); + } + } + // N1: the block this (KV head, row, slot) selected at the layer's previous decode step recurs ~2/3 of the time; + // pull its K/V page and scales into L2 while the select runs. The ids were written by this kernel one step back + // (complete before the graph's next replay) and the block-table row before the graph runs. Read and written by + // the same thread, so the record below cannot overtake this read. Out-of-range ids (stale rows) are skipped. + bool const n1 = p.n1_ids != nullptr && tok < p.n1_rows; + int* const n1_slot = n1 ? p.n1_ids + kvh * p.n1_h + tok * kTopk + rank : nullptr; + if (n1 && tid == 64) + { + int const nb = *n1_slot; + if (nb >= 0 && nb < nblocks) + { + int64_t const pg = __ldg(p.block_table + req * p.bt_stride + nb); + if (pg >= 0 && pg < p.n1_pages) + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.k + pg * p.k_blk + kvh * p.k_h), + "r"(kPage * 64) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.v + pg * p.v_blk + kvh * p.v_h), + "r"(kPage * 64) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.ksc + pg * p.ksc_blk + kvh * p.ksc_h), + "r"(kPage * 8) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.vsc + pg * p.vsc_blk + kvh * p.vsc_h), + "r"(kPage * 8) + : "memory"); + } + } + } + asm volatile("griddepcontrol.wait;" ::: "memory"); + if constexpr (EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } + int blk = static_cast(rank) < real_topk ? p.topk[kvh * p.t_h + tok * p.t_tok + rank * p.t_k] : -1; + if (n1 && tid == 64) + { + *n1_slot = blk; + } + bool const active = blk >= 0; + if constexpr (PREWAIT) + { + __syncthreads(); // s_bt complete + } + int64_t const page_pw = !active ? 0 : (bt_smem ? s_bt[blk] : p.block_table[req * p.bt_stride + blk]); + + float m_row[2] = {kNegInf, kNegInf}; + float acc[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + acc[j][i] = 0.F; + float l_row[2] = {0.F, 0.F}; + + // K stays packed in registers: the MMA B fragment of token n = lane / 4 at k-step s is the f16x2 pair of + // packed bytes 8s + q and 8s + 4 + q (q = lane % 4), both under scale column s. + uint32_t kw[4][16]; + uint2 ksc[4]; + int64_t page = 0; + if (active) + { + page = page_pw; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const n = warp * 32 + j * 8 + (lane >> 2); + bool const valid = blk * kPage + n < kv_len; + uint4 const* src = reinterpret_cast(p.k + page * p.k_blk + kvh * p.k_h + n * 64); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 const w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + kw[j][4 * c] = w.x; + kw[j][4 * c + 1] = w.y; + kw[j][4 * c + 2] = w.z; + kw[j][4 * c + 3] = w.w; + } + ksc[j] = valid ? __ldg(reinterpret_cast(p.ksc + page * p.ksc_blk + kvh * p.ksc_h + n * 8)) + : make_uint2(0, 0); + } + bool const valid = blk * kPage + tid < kv_len; + uint8_t vsc[8]; + if (valid) + { + uint8_t const* vb = p.vsc + page * p.vsc_blk + kvh * p.vsc_h + (tid >> 2) * 32 + (tid & 3); +#pragma unroll + for (int c = 0; c < 8; ++c) + vsc[c] = __ldg(vb + 4 * c); + } + dequant_row(s.vs + tid * kPitch, p.v + page * p.v_blk + kvh * p.v_h + tid * 64, vsc, valid); + } + __syncthreads(); + + if (active) + { + float const scale_log2 = p.sm_scale * 1.4426950409F * __ldg(p.kgs); + int const q4 = (lane & 3) * 8; + // S = q K^T for this warp's 32 tokens: rows g / g + 8, columns 2 * (lane % 4) + {0, 1} of each n-tile. + float sacc[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + sacc[j][i] = 0.F; +#pragma unroll + for (int ks = 0; ks < kDim / 16; ++ks) + { + uint32_t a[4]; + ldsm_x4(a, smem_u32(s.qs + (lane & 15) * kPitch + ks * 16 + (lane >> 4) * 8)); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + uint32_t const scw = ks < 4 ? ksc[j].x : ksc[j].y; + uint8_t const sc = static_cast(scw >> (8 * (ks & 3))); + uint32_t const sh = e4m3x2_to_f16x2(static_cast(sc | (sc << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); + uint32_t const r0 = e2m1x2_to_f16x2((kw[j][2 * ks] >> q4) & 0xFF); + uint32_t const r1 = e2m1x2_to_f16x2((kw[j][2 * ks + 1] >> q4) & 0xFF); + __half2 const h0 = __hmul2(*reinterpret_cast<__half2 const*>(&r0), scale); + __half2 const h1 = __hmul2(*reinterpret_cast<__half2 const*>(&r1), scale); + mma(sacc[j], a, *reinterpret_cast(&h0), *reinterpret_cast(&h1)); + } + } + int const g = lane >> 2; + float mloc[2] = {kNegInf, kNegInf}; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const t = warp * 32 + j * 8 + 2 * (lane & 3) + (i & 1); + float const v = blk * kPage + t < kv_len ? 0.F + sacc[j][i] * scale_log2 : kNegInf; + sacc[j][i] = v; + mloc[i >> 1] = fmaxf(mloc[i >> 1], v); + } +#pragma unroll + for (int r = 0; r < 2; ++r) + { + mloc[r] = fmaxf(mloc[r], __shfl_xor_sync(0xffffffff, mloc[r], 1)); + mloc[r] = fmaxf(mloc[r], __shfl_xor_sync(0xffffffff, mloc[r], 2)); + } + if ((lane & 3) == 0) + { + s.red[warp][g] = mloc[0]; + s.red[warp][g + 8] = mloc[1]; + } + __syncthreads(); + if (tid < kHeads) + s.rowm[tid] = fmaxf(fmaxf(s.red[0][tid], s.red[1][tid]), fmaxf(s.red[2][tid], s.red[3][tid])); + __syncthreads(); + m_row[0] = s.rowm[g]; + m_row[1] = s.rowm[g + 8]; + float lloc[2] = {0.F, 0.F}; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; i += 2) + { + float const p0 = ex2(sacc[j][i] - m_row[i >> 1]); + float const p1 = ex2(sacc[j][i + 1] - m_row[i >> 1]); + lloc[i >> 1] += p0 + p1; + int const row = g + (i >> 1) * 8; + int const t = warp * 32 + j * 8 + 2 * (lane & 3); + *reinterpret_cast<__half2*>(s.ps + row * kPitch + t) = __floats2half2_rn(p0, p1); + } +#pragma unroll + for (int r = 0; r < 2; ++r) + { + lloc[r] += __shfl_xor_sync(0xffffffff, lloc[r], 1); + lloc[r] += __shfl_xor_sync(0xffffffff, lloc[r], 2); + } + __syncthreads(); // rowm reads done before red is reused; ps complete + if ((lane & 3) == 0) + { + s.red[warp][g] = lloc[0]; + s.red[warp][g + 8] = lloc[1]; + } + // O[:, warp's 32 dims] = P V over all 128 tokens. +#pragma unroll + for (int kk = 0; kk < kPage; kk += 16) + { + uint32_t a[4]; + ldsm_x4(a, smem_u32(s.ps + (lane & 15) * kPitch + kk + (lane >> 4) * 8)); +#pragma unroll + for (int jp = 0; jp < 2; ++jp) + { + uint32_t b[4]; + int const t = kk + (lane & 7) + ((lane >> 3) & 1) * 8; + int const d = warp * 32 + jp * 16 + (lane >> 4) * 8; + ldsm_x4_t(b, smem_u32(s.vs + t * kPitch + d)); + mma(acc[2 * jp], a, b[0], b[1]); + mma(acc[2 * jp + 1], a, b[2], b[3]); + } + } + __syncthreads(); + if (tid < kHeads) + s.rowl[tid] = (s.red[0][tid] + s.red[1][tid]) + (s.red[2][tid] + s.red[3][tid]); + __syncthreads(); + l_row[0] = s.rowl[g]; + l_row[1] = s.rowl[g + 8]; + } + + // Normalize the partial and push head h's row with its LSE (log2) into CTA h's recv[rank] slot, as the Triton + // producer stores chunk `rank`; st.async stores that complete_tx on CTA h's barrier. + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + { + int const g = lane >> 2; + float const vgs = __ldg(p.vgs); +#pragma unroll + for (int r = 0; r < 2; ++r) + { + float scale, lse; + if constexpr (LINEAR) + { + scale = l_row[r] > 0.F ? 1.0F / l_row[r] : 0.F; + lse = l_row[r] > 0.F ? m_row[r] + lg2(l_row[r]) : kNegInf; + } + else + { + float const lse_i = active ? m_row[r] + lg2(ex2(kNegInf - m_row[r]) + l_row[r]) : kNegInf; + scale = lse_i > kNegInf ? ex2(m_row[r] - lse_i) : 0.F; + lse = lse_i; + } + scale = scale * vgs; + int const row = g + r * 8; + uint32_t const dst = map_cluster(&s.recv[rank][0], row); + uint32_t const rb = map_cluster(&s.rbar, row); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const d = warp * 32 + j * 8 + 2 * (lane & 3); + st_async_v2(dst + d * 4, acc[j][2 * r] * scale, acc[j][2 * r + 1] * scale, rb); + } + if (warp == 0 && (lane & 3) == 0) + st_async_f32(map_cluster(&s.recv_lse[rank], row), lse, rb); + } + } + { + // Every CTA's partial row and LSE for head `rank` landed (complete_tx on this CTA's barrier). + uint32_t done = 0; + while (!done) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(smem_u32(&s.rbar)) + : "memory"); + } + } + + // CTA `rank` merges head `rank` across the 16 partials; thread = output dim. The weights are the Triton merge's + // (exp2 of each chunk's LSE minus the largest, divided by their sum in chunk order), formed once per CTA: thread + // c < 16 takes chunk c's exponential, then divides it by the sum every one of those threads forms in the same + // order. + int const h = static_cast(rank); + if (tid < kTopk) + { + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + lse_max = fmaxf(lse_max, s.recv_lse[c]); + lse_max = lse_max == kNegInf ? 0.F : lse_max; + s.red[0][tid] = ex2(s.recv_lse[tid] - lse_max); + } + __syncthreads(); + if (tid < kTopk) + { + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + denom += s.red[0][c]; + denom = denom > 0.F ? denom : 1.F; + s.wgt[tid] = __fdiv_rn(s.red[0][tid], denom); + } + __syncthreads(); + float o = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + o += s.recv[c][tid] * s.wgt[c]; + __nv_bfloat16 const ob = __float2bfloat16_rn(o); + p.out[tok * p.out_stride + (kvh * kHeads + h) * kDim + tid] = ob; + + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + float const xq = x * qs; + uint16_t q2; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2) : "f"(0.F), "f"(xq)); + p.q8[static_cast(tok) * (p.num_sf_k_tiles * 128) + (kvh * kHeads + h) * kDim + tid] + = static_cast(q2 & 0xFF); + if (lane == 0) + { + int const k_sf = (kvh * kHeads + h) * 4 + warp; + int64_t const off = static_cast(tok / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (tok % 32) * 16 + ((tok % 128) / 32) * 4 + k_sf % 4; + p.sf[off] = static_cast(e); + } + } + if constexpr (!EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +cudaError_t configure() +{ + auto kernel = m3_lc_attn_kernel; + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes())); + if (err != cudaSuccess) + { + return err; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); +} + +template +cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) +{ + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs + // every variant eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + auto kernel = m3_lc_attn_kernel; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(num_tokens * kTopk, prm.num_kv_heads); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smemBytes(); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kTopk; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +} // namespace + +cudaError_t invokeMinimaxM3Sm107Tp2SparseDecodeLc(MinimaxM3Sm107Tp2SparseDecodeLcParams const& params, + int32_t numTokens, bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream) +{ + if (prewait) + { + if (linearSoftmax) + { + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); + } + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); + } + if (linearSoftmax) + { + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); + } + return mxfp8 ? launch(params, numTokens, usePdl, stream) + : launch(params, numTokens, usePdl, stream); +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h new file mode 100644 index 000000000000..b20a1894f255 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h @@ -0,0 +1,77 @@ +/* + * 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 +{ + +//! Operands of the MiniMax-M3 one-launch long-context sparse decode, SM107 TP2 geometry. Strides are in elements of +//! the pointed-to type. Geometry is fixed: 16 query heads per KV head, head_dim 128, page 128, top-k 16. +struct MinimaxM3Sm107Tp2SparseDecodeLcParams +{ + uint8_t const* q; // e4m3 [T, kv_heads * 16, 128] + int64_t q_tok, q_head; + uint8_t const* k; // packed e2m1 [pages, kv_heads, 128, 64] + uint8_t const* v; + int64_t k_blk, v_blk, k_h, v_h; + uint8_t const* ksc; // e4m3 [pages, kv_heads, 128, 8], token-major + uint8_t const* vsc; // e4m3, 4x4 token-quad order + int64_t ksc_blk, vsc_blk, ksc_h, vsc_h; + float const* kgs; + float const* vgs; + int const* topk; // [kv_heads, T, 16] + int64_t t_h, t_tok, t_k; + int const* block_table; + int64_t bt_stride; + int const* seq_lens; + int dql; + float sm_scale; + __nv_bfloat16* out; // [T, kv_heads * 2048] + int64_t out_stride; + uint8_t* q8; // e4m3 [rows, num_sf_k_tiles * 128] + uint8_t* sf; + int num_sf_k_tiles; + int num_kv_heads; + // N1 (optional): this layer's selected block id per (KV head, query row, slot) at its previous decode step, + // int32 [kv_heads, n1_rows, 16]. Before the grid dependency each CTA pulls into L2 the K/V page and scales of the + // block its slot selected last step, then records this step's block. A hint: no value changes. + int* n1_ids; + int64_t n1_h; + int n1_rows; + int64_t n1_pages; +}; + +//! One cluster of 16 CTAs per (query token, KV head): CTA r attends the token's r-th selected block of that head, +//! then the cluster merges the 16 partials over DSMEM into the BF16 rows (and, when mxfp8, the MXFP8 copy with +//! swizzled UE8M0 scales). prewait stages the request's block-table row in shared memory before the grid dependency +//! wait (rows longer than 4096 blocks read it from global memory); the arithmetic is the same in both variants. +//! The device code is SM100-family only; on other architectures the kernel compiles to an empty body. +cudaError_t invokeMinimaxM3Sm107Tp2SparseDecodeLc(MinimaxM3Sm107Tp2SparseDecodeLcParams const& params, + int32_t numTokens, bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu new file mode 100644 index 000000000000..22c073f74b3c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu @@ -0,0 +1,58 @@ +/* + * 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/minimaxM3Sm107Tp2/specOffsets.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr int kThreads = 256; + +__global__ void __launch_bounds__(kThreads) + specOffsetsKernel(int32_t* __restrict__ pos, int32_t const* __restrict__ posOff, int64_t nPos, + int32_t* __restrict__ kv, int32_t const* __restrict__ kvOff, int64_t nKv) +{ + for (int64_t i = threadIdx.x; i < nPos; i += kThreads) + { + pos[i] = static_cast(static_cast(pos[i]) + static_cast(posOff[i])); + } + for (int64_t j = threadIdx.x; j < nKv; j += kThreads) + { + kv[j] = static_cast(static_cast(kv[j]) + static_cast(kvOff[j])); + } +} + +} // namespace + +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream) +{ + specOffsetsKernel<<<1, kThreads, 0, stream>>>(pos, posOff, nPos, kv, kvOff, nKv); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h new file mode 100644 index 000000000000..1de212a37179 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h @@ -0,0 +1,41 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ + +//! The overlap scheduler's speculative step correction in one launch: pos[i] += posOff[i] for i < nPos and +//! kv[j] += kvOff[j] for j < nKv (int32, wrapping as the ATen adds do). A plain launch: it is a step-input writer that +//! pre-wait readers rely on (rubin-20261015/PREWAIT_AUDIT.md item b). +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu new file mode 100644 index 000000000000..9b6dbb42faed --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu @@ -0,0 +1,111 @@ +/* + * 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/minimaxM3Sm107Tp2/stepPrologue.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr int kThreads = 512; +constexpr int kWarps = kThreads / 32; + +// One warp per copy. Plain (coherent) loads: a device copy may read bytes another warp of this CTA just stored, which +// the CTA barrier between the two kinds makes visible; a non-coherent load could return the stale line. +__device__ __forceinline__ void copyWarp(uint8_t* dst, uint8_t const* src, int64_t bytes, int lane) +{ + auto const d = reinterpret_cast(dst); + auto const s = reinterpret_cast(src); + if (((d | s | static_cast(bytes)) & 15) == 0) + { + auto* d16 = reinterpret_cast(dst); + auto const* s16 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 16; i += 32) + { + d16[i] = s16[i]; + } + } + else if (((d | s | static_cast(bytes)) & 3) == 0) + { + auto* d4 = reinterpret_cast(dst); + auto const* s4 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 4; i += 32) + { + d4[i] = s4[i]; + } + } + else + { + for (int64_t i = lane; i < bytes; i += 32) + { + dst[i] = src[i]; + } + } +} + +// One CTA: a step's window is a few dozen copies of a few hundred bytes, so the launch is the cost; a CTA barrier +// orders the host copies before the device copies where the parallel scatter needed a second launch. +__global__ void __launch_bounds__(kThreads) stepPrologueKernel(uint8_t* window, int64_t* epoch, int64_t maxDesc) +{ + auto* header = reinterpret_cast(window); + int64_t const numHost = min(header[0], maxDesc); + int64_t const numDevice = min(header[1], maxDesc); + int64_t const* desc = header + 2; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (epoch != nullptr && threadIdx.x == 0) + { + *epoch += 1; + } + for (int64_t i = warp; i < numHost; i += kWarps) + { + int64_t const* c = desc + 3 * i; + copyWarp(reinterpret_cast(c[0]), window + c[1], c[2], lane); + } + __syncthreads(); + for (int64_t i = warp; i < numDevice; i += kWarps) + { + // The device copies' slots follow all header[0] host copies, as the parallel scatter indexes them. + int64_t const* c = desc + 3 * (header[0] + i); + copyWarp(reinterpret_cast(c[0]), reinterpret_cast(c[1]), c[2], lane); + } + // Every thread has read the header and the descriptors: consume the window. + __syncthreads(); + if (threadIdx.x < 2) + { + header[threadIdx.x] = 0; + } +} + +} // namespace + +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream) +{ + stepPrologueKernel<<<1, kThreads, 0, stream>>>(window, epoch, maxDesc); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h new file mode 100644 index 000000000000..44a7b7bcb44d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h @@ -0,0 +1,47 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ + +//! The decode graph's step-input prologue in one launch (step_inputs.StepInputPack with early_copy): the work of the +//! two parallel scatter launches and the header clear that follow them, in that order. +//! +//! The window is the pack's staging layout (kernels/stepInputsScatter.h): an int64 header `[numHostCopies, +//! numDeviceCopies]`, then one `(dst, src, bytes)` int64 triple per copy, host copies first; a host copy's `src` is a +//! byte offset into the window, a device copy's a device address. Up to `maxDesc` copies of each kind are written (the +//! slots the parallel scatter's grid covers), every host copy before any device copy, so a device copy may read a +//! destination a host copy filled; the copies of one kind must not overlap (the pack issues the window only then). +//! Afterwards the header is zeroed, so a replay with no new upload writes nothing. With `epoch` non-null, it is +//! incremented once. +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu new file mode 100644 index 000000000000..5da59fcec01a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp2 +{ + +namespace +{ + +constexpr int kThreads = 256; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves): the +// instruction sequence of the Triton kernel at limit 7, alpha 1.702, beta 1 (Indexer kdev17, bitwise on SM107): the +// gate clamp as min.bf16x2 before the widening, the up clamp as min.xorsign.abs, (u + 1), (g * -1.702) * log2(e), +// ex2.approx, + 1, div.full 1 / that, s * g, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t swigluOaiBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, %1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +template +__global__ void __launch_bounds__(kThreads) + swigluOaiKernel(uint32_t const* __restrict__ x, uint32_t* __restrict__ out, int64_t rows, int64_t halfD) +{ + if constexpr (kPdl) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int64_t const i = blockIdx.x * static_cast(kThreads) + threadIdx.x; + if (i >= rows * halfD) + { + return; + } + int64_t const row = i / halfD; + int64_t const col = i - row * halfD; + uint32_t const* xr = x + row * 2 * halfD; + out[i] = swigluOaiBf16x2(xr[col], xr[halfD + col]); +} + +} // namespace + +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream) +{ + int64_t const n = rows * halfD; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast((n + kThreads - 1) / kThreads)); + cfg.blockDim = dim3(kThreads); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return pdl ? cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD) + : cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD); +} + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h new file mode 100644 index 000000000000..c8db72a5efd3 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h @@ -0,0 +1,43 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp2 +{ + +//! The dense MLP's swigluoai activation: out[m, j] = bf16(min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + +//! 1)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both contiguous, d even): bit for bit +//! the Triton silu_and_mul kernel at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With `pdl` the launch allows +//! programmatic serialization and the grid waits on its dependency, then releases its dependents before it computes +//! (the down GEMV stages its weights during the activation and reads it only after its own grid wait). +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp2 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu new file mode 100644 index 000000000000..e173c91016c1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.cu @@ -0,0 +1,296 @@ +/* + * 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. + */ + +// The TP4 dense layers' QK norm + RoPE with FP8 output. The norm and RoPE are the generic fusedQKNormRopeKernel's at +// head_dim 128 without YaRN or mRoPE, expression for expression; the store rounds Q/K to BF16 and then to E4M3 +// (saturating), the two roundings of the stock BF16 kernel followed by its FP8 cast. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace common +{ +// tensorrt_llm/common/envUtils.h (not included: its CUDA utility headers do not build under the extension's +// no-half-conversion flags). +bool getEnvEnablePDL(); +} // namespace common + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace denseQkNormRopeFp8 +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kVecSize = kElemsPerThread * sizeof(__nv_bfloat16) / 4; +constexpr int kBlockSize = 256; +static_assert(kVecSize == 2, "a lane loads 4 BF16 values as one uint2"); + +// tensorrt_llm::common::warpReduceSum: the same butterfly, so the same sum. +__device__ __forceinline__ float warpReduceSum(float val) +{ +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + { + val = val + __shfl_xor_sync(0xffffffff, val, mask, 32); + } + return val; +} + +__device__ __forceinline__ void storeFp8ViaBf16(__nv_fp8_e4m3* out, long long offsetThread, float const (&elements)[4]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + float2 const rounded = __bfloat1622float2(__float22bfloat162_rn(make_float2(elements[i], elements[i + 1]))); + __nv_fp8x2_e4m3 const packed(rounded); + reinterpret_cast<__nv_fp8x2_storage_t*>(&out[offsetThread])[i / 2] = packed.__x; + } +} + +// The K (kv 0) or V (kv 1) row of one head of token tokenIdx in the HND paged pool, at the token's cache slot. +__device__ __forceinline__ __nv_fp8_e4m3* cacheRow(KvCacheView const& cache, int kv, int headIdx, int tokenIdx) +{ + long long const slot = static_cast(cache.outCacheLoc[tokenIdx]); + long long const page = slot / cache.tokensPerBlock; + long long const within = slot % cache.tokensPerBlock; + return cache.base + page * cache.pageStride + kv * cache.kvStride + headIdx * cache.headStride + + within * cache.tokenStride; +} + +template +__global__ void denseQkNormRopeFp8Kernel(__nv_bfloat16 const* qkvIn, __nv_fp8_e4m3* qkvOut, int numHeadsQ, + int numHeadsK, int numHeadsV, int rotaryDim, float eps, __nv_bfloat16 const* qWeight, __nv_bfloat16 const* kWeight, + float base, int const* positionIds, int numTokens, bool useGemma, KvCacheView cache) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + // PDL: every read of this step's q/k/v rows, positions and cache tables follows the wait. The dependent (the + // full-context decode) reads this kernel's outputs only after its own griddepcontrol.wait, so it is released at + // entry and stages its CTAs while this kernel runs. + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const warpsPerBlock = blockDim.x / 32; + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId; + int const totalQkHeads = numHeadsQ + numHeadsK; + int const totalHeads = totalQkHeads + numHeadsV; + int const tokenIdx = globalWarpIdx / totalHeads; + int const localHeadIdx = globalWarpIdx % totalHeads; + if (tokenIdx >= numTokens) + { + return; + } + bool const isQ = localHeadIdx < numHeadsQ; + bool const isV = localHeadIdx >= totalQkHeads; + int headIdx; + int segStart; + if (isQ) + { + headIdx = localHeadIdx; + segStart = 0; + } + else if (!isV) + { + headIdx = localHeadIdx - numHeadsQ; + segStart = numHeadsQ * kHeadDim; + } + else + { + headIdx = localHeadIdx - totalQkHeads; + segStart = totalQkHeads * kHeadDim; + } + + using vec_T = uint2; + int const offsetThread + = tokenIdx * totalHeads * kHeadDim + segStart + headIdx * kHeadDim + laneId * kElemsPerThread; + float elements[kElemsPerThread]; + float sumOfSquares = 0.0f; + { + vec_T vec = *reinterpret_cast(&qkvIn[offsetThread]); +#pragma unroll + for (int i = 0; i < kVecSize; i++) + { + float2 vals = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162*>(reinterpret_cast(&vec) + i)); + sumOfSquares += vals.x * vals.x; + sumOfSquares += vals.y * vals.y; + elements[2 * i] = vals.x; + elements[2 * i + 1] = vals.y; + } + } + if (isV) + { + // V is BF16 already: its single E4M3 rounding is the stock cast's. + storeFp8ViaBf16(qkvOut, offsetThread, elements); + if (kWriteCache && tokenIdx < cache.numLiveTokens && cache.outCacheLoc[tokenIdx] >= 0) + { + storeFp8ViaBf16(cacheRow(cache, 1, headIdx, tokenIdx), laneId * kElemsPerThread, elements); + } + return; + } + + sumOfSquares = warpReduceSum(sumOfSquares); + float rms_rcp = rsqrtf(sumOfSquares / static_cast(kHeadDim) + eps); + for (int i = 0; i < kElemsPerThread; i++) + { + int dim = laneId * kElemsPerThread + i; + float weight = isQ ? __bfloat162float(qWeight[dim]) : __bfloat162float(kWeight[dim]); + elements[i] *= rms_rcp * (useGemma ? (1.0f + weight) : weight); + } + + float elements2[kElemsPerThread]; + float cos_vals[kElemsPerThread]; + float sin_vals[kElemsPerThread]; + float const neg2_log2base_over_rd = -2.0f * __log2f(base) / static_cast(rotaryDim); + int const rd_mask = rotaryDim - 1; + bool const rd_is_pow2 = ((rotaryDim & rd_mask) == 0); + float const pos_id = static_cast(positionIds[tokenIdx]); + if constexpr (kInterleave) + { + for (int i = 0; i < kElemsPerThread; i++) + { + if (i % 2 == 0) + { + elements2[i] = -elements[i + 1]; + } + else + { + elements2[i] = elements[i - 1]; + } + int dim_idx = laneId * kElemsPerThread + i; + int half_dim = dim_idx / 2; + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); + } + } + else + { + __syncwarp(); + int pairOffset = (rotaryDim / 2) / kElemsPerThread; + for (int i = 0; i < kElemsPerThread; i++) + { + elements2[i] = __shfl_xor_sync(0xffffffff, elements[i], pairOffset); + if (laneId < pairOffset) + { + elements2[i] = -elements2[i]; + } + int dim_idx = laneId * kElemsPerThread + i; + dim_idx = rd_is_pow2 ? ((dim_idx * 2) & rd_mask) : ((dim_idx * 2) % rotaryDim); + int half_dim = dim_idx / 2; + float freq = exp2f(static_cast(half_dim) * neg2_log2base_over_rd); + float theta = pos_id * freq; + __sincosf(theta, &sin_vals[i], &cos_vals[i]); + } + __syncwarp(); + } + + // attention_factor is 1 without YaRN: (x * cos + x2 * sin) * 1 rounds as the generic kernel's product does. + float const attention_factor = 1.0f; + if (rotaryDim == kHeadDim) + { + for (int i = 0; i < kElemsPerThread; i++) + { + elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor; + } + } + else + { + for (int i = 0; i < kElemsPerThread; i++) + { + int dim_idx = laneId * kElemsPerThread + i; + if (dim_idx < rotaryDim) + { + elements[i] = (elements[i] * cos_vals[i] + elements2[i] * sin_vals[i]) * attention_factor; + } + } + } + storeFp8ViaBf16(qkvOut, offsetThread, elements); + if constexpr (kWriteCache) + { + // Rows past the live count (padded graph rows) own no slot, nor does a negative (sentinel) slot. + if (!isQ && tokenIdx < cache.numLiveTokens && cache.outCacheLoc[tokenIdx] >= 0) + { + storeFp8ViaBf16(cacheRow(cache, 0, headIdx, tokenIdx), laneId * kElemsPerThread, elements); + } + // The compact query [tokens, q heads, 128] the decode reads, beside the packed row. + if (isQ && cache.qCompact != nullptr) + { + storeFp8ViaBf16(cache.qCompact, + (static_cast(tokenIdx) * numHeadsQ + headIdx) * kHeadDim + laneId * kElemsPerThread, + elements); + } + } +} + +} // namespace + +cudaError_t launch(void const* qkvIn, void* qkvOut, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, + int rotaryDim, float eps, void const* qWeight, void const* kWeight, float base, bool interleave, + int const* positionIds, bool useGemma, KvCacheView const* cache, cudaStream_t stream) +{ + int const totalWarps = numTokens * (numHeadsQ + numHeadsK + numHeadsV); + int const warpsPerBlock = kBlockSize / 32; + dim3 const grid((totalWarps + warpsPerBlock - 1) / warpsPerBlock); + auto const* in = static_cast<__nv_bfloat16 const*>(qkvIn); + auto* out = static_cast<__nv_fp8_e4m3*>(qkvOut); + auto const* qw = static_cast<__nv_bfloat16 const*>(qWeight); + auto const* kw = static_cast<__nv_bfloat16 const*>(kWeight); + KvCacheView const view = cache != nullptr ? *cache : KvCacheView{}; + auto const run = [&](auto kernel) -> cudaError_t + { + // The maximum shared-memory carveout lets the dependent's large-smem CTAs co-reside on the SMs this small + // kernel occupies once it releases them. + cudaError_t const err = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (err != cudaSuccess) + { + return err; + } + cudaLaunchConfig_t config = {}; + config.gridDim = grid; + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL() ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + return cudaLaunchKernelEx(&config, kernel, in, out, numHeadsQ, numHeadsK, numHeadsV, rotaryDim, eps, qw, kw, + base, positionIds, numTokens, useGemma, view); + }; + if (cache != nullptr) + { + return interleave ? run(denseQkNormRopeFp8Kernel) : run(denseQkNormRopeFp8Kernel); + } + return interleave ? run(denseQkNormRopeFp8Kernel) : run(denseQkNormRopeFp8Kernel); +} + +} // namespace denseQkNormRopeFp8 +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h new file mode 100644 index 000000000000..1049033b07c9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h @@ -0,0 +1,68 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ +namespace denseQkNormRopeFp8 +{ + +constexpr int kHeadDim = 128; + +//! A layer's FP8 K/V pool in the HND paged layout [pages, 2 (K, V), kv heads, tokens per block, 128] (element +//! strides), with the step's per-token cache slots (page * tokensPerBlock + within) for its first numLiveTokens rows. +struct KvCacheView +{ + __nv_fp8_e4m3* base = nullptr; + int const* outCacheLoc = nullptr; + long long pageStride = 0; + long long kvStride = 0; + long long headStride = 0; + long long tokenStride = 0; + int tokensPerBlock = 1; + int numLiveTokens = 0; + //! Optional: every row's Q also stored compact, [tokens, q heads, 128]. + __nv_fp8_e4m3* qCompact = nullptr; +}; + +//! The MiniMax-M3 TP4 dense layers' per-head QK RMSNorm + partial RoPE of a packed BF16 [tokens, (q + k + v) * 128] +//! row, written out of place as FP8 E4M3 with V copy-cast. Q and K are rounded to BF16 before E4M3, as the stock path +//! (the BF16 fused_qk_norm_rope output cast to FP8) rounds them. With cache, the K and V rows of the live tokens are +//! also stored into the paged pool at their slots (the bytes of the qkvOut rows). +//! +//! PDL (with TRTLLM_ENABLE_PDL): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction, so +//! nothing is read or written before the producer has completed, and the kernel releases its dependent right after +//! it. The dependent may read only step-invariant data before its own griddepcontrol.wait. +cudaError_t launch(void const* qkvIn, void* qkvOut, int numTokens, int numHeadsQ, int numHeadsK, int numHeadsV, + int rotaryDim, float eps, void const* qWeight, void const* kWeight, float base, bool interleave, + int const* positionIds, bool useGemma, KvCacheView const* cache, cudaStream_t stream); + +} // namespace denseQkNormRopeFp8 +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu new file mode 100644 index 000000000000..2e67b3ae7ac1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu @@ -0,0 +1,667 @@ +/* + * 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/minimaxM3Sm107Tp4/draftStep.h" + +#include +#include +#include + +// One pass of the fused greedy Eagle3 draft loop of MiniMax-M3 (GB300, TP4, <= 8 requests per pass). +// +// One kernel replaces the draft lm_head GEMM, the logits all-gather / float cast / argmax of the in-tree greedy draft +// sampling and the draft loop's per-pass bookkeeping, at the in-tree precision: +// +// * The vocab-sharded BF16 lm_head weight is stored in mma.m16n8k16 A-fragment order (the same values): per 16-row +// tile and k16 step, 512 bytes; lane (g, j) holds 16 bytes = {row g k 2j..2j+1, row g+8 k 2j.., row g k 2j+8.., +// row g+8 k 2j+8..}. +// * Or, by default, in a lossless 12-bit coding of those fragments (75% of the bytes; the target head's coding, see +// targetHead.cu): per 16-row tile and unit of 4 k16 steps a lane reads three 16-byte words, the sign/mantissa bytes +// of steps 0-1, of steps 2-3, then 4 code words (exponent - base, 4 bits; 15 = an exception). The decoded values +// are the checkpoint's BF16 bits; an exception decodes to +0 and the segment's flush adds the exact product x * w +// from the tile's list (products of BF16 values are exact in FP32). +// * A persistent grid (one 16-warp CTA per SM) streams whole tiles over the full K with BF16 MMAs and FP32 +// accumulation; the BF16 head input rows sit in shared memory in B-fragment order, unchanged. +// * Each logit is rounded to BF16 (the in-tree logits are the BF16 GEMM output) and becomes a 64-bit key +// ordered(value) << 32 | ~global_index, with -0 taken as +0: the larger value wins and ties pick the lowest index, as +// torch.max / torch.argmax over the gathered logits do. Keys are max-reduced per row: warp shuffles, then shared +// memory, then atomicMax. +// * The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping +// (Lamport: 0 = empty, a key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the +// slots, and runs the pass's bookkeeping. Call n uses slot half n % 2. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace draftStep +{ +namespace +{ +constexpr int kUnroll = 8; // k16 chunks in flight per warp +constexpr int kChunkBytes = 512; // one tile's k16 step: 32 lanes x 16 bytes + +// Coded units (4 k16 steps, 3 x 16 bytes per lane) in flight per warp: 3 up to 4 rows, 2 at 8 rows. +template +constexpr int codedUnroll() +{ + return MP >= 8 ? 2 : 3; +} + +__device__ __forceinline__ void mma_bf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + 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};\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)); +} + +// BF16 pair p (values 2p, 2p + 1) of a lane's 8 coded values: sign/mantissa bytes sm (4 per word), codes word cw. +__device__ __forceinline__ uint32_t coded_decode_pair(uint32_t sm_word, int byte0, uint32_t cw, int p, uint32_t base) +{ + uint32_t const b = __byte_perm(sm_word, 0u, byte0 == 0 ? 0x4140 : 0x4342); // {b0, 0, b1, 0} + uint32_t const q = (cw >> (4 * p)) & 0x000f000fu; + uint32_t const e = q * 128u + base; // exponent field of both halves + // Code 15 (an exception, sign/mantissa byte 0) decodes to +0. + return ((b & 0x007f007fu) | ((b << 8) & 0x80008000u) | e) & __vcmpne2(q, 0x000f000fu); +} + +// One coded unit (4 k16 steps) against the BF16 B fragments: steps alternate between the two accumulators, as the +// k16 steps of the plain BF16 head do. +template +__device__ __forceinline__ void chunk_mma_coded(uint4 const& sm01, uint4 const& sm23, uint4 const& cw, uint32_t base, + uint2 const* xs, int unit, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint32_t const smw[8] = {sm01.x, sm01.y, sm01.z, sm01.w, sm23.x, sm23.y, sm23.z, sm23.w}; + uint32_t const cws[4] = {cw.x, cw.y, cw.z, cw.w}; +#pragma unroll + for (int ci = 0; ci < 4; ++ci) + { + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[((unit * 4 + ci) * 4 + j) * MP + g]; + uint32_t const w0 = coded_decode_pair(smw[2 * ci], 0, cws[ci], 0, base); + uint32_t const w1 = coded_decode_pair(smw[2 * ci], 2, cws[ci], 1, base); + uint32_t const w2 = coded_decode_pair(smw[2 * ci + 1], 0, cws[ci], 2, base); + uint32_t const w3 = coded_decode_pair(smw[2 * ci + 1], 2, cws[ci], 3, base); + if (ci & 1) + mma_bf16(db, w0, w1, w2, w3, b.x, b.y); + else + mma_bf16(da, w0, w1, w2, w3, b.x, b.y); + } +} + +// Exceptions of `tile` with k in [kb, ke): v += x * w (exact) for the owner lanes (row g or g + 8; tokens 2j, +// 2j + 1 < M), x from the staged B fragments. Warp-uniform (shuffles). +template +__device__ __forceinline__ void coded_exceptions( + CodedArgs const& cd, int tile, int kb, int ke, uint2 const* xs, int M, int lane, float (&v)[4]) +{ + int const g = lane >> 2, j = lane & 3; + int const e0 = __ldg(cd.exc_off + tile), e1 = __ldg(cd.exc_off + tile + 1); + for (int b0 = e0; b0 < e1; b0 += 32) + { + uint32_t pos = 0xffffffffu; + float2 wv = make_float2(0.f, 0.f); + if (b0 + lane < e1) + { + pos = __ldg(cd.exc_pos + b0 + lane); + wv = __ldg(cd.exc_w + b0 + lane); + } + int const cnt = min(32, e1 - b0); + for (int q = 0; q < cnt; ++q) + { + uint32_t const pq = __shfl_sync(0xffffffffu, pos, q); + float const w0 = __shfl_sync(0xffffffffu, wv.x, q); + int const r = static_cast(pq >> 16), k = static_cast(pq & 0xffffu); + if (k >= kb && k < ke && (r & 7) == g) + { + int const st = k >> 4, rr = k & 15; + bool const hi = r >= 8; +#pragma unroll + for (int q2 = 0; q2 < 2; ++q2) + { + int const n = 2 * j + q2; + if (n < M) + { + uint2 const bw = xs[(st * 4 + ((rr & 7) >> 1)) * MP + n]; + uint32_t const word = rr >= 8 ? bw.y : bw.x; + float const xv = __uint_as_float((rr & 1) ? (word & 0xffff0000u) : (word << 16)); + float const d = __fmul_rn(xv, w0); + if (hi) + v[2 + q2] += d; + else + v[q2] += d; + } + } + } + __syncwarp(); + } + } +} + +__device__ __forceinline__ uint4 ldg_stream(uint4 const* p) +{ + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.u32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +// The argmax value of a logit: its BF16 rounding (the in-tree logits' dtype), -0 as +0. +__device__ __forceinline__ float key_value(float v) +{ + float const r = __bfloat162float(__float2bfloat16(v)); + return r == 0.f ? 0.f : r; +} + +__device__ __forceinline__ unsigned long long make_key(float v, uint32_t index) +{ + uint32_t u = __float_as_uint(v); + u = (u & 0x80000000u) ? ~u : (u | 0x80000000u); + return (static_cast(u) << 32) | static_cast(0xFFFFFFFFu - index); +} + +__device__ __forceinline__ unsigned long long shfl_max(unsigned long long k, int o) +{ + unsigned long long const ok = __shfl_xor_sync(0xffffffffu, k, o); + return ok > k ? ok : k; +} + +// One k16 step of a tile: the lane's A fragment w against the activation B fragment of step `chunk` (rows g < MP). +// Even and odd steps accumulate separately (two independent MMA chains); the tile's sum is da + db. +template +__device__ __forceinline__ void chunk_mma( + uint4 const& w, uint2 const* xs, int chunk, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[(chunk * 4 + j) * MP + g]; + if (chunk & 1) + mma_bf16(db, w.x, w.y, w.z, w.w, b.x, b.y); + else + mma_bf16(da, w.x, w.y, w.z, w.w, b.x, b.y); +} + +// A warp's pieces of the tiles it shares with other warps (its first and last local tile), kept for the warp-order sum. +struct Pieces +{ + float4 head = make_float4(0.f, 0.f, 0.f, 0.f); + float4 tail = make_float4(0.f, 0.f, 0.f, 0.f); +}; + +// Whether the warp's chunk range [q_begin, q_end) covers the CTA's whole range of local tile tl (the CTA streams +// chunks [lo, lo + span) of its local tiles). +__device__ __forceinline__ bool covers_tile(int tl, int kchunks, int lo, int span, int q_begin, int q_end) +{ + return q_begin <= max(tl * kchunks, lo) && q_end >= min((tl + 1) * kchunks, lo + span); +} + +// Local tile t's sum (D-FIX): the accumulator when one warp streamed the CTA's whole range of it, else its pieces +// from the warps that streamed it, added in warp order from 0 (each warp's first tile is its head piece, any later +// one its tail piece). Deterministic: the order is fixed by the launch shape, and with two pieces it is the same +// sum as the atomics' (two addends from 0 commute). +__device__ __forceinline__ float4 ordered_tile( + float4 const* acc_t, float4 const* stage, int t, int kchunks, int lo, int span, int lane) +{ + int const r_lo = max(t * kchunks, lo); + int const r_last = min((t + 1) * kchunks, lo + span) - 1; + int w_lo = 0, w_hi = 0; + for (int w = 1; w < kWarps; ++w) + { + int const qb = lo + w * span / kWarps; + if (qb <= r_lo) + w_lo = w; + if (qb <= r_last) + w_hi = w; + } + if (w_lo == w_hi) + return acc_t[lane]; + float4 s = make_float4(0.f, 0.f, 0.f, 0.f); + for (int w = w_lo; w <= w_hi; ++w) + { + int const qb = lo + w * span / kWarps; + float4 const x = stage[(w * 2 + (qb / kchunks == t ? 0 : 1)) * 32 + lane]; + s.x += x.x; + s.y += x.y; + s.z += x.z; + s.w += x.w; + } + return s; +} + +// Add a warp's segment of a tile into the CTA's accumulator (with a coded head, the segment's exceptions first); +// with `fixed` (D-FIX) a shared tile's piece goes to `pc` instead (ordered_tile adds it). +template +__device__ __forceinline__ void flush_tile(float4* acc, float (&da)[4], float (&db)[4], Params const& p, int tile, + int kb, int ke, uint2 const* xs, int lane, Pieces& pc, bool fixed, bool whole, bool first) +{ + float v[4] = {da[0] + db[0], da[1] + db[1], da[2] + db[2], da[3] + db[3]}; + if constexpr (CODED) + coded_exceptions(p.cd, tile, kb, ke, xs, p.M, lane, v); + if (fixed) + { + // A piece covering the CTA's whole range of the tile is its only addend (0 + v, as the atomic gave); a + // warp's other pieces (at most its first and last tile) wait in registers for the warp-order sum. + if (whole) + *acc = make_float4(0.f + v[0], 0.f + v[1], 0.f + v[2], 0.f + v[3]); + else if (first) + pc.head = make_float4(v[0], v[1], v[2], v[3]); + else + pc.tail = make_float4(v[0], v[1], v[2], v[3]); + } + else + { + float* a = reinterpret_cast(acc); + atomicAdd(a + 0, v[0]); + atomicAdd(a + 1, v[1]); + atomicAdd(a + 2, v[2]); + atomicAdd(a + 3, v[3]); + } +#pragma unroll + for (int e = 0; e < 4; ++e) + { + da[e] = 0.f; + db[e] = 0.f; + } +} + +// Balanced streaming GEMV over the fragment-ordered weight, persistent grid of 16-warp CTAs. CTA c owns tiles +// [t_begin, t_end); its (tile, chunk) space is split into 16 equal contiguous ranges, one per warp, so every warp +// streams the same number of bytes. A tile straddles at most two warps; their partial sums are added into shared +// memory (two addends: the order does not change the sum). MP = rows padded to 1 / 2 / 4 / 8. A chunk is one k16 +// step of the BF16 head, or one coded unit (4 k16 steps) of the coded head (CODED: no shared-memory prologue). +template +__global__ void __launch_bounds__(kWarps * 32, 1) + draft_step_kernel(Params const p, int kchunks, int pf_chunks, int pf_off) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(16) uint2 xs[]; // [K/16 steps][4 (j)][MP (g)], then the weight prologue + __shared__ float4 acc_s[kMaxTilesPerCta][32]; + __shared__ unsigned long long red[kWarps][kMaxTokens]; + __shared__ __align__(8) unsigned long long pf_bar[kWarps]; + __shared__ bool s_last; + __shared__ int s_tok[kMaxTokens]; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const g = lane >> 2; + int const j = lane & 3; + constexpr int kChunkK = CODED ? 64 : 16; // k per chunk + constexpr int kWords = CODED ? 3 : 1; // 16-byte words per lane and chunk + constexpr int KU = CODED ? codedUnroll() : kUnroll; + int const K = kchunks * kChunkK; + int const M = p.M; + int const t_begin = static_cast((static_cast(blockIdx.x) * p.n_tiles) / gridDim.x); + int const t_end = static_cast((static_cast(blockIdx.x + 1) * p.n_tiles) / gridDim.x); + int const n_local = t_end - t_begin; + int const span = n_local * kchunks; + int const q_begin = warp * span / kWarps; + int const q_end = (warp + 1) * span / kWarps; + long long const base = static_cast(t_begin) * kchunks; // first chunk of this CTA + uint4 const* const wf = p.wf; + + // Weights are static: start this warp's stream before waiting on the producer of x. The warp's first pf_chunks + // chunks come in by one bulk copy into shared memory (dynamic offset pf_off), which fills the grid-dependency + // wait and the staging; the register ring starts after them. + int const n_pf = CODED ? 0 : min(pf_chunks, q_end - q_begin); + uint4* const wpf = reinterpret_cast(reinterpret_cast(xs) + pf_off) + warp * pf_chunks * 32; + if (n_pf > 0 && lane == 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + uint32_t const bytes = static_cast(n_pf) * kChunkBytes; + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + static_cast(__cvta_generic_to_shared(wpf))), + "l"(wf + (base + q_begin) * 32), "r"(bytes), "r"(bar) + : "memory"); + } + __syncwarp(); + int const q_reg = q_begin + n_pf; + uint4 buf[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q_reg + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = ldg_stream(wf + ((base + q_reg + u) * kWords + l) * 32 + lane); + } + } + for (int i = threadIdx.x; i < n_local * 32; i += blockDim.x) + acc_s[i / 32][i % 32] = make_float4(0.f, 0.f, 0.f, 0.f); + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // Stage the head input rows in B-fragment order: item (k16 step S, row gg) = two 16-byte loads of x, four + // fragment words {x[2jj], x[2jj + 1]}, {x[2jj + 8], x[2jj + 9]} (BF16 pairs, moved unchanged). Rows >= M are 0. + constexpr int kSU = MP >= 4 ? 4 : 2; + int const n_items = (K / 16) * MP; + for (int i0 = threadIdx.x; i0 < n_items; i0 += kSU * blockDim.x) + { + uint4 xv[kSU][2]; +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + int const gg = i % MP; + int const S = i / MP; + xv[u][0] = make_uint4(0u, 0u, 0u, 0u); + xv[u][1] = make_uint4(0u, 0u, 0u, 0u); + if (i < n_items && gg < M) + { + long long const xr = p.x_rows != nullptr ? p.x_rows[gg] : gg; + uint4 const* xp = reinterpret_cast(p.x + xr * K + 16 * S); + xv[u][0] = xp[0]; + xv[u][1] = xp[1]; + } + } +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + if (i >= n_items) + continue; + int const gg = i % MP; + int const S = i / MP; + uint32_t const lo[4] = {xv[u][0].x, xv[u][0].y, xv[u][0].z, xv[u][0].w}; + uint32_t const hi[4] = {xv[u][1].x, xv[u][1].y, xv[u][1].z, xv[u][1].w}; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) + xs[(S * 4 + jj) * MP + gg] = make_uint2(lo[jj], hi[jj]); + } + } + __syncthreads(); + + bool const b_active = g < MP; + float da[4] = {0.f, 0.f, 0.f, 0.f}; + float db[4] = {0.f, 0.f, 0.f, 0.f}; + int tl = q_begin / kchunks; // local tile of the current accumulator + int c = q_begin - tl * kchunks; // chunk within that tile + int cs = c; // first chunk of the warp's segment of that tile + int const tl0 = tl; // the warp's first local tile (its head piece) + Pieces pc; + if (n_pf > 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + asm volatile( + "{\n .reg .pred p;\n PF_WAIT%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n" + " @!p bra PF_WAIT%=;\n}\n" ::"r"(bar) + : "memory"); + for (int q = q_begin; q < q_reg; ++q) + { + chunk_mma(wpf[(q - q_begin) * 32 + lane], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, 0, span, q_begin, q_end), tl == tl0); + c = 0; + cs = 0; + ++tl; + } + } + } + for (int q = q_reg; q < q_end; q += KU) + { + uint4 nxt[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + KU + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + nxt[u][l] = ldg_stream(wf + ((base + q + KU + u) * kWords + l) * 32 + lane); + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + u < q_end) + { + if constexpr (CODED) + chunk_mma_coded(buf[u][0], buf[u][kWords > 1 ? 1 : 0], buf[u][kWords > 2 ? 2 : 0], p.cd.base, + xs, c, j, g, b_active, da, db); + else + chunk_mma(buf[u][0], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, 0, span, q_begin, q_end), tl == tl0); + c = 0; + cs = 0; + ++tl; + } + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = nxt[u][l]; + } + if (c != 0) + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, c * kChunkK, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, 0, span, q_begin, q_end), tl == tl0); + __syncthreads(); + // D-FIX: stage every warp's head / tail piece. xs and the weight prologue are dead after the barrier above (the + // launcher sizes them to hold the 16 x 2 x 32 pieces). + float4* const stage = reinterpret_cast(xs); + if (p.fixedOrder) + { + stage[(warp * 2) * 32 + lane] = pc.head; + stage[(warp * 2 + 1) * 32 + lane] = pc.tail; + __syncthreads(); + } + + // Accumulator lane (g, j) of a tile: rows g / g + 8 (x, y / z, w) of tokens 2j / 2j + 1 (x, z / y, w). + unsigned long long best0 = 0ull, best1 = 0ull; // tokens 2j, 2j + 1 + for (int t = warp; t < n_local; t += kWarps) + { + float4 const v = p.fixedOrder ? ordered_tile(acc_s[t], stage, t, kchunks, 0, span, lane) : acc_s[t][lane]; + int const r0 = (t_begin + t) * 16 + g; + uint32_t const i0 = static_cast(p.vocab_offset + r0); + unsigned long long k0 = make_key(key_value(v.x), i0); + unsigned long long const k0b = make_key(key_value(v.z), i0 + 8); + unsigned long long k1 = make_key(key_value(v.y), i0); + unsigned long long const k1b = make_key(key_value(v.w), i0 + 8); + k0 = k0b > k0 ? k0b : k0; + k1 = k1b > k1 ? k1b : k1; + best0 = k0 > best0 ? k0 : best0; + best1 = k1 > best1 ? k1 : best1; + } +#pragma unroll + for (int o = 4; o < 32; o <<= 1) + { + best0 = shfl_max(best0, o); + best1 = shfl_max(best1, o); + } + asm volatile("griddepcontrol.launch_dependents;"); + if (lane < 4) + { + red[warp][2 * lane] = best0; + red[warp][2 * lane + 1] = best1; + } + __syncthreads(); + if (threadIdx.x < M) + { + unsigned long long b = 0ull; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + b = red[w][threadIdx.x] > b ? red[w][threadIdx.x] : b; + if (b != 0ull) + atomicMax(p.keys + threadIdx.x, b); + } + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned int const prev = atomicAdd(p.counters, 1u); + s_last = prev == gridDim.x - 1; + } + __syncthreads(); + if (!s_last) + return; + + __threadfence(); + unsigned int const call = *reinterpret_cast(p.counters + 1); + int const half = static_cast(call & 1u); + if (threadIdx.x < M) + { + int const t = threadIdx.x; + unsigned long long const mine = atomicExch(p.keys + t, 0ull); + unsigned long long* slot_mc = p.ws_mc + (half * kMaxTokens + t) * p.world; + unsigned long long* slot_uc = p.ws_uc + (half * kMaxTokens + t) * p.world; + slot_mc[p.rank] = mine; + unsigned long long best = 0ull; + for (int r = 0; r < p.world; ++r) + { + unsigned long long v; + do + { + v = *reinterpret_cast(slot_uc + r); + } while (v == 0ull); + best = v > best ? v : best; + slot_uc[r] = 0ull; + } + s_tok[t] = static_cast(0xFFFFFFFFu - static_cast(best & 0xFFFFFFFFull)); + if (p.out_keys != nullptr) + p.out_keys[t] = best; + } + __syncthreads(); + if (threadIdx.x == 0) + { + p.counters[0] = 0u; + p.counters[1] = call + 1u; + } + if (threadIdx.x < M) + { + StepArgs const& st = p.st; + int const b = threadIdx.x; + int const tok = s_tok[b]; + st.draft_out[b * st.draft_stride + st.step] = tok; + st.ids_next[b] = tok; + long long const r = p.x_rows != nullptr ? p.x_rows[b] : b; + st.pos_next[b] = st.pos_src[r] + 1; + if (st.step == 0 && st.num_acc != nullptr) + { + st.kv_lens[b] -= st.draft_len - st.num_acc[b]; + st.seq_lens[b] = 1; + } + else if (st.step > 0) + { + st.kv_lens[b] += 1; + } + } + if (p.hasFold) + { + // The fold tail (H-R2-3) after this pass's tokens and bookkeeping: the barrier makes this CTA's writes above + // visible to warp 0; every earlier operand was written by grids this one waited on. All CTAs triggered their + // dependents before the exchange, which is safe: a dependent's griddepcontrol.wait waits for this whole grid + // (this tail included) to complete, so only a dependent that read the tail's outputs before its own wait + // could race, and none does (the next step's readers run in the next graph replay). + __syncthreads(); + if (threadIdx.x < 32) + { + foldTailWarp(p.fold, static_cast(threadIdx.x)); + } + } +#endif +} + +template +cudaError_t launch(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + auto kern = draft_step_kernel; + size_t const xs_bytes = static_cast(K / 16) * 4 * MP * sizeof(uint2); + // Shared-memory weight prologue: chunks per warp that fit next to the kernel's static + dynamic shared memory. + static int const sStaticSmem = [] + { + cudaFuncAttributes fa{}; + return cudaFuncGetAttributes(&fa, draft_step_kernel) == cudaSuccess + ? static_cast(fa.sharedSizeBytes) + : -1; + }(); + static int const sMaxSmemPerBlock = [] + { + int dev = 0, maxSmem = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + return maxSmem; + }(); + int pf = 0; + if (!CODED && sStaticSmem >= 0) + { + long long const room + = static_cast(sMaxSmemPerBlock) - sStaticSmem - static_cast(xs_bytes) - 1024; + long long const want = static_cast(kPrefetchKb) * 1024; + pf = static_cast(std::max(0ll, std::min(room, want)) / (kWarps * kChunkBytes)); + } + // At least the D-FIX piece staging (16 warps x head / tail x 32 lanes of float4), which reuses this space. + size_t const smem = std::max(xs_bytes + static_cast(pf) * kWarps * kChunkBytes, + static_cast(kWarps) * 2 * 32 * sizeof(float4)); + cudaError_t err = cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + return err; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(grid)); + cfg.blockDim = dim3(kWarps * 32); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, kern, p, K / (CODED ? 64 : 16), pf, static_cast(xs_bytes)); +} + +template +cudaError_t dispatchRows(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + if (p.M <= 1) + return launch<1, CODED>(p, K, grid, pdl, stream); + if (p.M <= 2) + return launch<2, CODED>(p, K, grid, pdl, stream); + if (p.M <= 4) + return launch<4, CODED>(p, K, grid, pdl, stream); + return launch<8, CODED>(p, K, grid, pdl, stream); +} + +} // namespace + +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + bool const coded = p.cd.exc_off != nullptr; + if (p.M < 1 || p.M > kMaxTokens || K % (coded ? 64 : 16) != 0 || grid < 1) + return cudaErrorInvalidValue; + return coded ? dispatchRows(p, K, grid, pdl, stream) : dispatchRows(p, K, grid, pdl, stream); +} + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h new file mode 100644 index 000000000000..85296a3fd6d7 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h @@ -0,0 +1,107 @@ +/* + * 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 "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +//! One pass of the fused greedy Eagle3 draft loop of the MiniMax-M3 NVFP4 TP4 decode path (sm_103a): the BF16 +//! vocab-sharded draft head, the cross-rank argmax and the draft loop's per-pass bookkeeping in one kernel. See +//! draftStep.cu. +namespace draftStep +{ + +constexpr int kMaxTokens = 8; // head rows (requests) of one launch +constexpr int kWarps = 16; // warps per CTA +constexpr int kMaxTilesPerCta = 32; // 16-row weight tiles one CTA can own +//! Shared-memory weight prologue per CTA, in KB (capped by the free shared memory). +constexpr int kPrefetchKb = 112; + +// Draft-pass epilogue (last CTA, one thread per request row b): stores the token as draft token `step`, stages the +// next pass's input id and position, and advances the draft KV lengths the way the in-tree draft loop does between +// passes (pass 0 with num_acc: drop the rejected verify rows and switch to one token per request; later passes: +1). +struct StepArgs +{ + int* draft_out; // [nreq, draft_stride]: column `step` receives the token + int draft_stride; + int step; + int* ids_next; // [nreq] + int const* pos_src; // positions of the rows the head read (indexed like x) + int* pos_next; // [nreq] = pos_src[x_rows[b] or b] + 1 + int* kv_lens; // [nreq] draft KV lengths, updated in place + int const* num_acc; // pass 0: accepted counts (nullptr: no rollback) + int* seq_lens; // pass 0: set to 1 + int draft_len; +}; + +// The 12-bit coded head (see draftStep.cu): the per-tile exception lists and the exponent base. exc_off == nullptr: +// the head is the plain BF16 fragment-ordered copy. +struct CodedArgs +{ + int const* exc_off; // [n_tiles + 1] exception list offsets + uint32_t const* exc_pos; // row in tile << 16 | k, per exception + float2 const* exc_w; // {w, 0} per exception: the exact weight + uint32_t base; // exponent base << 7 in both BF16 halves (base << 7 | base << 23) +}; + +// Everything a draft-step launch takes (see invokeDraftStep). +struct Params +{ + __nv_bfloat16 const* x; // [*, K] head input rows + long long const* x_rows; // optional [M]: head input row of request b (pass 0: the last accepted verify row) + int M; // requests, <= kMaxTokens + uint4 const* wf; // [n_tiles][K / 16][32 lanes] 16-byte BF16 A fragments, or the 12-bit coded stream + // [n_tiles][K / 64][3][32 lanes] when cd.exc_off is set (see draftStep.cu) + int n_tiles; + int vocab_offset; // global id of this rank's first head row + unsigned long long* keys; + unsigned int* counters; + unsigned long long* ws_uc; // this rank's exchange slots [2][kMaxTokens][world] + unsigned long long* ws_mc; // their multicast address + int rank; + int world; + unsigned long long* out_keys; // optional [M]: the winning global key of each row + StepArgs st; + CodedArgs cd; + // The step's forward-tail store + next-step staging (foldTail.h), run by the last CTA after this pass's + // tokens and bookkeeping: set on the drafter's last pass only (dispatch.DRAFT_STEP_FOLD_TAIL). + FoldTailParams fold{}; + bool hasFold = false; + // D-FIX: sum the pieces of a tile that several warps stream in warp order (deterministic) instead of by + // shared-memory atomics in arrival order (dispatch.DRAFT_STEP_FIXED_ORDER). + bool fixedOrder = false; +}; + +//! Launch the draft step over `grid` 512-thread CTAs (K = the hidden size, a multiple of 16; of 64 for a coded head). +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream); + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu new file mode 100644 index 000000000000..a506548c1aec --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu @@ -0,0 +1,189 @@ +/* + * 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. + */ + +// The TP4 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's GPT-NeoX +// path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin negated, +// dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale E4M3 +// store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query +// saturates the same way (448, not NaN). + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace drafterRopeFp8Kv +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kBlockSize = 256; + +template +__device__ __forceinline__ float rotate(float c, float x, float s, float pair) +{ + if constexpr (kRopeMode == 0) + { + return __fadd_rn(__fmul_rn(c, x), __fmul_rn(s, pair)); + } + else if constexpr (kRopeMode == 1) + { + return __fmaf_rn(c, x, __fmul_rn(s, pair)); + } + else + { + return __fmaf_rn(s, pair, __fmul_rn(c, x)); + } +} + +__device__ __forceinline__ void storeE4m3Pairs(__nv_fp8_e4m3* dst, __nv_bfloat16 const (&v)[kElemsPerThread]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + __nv_bfloat162 pair; + pair.x = v[i]; + pair.y = v[i + 1]; + __nv_fp8x2_e4m3 const packed(pair); + reinterpret_cast<__nv_fp8x2_storage_t*>(dst)[i / 2] = packed.__x; + } +} + +template +__global__ void __launch_bounds__(kBlockSize) drafterRopeFp8KvKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const lane = threadIdx.x % 32; + int const warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; + int const slots = p.numHeadsQ + p.numHeadsKv; + int const token = warp / slots; + if (token >= p.numTokens) + { + return; + } + int const slot = warp % slots; + bool const isQ = slot < p.numHeadsQ; + int const head = isQ ? slot : slot - p.numHeadsQ; + int const request = token / p.qLen; + int const row = token % p.qLen; + int const cacheLen = p.seqLens[request]; + int const kvPos = cacheLen - p.qLen + row; + int const ropePos + = p.specPositionOffsets != nullptr ? p.specPositionOffsets[request * p.qLen + row] + cacheLen - p.qLen : kvPos; + int const rowWidth = (p.numHeadsQ + 2 * p.numHeadsKv) * kHeadDim; + __nv_bfloat16 const* src + = p.qkv + static_cast(token) * rowWidth + (isQ ? head * kHeadDim : (p.numHeadsQ + head) * kHeadDim); + int const half = p.rotaryDim / 2; + float2 const* coef = p.cosSin + static_cast(ropePos) * half; + + __nv_bfloat16 rotated[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + int const d = lane * kElemsPerThread + i; + bool const firstHalf = d < half; + float const x = __bfloat162float(src[d]); + float const pair = __bfloat162float(src[firstHalf ? d + half : d - half]); + float2 const cs = d < p.rotaryDim ? coef[firstHalf ? d : d - half] : make_float2(1.0f, 0.0f); + float const s = firstHalf ? -cs.y : cs.y; + rotated[i] = __float2bfloat16(rotate(cs.x, x, s, pair)); + } + + if (isQ) + { + storeE4m3Pairs( + p.qOut + (static_cast(token) * p.numHeadsQ + head) * kHeadDim + lane * kElemsPerThread, rotated); + return; + } + + int const block = kvPos / p.tokensPerBlock; + long long const inPage = (static_cast(head) * p.tokensPerBlock + kvPos % p.tokensPerBlock) * kHeadDim + + lane * kElemsPerThread; + long long const pageElems = static_cast(p.numHeadsKv) * p.tokensPerBlock * kHeadDim; + int const* table = p.blockTable + request * p.blockTableRequestStride; + long long const kPage = table[block]; + long long const vPage = table[p.blockTableKvStride + block]; + storeE4m3Pairs(p.pool + kPage * pageElems + inPage, rotated); + + __nv_bfloat16 const* vSrc = p.qkv + static_cast(token) * rowWidth + + (p.numHeadsQ + p.numHeadsKv + head) * kHeadDim + lane * kElemsPerThread; + __nv_bfloat16 v[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + v[i] = vSrc[i]; + } + storeE4m3Pairs(p.pool + vPage * pageElems + inPage, v); +} + +//! With PDL and maxSharedCarveout, the maximum shared-memory carveout lets the full-context decode's large-smem CTAs +//! (the dependent this kernel releases right after its wait) co-reside on the SMs this small kernel occupies. Set once +//! per instantiation. +template +cudaError_t launchWithMaxSharedCarveout( + cudaLaunchConfig_t const& config, Params const& params, bool pdl, bool maxSharedCarveout) +{ + auto const kernel = drafterRopeFp8KvKernel; + if (pdl && maxSharedCarveout) + { + static cudaError_t const carveout = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (carveout != cudaSuccess) + { + return carveout; + } + } + return cudaLaunchKernelEx(&config, kernel, params); +} + +} // namespace + +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream) +{ + int const warps = params.numTokens * (params.numHeadsQ + params.numHeadsKv); + int const warpsPerBlock = kBlockSize / 32; + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((warps + warpsPerBlock - 1) / warpsPerBlock); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + switch (params.ropeMode) + { + case 0: return launchWithMaxSharedCarveout<0>(config, params, pdl, maxSharedCarveout); + case 1: return launchWithMaxSharedCarveout<1>(config, params, pdl, maxSharedCarveout); + case 2: return launchWithMaxSharedCarveout<2>(config, params, pdl, maxSharedCarveout); + default: return cudaErrorInvalidValue; + } +} + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h new file mode 100644 index 000000000000..2a39d03d9d76 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h @@ -0,0 +1,79 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ +namespace drafterRopeFp8Kv +{ + +constexpr int kHeadDim = 128; + +struct Params +{ + //! [tokens, (numHeadsQ + 2 * numHeadsKv) * 128] BF16: the drafter layer's packed Q, K, V rows. + __nv_bfloat16 const* qkv; + //! [positions, rotaryDim / 2] (cos, sin) pairs: the layer's rotary_cos_sin table. + float2 const* cosSin; + //! [requests] cache lengths including this step's rows. + int const* seqLens; + //! [requests, qLen] spec-decoding position offsets of a multi-token pass, or null. + int const* specPositionOffsets; + //! [tokens, numHeadsQ, 128] E4M3 query out. + __nv_fp8_e4m3* qOut; + //! Base of the layer's K / V sub-pages, [*, numHeadsKv, tokensPerBlock, 128] E4M3. + __nv_fp8_e4m3* pool; + //! [requests, 2, maxBlocks] sub-page indices: K in row 0, V in row 1. + int const* blockTable; + long long blockTableRequestStride; + long long blockTableKvStride; + int numTokens; + int qLen; + int numHeadsQ; + int numHeadsKv; + int rotaryDim; + int tokensPerBlock; + //! How the rotation's two products are summed: 0 two products then the add, 1 fma(cos, x, sin * pair), + //! 2 fma(sin, pair, cos * x). + int ropeMode; +}; + +//! The MiniMax-M3 TP4 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, +//! rounded to BF16 (the in-tree preprocess's q output and cache input), then Q, K and V to saturating E4M3: K / V +//! as the in-tree preprocess stores them at unit KV scale, Q as torch's CUDA cast of the BF16 query. Rows map to +//! (request token / qLen, row token % qLen) at cache position seqLen - qLen + row. +//! +//! PDL (pdl): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction and the kernel +//! releases its dependent right after it. The dependent may read only step-invariant data before its own wait. +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream); + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu new file mode 100644 index 000000000000..e6a4f0da2039 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu @@ -0,0 +1,64 @@ +/* + * 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/minimaxM3Sm107Tp4/foldTail.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +namespace +{ + +// Row i of the step (lane c = column), as the Triton fold tail's program i: the store's padded rows, then the next +// step's inputs read from them (zero past each source width). +__global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) +{ + int const i = blockIdx.x; + if (i >= *p.ctrl) + { + return; + } + foldTailRow(p, i, static_cast(threadIdx.x)); +} + +__global__ void __launch_bounds__(32) foldTailOneWarpKernel(FoldTailParams const p) +{ + foldTailWarp(p, static_cast(threadIdx.x)); +} + +} // namespace + +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream) +{ + foldTailKernel<<>>(p); + return cudaGetLastError(); +} + +cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream) +{ + foldTailOneWarpKernel<<<1, 32, 0, stream>>>(p); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h new file mode 100644 index 000000000000..54a58b0569e0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h @@ -0,0 +1,129 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ + +//! Operands of the speculative step's forward-tail store and next-step staging (kernels/m3_fold_tail.py). Every tensor +//! is int32 except `slots` (int64); strides are in elements. The store views are [slots, width], the sources +//! [rows, width]. +struct FoldTailParams +{ + int32_t const* ctrl; // ctrl[0]: the rows the step armed (0: write nothing) + int64_t const* slots; + int32_t const* srcNew; + int64_t sn0, sn1; + int32_t snW; + int32_t* dstNew; + int64_t dn0, dn1; + int32_t dnW; + int32_t const* srcNext; + int64_t sx0, sx1; + int32_t sxW; + int32_t* dstNext; + int64_t dx0, dx1; + int32_t dxW; + int32_t const* srcDraft; + int64_t sd0, sd1; + int32_t sdW; + int32_t* dstDraft; + int64_t dd0, dd1; + int32_t ddW; + int32_t const* srcLens; + int32_t* dstLens; + int32_t* ids; // [rows * T] + int32_t* draftIds; // [rows * (T - 1)] + int32_t* posOff; // [rows * T] + int32_t* kvOff; // [rows] + int32_t tokensPerStep; +}; + +//! One CTA of 32 lanes per row, `rows` CTAs; the first ctrl[0] rows write. Widths and tokensPerStep must be <= 32. +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream); + +//! The same writes from ONE warp (one CTA of 32 lanes looping over the rows, foldTailWarp): the in-kernel form a caller +//! such as the drafter's last-pass draft_step runs in its last CTA; launched alone it is that form's oracle. +cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream); + +#ifdef __CUDACC__ +//! Row i of the tail at lane c (0..31), the work of foldTailKernel's CTA i, for a caller that runs the tail inside its +//! own kernel (e.g. the drafter's last-pass draft_step in its last CTA, after its final tokens are written and fenced). +//! The caller runs rows i < *p.ctrl (one warp per row, or one warp looping over the rows); same code as the kernel, so +//! the writes are bit for bit the standalone tail's. +__device__ __forceinline__ void foldTailRow(FoldTailParams const& p, int i, int c) +{ + int64_t const s = p.slots[i]; + int32_t const nw = c < p.snW ? p.srcNew[i * p.sn0 + c * p.sn1] : 0; + if (c < p.dnW) + { + p.dstNew[s * p.dn0 + c * p.dn1] = nw; + } + int32_t const nxt = c < p.sxW ? p.srcNext[i * p.sx0 + c * p.sx1] : 0; + if (c < p.dxW) + { + p.dstNext[s * p.dx0 + c * p.dx1] = nxt; + } + int32_t const drf = c < p.sdW ? p.srcDraft[i * p.sd0 + c * p.sd1] : 0; + if (c < p.ddW) + { + p.dstDraft[s * p.dd0 + c * p.dd1] = drf; + } + int32_t const ln = p.srcLens[i]; + int32_t const t = p.tokensPerStep; + int32_t const d = t - 1; + if (c == 0) + { + p.dstLens[s] = ln; + p.kvOff[i] = ln - t; + } + if (c < t) + { + p.ids[i * t + c] = c < p.dxW ? nxt : 0; + p.posOff[i * t + c] = ln; + } + if (c < d) + { + p.draftIds[i * d + c] = c < p.ddW ? drf : 0; + } +} + +//! The whole tail from one warp of a running kernel (lane = threadIdx.x % 32): rows i < *p.ctrl in order. +__device__ __forceinline__ void foldTailWarp(FoldTailParams const& p, int lane) +{ + int const rows = *p.ctrl; + for (int i = 0; i < rows; ++i) + { + foldTailRow(p, i, lane); + } +} +#endif + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu new file mode 100644 index 000000000000..2937d461854c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu @@ -0,0 +1,268 @@ +/* + * 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. + */ + +// The split-KV merge for the MiniMax-M3 NVFP4 TP2 (Rubin) full-context decode. The producer is the PrimTS split-KV +// decode with its separate-reduction scratch: per (request, KV head, split, row) one normalized 16-bit partial O and +// one fp32 log2-LSE. Active splits per row follow the producer's split ranges (the causal union of the row's grouped +// Q CTA, instruction-group aligned): splits past them hold nothing and are not read. +// +// One CTA per (request, KV head, kRowsPerCta rows); per row kGroups split groups x kThreadsPerRow threads x 8 dims. +// A thread issues the loads of every split it owns (strided over the groups) before using any, folds them under its +// group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to +// BF16 once. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace fullCtxReduce +{ +namespace +{ + +constexpr int kGroups = 8; +constexpr int kThreadsPerRow = kHeadDim / 8; +constexpr int kRowsPerCta = 1; +constexpr int kSplitsPerGroup = kMaxSplits / kGroups; +constexpr int kThreads = kGroups * kThreadsPerRow * kRowsPerCta; + +__device__ int activeSplits(Params const& p, int b, int row) +{ + int const seqKv = p.seqLensKv[b]; + int visible = seqKv; + if (p.causal) + { + int const token = row / p.headsPerKv; + int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; + int const end = min(base + p.qTokensPerCta, p.seqLenQ); + visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); + } + int const tiles = (visible + p.tileKv - 1) / p.tileKv; + int const tilesPerGroup = p.splits * p.instsKv; + int const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; + int const local = max(groups * p.instsKv, p.instsKv); + return (tiles + local - 1) / local; +} + +template +__device__ float2 unpack(uint32_t word) +{ + if constexpr (kBf16Partial) + { + return __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&word)); + } + else + { + return __half22float2(*reinterpret_cast<__half2 const*>(&word)); + } +} + +template +__global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); + // The dependent (the o_proj activation quantization / GEMV) reads `out` only after its own griddepcontrol.wait, + // which waits for this whole grid, so it is released here and runs its weight-only prologue under the merge. + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const b = blockIdx.z; + int const h = blockIdx.y; + int const localRow = threadIdx.x / (kGroups * kThreadsPerRow); + int const row = blockIdx.x * kRowsPerCta + localRow; + int const g = (threadIdx.x / kThreadsPerRow) % kGroups; + int const t = threadIdx.x % kThreadsPerRow; + bool const validRow = row < p.rows; + int const active = validRow ? activeSplits(p, b, row) : 0; + + __shared__ float sMax[kRowsPerCta][kGroups]; + __shared__ float sSum[kRowsPerCta][kGroups]; + __shared__ float sAcc[kRowsPerCta][kGroups][kHeadDim]; + + uint4 const* partialO = static_cast(p.partialO); + size_t const kvIdx = static_cast(b) * p.numKvHeads + h; + float lse[kSplitsPerGroup]; + uint4 vec[kSplitsPerGroup]; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + int const s = g + i * kGroups; + if (s < active) + { + size_t const r = (kvIdx * p.splits + s) * p.rows + row; + lse[i] = __ldcg(p.partialStats + r); + vec[i] = __ldcg(partialO + r * kThreadsPerRow + t); + } + } + float groupMax = -CUDART_INF_F; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active) + { + groupMax = fmaxf(groupMax, lse[i]); + } + } + // A split of the CTA's causal union can hold no key of this row (log2-LSE -inf). Such a split, and a group of + // only such splits, is skipped rather than weighted 0: its partial O carries no defined value. + float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + float sum = 0.f; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active && lse[i] != -CUDART_INF_F) + { + float const w = exp2f(lse[i] - groupMax); + sum += w; + uint32_t const* words = reinterpret_cast(&vec[i]); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = unpack(words[j]); + acc[2 * j] += w * f.x; + acc[2 * j + 1] += w * f.y; + } + } + } + if (t == 0) + { + sMax[localRow][g] = groupMax; + sSum[localRow][g] = sum; + } +#pragma unroll + for (int j = 0; j < 8; ++j) + { + sAcc[localRow][g][t * 8 + j] = acc[j]; + } + __syncthreads(); + if (g != 0 || !validRow) + { + return; + } + float rowMax = -CUDART_INF_F; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + rowMax = fmaxf(rowMax, sMax[localRow][k]); + } + float total = 0.f; + float o[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + if (sMax[localRow][k] == -CUDART_INF_F) + { + continue; + } + float const scale = exp2f(sMax[localRow][k] - rowMax); + total += scale * sSum[localRow][k]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + o[j] += scale * sAcc[localRow][k][t * 8 + j]; + } + } + float const inv = total > 0.f ? 1.f / total : 0.f; + int const token = row / p.headsPerKv; + int const head = row % p.headsPerKv; + __nv_bfloat162 packed[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + packed[j] = __floats2bfloat162_rn(o[2 * j] * inv, o[2 * j + 1] * inv); + } + size_t const outRow + = (static_cast(b) * p.seqLenQ + token) * (p.numKvHeads * p.headsPerKv) + h * p.headsPerKv + head; + reinterpret_cast(p.out + outRow * kHeadDim)[t] = *reinterpret_cast(packed); + if (p.q8 == nullptr) + { + return; + } + // MXFP8 copy of the BF16 row: a 32-column block is 4 neighbouring threads x 8 columns (lanes 0..15 of warp 0 + // hold the row). Block scale = the UE8M0 exponent of ceil(amax / 448), values = satfinite E4M3 of x * 2^-scale. + float x[8]; + float amax = 0.f; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = __bfloat1622float2(packed[j]); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + amax = fmaxf(amax, fmaxf(fabsf(f.x), fabsf(f.y))); + } + amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 2)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + uint16_t q2[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2[j]) : "f"(x[2 * j + 1] * qs), "f"(x[2 * j] * qs)); + } + int const tok = b * p.seqLenQ + token; + int const col = (h * p.headsPerKv + head) * kHeadDim + t * 8; + *reinterpret_cast(p.q8 + static_cast(tok) * (p.numSfKTiles * 128) + col) + = *reinterpret_cast(q2); + if (t % 4 == 0) + { + int const kSf = col / 32; + int64_t const off = static_cast(tok / 128) * p.numSfKTiles * 512 + (kSf / 4) * 512 + (tok % 32) * 16 + + ((tok % 128) / 32) * 4 + kSf % 4; + p.sf[off] = static_cast(e); + } +} + +} // namespace + +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) +{ + if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((params.rows + kRowsPerCta - 1) / kRowsPerCta, params.numKvHeads, params.batch); + config.blockDim = dim3(kThreads); + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + return params.bf16Partial ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) + : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); +} + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h new file mode 100644 index 000000000000..95af79e2b2ce --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h @@ -0,0 +1,77 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ +//! The split-KV merge of the MiniMax-M3 TP2 full-context decode (PrimTS split-KV decode, separate-reduction scratch) +//! straight to BF16. See fullCtxReduce.cu. +namespace fullCtxReduce +{ + +constexpr int kHeadDim = 128; +constexpr int kMaxSplits = 128; + +struct Params +{ + void const* partialO; // [B, Hkv, S, R, 128] normalized 16-bit partial O per split + float const* partialStats; // [B, Hkv, S, R] log2-LSE per split + int const* seqLensKv; // [B] + __nv_bfloat16* out; // [B, SQ, Hkv * hR, 128] + int batch; + int numKvHeads; + int splits; // S, the planned split capacity (<= kMaxSplits) + int rows; // R = hR * SQ, rows [token, local head] + int seqLenQ; // SQ + int headsPerKv; // hR + int qTokensPerCta; // query tokens one producer CTA owns (grouped Q) + int tileKv; // producer K/V tile + int instsKv; // producer K/V instructions per loop step + bool causal; + bool bf16Partial; + uint8_t* q8 = nullptr; // optional E4M3 copy of out [B * SQ, Hkv * hR * 128] (the o_proj MXFP8 hand-off) + uint8_t* sf = nullptr; // its UE8M0 block scales, 128x4-swizzled as trtllm::mxfp8_quantize(x, True) lays them out + int numSfKTiles = 0; // Hkv * hR * 128 / 128 +}; + +//! Merges the active splits of every row in fp32 with one max rescale and writes BF16; with q8 set, also the MXFP8 +//! copy of the BF16 rows (bit for bit trtllm::mxfp8_quantize of out). With pdl the launch is +//! programmatic: the grid waits on the producer (griddepcontrol.wait) before reading its scratch. +//! +//! PDL contract: griddepcontrol.wait is the kernel's first instruction, so nothing (scratch, lengths, out) is read or +//! written before every prerequisite grid has completed and flushed; a producer's early trigger +//! (griddepcontrol.launch_dependents) only lets this grid launch and wait sooner. The kernel triggers its own +//! dependents right after its wait: the next kernel's griddepcontrol.wait still waits for this whole grid, so its +//! post-wait code sees all of out, and its pre-wait prologue must not read out. +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream); + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu new file mode 100644 index 000000000000..1942082c6df9 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu @@ -0,0 +1,1264 @@ +/* + * 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/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// SM107 (Rubin) TP4 copy of minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu for an sm_107 TP4 target: the row kernel at 4 +// ranks. Two changes in the arithmetic: the rank sum rounds to BF16 before the residual add, the rounding points of the +// stock all-reduce + fused residual RMSNorm on this target (bf16 all-reduce output, then bf16(residual + sum)); and +// every call but the MoE finalize runs the row kernel below, whose RMS statistic is the stock fused norm's reduction +// tree (bitwise the stock outputs). +// +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions): they are compiled for every sm_10x target of the build, the sm_100f family target +// included (which runs on SM107). A device pass for any other architecture compiles them to a trap, never to a silent +// empty body, and the host entry point checks that the kernel binary it launches holds sm_10x code (see +// checkBinary). +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace cg = cooperative_groups; + +namespace +{ + +constexpr int kEmpty = static_cast(0x80000000u); +constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k + +// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], +// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c +// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp +// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its +// arithmetic is the same source. (The TP2 target's mcAllreduceNormTp2.cu has the same epilogue.) +constexpr int kCapThreads = 256; +constexpr int kCapIters = 3; +constexpr int kCapWarps = kCapThreads / 32; +constexpr int kCapHidden = kCapThreads * kCapIters * 8; + +// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). +__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) +{ + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int q = 0; q < 8; ++q) + { + float const v = __bfloat162float(e[q]); + ss += v * v; + } + return ss; +} + +__device__ __forceinline__ float cap_warp_sum(float v) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. +__device__ __forceinline__ void cap_store( + float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) +{ + uint4 const wr = *reinterpret_cast(w); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int q = 0; q < 8; ++q) + oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); + *reinterpret_cast(cap) = o; +} + +__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) +{ + if ((w & 0xFFFFu) == 0x8000u) + w &= 0xFFFF0000u; + if ((w >> 16) == 0x8000u) + w &= 0x0000FFFFu; + return w; +} + +__device__ __forceinline__ float bf16_lo(uint32_t w) +{ + return __uint_as_float(w << 16); +} + +// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's +// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals +// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the +// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. +__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) +{ + float value; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); + uint32_t result; + asm("{\n" + " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" + " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" + " setp.le.f32 p_zero, %1, 0f00000000;\n" + " mov.b32 bits, %1;\n" + " shr.b32 exp_biased, bits, 23;\n" + " and.b32 exp_biased, exp_biased, 255;\n" + " and.b32 mantissa, bits, 0x7FFFFF;\n" + " setp.ne.u32 p_has_mant, mantissa, 0;\n" + " selp.u32 bump, 1, 0, p_has_mant;\n" + " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" + " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" + " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" + " @p_tiny_sub mov.u32 bump, 0;\n" + " add.u32 r, exp_biased, bump;\n" + " setp.gt.u32 p_ovf, r, 254;\n" + " selp.u32 r, 254, r, p_ovf;\n" + " selp.u32 %0, 0, r, p_zero;\n" + "}\n" + : "=r"(result) + : "f"(value)); + return result; +} + +__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) +{ + float inv; + asm("{\n" + " .reg .s32 new_exp;\n" + " .reg .b32 float_bits;\n" + " .reg .pred p_zero;\n" + " setp.eq.u32 p_zero, %1, 0;\n" + " sub.s32 new_exp, 254, %1;\n" + " max.s32 new_exp, new_exp, 0;\n" + " shl.b32 float_bits, new_exp, 23;\n" + " mov.b32 %0, float_bits;\n" + " @p_zero mov.b32 %0, 0;\n" + "}\n" + : "=f"(inv) + : "r"(ue8m0)); + return inv; +} + +// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. +__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) +{ + uint16_t pair; + asm("{\n" + " .reg .f32 f0, f1;\n" + " mul.f32 f0, %1, %3;\n" + " mul.f32 f1, %2, %3;\n" + " min.f32 f0, f0, 0f43E00000;\n" + " max.f32 f0, f0, 0fC3E00000;\n" + " min.f32 f1, f1, 0f43E00000;\n" + " max.f32 f1, f1, 0fC3E00000;\n" + " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" + "}\n" + : "=h"(pair) + : "f"(lo), "f"(hi), "f"(inv)); + return pair; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ float bf16_hi(uint32_t w) +{ + return __uint_as_float(w & 0xFFFF0000u); +} + +__device__ __forceinline__ int4 ld_volatile_v4(int const* p) +{ + int4 v; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +struct FinalizeArgs +{ + __nv_bfloat16 const* fc2; // [permuted rows, H] routed-expert FC2 output (null: plain input x) + int const* perm; // [M, topk] expanded index -> permuted row, -1 = dropped + void const* scales; // [M, topk] routing weights (fp32 or bf16) + int scale_fp32; + int topk; +}; + +// grid = min(M, MAXT) * CLUSTER CTAs, cluster (CLUSTER, 1, 1), THREADS threads. +// OPTS (each bit-identical): 1 residual + norm weight loaded at entry (before the poll); 2 cluster RMS exchange by +// st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated right after the poll; 8 max-smem +// carveout so the next kernel's CTAs co-reside early. +// OPTS 16 (NOT bit-identical with the others; a different boundary): the Eagle3 drafter's boundary, whose stock is a +// plain BF16 all-reduce, then FlashInfer fused_add_rmsnorm: the RMS statistic and the normalization take the +// unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded to BF16). +template +__global__ void __launch_bounds__(THREADS) mc_allreduce_norm_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, FinalizeArgs fin, int M, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int world, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ alignas(16) float cluster_sq[CLUSTER]; + extern __shared__ __align__(16) uint4 cap_slice[]; // [cols / 8] this CTA's updated chunks (cap_out only) + // Option 32, ONESHOT order (H 6144, CLUSTER 8, not the drafter option 16): every CTA's 3 active warp sums gathered + // in every CTA (global warp 3c + w), then reduced as the stock TP4 oneshot fusion does (see the consumer below). + constexpr bool kOneshotOrder = (OPTS & 32) != 0; + static_assert(!kOneshotOrder || (CLUSTER == 8 && THREADS == 128 && (OPTS & 16) == 0), "option 32: CLUSTER 8 x 128"); + constexpr int kAW = 3; + __shared__ alignas(16) float cluster_wsq[CLUSTER * kAW]; + __shared__ float oneshot_total; + __shared__ alignas(8) unsigned long long xbar; + cg::cluster_group cluster = cg::this_cluster(); + if constexpr ((OPTS & 2) != 0) + { + if (threadIdx.x == 0) + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(b) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + } + int const crank = static_cast(cluster.block_rank()); + int const t = blockIdx.x / CLUSTER; + int const cols = H / CLUSTER; + int const c0 = crank * cols + threadIdx.x * 8; + bool const active = t < M && threadIdx.x * 8 < cols; + int64_t const half_words = static_cast(MAXT) * world * (H / 2); + + // prepushed 1: the producer (m3_mx_gemv) already pushed this rank's rows into the exchange slots; 2: and it + // launched this grid only after its pushes (its grid wait passed), so the slots are polled without a grid wait; + // 3: as 2, with x in set A and x2 in set B. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const parity = *reinterpret_cast(flags + blockIdx.x) & 1; + uint4 rraw = make_uint4(0u, 0u, 0u, 0u); + uint4 wraw = make_uint4(0u, 0u, 0u, 0u); + if constexpr ((OPTS & 1) != 0) + { + // the residual row was written grids earlier (prepushed >= 2: the producer launched this grid only after its + // own grid wait; otherwise the grid wait above) + if (active) + { + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + wraw = *reinterpret_cast(norm_w + c0); + } + } + // Push this rank's slice into every rank's buffer. + if (active && !prepushed) + { + uint4 raw; + if (fin.fc2 != nullptr) + { + // MoE finalize: fp32 sum of the routing-weighted top-k expert rows plus the shared expert, rounded once. + float f[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + if (fin.topk <= kFinMaxTopk) + { + // All top-k indices and scales, then every expert-row load in flight at once, then the same in-order + // FP32 accumulation as the loop below (one memory round trip instead of topk; bit-identical). + int prow[kFinMaxTopk]; + float sc[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + prow[k] = -1; + sc[k] = 0.f; + if (k < fin.topk) + { + int const e = t * fin.topk + k; + prow[k] = fin.perm[e]; + sc[k] = fin.scale_fp32 ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + } + } + uint4 rv[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] >= 0) + rv[k] = *reinterpret_cast(fin.fc2 + static_cast(prow[k]) * H + c0); + } +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] < 0) + continue; + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv[k]); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc[k]; + f[2 * q + 1] += v.y * sc[k]; + } + } + } + else + for (int k = 0; k < fin.topk; ++k) + { + int const e = t * fin.topk + k; + int const prow = fin.perm[e]; + if (prow < 0) + continue; + float const sc = fin.scale_fp32 + ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + uint4 const rv = *reinterpret_cast(fin.fc2 + static_cast(prow) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc; + f[2 * q + 1] += v.y * sc; + } + } + if (x2 != nullptr) + { + // As stock: the routed output is rounded to BF16 first, then the shared expert is added (one more + // BF16 rounding of the FP32 sum). + uint4 const rv = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const routed = __bfloat1622float2(__floats2bfloat162_rn(f[2 * q], f[2 * q + 1])); + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] = routed.x + v.x; + f[2 * q + 1] = routed.y + v.y; + } + } + __nv_bfloat162* o = reinterpret_cast<__nv_bfloat162*>(&raw); +#pragma unroll + for (int q = 0; q < 4; ++q) + o[q] = __floats2bfloat162_rn(f[2 * q], f[2 * q + 1]); + } + else + { + raw = *reinterpret_cast(x + static_cast(t) * H + c0); + } + if (fin.fc2 == nullptr && x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * world + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float s[8] = {}; + float sq = 0.f; + if (active) + { + // The residual row (written grids earlier) and every rank's slot are loaded together; slots still empty are + // re-polled. Ranks are then summed in rank order, as before. + if constexpr ((OPTS & 1) == 0) + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + constexpr int kMaxWorld = 8; + int* const p0 = ws_uc + parity * half_words + static_cast(t) * world * (H / 2) + c0 / 2; + int4 v[kMaxWorld]; + int4 v2[kMaxWorld]; + bool const two = prepushed == 3; + int* const p1 = p0 + 2 * half_words; // set B + uint32_t pending = (1u << world) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (r < world) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr ((OPTS & 16) != 0) + { + float const x0 = s[2 * q] + rf.x, x1 = s[2 * q + 1] + rf.y; + s[2 * q] = x0; + s[2 * q + 1] = x1; + sq += x0 * x0 + x1 * x1; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + if constexpr (kOneshotOrder) + { + // stock oneshot: acc += v * v, one column at a time + sq += uf.x * uf.x; + sq += uf.y * uf.y; + } + else + sq += uf.x * uf.x + uf.y * uf.y; + } + } + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr ((OPTS & 16) == 0) + { + if (cap_out != nullptr) + cap_slice[threadIdx.x] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + if constexpr ((OPTS & 4) != 0) + { + // every CTA (active or not) rotates its call parity once its own poll is done (read only by later grids) + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffffu, sq, off); + if ((threadIdx.x & 31) == 0) + warp_sq[threadIdx.x / 32] = sq; + __syncthreads(); + if constexpr ((OPTS & 2) != 0) + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); // peers' xbar initialized (long since) + if (threadIdx.x == 0) + { + if constexpr (kOneshotOrder) + { + // Every CTA gets this CTA's kAW warp sums at [kAW * crank + w]. + if constexpr ((OPTS & 2) == 0) + { + for (int c = 0; c < CLUSTER; ++c) + for (int w = 0; w < kAW; ++w) + *cluster.map_shared_rank(&cluster_wsq[kAW * crank + w], c) = warp_sq[w]; + } + else + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * kAW * 4) + : "memory"); + for (int c = 0; c < CLUSTER; ++c) + { + uint32_t rb; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); + for (int w = 0; w < kAW; ++w) + { + uint32_t rd; + uint32_t const d + = static_cast(__cvta_generic_to_shared(&cluster_wsq[kAW * crank + w])); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), + "r"(__float_as_uint(warp_sq[w])), "r"(rb) + : "memory"); + } + } + } + } + else + { + float total = 0.f; +#pragma unroll + for (int w = 0; w < THREADS / 32; ++w) + total += warp_sq[w]; + if constexpr ((OPTS & 2) == 0) + for (int c = 0; c < CLUSTER; ++c) + *cluster.map_shared_rank(&cluster_sq[crank], c) = total; + else + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t const d = static_cast(__cvta_generic_to_shared(&cluster_sq[crank])); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * 4) + : "memory"); + for (int c = 0; c < CLUSTER; ++c) + { + uint32_t rd, rb; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), + "r"(__float_as_uint(total)), "r"(rb) + : "memory"); + } + } + } + } + if constexpr ((OPTS & 2) == 0) + cluster.sync(); + else + { + // every peer's sum landed (st.async complete_tx on this CTA's barrier); no cluster-wide release-arrive + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t done = 0; + while (!done) + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(b) + : "memory"); + } + if constexpr ((OPTS & 16) == 0) + { + if (cap_out != nullptr) + { + // The row's chunks are spread over the cluster (CTA c holds chunks [c * per, (c + 1) * per)); every CTA + // forms capture_norm's statistic itself, its 128 threads standing in for capture threads ct and ct + 128. + __shared__ float cap_red[kCapWarps]; + int const per = cols / 8; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + float ca = 0.f, cb = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + { + int const ka = it * kCapThreads + static_cast(threadIdx.x); + int const kb = ka + kCapThreads / 2; + ca = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[ka % per], ka / per), ca); + cb = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[kb % per], kb / per), cb); + } + // Every peer read is done before any CTA moves on (and later exits). + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + ca = cap_warp_sum(ca); + cb = cap_warp_sum(cb); + if ((threadIdx.x & 31) == 0) + { + cap_red[threadIdx.x / 32] = ca; + cap_red[threadIdx.x / 32 + kCapWarps / 2] = cb; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + if (active) + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + if constexpr (kOneshotOrder) + { + // The stock TP4 oneshot fusion at H 6144 (4 CTAs x 192 threads): blockReduceSumV2 per CTA = a 32-lane + // butterfly over its 6 warp sums (lanes >= 6 zero), lane 0's value; then thread 0 adds the 4 CTA sums in + // order from 0.f. Global warps 6g..6g+5 = this cluster's CTAs 2g, 2g+1. + if (threadIdx.x < 32) + { + int const lane = threadIdx.x; + float acc = 0.f; + for (int g = 0; g < 4; ++g) + { + float val = lane < 6 ? cluster_wsq[6 * g + lane] : 0.f; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + val += __shfl_xor_sync(0xffffffffu, val, off); + acc += __shfl_sync(0xffffffffu, val, 0); + } + if (lane == 0) + oneshot_total = acc; + } + __syncthreads(); + } + if (active) + { + float total = 0.f; + if constexpr (kOneshotOrder) + total = oneshot_total; + else + { +#pragma unroll + for (int c = 0; c < CLUSTER; ++c) + total += cluster_sq[c]; + } + float const inv = rsqrtf(total / static_cast(H) + eps); + if constexpr ((OPTS & 1) == 0) + wraw = *reinterpret_cast(norm_w + c0); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + } + if (out_q != nullptr) + { + // MXFP8 copy of the normed row for the next GEMM, as the CuTe DSL mxfp8 quantize of `normed` (the tree's + // decode quantize, MX-TAC; see mxfp8_ue8m0): four threads hold one 32-column scale block. + float amax = 0.f; + if (active) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + if (active) + { + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((threadIdx.x & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } + } + if constexpr ((OPTS & 4) == 0) + { + __syncthreads(); + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +// Refuses (raises) to launch a kernel whose loaded binary holds no sm_10x code (PTX and SASS version >= 100): the +// device passes for other architectures compile the kernels to a trap, and a build that left out every sm_10x target +// must not reach the GPU at all. +template +bool checkBinary(Kernel kernel, char const* name) +{ + cudaFuncAttributes attr{}; + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, reinterpret_cast(kernel))); + TLLM_CHECK_WITH_INFO(attr.ptxVersion >= 100 && attr.binaryVersion >= 100, + "%s: the loaded binary holds no sm_10x code (PTX %d, SASS %d): rebuild with an sm_100f / sm_107a target", name, + attr.ptxVersion, attr.binaryVersion); + return true; +} + +template +void launch(MinimaxM3Sm107Tp4McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) +{ + int const M = p.numTokens, H = p.hidden; + TLLM_CHECK_WITH_INFO(H % (CLUSTER * 8) == 0 && H / CLUSTER / 8 <= THREADS, "unsupported hidden size %d", H); + TLLM_CHECK_WITH_INFO(M >= 1 && M <= MAXT, "%d tokens do not fit the %d-token bucket", M, MAXT); + auto kernel = mc_allreduce_norm_kernel; + static bool const binaryOk = checkBinary(kernel, "mc_allreduce_norm_kernel"); + (void) binaryOk; + static bool const configured = [&]() + { + bool ok = CLUSTER <= 8 + || cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess; + if ((OPTS & 8) != 0) + ok = ok + && cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared) + == cudaSuccess; + if (!ok) + cudaGetLastError(); + return ok; + }(); + TLLM_CHECK_WITH_INFO(configured, "cluster size not supported"); + FinalizeArgs const fin{p.fc2, p.perm, p.scales, p.scaleFp32 ? 1 : 0, p.topk}; + cudaLaunchConfig_t cfg = {}; + // The grid is sized to the call's rows: min(M, MAXT) clusters. Each (row, slice) slot keeps its own parity flag, + // read by the producers before they push and rotated only by the slot's CTA after it polled: a slot that a call + // does not use keeps its parity, and every rank runs the same calls, so the parities stay consistent when M + // changes between calls. + cfg.gridDim = dim3((M < MAXT ? M : MAXT) * CLUSTER); + cfg.blockDim = dim3(THREADS); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO((OPTS & 16) == 0 && H == kCapHidden && 2 * THREADS == kCapThreads, + "the capture epilogue takes H %d on the 128-thread cluster kernel, not the drafter boundary", kCapHidden); + cfg.dynamicSmemBytes = static_cast(H / CLUSTER / 8) * sizeof(uint4); + } + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = CLUSTER; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, fin, M, H, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, p.world, p.outNormed, p.outResidual, p.outQuant, p.outSf, (H / 32 + 3) / 4, + p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +// ws_uc / ws_mc are this rank's unicast and multicast views of ONE exchange buffer (they alias): no __restrict__. +// prepushed >= 2 contract: the producer (m3_mx_gemv's o_proj push, the m3_moe engine's finalize push) passes its own +// grid dependency, pushes its rows into the exchange, and only then releases this grid (griddepcontrol +// .launch_dependents); this kernel therefore skips griddepcontrol.wait and reads only the pushed (Lamport-polled) +// words, its parity flags and the residual / norm weight, which were written two or more grids earlier. +// +// The TP4 row kernel (4 ranks): one CTA of H / 8 = 768 threads per token row, used for every call but the MoE finalize. +// Its residual add + RMSNorm is the stock TP4 all-reduce fusion (allReduceFusionKernels.cu, oneshot Lamport), the fused +// norm after the all-reduce on this target, operation for operation at H 6144: thread j owns columns [8j, 8j + 8), its +// sum of squares is acc += v * v over them in order, then the warp butterfly 16..1, each group of 6 warps' sums through +// shared memory into a second 32-lane butterfly (lanes >= 6 zero), the 4 group sums added in order from 0.f (the stock +// oneshot kernel's 4-CTA cluster), rsqrtf(acc / H + eps) and (v * denom) * w. With the BF16 rank sum before the +// residual add, the normed rows, the updated residual and the MXFP8 copy are bitwise the stock ones. The exchange slots +// and the per-(row, slice) parity flags keep the cluster kernel's layout (slice = 96 threads), so producers that push +// (the o_proj GEMV) are unchanged. +constexpr int kSlices = 8; +constexpr int kSliceThreads = 96; +constexpr int kRowThreads = 768; +// The stock TP4 fused norm's launch at H 6144 (allreduce_fusion_kernel_launcher): 768 threads per token start as 8 +// CTAs of 96, then widen to CTAs of at least 128 threads: 4 CTAs of 192 per row cluster. +constexpr int kOneshotBlocks = 4; +constexpr int kOneshotThreads = kRowThreads / kOneshotBlocks; +static_assert(kOneshotThreads % 32 == 0 && kOneshotThreads / 32 <= 32, "one second-level butterfly per CTA"); + +// grid = M CTAs (one per row), THREADS = H / 8 threads. WORLD ranks (compile time). DRAFTER (opts bit 16): the Eagle3 +// drafter boundary, whose stock is a plain BF16 all-reduce then FlashInfer fused_add_rmsnorm: the RMS statistic and +// the normalization take the unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded). +// Its reduction is FlashInfer's CuTe-DSL FusedAddRMSNormKernel at H 6144 (what flashinfer.norm.fused_add_rmsnorm runs +// on SM 10.7 in this image: 128-thread CTAs, 64 threads per row, cluster 1): FI thread t sums its 96 columns +// 512 b + 8 t + e in (b, e) order from 0.f, a butterfly over its warp with offsets 1, 2, .., 16, the row's 2 warp sums +// through a second ascending butterfly (lanes >= 2 zero); rstd = rsqrt.approx(div.rn(sum, H) + eps), y = (h * rstd) * +// w. Bitwise to the stock boundary (oracle 777774: normed and residual, eager and graph, 4 ranks, 128 calls). CARVE +// (opts bit 64): the same kernel, its own function with the max-smem carveout preference set, so the next kernel's CTAs +// (the QKV GEMV's big staged tiles) can co-reside on the SM while it runs (the cluster kernel's option 8). ws_uc / +// ws_mc alias (one exchange buffer): no __restrict__. +template +__global__ void __launch_bounds__(THREADS, 1) mc_allreduce_norm_row_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int maxt, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ float group_sq[kOneshotBlocks]; + extern __shared__ __align__(16) uint4 cap_row[]; // [THREADS] the updated row's chunks (cap_out, !DRAFTER only) + // DRAFTER: the row's FP32 sums h, regrouped into FlashInfer's per-thread column sets. + __shared__ float h_row[DRAFTER ? THREADS * 8 : 1]; + __shared__ float total_sh; + int const t = blockIdx.x; + int const j = threadIdx.x; + int const c0 = j * 8; + int const slice = j / kSliceThreads; + int64_t const half_words = static_cast(maxt) * WORLD * (H / 2); + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); + int const parity = *reinterpret_cast(flags + t * kSlices + slice) & 1; + uint4 const rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + uint4 const wraw = *reinterpret_cast(norm_w + c0); + if (!prepushed) + { + uint4 raw = *reinterpret_cast(x + static_cast(t) * H + c0); + if (x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * WORLD + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } + asm volatile("griddepcontrol.launch_dependents;"); + float s[8] = {}; + float sq = 0.f; + { + int* const p0 = ws_uc + parity * half_words + static_cast(t) * WORLD * (H / 2) + c0 / 2; + int* const p1 = p0 + 2 * half_words; // set B + bool const two = prepushed == 3; + int4 v[WORLD]; + int4 v2[WORLD]; + uint32_t pending = (1u << WORLD) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr (DRAFTER) + { + s[2 * q] += rf.x; + s[2 * q + 1] += rf.y; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + } + } +#pragma unroll + for (int e = 0; e < 8; ++e) + sq += s[e] * s[e]; + if constexpr (DRAFTER) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + h_row[c0 + e] = s[e]; + } + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr (!DRAFTER) + { + if (cap_out != nullptr) + cap_row[j] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + // The stock TP4 all-reduce's fused RMSNorm (allReduceFusionKernels.cu, allreduce_fusion_kernel_oneshot_lamport at + // H 6144): a cluster of kOneshotBlocks CTAs x kOneshotThreads threads per row; each CTA reduces its warps with + // blockReduceSumV2 (warp butterfly 16..1, the CTA's warp sums into a second 32-lane butterfly, lanes past its warp + // count zero), then thread 0 sums the CTAs' totals in cluster-rank order from 0.f. + sq += __shfl_xor_sync(~0, sq, 16); + sq += __shfl_xor_sync(~0, sq, 8); + sq += __shfl_xor_sync(~0, sq, 4); + sq += __shfl_xor_sync(~0, sq, 2); + sq += __shfl_xor_sync(~0, sq, 1); + if ((j & 31) == 0) + warp_sq[j / 32] = sq; + __syncthreads(); + if constexpr (DRAFTER) + { + constexpr int kFiThreads = 64; // FlashInfer CuTe threads per row at H 6144 + constexpr int kFiBlocks = THREADS * 8 / (kFiThreads * 8); // 12 column blocks of 512 + if (j < kFiThreads) + { + float acc = 0.f; +#pragma unroll + for (int b = 0; b < kFiBlocks; ++b) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + { + float const v = h_row[b * kFiThreads * 8 + j * 8 + e]; + acc = __fmaf_rn(v, v, acc); // the DSL's h * h + reduce contracts to FMA + } + } +#pragma unroll + for (int off = 1; off < 32; off <<= 1) + acc += __shfl_xor_sync(~0, acc, off); + if ((j & 31) == 0) + group_sq[j / 32] = acc; + } + __syncthreads(); + if (j < 32) + { + float val = j < kFiThreads / 32 ? group_sq[j] : 0.f; +#pragma unroll + for (int off = 1; off < 32; off <<= 1) + val += __shfl_xor_sync(~0, val, off); + if (j == 0) + total_sh = val; + } + } + else if (j < kOneshotBlocks * 32) + { + constexpr int kWarpsPerBlock = kOneshotThreads / 32; + int const g = j / 32; + int const lane = j & 31; + float val = lane < kWarpsPerBlock ? warp_sq[g * kWarpsPerBlock + lane] : 0.f; + val += __shfl_xor_sync(~0, val, 16); + val += __shfl_xor_sync(~0, val, 8); + val += __shfl_xor_sync(~0, val, 4); + val += __shfl_xor_sync(~0, val, 2); + val += __shfl_xor_sync(~0, val, 1); + if (lane == 0) + group_sq[g] = val; + } + __syncthreads(); + if (!DRAFTER && j == 0) + { + float acc = 0.f; +#pragma unroll + for (int g = 0; g < kOneshotBlocks; ++g) + acc += group_sq[g]; + total_sh = acc; + } + // The parity of this row's slices rotates once every thread has read the exchange (later grids only). + if (j % kSliceThreads == 0) + flags[t * kSlices + slice] = parity ^ 1; + __syncthreads(); + if constexpr (!DRAFTER) + { + if (cap_out != nullptr) + { + // The whole row is in this CTA: threads 0..255 run capture_norm's chains over the staged chunks. + __shared__ float cap_red[kCapWarps]; + if (j < kCapThreads) + { + float css = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + css = cap_chunk_sq(cap_row[it * kCapThreads + j], css); + css = cap_warp_sum(css); + if ((j & 31) == 0) + cap_red[j / 32] = css; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + float inv; + if constexpr (DRAFTER) + { + float const mean = __fdiv_rn(total_sh, static_cast(H)); + asm("rsqrt.approx.f32 %0, %1;" : "=f"(inv) : "f"(__fadd_rn(mean, eps))); + } + else + inv = rsqrtf(total_sh / H + eps); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + if (out_q != nullptr) + { + float amax = 0.f; +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((j & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +template +auto rowKernel() +{ + auto kernel = mc_allreduce_norm_row_kernel<4, kRowThreads, DRAFTER, CARVE>; + static bool const ok = [&]() + { + checkBinary(kernel, "mc_allreduce_norm_row_kernel"); + if constexpr (CARVE) + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared)); + return true; + }(); + (void) ok; + return kernel; +} + +void launchRow(MinimaxM3Sm107Tp4McAllreduceNormParams const& p, int32_t maxTokens, bool drafter, bool carve, bool pdl, + cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(p.hidden == kRowThreads * 8 && p.world == 4, "the row kernel takes H 6144 at 4 ranks"); + TLLM_CHECK_WITH_INFO(p.numTokens >= 1 && p.numTokens <= maxTokens, "%d tokens do not fit the %d-token bucket", + p.numTokens, maxTokens); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(p.numTokens); + cfg.blockDim = dim3(kRowThreads); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO( + !drafter && p.hidden == kCapHidden, "the capture epilogue takes H %d, not the drafter", kCapHidden); + cfg.dynamicSmemBytes = static_cast(kRowThreads) * sizeof(uint4); + } + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + auto kernel = drafter ? (carve ? rowKernel() : rowKernel()) + : (carve ? rowKernel() : rowKernel()); + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, p.hidden, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, maxTokens, p.outNormed, p.outResidual, p.outQuant, p.outSf, + (p.hidden / 32 + 3) / 4, p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +constexpr int kCluster = 8; +constexpr int kThreads = 128; + +} // namespace + +bool isMinimaxM3Sm107Tp4McAllreduceNormSupported(int32_t maxTokens, int32_t opts) +{ + opts &= ~(64 | 128); // the row kernel's carveout instance; 128 only routes the call (invoke) + return (maxTokens == 8 && (opts == 7 || opts == 15 || opts == 23 || opts == 31 || opts == 39 || opts == 47)) + || ((maxTokens == 16 || maxTokens == 48) && (opts == 15 || opts == 47)); +} + +// Option 128 routes a world-4 call to the 8-CTA cluster kernel instead of the row kernel: the drafter boundary's +// cluster form (opts 23 / 31 = 16 | 7 / 16 | 15, dispatch.DRAFT_MC_AR_CLUSTER). Deterministic and exact-class to the +// row kernel's FlashInfer order: the updated residual is bitwise, the RMS sum of squares is the cluster's (per-CTA +// warps, then the 8-CTA total), so a normed element can differ by one BF16 rounding step. +void invokeMinimaxM3Sm107Tp4McAllreduceNorm(MinimaxM3Sm107Tp4McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream) +{ + if (params.fc2 == nullptr && params.world == 4 && params.hidden == kRowThreads * 8 && (opts & (32 | 128)) == 0) + { + launchRow(params, maxTokens, (opts & 16) != 0, (opts & 64) != 0, pdl, stream); + return; + } + opts &= ~(64 | 128); + if (maxTokens == 8 && opts == 39) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 47) + { + launch(params, pdl, stream); + } + else if (maxTokens == 16 && opts == 47) + { + launch(params, pdl, stream); + } + else if (maxTokens == 48 && opts == 47) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 7) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 23) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 31) + { + launch(params, pdl, stream); + } + else if (maxTokens == 16 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 48 && opts == 15) + { + launch(params, pdl, stream); + } + else + { + TLLM_THROW("unsupported (bucket %d, opts %d)", maxTokens, opts); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h new file mode 100644 index 000000000000..85938182c098 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h @@ -0,0 +1,83 @@ +/* + * 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 +{ + +//! Operands of the MiniMax-M3 multicast all-reduce + residual + RMSNorm kernel (decode, at most 48 tokens). +//! +//! Per rank the kernel forms the partial `bf16(x + x2)` (or, in finalize mode, +//! `bf16(bf16(sum_k scale * fc2[perm]) + x2)`, the routed sum rounded before the shared expert is added, as stock), +//! pushes it through the multicast mapping `wsMc` into every rank's exchange, polls its own unicast view `wsUc`, +//! sums the ranks in FP32 in rank order, adds the residual (BF16 round) and applies RMSNorm. With `outQuant` it +//! also writes the MXFP8 copy of the normed rows (E4M3 [numTokens, hidden] and swizzled UE8M0 scales). +struct MinimaxM3Sm107Tp4McAllreduceNormParams +{ + __nv_bfloat16 const* x{nullptr}; //!< [numTokens, hidden] partial; in finalize mode the FC2 rows [rows, hidden]. + __nv_bfloat16 const* x2{nullptr}; //!< Optional second addend [numTokens, hidden]. + __nv_bfloat16 const* fc2{nullptr}; //!< Finalize mode: the permuted FC2 output (same buffer as x); null otherwise. + int32_t const* perm{nullptr}; //!< Finalize mode: [numTokens, topk] expanded index -> permuted row, -1 dropped. + void const* scales{nullptr}; //!< Finalize mode: [numTokens, topk] routing weights (fp32 or bf16). + bool scaleFp32{false}; + int32_t topk{0}; + int32_t numTokens{0}; + int32_t hidden{0}; + __nv_bfloat16 const* residual{nullptr}; //!< [numTokens, hidden] + __nv_bfloat16 const* normWeight{nullptr}; //!< [hidden] + float eps{0.F}; + int32_t* wsUc{nullptr}; //!< This rank's unicast view of the exchange (int32 words, re-armed after reading). + int32_t* wsMc{nullptr}; //!< The multicast mapping of the exchange. + int32_t* flags{nullptr}; //!< Per-CTA call-parity counters, [maxTokens * 8]. + int32_t rank{0}; + int32_t world{0}; + __nv_bfloat16* outNormed{nullptr}; + __nv_bfloat16* outResidual{nullptr}; + uint8_t* outQuant{nullptr}; //!< Optional E4M3 [numTokens, hidden]. + uint8_t* outSf{nullptr}; //!< Optional UE8M0 scales, 128x4-swizzled. + //! 0: push x (+ x2) here; 2: a producer already pushed x and launched this grid after its pushes; 3: as 2, with x + //! in exchange set A and x2 in set B. + int32_t prepushed{0}; + //! Optional Eagle3 capture of a capture layer's boundary: capOut[t * capStride + c] = the m3_spec_tail + //! capture_norm of the updated residual row with weight capWeight [hidden] and capEps, bit for bit (hidden 6144, + //! not the drafter boundary). + __nv_bfloat16* capOut{nullptr}; + int64_t capStride{0}; + __nv_bfloat16 const* capWeight{nullptr}; + float capEps{0.F}; +}; + +//! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). +bool isMinimaxM3Sm107Tp4McAllreduceNormSupported(int32_t maxTokens, int32_t opts); + +//! Launch with 8-CTA clusters of 128 threads, one cluster per token. `maxTokens` selects the exchange bucket; +//! `opts` the epilogue variant (every variant is bit-identical). +void invokeMinimaxM3Sm107Tp4McAllreduceNorm(MinimaxM3Sm107Tp4McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu new file mode 100644 index 000000000000..a0d5ac50025a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu @@ -0,0 +1,340 @@ +/* + * 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/minimaxM3Sm107Tp4/producerPw4.h" + +#include +#include +#include + +// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP4 heads. +// +// The SM107 TP4 target's copy of the SM107 TP2 producerPw2.cu (itself a copy of minimaxM3/producerPw.cu, which the +// sm_103 targets keep), for its pre-wait step inputs. Same math and +// outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with --use_fast_math like +// the rest of the library), specialized to 16 query, 1 KV and 1 index head per rank (one index-K head), head_dim 128, +// rotary_dim 64, page 128, NVFP4 main KV and FP8 index-K. One warp per (token, packed head); K and V rows go to their +// head's block of the HND pools (the native kernel's head strides). +// +// The decode producer is latency bound: after griddepcontrol.wait the native kernel runs a chain of dependent global +// loads (packed row -> norm weight -> position id -> RoPE row -> cache slot -> quant scale). This kernel reads the +// constant norm weights and KV quant scale before griddepcontrol.wait and issues the packed row, position id and cache +// slot loads together after it, followed by the RoPE row; with `prewaitStepInputs` the position id, the slot and the +// RoPE row are read before the wait too, leaving only the packed row after it. With `early` it also calls +// griddepcontrol.launch_dependents on entry: consumers read its outputs only after their own griddepcontrol.wait, which +// waits for this grid to complete, so the early trigger only lets the next kernel launch and run its pre-wait prologue +// while this one runs. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace producerPw4 +{ +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotaryDim = 64; +constexpr int kPageSize = 128; +constexpr int kElems = kHeadDim / 32; +constexpr int kBlock = 256; + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ uint32_t sf_group_mask4() +{ + uint32_t laneId; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return 0xFu << (laneId & ~3u); +} + +__device__ __forceinline__ uint32_t fp32x8_to_e2m1(float (&a)[8]) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(a[0]), "f"(a[1]), "f"(a[2]), "f"(a[3]), "f"(a[4]), "f"(a[5]), "f"(a[6]), "f"(a[7])); + return val; +#else + return 0; +#endif +} + +__device__ __forceinline__ void store_fp8x4(__nv_fp8_e4m3* out, int64_t off, float const (&e)[kElems]) +{ + __nv_fp8x2_e4m3 const low(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const high(make_float2(e[2], e[3])); + *reinterpret_cast(out + off) = static_cast(low.__x) | (static_cast(high.__x) << 16); +} + +__device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(e[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = sf_group_mask4(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + *scaleOut = encodedScale.__x; + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float q[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + q[i] = __bfloat162float(rounded[i]) * outputScale; + return static_cast(fp32x8_to_e2m1(q)); +} + +template +__global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (kEarly) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * (kBlock / 32) + warpId; + int const tokenIdx = globalWarp / kTotalHeads; + int const localHead = globalWarp % kTotalHeads; + if (tokenIdx >= p.numTokens) + return; + + constexpr int kBegin = kHeadsQ; + constexpr int vBegin = kBegin + kHeadsKV; + constexpr int indexQBegin = vBegin + kHeadsKV; + constexpr int indexKHead = indexQBegin + kHeadsIndex; + bool const isQ = localHead < kBegin; + bool const isK = localHead >= kBegin && localHead < vBegin; + bool const isV = localHead >= vBegin && localHead < indexQBegin; + bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; + bool const isIndexK = localHead == indexKHead; + + // Constant data read before the wait: norm weights and the KV quant scale. + float w[kElems] = {}; + float quantScale = 0.0F; + __nv_bfloat16 const* normWeight + = isQ ? p.qWeight : (isK ? p.kWeight : (isIndexQ ? p.indexQWeight : p.indexKWeight)); + if (!isV) + { + uint2 const wv = *reinterpret_cast(normWeight + laneId * kElems); + float2 const w01 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.x)); + float2 const w23 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.y)); + w[0] = w01.x; + w[1] = w01.y; + w[2] = w23.x; + w[3] = w23.y; + } + if (isK || isV) + { + quantScale = p.kvQuantScale[isV ? 2 : 1]; + } + + int64_t const inputOffset = (static_cast(tokenIdx) * kTotalHeads + localHead) * kHeadDim + laneId * kElems; + bool const rotates = !isV && laneId * kElems < kRotaryDim; + bool const slotted = isK || isV || isIndexK; + int const coefficient = (laneId * kElems) % (kRotaryDim / 2); + // kPrewait: the position id, the cache slot and the RoPE row are read before the wait. They are per-step metadata + // written before the replay (host-to-device) or by plain-launched kernels earlier in the step (position_ids by the + // step's input preprocessing, the slots by msa_kv_update or its torch fallback), so they are final before this + // grid can start whatever its predecessors trigger; the RoPE table is constant. + int positionId = 0; + int slot = 0; + float4 c = make_float4(0.0F, 0.0F, 0.0F, 0.0F); + float4 s = c; + auto const loadStepInputs = [&]() + { + positionId = rotates ? p.positionIds[tokenIdx] : 0; + slot = slotted ? p.outCacheLoc[tokenIdx] : 0; + if (rotates) + { + int64_t const ropeRow = static_cast(positionId) * kRotaryDim; + c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); + s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); + } + }; + if constexpr (kPrewait) + { + loadStepInputs(); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + // The packed row comes from the predecessor grid; without kPrewait the step inputs are issued with it, so the + // native kernel's serial chain of dependent loads becomes two round trips. + uint2 const packed = *reinterpret_cast(p.packed + inputOffset); + if constexpr (!kPrewait) + { + loadStepInputs(); + } + float cosine[kElems] = {c.x, c.y, c.z, c.w}; + float sine[kElems] = {s.x, s.y, s.z, s.w}; + + float e[kElems]; + float sumSquares = 0.0F; +#pragma unroll + for (int pair = 0; pair < 2; ++pair) + { + float2 const v = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); + e[2 * pair] = v.x; + e[2 * pair + 1] = v.y; + if (!isV) + sumSquares += v.x * v.x + v.y * v.y; + } + + if (!isV) + { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + sumSquares += __shfl_xor_sync(0xffffffff, sumSquares, mask, 32); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kHeadDim) + p.eps); +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + constexpr int kPairOffset = (kRotaryDim / 2) / kElems; +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + int const dim = laneId * kElems + i; + float paired = __shfl_xor_sync(0xffffffff, e[i], kPairOffset); + if (dim < kRotaryDim) + { + if (dim < kRotaryDim / 2) + paired = -paired; + // The native kernel's contraction of e * cos + paired * sin; spelled out so this build matches it. + e[i] = __fmaf_rn(paired, sine[i], __fmul_rn(e[i], cosine[i])); + } + } + } + + if (isQ) + { + store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); + return; + } + if (isIndexQ) + { + store_fp8x4(p.indexQOut, + (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); + return; + } + + if (slot < 0) + return; + int const page = slot / kPageSize; + if (page >= p.numPages) + return; + int const withinPage = slot & (kPageSize - 1); + if (isIndexK) + { + store_fp8x4(p.indexK, + static_cast(page) * p.indexPageStride + static_cast(withinPage) * p.indexTokenStride + + laneId * kElems, + e); + return; + } + + int const head = isK ? localHead - kBegin : localHead - vBegin; + int const plane = isV ? 1 : 0; + int const scaleColumn = laneId / 4; + int64_t const scaleWithinPage = isV + ? static_cast(withinPage / 4) * 32 + static_cast(scaleColumn) * 4 + withinPage % 4 + : static_cast(withinPage) * 8 + scaleColumn; + int64_t const scaleOffset = static_cast(page) * p.scalePageStride + + static_cast(plane) * p.scalePlaneStride + static_cast(head) * p.scaleHeadStride + + scaleWithinPage; + uint8_t* scaleOut = laneId % 4 == 0 ? p.kvScale + scaleOffset : nullptr; + uint16_t const packedFp4 = quantize_fp4x4(e, quantScale, scaleOut); + int64_t const dataOffset = static_cast(page) * p.dataPageStride + + static_cast(plane) * p.dataPlaneStride + static_cast(head) * p.dataHeadStride + + static_cast(withinPage) * p.dataTokenStride + laneId * 2; + *reinterpret_cast(p.kvData + dataOffset) = packedFp4; +#endif +} + +template +cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) +{ + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; + int const warps = numTokens * kTotalHeads; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); + cfg.blockDim = dim3(kBlock); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); +} + +} // namespace + +cudaError_t invokeProducerPw4( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream) +{ + if (numTokens <= 0) + return cudaSuccess; + if (prewaitStepInputs) + { + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); + } + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); +} + +} // namespace producerPw4 +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h new file mode 100644 index 000000000000..6f8aa54a812b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h @@ -0,0 +1,76 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP4 heads (16 query, 1 KV and +//! 1 index head per rank) with its constant loads before the PDL grid-dependency wait and an optional early dependents +//! trigger. Same outputs, bit for bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw4.cu. +namespace producerPw4 +{ + +struct Params +{ + __nv_bfloat16 const* packed; + __nv_fp8_e4m3* qOut; + __nv_fp8_e4m3* indexQOut; + uint8_t* kvData; + uint8_t* kvScale; + __nv_fp8_e4m3* indexK; + int const* outCacheLoc; + float const* kvQuantScale; + int64_t dataPageStride, dataPlaneStride, dataHeadStride, dataTokenStride; + int64_t scalePageStride, scalePlaneStride, scaleHeadStride; + int64_t indexPageStride, indexTokenStride, numPages; + int numTokens; + float eps; + __nv_bfloat16 const* qWeight; + __nv_bfloat16 const* kWeight; + __nv_bfloat16 const* indexQWeight; + __nv_bfloat16 const* indexKWeight; + float const* rotaryCosSin; + int const* positionIds; +}; + +//! The heads per rank this producer is built for. +constexpr int kHeadsQ = 16; +constexpr int kHeadsKV = 1; +constexpr int kHeadsIndex = 1; + +//! Launch over `numTokens` tokens, one warp per (token, packed head). `prewaitStepInputs` reads the position ids, the +//! cache slots and the RoPE rows before the grid-dependency wait (they must not be written by a PDL predecessor). +cudaError_t invokeProducerPw4( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream); + +} // namespace producerPw4 +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu new file mode 100644 index 000000000000..3d041421a92b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu @@ -0,0 +1,944 @@ +/* + * 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/minimaxM3Sm107Tp4/sampledAccept.h" + +#include "tensorrt_llm/common/assert.h" + +#include +#include + +// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). +// +// sm_107 (Rubin) TP4 copy of kernels/minimaxM3Sm107Tp2/sampledAccept.cu: identical algorithm and arithmetic as the +// shared kernels/minimaxM3/sampledAccept.cu, with the TP2 copy's launch (a row's CTAs as one thread-block cluster) and +// bounded record polls. The vocab shard limit stays the shared kernel's (kMaxVs 51200 >= the TP4 shard 50016). +// +// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +// active must stay on the generic path; the caller decides. +// +// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens +// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the +// nucleus with probability >= top_p, so most rows decide on their second exchange. +// +// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based +// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one +// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and +// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the +// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing +// and two call halves keep a rank one call ahead from overwriting a slot still being read. +// +// SA-1R (Args.one_round_k 9 / 82; MoE prototype moe/sa1r/6c753632ec, exact on 39,968 rows, eager + graph): each CTA +// also multicasts its largest bf16 logits with the round-0 record (3 per extra word). After round 0 every CTA knows M, +// the candidate logit cv and wc = __expf(cv - M); round 1 would accept iff S <= target, S = the fp32 sum (its tree) of +// w = __expf(__fmul_rn(x, tInv) - M) over every token with w > wc. x -> w is monotone (__fmul_rn by tInv > 0, fsub, +// __expf: checked exhaustively over all non-positive fp32 inputs), so if no CTA's unlisted value can beat wc (82: the +// bound tau; 9: the 9th listed), the listed w > wc are exactly round 1's summands, bitwise. Two fp32 sums of the same +// m non-negative terms differ by < m * 2^-23 * S, so accepting iff S_1r + m * 2^-21 * S_1r (round up) <= target +// implies round 1 accepts; m <= 1 is exact. Otherwise round 1 runs as before. An accepted row outputs the round-0 +// candidate, as the two-round kernel does, and consumes no Philox draw it would not. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ +namespace sampledAccept +{ +namespace +{ +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kMaxRec = kMaxWorld * kMaxSplits; +constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) +constexpr int kTagBits = 64 - 32 - kIdBits; +constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard +// Record polls trap after this many spins (the multicast all-reduce's bound): a missing writer becomes a CUDA error, +// never a hang. +constexpr long long kMaxSpins = 1LL << 31; + +// SA-1R variant kTopK: 0 none; K exact top-K list; 10 * K + P pooled (each thread keeps its top P plus its (P + 1)-th +// as a bound; the CTA lists the pool's top K plus tau = max(K-th listed, max thread bound)). +__host__ __device__ constexpr int mode_k(int m) +{ + return m >= 10 ? m / 10 : m; +} + +__host__ __device__ constexpr int mode_pool(int m) +{ + return m >= 10 ? m % 10 : 0; +} + +__host__ __device__ constexpr int mode_vals(int m) +{ + return mode_k(m) + (mode_pool(m) ? 1 : 0); +} + +// Words per record: key, logit of key, f2, f3, then the round-0 list (3 bf16 + a 16-bit tag per word). +template +struct Rec +{ + static constexpr int kXw = (mode_vals(kTopK) + 2) / 3; + static constexpr int kEx = 4 + kXw; +}; + +// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. +template +__host__ __device__ inline long long half_words(int world) +{ + return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * Rec::kEx; +} + +__device__ __forceinline__ uint32_t ordered(float v) +{ + uint32_t const b = __float_as_uint(v); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +// bf16 bits <-> an order-preserving 16-bit key. +__device__ __forceinline__ uint32_t ord16(uint32_t b) +{ + return (b & 0x8000u) ? (~b & 0xFFFFu) : (b | 0x8000u); +} + +__device__ __forceinline__ uint32_t unord16(uint32_t k) +{ + return (k & 0x8000u) ? (k & 0x7FFFu) : (~k & 0xFFFFu); +} + +// Philox4x32-10. +__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) +{ +#pragma unroll + for (int i = 0; i < 10; ++i) + { + uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; + uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; + c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); + k.x += 0x9E3779B9u; + k.y += 0xBB67AE85u; + } + return c; +} + +// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. +__device__ __forceinline__ float uniform(uint32_t bits) +{ + return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); +} + +// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. +__device__ __forceinline__ unsigned long long gkey(float lg, int id) +{ + return (static_cast(ordered(lg)) << kIdBits) + | static_cast((1u << kIdBits) - 1u - static_cast(id)); +} + +__device__ __forceinline__ int key_token(unsigned long long k) +{ + return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); +} + +__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) +{ + asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); +} + +__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) +{ + unsigned long long v; + asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); + return v; +} + +template +struct Shared +{ + __nv_bfloat16 row[kMaxVs / kSplits]; + float redf[kWarps]; + unsigned long long redk[kWarps]; + unsigned long long rk[kMaxRec]; + float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; + float wsum[kWarps]; + int owner; + int pick; + float pickx; + bool last; + float dS[kWarps]; // SA-1R: per-warp partial mass above the candidate + int dM[kWarps]; // SA-1R: per-warp count | incomplete << 16 + uint32_t wtop[kWarps][kTopK > 0 ? mode_k(kTopK) : 1]; // SA-1R: per-warp top-K keys + uint32_t wbound[kWarps]; // SA-1R pooled: per-warp max thread bound + uint32_t ctop[kTopK > 0 ? mode_vals(kTopK) : 1]; // SA-1R: the CTA's top-K keys (descending), by rank + unsigned long long rx[kMaxRec][Rec::kXw > 0 ? Rec::kXw : 1]; +}; + +template +__device__ __forceinline__ float block_max(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = s.redf[0]; + for (int i = 1; i < kWarps; ++i) + r = fmaxf(r, s.redf[i]); + __syncthreads(); + return r; +} + +// Fixed-order block sum: identical result for identical inputs. +template +__device__ __forceinline__ float block_sum(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = 0.f; + for (int i = 0; i < kWarps; ++i) + r += s.redf[i]; + __syncthreads(); + return r; +} + +template +__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + if ((threadIdx.x & 31) == 0) + s.redk[threadIdx.x >> 5] = v; + __syncthreads(); + unsigned long long r = s.redk[0]; + for (int i = 1; i < kWarps; ++i) + r = s.redk[i] > r ? s.redk[i] : r; + __syncthreads(); + return r; +} + +// Exclusive prefix of v over the threads in thread order, and the total. +template +__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + float inc = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + float const t = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) + inc += t; + } + if (lane == 31) + s.wsum[warp] = inc; + __syncthreads(); + float base = 0.f; + total = 0.f; + for (int i = 0; i < kWarps; ++i) + { + base += i < warp ? s.wsum[i] : 0.f; + total += s.wsum[i]; + } + __syncthreads(); + return base + inc - v; +} + +// Warp-wide integer max / sum / or (the redux instructions from sm_80; a butterfly otherwise, same results). +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_max_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + uint32_t const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_add_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_add_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_or_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_or_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v |= __shfl_xor_sync(0xffffffffu, v, o); + return v; +#endif +} + +// Insert key into a descending register list of K (fully unrolled: no local memory). +template +__device__ __forceinline__ void topk_insert(uint32_t (&a)[K], uint32_t key) +{ +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const hi = key > a[j] ? key : a[j]; + key = key > a[j] ? a[j] : key; + a[j] = hi; + } +} + +template +__device__ __forceinline__ void topk_pop(uint32_t (&a)[K]) +{ +#pragma unroll + for (int j = 0; j + 1 < K; ++j) + a[j] = a[j + 1]; + a[K - 1] = 0u; +} + +// Per-thread insertion with an early reject: after the first few elements almost every key fails the K-th test. +template +__device__ __forceinline__ void thread_insert(uint32_t (&a)[K], uint32_t key) +{ + if (key > a[K - 1]) + topk_insert(a, key); +} + +// The warp's K largest keys (descending) into s.wtop[warp]; published by the caller's next block barrier. +template +__device__ __forceinline__ void warp_topk(uint32_t (&a)[K], S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const m = warp_max_u32(a[0]); + if (a[0] == m && m != 0u) + topk_pop(a); + if (lane == 0) + s.wtop[warp][j] = m; + } +} + +// Pooled: K rounds over each lane's P pool keys (descending) and the warp's max thread bound. +template +__device__ __forceinline__ void warp_pool(uint32_t (&a)[P + 1], S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + uint32_t const bound = warp_max_u32(a[P]); + uint32_t pl[P]; +#pragma unroll + for (int j = 0; j < P; ++j) + pl[j] = a[j]; +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const m = warp_max_u32(pl[0]); + if (pl[0] == m && m != 0u) + topk_pop

(pl); + if (lane == 0) + s.wtop[warp][j] = m; + } + if (lane == 0) + s.wbound[warp] = bound; +} + +// A row's kSplits CTAs wait for each other's records, so they launch as one thread-block cluster: the hardware +// schedules a cluster's CTAs together, whatever other streams hold. Rows never wait on other rows. +template +__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) +{ + using R = Rec; + extern __shared__ __align__(16) unsigned char smem[]; + Shared& s = *reinterpret_cast*>(smem); + float constexpr kNegInf = -__builtin_huge_valf(); + int const tid = threadIdx.x; + int const row = blockIdx.x / kSplits; + int const split = blockIdx.x % kSplits; + int const world = a.world; + int const nrec = world * kSplits; + // The row's sampling parameters are step inputs (the spec metadata, written before the target forward), not the + // predecessor head GEMM's output, so with prewait_params they load before the dependency wait and overlap the + // head's tail; the same values either way. + float temperature = 0.f, topP = 0.f; + long long seed_in = 0, offset_in = 0; + auto load_params = [&]() + { + temperature = a.temperatures[row]; + topP = a.top_ps[row]; + seed_in = a.seeds[row]; + offset_in = a.offsets[row]; + }; + if (a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (!a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // early_trigger: the dependent grid (the first draft pass's input) may launch now and park at its own + // griddepcontrol.wait, which still returns only once this grid has completed and flushed; it skips the launch + // latency after this kernel ends. Otherwise dependents are released at the acceptance epilogue (the CTAs of + // splits 1.. exit without releasing). + if (a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const tempInv = 1.0f / (temperature + kTempEpsilon); + unsigned long long const seed = static_cast(seed_in); + unsigned long long const offset = static_cast(offset_in); + unsigned int const call = *reinterpret_cast(a.counters + 1); + // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of + // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round + // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten + // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. + auto ktag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; + }; + auto ftag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; + }; + // SA-1R round-0 list words: a 16-bit tag, never 0. + unsigned long long const xtag = (((call >> 1) << 1) & 0x7FFFu) | 0x8000u; + long long const rec_stride = static_cast(world) * kMaxSplits * R::kEx; + unsigned long long* const xu = a.ws_uc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride; + unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * R::kEx; + long long const round_stride = static_cast(kMaxRows) * rec_stride; + + // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- + __nv_bfloat16 const* src = a.logits + row * a.ld; + int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; + int const s0 = min(split * sl_len, a.Vs); + int const n = min(s0 + sl_len, a.Vs) - s0; + int const id0 = a.vocab_offset + s0; + float tmax = kNegInf; + // SA-1R: each thread's largest (ord16(bf16) << 16 | ~index) keys while staging (only rows with a nucleus cut). + bool const want_x = kTopK > 0 && topP < 1.0f; + constexpr int kK = mode_k(kTopK), kPool = mode_pool(kTopK), kVals = mode_vals(kTopK); + constexpr int kTk = kPool ? kPool + 1 : (kK > 0 ? kK : 1); + uint32_t tk[kTk]; +#pragma unroll + for (int j = 0; j < kTk; ++j) + tk[j] = 0u; + bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; + int const nvec = vec ? n / 8 : 0; + for (int v = tid; v < nvec; v += kThreads) + { + uint4 const w = *reinterpret_cast(src + s0 + 8 * v); + *reinterpret_cast(&s.row[8 * v]) = w; + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); +#pragma unroll + for (int j = 0; j < 8; ++j) + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); + if constexpr (kTopK > 0) + { + if (want_x) + { +#pragma unroll + for (int j = 0; j < 8; ++j) + { + uint32_t const key + = (ord16(__bfloat16_as_ushort(e[j])) << 16) | (0xFFFFu - static_cast(8 * v + j)); + if constexpr (kPool > 0) + topk_insert(tk, key); // branchless: P + 1 compare-swaps + else + thread_insert(tk, key); + } + } + } + } + for (int i = nvec * 8 + tid; i < n; i += kThreads) + { + s.row[i] = src[s0 + i]; + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); + if constexpr (kTopK > 0) + { + if (want_x) + topk_insert( + tk, (ord16(__bfloat16_as_ushort(src[s0 + i])) << 16) | (0xFFFFu - static_cast(i))); + } + } + if constexpr (kTopK > 0) + { + if (tid < kVals) + s.ctop[tid] = 0u; // published with s.wtop by block_max's barrier + if (want_x) + { + if constexpr (kPool > 0) + warp_pool(tk, s); // s.wtop / s.wbound published by block_max's barrier + else + warp_topk(tk, s); // s.wtop published by block_max's barrier + } + } + float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier + if constexpr (kTopK > 0) + { + // SA-1R: the CTA's top K of the warp lists by rank (keys are unique): a lane quad per key, no serial merge; + // s.ctop is read after the round-0 draw's first barrier + constexpr int kN = kWarps * kK; + constexpr int kSp = kN * 4 <= kThreads ? 4 : 1; // threads per key + if (want_x && tid < kN * kSp) + { + uint32_t const* const wl = &s.wtop[0][0]; + int const part = tid % kSp; + uint32_t const key = wl[tid / kSp]; + int rank = 0; +#pragma unroll + for (int j = part; j < kN; j += kSp) + rank += wl[j] > key ? 1 : 0; + if constexpr (kSp == 4) + { + rank += __shfl_xor_sync(0xffffffffu, rank, 1); + rank += __shfl_xor_sync(0xffffffffu, rank, 2); + } + if (part == 0 && key != 0u && rank < kK) + s.ctop[rank] = key; + } + } + float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z + unsigned long long ckey = 0ULL; + int const gid = a.rank * kMaxSplits + split; + for (int r = 0;; ++r) + { + // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) + float const ref = r == 0 ? smax : M; + float const wmin = r == 0 ? 0.f : wc; + float ts = 0.f; + for (int i = tid; i < n; i += kThreads) + { + float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); + ts += w > wmin ? w : 0.f; + } + float total; + float const excl = block_excl_scan(ts, total, s); + if constexpr (kTopK > 0) + { + // SA-1R round 0: the list words depend only on the staged slice, so they go out now (s.ctop / s.wbound are + // published by the barrier above), one lane per word, ahead of the main record: off the readers' path + if (r == 0 && want_x && tid < R::kXw) + { + uint32_t tau = 0u; + if constexpr (kPool > 0) + { + tau = s.ctop[kK - 1]; // the largest unlisted pool key is <= the K-th listed +#pragma unroll + for (int wq = 0; wq < kWarps; ++wq) + tau = s.wbound[wq] > tau ? s.wbound[wq] : tau; + } + unsigned long long w = xtag; +#pragma unroll + for (int t = 0; t < 3; ++t) + { + int const e = 3 * tid + t; + uint32_t const kv = e < kK ? s.ctop[e] : (kPool > 0 && e == kK ? tau : 0u); + // missing entries: -inf (bf16 0xFF80), w = 0 + uint32_t const b = kv != 0u ? unord16(kv >> 16) : 0xFF80u; + w |= static_cast(b) << (16 * (3 - t)); + } + st_sys(xm + 4 + tid, w); // round 0 = slot 0 + } + } + uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), + static_cast(offset), static_cast(offset >> 32)), + make_uint2(static_cast(seed), static_cast(seed >> 32))); + float const tgt = uniform(rnd.x) * total; + if (tid == 0) + s.owner = -1; + __syncthreads(); + if (ts > 0.f && excl <= tgt) + atomicMax(&s.owner, tid); + __syncthreads(); + if (tid == s.owner) + { + float acc = excl; + int pick = -1; + float px = 0.f; + if constexpr (kIlpWalk) + { + // SA-PW: the owner's strided tokens 8 at a time. The 8 x and w (the same expressions) are + // independent, so their MUFU work overlaps; the serial sum then adds them in the same order with + // the same early exit: the same acc sequence, pick and px as the one-token walk, bit for bit. + constexpr int kW = 8; + bool done = false; + for (int i0 = tid; i0 < n && !done; i0 += kW * kThreads) + { + float xs[kW], wv[kW]; +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + xs[k] = i < n ? __fmul_rn(__bfloat162float(s.row[i]), tempInv) : 0.f; + wv[k] = __expf(xs[k] - ref); + } +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + if (!done && i < n && wv[k] > wmin) + { + acc += wv[k]; + pick = i; + px = xs[k]; + if (acc > tgt) + done = true; + } + } + } + } + else + { + for (int i = tid; i < n; i += kThreads) + { + float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); + float const w = __expf(x - ref); + if (w > wmin) + { + acc += w; + pick = i; + px = x; + if (acc > tgt) + break; + } + } + } + s.pick = pick; + s.pickx = px; + } + __syncthreads(); + unsigned long long sb = 0ULL; + if (total > 0.f && s.owner >= 0) + sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); + float const smab = total; + unsigned long long const ktag = ktag_of(r); + unsigned long long const ftag = ftag_of(r); + unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; + bool const with_x = want_x && r == 0; + if (tid == 0) + { + unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; + st_sys(xw + 0, (sb << kTagBits) | ktag); + st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); + st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); + st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); + } + // ---- every record of the row, every rank --------------------------------------------------------------- + int const nw = with_x ? R::kEx : 4; + for (int t = tid; t < nrec * nw; t += kThreads) + { + int const q = t / nw, f = t % nw; + unsigned long long const* const p = xr + (q / kSplits * kMaxSplits + q % kSplits) * R::kEx + f; + unsigned long long v; + long long spins = 0; + if (f == 0) + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); + s.rk[q] = v >> kTagBits; + } + else if (f < 4) + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFFFFF00000000ULL) != ftag); + float const fv = __uint_as_float(static_cast(v)); + (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; + } + else + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFULL) != xtag); + s.rx[q][f - 4] = v; + } + } + __syncthreads(); + // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- + unsigned long long gk = 0ULL; + float gv = 0.f; + for (int q = 0; q < nrec; ++q) + { + if (s.rk[q] > gk) + { + gk = s.rk[q]; + gv = s.rv[q]; + } + } + bool accept; + if (r == 0) + { + M = kNegInf; + for (int q = 0; q < nrec; ++q) + M = fmaxf(M, s.r2[q]); + float Z = 0.f; + for (int q = 0; q < nrec; ++q) + Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); + target = topP * Z; + accept = !(topP < 1.0f); // no nucleus cut: the first draw stands + ckey = gk; + cv = gv; + if constexpr (kTopK > 0) + { + // SA-1R: every thread takes (record, entry) items; the per-warp partials are published by the round's + // barrier below and every thread combines them in warp order after it (no new barrier) + if (with_x) + { + float const wcand = __expf(cv - M); // = round 1's wmin + float part = 0.f; + int cnt = 0, inc = 0; + for (int t = tid; t < nrec * kVals; t += kThreads) + { + int const q = t / kVals, e = t % kVals; + uint32_t const b = static_cast(s.rx[q][e / 3] >> (16 * (3 - e % 3))) & 0xFFFFu; + float const w = __expf( + __fmul_rn(__bfloat162float(__ushort_as_bfloat16(static_cast(b))), tempInv) + - M); // w bitwise as round 1 computes it + if (w > wcand) + { + if (kPool > 0 && e == kK) + inc = 1; // tau (the largest unlisted value) is heavier than the candidate + else + { + part += w; + ++cnt; + if (kPool == 0 && e == kK - 1) + inc = 1; // the K-th is heavier: unlisted heavier tokens may exist + } + } + } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + part += __shfl_xor_sync(0xffffffffu, part, o); // any order: the margin covers it + cnt = static_cast(warp_add_u32(static_cast(cnt))); + inc = static_cast(warp_or_u32(static_cast(inc))); + if ((tid & 31) == 0) + { + s.dS[tid >> 5] = part; + s.dM[tid >> 5] = cnt | (inc << 16); + } + } + } + } + else + { + float S = 0.f; // mass strictly above the current candidate + for (int q = 0; q < nrec; ++q) + S += s.r2[q]; + accept = S <= target || gk == 0ULL; + if (!accept) + { + ckey = gk; + cv = gv; + } + } + wc = __expf(cv - M); + __syncthreads(); // records consumed before the next round's writes to s + bool decided = false; + if constexpr (kTopK > 0) + { + if (r == 0 && with_x) + { + float S1 = 0.f; + int m = 0, inc = 0; + for (int wq = 0; wq < kWarps; ++wq) + { + S1 += s.dS[wq]; + m += s.dM[wq] & 0xFFFF; + inc |= s.dM[wq] >> 16; + } + if (!inc) + { + float const margin = m <= 1 ? 0.f : S1 * static_cast(m) * (1.0f / 2097152.0f); + decided = __fadd_ru(S1, margin) <= target; + accept = accept || decided; + } + } + } + if (accept) + { + if (split != 0) + return; + if (tid == 0) + { + a.tokens[row] = key_token(ckey); + if (a.extra_attempts != nullptr && r > 1) + atomicAdd(a.extra_attempts, r - 1); + if (a.decided != nullptr && decided) + atomicAdd(a.decided, 1); + } + break; + } + } + + // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (!a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (tid == 0) + s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; + __syncthreads(); + if (!s.last) + return; + __threadfence(); + if (tid == 0) + { + a.counters[0] = 0u; + a.counters[1] = call + 1u; + } + if (tid < a.nreq) + { + int const* tok = a.tokens; + int const b = tid; + int const k1 = a.draft_len + 1; + int nacc = 1; + bool run = true; + for (int k = 0; k < a.draft_len; ++k) + { + int const tt = *reinterpret_cast(tok + b * k1 + k); + a.acc_tokens[b * k1 + k] = tt; + run = run && a.draft[b * a.draft_len + k] == tt; + nacc += run ? 1 : 0; + } + a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); + if (a.force_vec != nullptr) + nacc = a.force_vec[b]; + else if (a.force_total > 0) + nacc = a.force_total; + a.num_acc[b] = nacc; + a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; + if (a.next_new != nullptr) + a.next_new[b * a.next_stride] + = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); + } +} + +template +cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) +{ + static int max_clusters = 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(a.rows * kSplits); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Shared); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + if (max_clusters == 0) + { + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } + } + TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); + return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); +} + +template +cudaError_t launchRows(Args const& a, bool use_pdl, cudaStream_t stream) +{ + // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); + // more rows: 4 CTAs per row. + if (a.ilp_walk) + { + return a.rows <= 32 ? launch<8, kTopK, true>(a, use_pdl, stream) : launch<4, kTopK, true>(a, use_pdl, stream); + } + return a.rows <= 32 ? launch<8, kTopK, false>(a, use_pdl, stream) : launch<4, kTopK, false>(a, use_pdl, stream); +} + +} // namespace + +long long sampledAcceptWsWords(int world, int oneRoundK) +{ + TLLM_CHECK_WITH_INFO( + oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "SA-1R variant %d is not built", oneRoundK); + return 2 * (oneRoundK == 9 ? half_words<9>(world) : oneRoundK == 82 ? half_words<82>(world) : half_words<0>(world)); +} + +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) +{ + switch (a.one_round_k) + { + case 0: return launchRows<0>(a, usePdl, stream); + case 9: return launchRows<9>(a, usePdl, stream); + case 82: return launchRows<82>(a, usePdl, stream); + default: TLLM_CHECK_WITH_INFO(false, "SA-1R variant %d is not built", a.one_round_k); + } + return cudaErrorInvalidValue; +} + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h new file mode 100644 index 000000000000..0ebd153992d0 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h @@ -0,0 +1,87 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ +//! Fused sampled (temperature + top-p) verification of the MiniMax-M3 Eagle3 target head over the TP-sharded +//! vocabulary, decode rows <= kMaxRows; see sampledAccept.cu. +namespace sampledAccept +{ + +constexpr int kMaxSplits = 8; // CTAs per verify row (at most) +constexpr int kMaxRows = 48; // verify rows per call +constexpr int kMaxWorld = 8; // TP ranks +constexpr int kMaxVs = 51200; // vocab shard width (TP4: 200064 / 4 = 50016) +constexpr int kIdBits = 18; // token ids < 2^18 + +struct Args +{ + __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard + long long ld; + int rows, Vs, vocab_offset; + float const* temperatures; + float const* top_ps; + long long const* seeds; + long long const* offsets; + unsigned long long* ws_uc; + unsigned long long* ws_mc; + unsigned int* counters; // [0] finished rows, [1] call + int* tokens; // [rows] sampled token per row + int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows + int rank, world; + int const* draft; + int const* batch_idx; + int* acc_tokens; + int* num_acc; + long long* gather_ids; + int draft_len, nreq, force_total; + int* next_new; + int next_stride; + int const* force_vec; + int prewait_params; // nonzero: the row's sampling parameters load before the PDL dependency wait + int early_trigger; // nonzero: dependents are released right after the dependency wait, not at the epilogue + // SA-1R (one-round decision): 0 the two-round kernel; 9 each CTA's top-9 logits ride its round-0 record; 82 pooled + // (each thread's top 2, the CTA lists the pool's top 8 plus the bound of everything unlisted). A row whose first + // candidate is provably in the nucleus then decides in round 0, with the same token as the two-round kernel. + int one_round_k; + int* decided; // optional [1]: rows decided in round 0 (SA-1R), summed + int ilp_walk; // nonzero: the owner walk takes its tokens 8 at a time (same weights, same summation order) +}; + +//! Words (uint64) of the multicast exchange workspace for `world` ranks and SA-1R variant `oneRoundK`: two call halves. +long long sampledAcceptWsWords(int world, int oneRoundK = 0); + +//! Launch the sampled verification over a.rows * (8 or 4) co-resident CTAs (8 per row up to 32 rows). +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream); + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.cu new file mode 100644 index 000000000000..c93f5089bf2b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.cu @@ -0,0 +1,312 @@ +/* + * 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. + */ + +// Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. +// +// One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +// thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups `block % 16` is a lower +// bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared +// memory and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every +// candidate against the whole row (re-read). Rows of at most 64 candidates (the common case) skip the all-thread rank: +// warp 0 sorts them in registers (a 64-key bitonic network), keeps the 16 smallest (key, block) pairs and emits them, +// and the other warps exit. The row's valid count is a step input and is read before the grid dependency wait. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +namespace +{ + +constexpr float kInitScore = 1.0e30f; +constexpr float kLocalScore = 1.0e29f; +constexpr uint32_t kNegInfKey = 0xff800000u; +constexpr int kItems = 9; // keys per thread: rows up to 9216 blocks (1M tokens + draft tokens) +constexpr int kWarpCandidates = 64; // candidates warp 0 ranks alone (2 per lane) + +// Descending order key: a smaller key is a larger score (NaN above +inf), as in the reference op. +__device__ __forceinline__ uint32_t orderedKey(float x) +{ + uint32_t const b = __float_as_uint(x); + return (b & 0x80000000u) ? b : (~b & 0x7fffffffu); +} + +__device__ __forceinline__ float applyForcing(float raw, int block, int initBlocks, long long localStart) +{ + if (block >= localStart) + return kLocalScore; + if (block < initBlocks) + return kInitScore; + return raw; +} + +struct SelectSmem +{ + uint32_t groupMin[16]; + unsigned long long cand[1024]; + int numCand; + unsigned long long picked[16]; +}; + +// One row with up to kItems * 1024 valid blocks; the caller picks kItems from the row's valid count. +// Returns -1 when sm.picked holds the row's picks, else the candidate count (at most kWarpCandidates, compacted to +// sm.cand and not yet ranked). The value is uniform across the CTA. +template +__device__ __forceinline__ int selectRow(SelectSmem& sm, float const* rowBase, long long blockStride, int valid, + int initBlocks, long long localStart, bool loadedTrigger) +{ + int const tid = threadIdx.x, lane = tid & 31; + uint32_t key[kItems]; + uint32_t localMin = 0xffffffffu; +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + key[s] = block < valid ? orderedKey(applyForcing(rowBase[block * blockStride], block, initBlocks, localStart)) + : 0xffffffffu; + localMin = min(localMin, key[s]); + } + if (valid <= 16) + { + if (tid < valid) + sm.picked[tid] = (static_cast(key[0]) << 32) | static_cast(tid); + if (loadedTrigger) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + return -1; + } + // Group g = block % 16 = tid % 16: lanes l and l ^ 16 of every warp. + localMin = min(localMin, __shfl_xor_sync(0xffffffffu, localMin, 16)); + if (lane < 16) + atomicMin(&sm.groupMin[lane], localMin); + __syncthreads(); + // Every score of the row is loaded: release the sparse decode, so its launch and its pre-wait prologue overlap + // the rank below. This comes after this grid's own dependency wait, which the sparse decode relies on: it reads + // q and the seq_lens / block-table row before its own wait (sparseDecodeLc2.cu), safe only while every kernel + // between the producer and it releases its dependents after its own wait. + if (loadedTrigger) + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + uint32_t bound = 0; +#pragma unroll + for (int g = 0; g < 16; ++g) + bound = max(bound, sm.groupMin[g]); +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + bool const isCand = block < valid && key[s] <= bound; + unsigned const ballot = __ballot_sync(0xffffffffu, isCand); + int base = 0; + if (lane == 0 && ballot) + base = atomicAdd(&sm.numCand, __popc(ballot)); + base = __shfl_sync(0xffffffffu, base, 0); + if (isCand) + { + int const i = base + __popc(ballot & ((1u << lane) - 1)); + if (i < 1024) + sm.cand[i] = (static_cast(key[s]) << 32) | static_cast(block); + } + } + __syncthreads(); + int const n = sm.numCand; + if (kFast && n <= kWarpCandidates) + return n; + if (n <= 1024) + { + for (int i = tid; i < n; i += 1024) + { + unsigned long long const mine = sm.cand[i]; + int rank = 0; + for (int j = 0; j < n; ++j) + rank += sm.cand[j] < mine; + if (rank < 16) + sm.picked[rank] = mine; + } + return -1; + } + // More than 1024 keys tied at the bound (degenerate rows): rank every candidate against the whole row. +#pragma unroll + for (int s = 0; s < kItems; ++s) + { + int const block = s * 1024 + tid; + if (block < valid && key[s] <= bound) + { + int rank = 0; + for (int j = 0; j < valid; ++j) + { + uint32_t const other = orderedKey(applyForcing(rowBase[j * blockStride], j, initBlocks, localStart)); + rank += other < key[s] || (other == key[s] && j < block); + } + if (rank < 16) + sm.picked[rank] = (static_cast(key[s]) << 32) | static_cast(block); + } + } + return -1; +} + +// Warp 0: the n <= kWarpCandidates candidates of sm.cand sorted ascending as 64 keys (element r * 32 + lane in +// register r); returns this lane's element (lanes 0..15 then hold the 16 smallest). +__device__ __forceinline__ unsigned long long warpSmallest(SelectSmem const& sm, int n, int lane) +{ + unsigned long long a = lane < n ? sm.cand[lane] : ~0ull; + unsigned long long b = lane + 32 < n ? sm.cand[lane + 32] : ~0ull; +#pragma unroll + for (int size = 2; size <= 64; size <<= 1) + { +#pragma unroll + for (int stride = size >> 1; stride > 0; stride >>= 1) + { + if (stride == 32) + { + unsigned long long const lo = min(a, b), hi = max(a, b); + a = lo; + b = hi; + } + else + { + unsigned long long const oa = __shfl_xor_sync(0xffffffffu, a, stride); + unsigned long long const ob = __shfl_xor_sync(0xffffffffu, b, stride); + bool const lower = (lane & stride) == 0; + bool const ascA = (lane & size) == 0, ascB = ((lane + 32) & size) == 0; + a = (lower == ascA) ? min(a, oa) : max(a, oa); + b = (lower == ascB) ? min(b, ob) : max(b, ob); + } + } + } + return a; +} + +// Warp 0: the picks (lane < 16 holds pick `lane`, ~0ull when absent) as ascending block ids, -1 for forced -inf picks +// and missing picks, padded at the tail. +__device__ __forceinline__ void emitPicks(unsigned long long c, int lane, int* outRow) +{ + uint32_t blockKey = ~0u; + if (lane < 16 && c != ~0ull && static_cast(c >> 32) != kNegInfKey) + blockKey = static_cast(c & 0xffffffffu); +#pragma unroll + for (int size = 2; size <= 32; size <<= 1) + { +#pragma unroll + for (int stride = size >> 1; stride > 0; stride >>= 1) + { + uint32_t const other = __shfl_xor_sync(0xffffffffu, blockKey, stride); + bool const ascending = (lane & size) == 0, lower = (lane & stride) == 0; + blockKey = (lower == ascending) ? min(blockKey, other) : max(blockKey, other); + } + } + if (lane < 16) + outRow[lane] = blockKey == ~0u ? -1 : static_cast(blockKey); +} + +// kFast = false is the select before the warp-0 tail (dispatch.SELECT_LONG_SL3 = False, the kill value). +template +__global__ void __launch_bounds__(1024) m3_select_long2_kernel(float const* __restrict__ scores, long long headStride, + long long blockStride, long long queryStride, int const* __restrict__ nValid, int* __restrict__ out, + long long outQueryStride, long long outHeadStride, int numKvHeads, int numBlocks, int initBlocks, int localBlocks, + bool earlyTrigger, bool loadedTrigger) +{ + __shared__ SelectSmem sm; + int const row = blockIdx.x, query = row / numKvHeads, head = row % numKvHeads; + int const tid = threadIdx.x, lane = tid & 31; + if (tid < 16) + { + sm.groupMin[tid] = 0xffffffffu; + sm.picked[tid] = ~0ull; + } + if (tid == 0) + sm.numCand = 0; + // With kFast the row's valid count is read before this grid's dependency wait: a step input (written between + // replays or at the step head), final before the wait, as for the sparse decode's pre-wait seq_lens / block-table + // reads. + int raw = 0; + if (kFast) + { + raw = nValid[query]; + __syncthreads(); + } +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + cudaGridDependencySynchronize(); + // The sparse decode that follows stages its table-independent inputs before its own grid wait; past this grid's + // wait those inputs are final, so it may launch now rather than at this grid's exit. + if (earlyTrigger) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + if (!kFast) + { + raw = nValid[query]; + __syncthreads(); + } + int const valid = max(0, min(raw, numBlocks)); + long long const localStart = max(static_cast(raw) - localBlocks, 0LL); + float const* rowBase = scores + head * headStride + query * queryStride; + int n; + if (valid <= 2048) + n = selectRow<2, kFast>(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + else if (valid <= 4096) + n = selectRow<4, kFast>(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + else + n = selectRow(sm, rowBase, blockStride, valid, initBlocks, localStart, loadedTrigger); + int* const outRow = out + query * outQueryStride + head * outHeadStride; + if (n >= 0) + { + if (tid >= 32) + return; + emitPicks(warpSmallest(sm, n, lane), lane, outRow); + return; + } + __syncthreads(); + if (tid >= 32) + return; + emitPicks(lane < 16 ? sm.picked[lane] : ~0ull, lane, outRow); +} + +} // namespace + +cudaError_t invokeSelectLong2(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* out, int64_t outQueryStride, int64_t outHeadStride, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQ, int32_t initBlocks, int32_t localBlocks, bool usePdl, bool earlyTrigger, + int32_t triggerPhase, bool fastTail, cudaStream_t stream) +{ + if (numBlocks > kSelectLong2MaxBlocks || triggerPhase < 0 || triggerPhase > 1) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(totalQ * numKvHeads); + cfg.blockDim = dim3(1024); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = usePdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, fastTail ? m3_select_long2_kernel : m3_select_long2_kernel, scores, + static_cast(headStride), static_cast(blockStride), static_cast(queryStride), + nValidBlocks, out, static_cast(outQueryStride), static_cast(outHeadStride), numKvHeads, + numBlocks, initBlocks, localBlocks, earlyTrigger, !earlyTrigger && triggerPhase == 1); +} + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.h new file mode 100644 index 000000000000..1e3683a3e7d1 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/selectLong2.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 + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +//! Largest row (blocks per query) the long-row top-16 select covers: 9 keys per thread of one 1024-thread CTA. +constexpr int32_t kSelectLong2MaxBlocks = 9216; + +//! MiniMax-M3 top-16 block select over per-KV-head scores [numKvHeads, numBlocks, totalQ] (element strides given), +//! one 1024-thread CTA per (query, KV head). Same contract as minimax_m3_select_blocks: init / local forcing, the +//! per-query valid-block count, ties to the smaller block, forced -inf picks emitted as -1, ascending ids with -1 +//! tail padding. Writes out[query * outQueryStride + head * outHeadStride + 0..15]. The SM107 TP4 copy of +//! minimaxM3/selectLong with a choice of where the dependent grid (the sparse decode) is released: at this grid's exit +//! (triggerPhase 0), right after every score of the row is loaded (1), or, with earlyTrigger, right after this grid's +//! dependency wait. Every release point follows this grid's own dependency wait. fastTail ranks rows of at most 64 +//! candidates on warp 0 and reads the valid count before the wait (the same picks); false keeps the all-thread rank +//! (dispatch.SELECT_LONG_SL3 = False). +cudaError_t invokeSelectLong2(float const* scores, int64_t headStride, int64_t blockStride, int64_t queryStride, + int32_t const* nValidBlocks, int32_t* out, int64_t outQueryStride, int64_t outHeadStride, int32_t numKvHeads, + int32_t numBlocks, int32_t totalQ, int32_t initBlocks, int32_t localBlocks, bool usePdl, bool earlyTrigger, + int32_t triggerPhase, bool fastTail, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.cu new file mode 100644 index 000000000000..970b2aa21cf4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.cu @@ -0,0 +1,89 @@ +/* + * 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/minimaxM3Sm107Tp4/siluAndMul.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +namespace +{ + +constexpr int kThreads = 128; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves). The +// instruction sequence of the Triton kernel's PTX at its defaults (host/silu/silu_and_mul_plain.ptx): g * -log2(e) as +// one product, ex2.approx, + 1, div.full 1 / that, g * s, (u + 0.0) as an IEEE add, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t siluMulBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 gl, gh, ul, uh, t0, t1, e0, e1, s0, s1, p0, p1, q0, q1, r0, r1, one, zero;\n\t" + ".reg .b16 hl, hh;\n\t" + "mov.b32 {hl, hh}, %1;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mul.rn.f32 t0, gl, 0fBFB8AA3B;\n\tmul.rn.f32 t1, gh, 0fBFB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 p0, gl, s0;\n\tmul.rn.f32 p1, gh, s1;\n\t" + "mov.b32 zero, 0f00000000;\n\tadd.rn.f32 q0, ul, zero;\n\tadd.rn.f32 q1, uh, zero;\n\t" + "mul.rn.f32 r0, q0, p0;\n\tmul.rn.f32 r1, q1, p1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +// Eight outputs per thread: 16-byte loads of the gate and the up half, one 16-byte store. +__global__ void __launch_bounds__(kThreads) + siluAndMulKernel(uint4 const* __restrict__ x, uint4* __restrict__ out, int64_t rows, int64_t vecs) +{ + int64_t const row = blockIdx.y; + int64_t const v = static_cast(blockIdx.x) * kThreads + threadIdx.x; + if (row >= rows || v >= vecs) + { + return; + } + uint4 const g = x[row * 2 * vecs + v]; + uint4 const u = x[row * 2 * vecs + vecs + v]; + uint4 o; + o.x = siluMulBf16x2(g.x, u.x); + o.y = siluMulBf16x2(g.y, u.y); + o.z = siluMulBf16x2(g.z, u.z); + o.w = siluMulBf16x2(g.w, u.w); + out[row * vecs + v] = o; +} + +} // namespace + +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream) +{ + int64_t const vecs = d / 8; + dim3 const grid(static_cast((vecs + kThreads - 1) / kThreads), static_cast(rows)); + siluAndMulKernel<<>>( + reinterpret_cast(x), reinterpret_cast(out), rows, vecs); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h new file mode 100644 index 000000000000..dccbb44572ef --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h @@ -0,0 +1,40 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ + +//! out[m, j] = bf16(silu(g) * (u + 0)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both +//! contiguous, d a multiple of 8, 16-byte aligned rows): bit for bit the Triton silu_and_mul kernel at its defaults +//! (no limit, alpha 1, beta 0), the plain SwiGLU of the Eagle3 drafter MLP. +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu new file mode 100644 index 000000000000..e26804fb47fa --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu @@ -0,0 +1,743 @@ +/* + * 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. + */ + +// MiniMax-M3 long-context sparse decode attention in one launch. +// +// Grid: one cluster of 16 CTAs per query token; CTA `r` attends the token's `r`-th selected block (the Triton +// producer's chunk `r` at 16 chunks), keeps the partial in shared memory, and after a cluster barrier merges head +// `r` from all 16 partials over DSMEM, writing the BF16 row and, optionally, its MXFP8 copy and swizzled UE8M0 +// scales exactly as the Triton merge's `_store_mxfp8_head` lays them out. +// +// Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +// PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +// log-domain recurrence. +// +// ONE_DIV merge (MSA_ONE_DIV_MERGE): each CTA pushes its unnormalized partial with its row max m and sum l, and the +// merge forms o = sum_c w_c acc_c / sum_c w_c l_c with w_c = 2^(m_c - max m), one division per output, instead of +// normalizing every partial and weighting by its LSE. Deterministic (fixed chunk order, no atomics); a different +// rounding sequence from the Triton merge, inside the same band of the float64 reference. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +namespace +{ + +constexpr int kHeads = 16; +constexpr int kDim = 128; +constexpr int kPage = 128; +constexpr int kTopk = 16; +constexpr int kThreads = 128; +constexpr int kPitch = kDim + 8; // halves per smem row: 272 B, conflict-free ldmatrix +constexpr float kNegInf = -__builtin_huge_valf(); + +using Params = MinimaxM3SparseDecodeLcParams; + +struct Smem +{ + __half qs[kHeads * kPitch]; + __half vs[kPage * kPitch]; + __half ps[kHeads * kPitch]; + float recv[kTopk][kDim]; // head `rank` partials pushed by the 16 CTAs, chunk-major + float recv_lse[kTopk]; // the chunk LSEs (log2); with ONE_DIV, the chunk row maxima + float recv_l[kTopk]; // ONE_DIV: the chunk row sums + float wgt[kTopk]; // the 16 merge weights, computed once per CTA + float red[4][kHeads]; + alignas(8) unsigned long long rbar; // the split-K partials of head `rank` landed (complete_tx) +}; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t e2m1x2_to_f16x2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 s;\n cvt.u8.u32 s, %1;\n cvt.rn.f16x2.e2m1x2 %0, s;\n}" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t e4m3x2_to_f16x2(uint16_t pair) +{ + uint32_t out; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(pair)); + return out; +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ void ldsm_x4(uint32_t (&r)[4], uint32_t addr) +{ + 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"(addr)); +} + +__device__ __forceinline__ void ldsm_x4_t(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma(float (&c)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.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"(b0), "r"(b1)); +} + +__device__ __forceinline__ uint32_t map_cluster(void const* local, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smem_u32(local)), "r"(rank)); + return remote; +} + +// Posted DSMEM stores that complete_tx on the receiving CTA's mbarrier `bar`. +__device__ __forceinline__ void st_async_v2(uint32_t addr, float a, float b, uint32_t bar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.v2.f32 [%0], {%1, %2}, [%3];" ::"r"(addr), + "f"(a), "f"(b), "r"(bar) + : "memory"); +} + +__device__ __forceinline__ void st_async_f32(uint32_t addr, float a, uint32_t bar) +{ + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(addr), + "r"(__float_as_uint(a)), "r"(bar) + : "memory"); +} + +// One packed NVFP4 row (64 bytes, 128 values) with its 8 E4M3 scales -> 128 fp16 values in smem. +__device__ __forceinline__ void dequant_row(__half* dst, uint8_t const* packed, uint8_t const (&sc)[8], bool valid) +{ + uint4 const* src = reinterpret_cast(packed); + uint32_t* d = reinterpret_cast(dst); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + uint32_t words[4] = {w.x, w.y, w.z, w.w}; + // 16 bytes = 32 values = scale columns 2c and 2c + 1. +#pragma unroll + for (int h = 0; h < 2; ++h) + { + uint8_t const s = valid ? sc[2 * c + h] : 0; + uint32_t const sh = e4m3x2_to_f16x2(static_cast(s | (s << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); +#pragma unroll + for (int j = 0; j < 2; ++j) + { + uint32_t const word = words[2 * h + j]; +#pragma unroll + for (int b = 0; b < 4; ++b) + { + uint32_t const v2 = e2m1x2_to_f16x2((word >> (8 * b)) & 0xFF); + __half2 const x = __hmul2(*reinterpret_cast<__half2 const*>(&v2), scale); + d[c * 16 + h * 8 + j * 4 + b] = *reinterpret_cast(&x); + } + } + } + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +// This CTA's 16-byte-aligned slice of [base, base + bytes) into L2, at most 16 KiB per bulk prefetch. +__device__ __forceinline__ void prefetch_l2_slice(uint8_t const* base, int64_t bytes) +{ + if (base == nullptr || bytes <= 0) + return; + int64_t const n = gridDim.x; + int64_t const slice = ((bytes + n - 1) / n + 15) & ~int64_t{15}; + int64_t const off = static_cast(blockIdx.x) * slice; + if (off >= bytes) + return; + int64_t const len = min(slice, bytes - off) & ~int64_t{15}; + for (int64_t o = 0; o < len; o += 16384) + { + uint32_t const sz = static_cast(min(int64_t{16384}, len - o)); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(base + off + o), "r"(sz) : "memory"); + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +// BT_PREWAIT: the request's block-table row is staged in dynamic shared memory before the grid dependency (at most +// kBtCap entries, 4096 x 128 = 512K tokens; longer rows keep the global load), so the post-wait page lookup is a shared +// memory hit instead of a dependent L2 load. 52 + 16 KB per CTA still fits beside the o_proj GEMV's 154 KB on one SM. +// The table is written only before the step's graph runs or by the step-input prologue (a plain launch, PRE-WAIT +// WRITER CONTRACT), so it is complete when this grid launches; the values are the same, so is the output. +constexpr int kBtCap = 4096; + +template +constexpr size_t smemBytes() +{ + return BTPW ? sizeof(Smem) + kBtCap * sizeof(int) : sizeof(Smem); +} + +template +__global__ void __launch_bounds__(kThreads) m3_lc_attn_kernel(__grid_constant__ Params const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(128) uint8_t smem_raw[]; + Smem& s = *reinterpret_cast(smem_raw); + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const tok = blockIdx.x / kTopk; + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + + if (tid == 0) + { + // One arrival (here) plus every cluster CTA's partial row (kDim floats) and LSE (ONE_DIV: max and sum) for + // head `rank`. + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(smem_u32(&s.rbar)) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(smem_u32(&s.rbar)), + "r"(kTopk * (kDim + (ONE_DIV ? 2 : 1)) * 4) + : "memory"); + } + // Every peer must be running before its shared memory takes a DSMEM store: arrive now, wait before pushing. + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + // The o_proj weights are read-only for the whole step: start pulling them into L2 now, while HBM is idle in the + // latency-bound attention phase (one lane of warps 0 and 1; nothing waits on these). + if (tid == 0) + prefetch_l2_slice(p.pf0, p.pf0_bytes); + else if (tid == 32) + prefetch_l2_slice(p.pf1, p.pf1_bytes); + + // Everything that does not depend on the top-16 table is read before the grid dependency: the request's length + // (written before the step's graph runs) and, behind a select, the query rows (written by the producer, which + // completed before the select ahead of this grid could release it: the select triggers its dependents only after + // its own grid wait). Identity steps have no select (see the q load below). + int const req = tok / p.dql; + int const qoff = tok - req * p.dql; + int const seq = p.seq_lens[req]; + int const kv_len = max(seq - p.dql + qoff + 1, 0); + int const nblocks = (kv_len + kPage - 1) / kPage; + int const real_topk = min(kTopk, nblocks); + // The next sparse layer's block that its CTA (tok, rank) selected at the previous step (a decode step selects + // mostly the blocks it selected one step earlier): its K/V pages and scales into L2 for that layer, from one lane + // of warp 2. The id buffer was written by a kernel of the previous step, long complete; a recorded id at or past + // this request's block count (another request's row) or a page outside the pool is skipped. + if (p.n1_next != nullptr && tid == 64) + { + int const nb = p.n1_next[tok * kTopk + rank]; + if (nb >= 0 && nb < nblocks) + { + int64_t const pg = p.block_table[req * p.bt_stride + nb]; + if (pg >= 0 && pg < p.n1_pages) + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_k + pg * p.k_blk), "r"(kPage * 64) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_v + pg * p.v_blk), "r"(kPage * 64) + : "memory"); + asm volatile( + "cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_ksc + pg * p.ksc_blk), "r"(kPage * 8) + : "memory"); + asm volatile( + "cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_vsc + pg * p.vsc_blk), "r"(kPage * 8) + : "memory"); + } + } + } + // q: 16 heads x 128 e4m3 -> fp16; thread = (head, 16-value chunk). + auto const load_q = [&]() + { + int const h = tid >> 3; + int const c = (tid & 7) * 16; + uint4 const w = *reinterpret_cast(p.q + tok * p.q_tok + h * p.q_head + c); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; + uint32_t* d = reinterpret_cast(s.qs + h * kPitch + c); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + d[2 * i] = e4m3x2_to_f16x2(static_cast(words[i] & 0xFFFF)); + d[2 * i + 1] = e4m3x2_to_f16x2(static_cast(words[i] >> 16)); + } + }; + // Before the wait only behind the select: an identity step has no scorer / select, so this grid's predecessor is + // the producer itself, which releases its dependents at entry, before it writes q. There q is read after the wait, + // whether the blocks come from the lengths (identity) or from the staged table. + if (!p.identity && !p.identity_step) + load_q(); + int const* const s_bt = reinterpret_cast(smem_raw + sizeof(Smem)); + bool bt_smem = false; + if constexpr (BTPW) + { + int const nbt = (max(seq, 0) + kPage - 1) / kPage; + bt_smem = nbt <= kBtCap; + if (bt_smem) + { + int* const dst = reinterpret_cast(smem_raw + sizeof(Smem)); + for (int i = tid; i < nbt; i += kThreads) + { + dst[i] = __ldg(p.block_table + req * p.bt_stride + i); + } + } + __syncthreads(); // the staged row is complete before any thread's lookup + } + + asm volatile("griddepcontrol.wait;" ::: "memory"); + if constexpr (EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } + if (p.identity || p.identity_step) + load_q(); + int blk = static_cast(rank) >= real_topk ? -1 + : p.identity ? static_cast(rank) + : p.topk[tok * p.t_tok + rank * p.t_k]; + // Fail closed on an id past the row (a stale or corrupt table), as the stock Triton producer does: the block is + // skipped. Read through the block table it would be another request's page, every token masked, and its softmax + // ex2(-inf - (-inf)) = NaN. + if (blk >= nblocks) + blk = -1; + bool const active = blk >= 0; + // Record this step's selection for this layer's next-step prefetch (read by the previous sparse layer's CTAs). + if (p.n1_self != nullptr && tid == 0) + p.n1_self[tok * kTopk + rank] = blk; + int64_t const page_pw = !active ? 0 : (bt_smem ? s_bt[blk] : p.block_table[req * p.bt_stride + blk]); + + float m_row[2] = {kNegInf, kNegInf}; + float acc[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + acc[j][i] = 0.F; + float l_row[2] = {0.F, 0.F}; + + // K stays packed in registers: the MMA B fragment of token n = lane / 4 at k-step s is the f16x2 pair of + // packed bytes 8s + q and 8s + 4 + q (q = lane % 4), both under scale column s. + uint32_t kw[4][16]; + uint2 ksc[4]; + int64_t page = 0; + if (active) + { + page = page_pw; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const n = warp * 32 + j * 8 + (lane >> 2); + bool const valid = blk * kPage + n < kv_len; + uint4 const* src = reinterpret_cast(p.k + page * p.k_blk + n * 64); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 const w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + kw[j][4 * c] = w.x; + kw[j][4 * c + 1] = w.y; + kw[j][4 * c + 2] = w.z; + kw[j][4 * c + 3] = w.w; + } + ksc[j] = valid ? __ldg(reinterpret_cast(p.ksc + page * p.ksc_blk + n * 8)) : make_uint2(0, 0); + } + bool const valid = blk * kPage + tid < kv_len; + uint8_t vsc[8]; + if (valid) + { + uint8_t const* vb = p.vsc + page * p.vsc_blk + (tid >> 2) * 32 + (tid & 3); +#pragma unroll + for (int c = 0; c < 8; ++c) + vsc[c] = __ldg(vb + 4 * c); + } + dequant_row(s.vs + tid * kPitch, p.v + page * p.v_blk + tid * 64, vsc, valid); + } + __syncthreads(); + + if (active) + { + float const scale_log2 = p.sm_scale * 1.4426950409F * __ldg(p.kgs); + int const q4 = (lane & 3) * 8; + // S = q K^T for this warp's 32 tokens: rows g / g + 8, columns 2 * (lane % 4) + {0, 1} of each n-tile. + float sacc[4][4]; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + sacc[j][i] = 0.F; +#pragma unroll + for (int ks = 0; ks < kDim / 16; ++ks) + { + uint32_t a[4]; + ldsm_x4(a, smem_u32(s.qs + (lane & 15) * kPitch + ks * 16 + (lane >> 4) * 8)); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + uint32_t const scw = ks < 4 ? ksc[j].x : ksc[j].y; + uint8_t const sc = static_cast(scw >> (8 * (ks & 3))); + uint32_t const sh = e4m3x2_to_f16x2(static_cast(sc | (sc << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); + uint32_t const r0 = e2m1x2_to_f16x2((kw[j][2 * ks] >> q4) & 0xFF); + uint32_t const r1 = e2m1x2_to_f16x2((kw[j][2 * ks + 1] >> q4) & 0xFF); + __half2 const h0 = __hmul2(*reinterpret_cast<__half2 const*>(&r0), scale); + __half2 const h1 = __hmul2(*reinterpret_cast<__half2 const*>(&r1), scale); + mma(sacc[j], a, *reinterpret_cast(&h0), *reinterpret_cast(&h1)); + } + } + int const g = lane >> 2; + float mloc[2] = {kNegInf, kNegInf}; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const t = warp * 32 + j * 8 + 2 * (lane & 3) + (i & 1); + float const v = blk * kPage + t < kv_len ? 0.F + sacc[j][i] * scale_log2 : kNegInf; + sacc[j][i] = v; + mloc[i >> 1] = fmaxf(mloc[i >> 1], v); + } +#pragma unroll + for (int r = 0; r < 2; ++r) + { + mloc[r] = fmaxf(mloc[r], __shfl_xor_sync(0xffffffff, mloc[r], 1)); + mloc[r] = fmaxf(mloc[r], __shfl_xor_sync(0xffffffff, mloc[r], 2)); + } + if ((lane & 3) == 0) + { + s.red[warp][g] = mloc[0]; + s.red[warp][g + 8] = mloc[1]; + } + __syncthreads(); + // Every thread reduces its two rows' four warp maxima itself (the same expression, so the same bits): no + // shared row-max stage and its second barrier. +#pragma unroll + for (int r = 0; r < 2; ++r) + { + int const row = g + r * 8; + m_row[r] = fmaxf(fmaxf(s.red[0][row], s.red[1][row]), fmaxf(s.red[2][row], s.red[3][row])); + } + // A row with no valid token has m = -inf: its probabilities are ex2(-inf - 0) = 0, never NaN. + float const m_use[2] = {m_row[0] == kNegInf ? 0.F : m_row[0], m_row[1] == kNegInf ? 0.F : m_row[1]}; + float lloc[2] = {0.F, 0.F}; +#pragma unroll + for (int j = 0; j < 4; ++j) +#pragma unroll + for (int i = 0; i < 4; i += 2) + { + float const p0 = ex2(sacc[j][i] - m_use[i >> 1]); + float const p1 = ex2(sacc[j][i + 1] - m_use[i >> 1]); + lloc[i >> 1] += p0 + p1; + int const row = g + (i >> 1) * 8; + int const t = warp * 32 + j * 8 + 2 * (lane & 3); + *reinterpret_cast<__half2*>(s.ps + row * kPitch + t) = __floats2half2_rn(p0, p1); + } +#pragma unroll + for (int r = 0; r < 2; ++r) + { + lloc[r] += __shfl_xor_sync(0xffffffff, lloc[r], 1); + lloc[r] += __shfl_xor_sync(0xffffffff, lloc[r], 2); + } + __syncthreads(); // row-max reads of red done before red is reused; ps complete + if ((lane & 3) == 0) + { + s.red[warp][g] = lloc[0]; + s.red[warp][g + 8] = lloc[1]; + } + // O[:, warp's 32 dims] = P V over all 128 tokens. +#pragma unroll + for (int kk = 0; kk < kPage; kk += 16) + { + uint32_t a[4]; + ldsm_x4(a, smem_u32(s.ps + (lane & 15) * kPitch + kk + (lane >> 4) * 8)); +#pragma unroll + for (int jp = 0; jp < 2; ++jp) + { + uint32_t b[4]; + int const t = kk + (lane & 7) + ((lane >> 3) & 1) * 8; + int const d = warp * 32 + jp * 16 + (lane >> 4) * 8; + ldsm_x4_t(b, smem_u32(s.vs + t * kPitch + d)); + mma(acc[2 * jp], a, b[0], b[1]); + mma(acc[2 * jp + 1], a, b[2], b[3]); + } + } + __syncthreads(); + // Likewise the row sums, in the same order. red is next written by the merge, after this CTA's rbar wait, + // which completes only once every warp here has pushed its partial (formed from these sums). +#pragma unroll + for (int r = 0; r < 2; ++r) + { + int const row = g + r * 8; + l_row[r] = (s.red[0][row] + s.red[1][row]) + (s.red[2][row] + s.red[3][row]); + } + } + + // Normalize the partial and push head h's row with its LSE (log2) into CTA h's recv[rank] slot, as the Triton + // producer stores chunk `rank`; st.async stores that complete_tx on CTA h's barrier. + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + { + int const g = lane >> 2; + float const vgs = __ldg(p.vgs); +#pragma unroll + for (int r = 0; r < 2; ++r) + { + float scale, lse; + if constexpr (ONE_DIV) + { + // Unnormalized: the merge divides once. An inactive CTA pushes m = -inf, l = 0 (weight 0). + scale = 1.F; + lse = m_row[r]; + } + else if constexpr (LINEAR) + { + scale = l_row[r] > 0.F ? 1.0F / l_row[r] : 0.F; + lse = l_row[r] > 0.F ? m_row[r] + lg2(l_row[r]) : kNegInf; + } + else + { + float const lse_i + = active && m_row[r] > kNegInf ? m_row[r] + lg2(ex2(kNegInf - m_row[r]) + l_row[r]) : kNegInf; + scale = lse_i > kNegInf ? ex2(m_row[r] - lse_i) : 0.F; + lse = lse_i; + } + scale = scale * vgs; + int const row = g + r * 8; + uint32_t const dst = map_cluster(&s.recv[rank][0], row); + uint32_t const rb = map_cluster(&s.rbar, row); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const d = warp * 32 + j * 8 + 2 * (lane & 3); + st_async_v2(dst + d * 4, acc[j][2 * r] * scale, acc[j][2 * r + 1] * scale, rb); + } + if (warp == 0 && (lane & 3) == 0) + { + st_async_f32(map_cluster(&s.recv_lse[rank], row), lse, rb); + if constexpr (ONE_DIV) + { + st_async_f32(map_cluster(&s.recv_l[rank], row), l_row[r], rb); + } + } + } + } + { + // Every CTA's partial row and LSE for head `rank` landed (complete_tx on this CTA's barrier). + uint32_t done = 0; + while (!done) + { + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(smem_u32(&s.rbar)) + : "memory"); + } + } + + // CTA `rank` merges head `rank` across the 16 partials; thread = output dim. + int const h = static_cast(rank); + float o = 0.F; + if constexpr (ONE_DIV) + { + // Thread c < 16 takes chunk c's weight 2^(m_c - max m) and its weighted row sum; every thread then sums the + // weighted partials and the weighted row sums in chunk order and divides once. + if (tid < kTopk) + { + float m_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + m_max = fmaxf(m_max, s.recv_lse[c]); + m_max = m_max == kNegInf ? 0.F : m_max; + float const w = ex2(s.recv_lse[tid] - m_max); + s.wgt[tid] = w; + s.red[0][tid] = w * s.recv_l[tid]; + } + __syncthreads(); + float num = 0.F; + float den = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + { + num += s.recv[c][tid] * s.wgt[c]; + den += s.red[0][c]; + } + o = den > 0.F ? __fdiv_rn(num, den) : 0.F; + } + else + { + // The weights are the Triton merge's (exp2 of each chunk's LSE minus the largest, divided by their sum in + // chunk order), formed once per CTA: thread c < 16 takes chunk c's exponential, then divides it by the sum + // every one of those threads forms in the same order. + if (tid < kTopk) + { + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + lse_max = fmaxf(lse_max, s.recv_lse[c]); + lse_max = lse_max == kNegInf ? 0.F : lse_max; + s.red[0][tid] = ex2(s.recv_lse[tid] - lse_max); + } + __syncthreads(); + if (tid < kTopk) + { + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + denom += s.red[0][c]; + denom = denom > 0.F ? denom : 1.F; + s.wgt[tid] = __fdiv_rn(s.red[0][tid], denom); + } + __syncthreads(); +#pragma unroll + for (int c = 0; c < kTopk; ++c) + o += s.recv[c][tid] * s.wgt[c]; + } + __nv_bfloat16 const ob = __float2bfloat16_rn(o); + p.out[tok * p.out_stride + h * kDim + tid] = ob; + + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + float const xq = x * qs; + uint16_t q2; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2) : "f"(0.F), "f"(xq)); + p.q8[static_cast(tok) * (p.num_sf_k_tiles * 128) + h * kDim + tid] = static_cast(q2 & 0xFF); + if (lane == 0) + { + int const k_sf = h * 4 + warp; + int64_t const off = static_cast(tok / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (tok % 32) * 16 + ((tok % 128) / 32) * 4 + k_sf % 4; + p.sf[off] = static_cast(e); + } + } + if constexpr (!EARLY) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +cudaError_t configure() +{ + auto kernel = m3_lc_attn_kernel; + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes())); + if (err != cudaSuccess) + { + return err; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); +} + +template +cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) +{ + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs + // every variant eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + auto kernel = m3_lc_attn_kernel; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(num_tokens * kTopk); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smemBytes(); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kTopk; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +template +cudaError_t launchBt(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) +{ + return prm.bt_prewait ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); +} + +} // namespace + +cudaError_t invokeSparseDecodeLc4(MinimaxM3SparseDecodeLcParams const& params, int32_t numTokens, bool linearSoftmax, + bool mxfp8, bool earlyTrigger, bool usePdl, cudaStream_t stream) +{ + if (!earlyTrigger) + { + return cudaErrorNotSupported; + } + if (params.one_div) + { + if (linearSoftmax) + { + return mxfp8 ? launchBt(params, numTokens, usePdl, stream) + : launchBt(params, numTokens, usePdl, stream); + } + return mxfp8 ? launchBt(params, numTokens, usePdl, stream) + : launchBt(params, numTokens, usePdl, stream); + } + if (linearSoftmax) + { + return mxfp8 ? launchBt(params, numTokens, usePdl, stream) + : launchBt(params, numTokens, usePdl, stream); + } + return mxfp8 ? launchBt(params, numTokens, usePdl, stream) + : launchBt(params, numTokens, usePdl, stream); +} + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h new file mode 100644 index 000000000000..6fea5ee87498 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h @@ -0,0 +1,104 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ + +//! SM107 TP4 copy of kernels/minimaxM3/sparseDecodeLc (the sm_103 targets keep that one), with the split-K merge +//! weights formed once per CTA. Operands of the MiniMax-M3 one-launch long-context sparse decode. Strides are in +//! elements of the pointed-to type. Geometry is fixed: 16 query heads on one KV head, head_dim 128, page 128, top-k 16. +struct MinimaxM3SparseDecodeLcParams +{ + uint8_t const* q; // e4m3 [T, 16, 128] + int64_t q_tok, q_head; + uint8_t const* k; // packed e2m1 [pages, 1, 128, 64] + uint8_t const* v; + int64_t k_blk, v_blk; + uint8_t const* ksc; // e4m3 [pages, 1, 128, 8], token-major + uint8_t const* vsc; // e4m3, 4x4 token-quad order + int64_t ksc_blk, vsc_blk; + float const* kgs; + float const* vgs; + int const* topk; // [T, 16] view (kv head 0) + int64_t t_tok, t_k; + int const* block_table; + int64_t bt_stride; + int const* seq_lens; + int dql; + float sm_scale; + __nv_bfloat16* out; // [T, 2048] + int64_t out_stride; + uint8_t* q8; // e4m3 [rows, num_sf_k_tiles * 128] + uint8_t* sf; + int num_sf_k_tiles; + // Optional L2 prefetch of two read-only ranges (the layer's o_proj weight and scales), split across the grid's + // CTAs and issued before the grid dependency: a hint, so the output does not depend on it. nullptr = none. + uint8_t const* pf0; + int64_t pf0_bytes; + uint8_t const* pf1; + int64_t pf1_bytes; + // Optional next-layer KV prefetch: CTA (token, r) records its selected block id in this layer's [rows, 16] id + // buffer (n1_self) and pulls into L2 the K/V pages and scales of the block the next sparse layer's CTA (token, r) + // selected at the previous step (n1_next, from that layer's pools n1_k / n1_v / n1_ksc / n1_vsc, this layer's + // strides). Hints only. nullptr = off. + int* n1_self; + int const* n1_next; + uint8_t const* n1_k; + uint8_t const* n1_v; + uint8_t const* n1_ksc; + uint8_t const* n1_vsc; + int64_t n1_pages; // pages in the next layer's pool; a recorded block outside it is skipped + // Identity selection (short-sequence steps, every request within 16 blocks): CTA r attends block r of its + // token's request when r < min(16, blocks), from the device sequence length; topk is not read. + int identity; + // The step is an identity step (short sequences: no scorer / select ran this step), whether or not this kernel + // derives its blocks from the lengths (identity) or reads the staged table: q is then read after the grid wait, + // because the predecessor is the producer, which releases its dependents at entry, before it writes q. + int identity_step; + // One-division merge (MSA_ONE_DIV_MERGE): the CTAs push unnormalized partials with their row max and sum, and the + // merge divides once per output. Deterministic, a different rounding sequence from the per-chunk normalization. + int one_div; + // BT_PREWAIT: stage the request's block-table row (up to 4096 entries) in shared memory before the grid + // dependency, so the post-wait page lookup reads shared memory. Same output. + int bt_prewait; +}; + +//! One cluster of 16 CTAs per query token: CTA r attends the token's r-th selected block, then the cluster merges +//! the 16 partials over DSMEM into the BF16 row (and, when mxfp8, its MXFP8 copy with swizzled UE8M0 scales). +//! +//! Instantiated variants: linearSoftmax in {false, true} x mxfp8 in {false, true}, with earlyTrigger (release PDL +//! dependents at entry). Returns cudaErrorNotSupported for any other variant. The device code is SM100-family only; +//! on other architectures the kernel compiles to an empty body, so callers check the device first. +cudaError_t invokeSparseDecodeLc4(MinimaxM3SparseDecodeLcParams const& params, int32_t numTokens, bool linearSoftmax, + bool mxfp8, bool earlyTrigger, bool usePdl, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu new file mode 100644 index 000000000000..679bcb002d1d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu @@ -0,0 +1,58 @@ +/* + * 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/minimaxM3Sm107Tp4/specOffsets.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +namespace +{ + +constexpr int kThreads = 256; + +__global__ void __launch_bounds__(kThreads) + specOffsetsKernel(int32_t* __restrict__ pos, int32_t const* __restrict__ posOff, int64_t nPos, + int32_t* __restrict__ kv, int32_t const* __restrict__ kvOff, int64_t nKv) +{ + for (int64_t i = threadIdx.x; i < nPos; i += kThreads) + { + pos[i] = static_cast(static_cast(pos[i]) + static_cast(posOff[i])); + } + for (int64_t j = threadIdx.x; j < nKv; j += kThreads) + { + kv[j] = static_cast(static_cast(kv[j]) + static_cast(kvOff[j])); + } +} + +} // namespace + +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream) +{ + specOffsetsKernel<<<1, kThreads, 0, stream>>>(pos, posOff, nPos, kv, kvOff, nKv); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h new file mode 100644 index 000000000000..878ca0197e9e --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h @@ -0,0 +1,41 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ + +//! The overlap scheduler's speculative step correction in one launch: pos[i] += posOff[i] for i < nPos and +//! kv[j] += kvOff[j] for j < nKv (int32, wrapping as the ATen adds do). A plain launch: it is a step-input writer that +//! pre-wait readers rely on (rubin-20261015/PREWAIT_AUDIT.md item b). +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu new file mode 100644 index 000000000000..42361b5ef47f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu @@ -0,0 +1,111 @@ +/* + * 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/minimaxM3Sm107Tp4/stepPrologue.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +namespace +{ + +constexpr int kThreads = 512; +constexpr int kWarps = kThreads / 32; + +// One warp per copy. Plain (coherent) loads: a device copy may read bytes another warp of this CTA just stored, which +// the CTA barrier between the two kinds makes visible; a non-coherent load could return the stale line. +__device__ __forceinline__ void copyWarp(uint8_t* dst, uint8_t const* src, int64_t bytes, int lane) +{ + auto const d = reinterpret_cast(dst); + auto const s = reinterpret_cast(src); + if (((d | s | static_cast(bytes)) & 15) == 0) + { + auto* d16 = reinterpret_cast(dst); + auto const* s16 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 16; i += 32) + { + d16[i] = s16[i]; + } + } + else if (((d | s | static_cast(bytes)) & 3) == 0) + { + auto* d4 = reinterpret_cast(dst); + auto const* s4 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 4; i += 32) + { + d4[i] = s4[i]; + } + } + else + { + for (int64_t i = lane; i < bytes; i += 32) + { + dst[i] = src[i]; + } + } +} + +// One CTA: a step's window is a few dozen copies of a few hundred bytes, so the launch is the cost; a CTA barrier +// orders the host copies before the device copies where the parallel scatter needed a second launch. +__global__ void __launch_bounds__(kThreads) stepPrologueKernel(uint8_t* window, int64_t* epoch, int64_t maxDesc) +{ + auto* header = reinterpret_cast(window); + int64_t const numHost = min(header[0], maxDesc); + int64_t const numDevice = min(header[1], maxDesc); + int64_t const* desc = header + 2; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (epoch != nullptr && threadIdx.x == 0) + { + *epoch += 1; + } + for (int64_t i = warp; i < numHost; i += kWarps) + { + int64_t const* c = desc + 3 * i; + copyWarp(reinterpret_cast(c[0]), window + c[1], c[2], lane); + } + __syncthreads(); + for (int64_t i = warp; i < numDevice; i += kWarps) + { + // The device copies' slots follow all header[0] host copies, as the parallel scatter indexes them. + int64_t const* c = desc + 3 * (header[0] + i); + copyWarp(reinterpret_cast(c[0]), reinterpret_cast(c[1]), c[2], lane); + } + // Every thread has read the header and the descriptors: consume the window. + __syncthreads(); + if (threadIdx.x < 2) + { + header[threadIdx.x] = 0; + } +} + +} // namespace + +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream) +{ + stepPrologueKernel<<<1, kThreads, 0, stream>>>(window, epoch, maxDesc); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h new file mode 100644 index 000000000000..c934decff161 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h @@ -0,0 +1,47 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ + +//! The decode graph's step-input prologue in one launch (step_inputs.StepInputPack with early_copy): the work of the +//! two parallel scatter launches and the header clear that follow them, in that order. +//! +//! The window is the pack's staging layout (kernels/stepInputsScatter.h): an int64 header `[numHostCopies, +//! numDeviceCopies]`, then one `(dst, src, bytes)` int64 triple per copy, host copies first; a host copy's `src` is a +//! byte offset into the window, a device copy's a device address. Up to `maxDesc` copies of each kind are written (the +//! slots the parallel scatter's grid covers), every host copy before any device copy, so a device copy may read a +//! destination a host copy filled; the copies of one kind must not overlap (the pack issues the window only then). +//! Afterwards the header is zeroed, so a replay with no new upload writes nothing. With `epoch` non-null, it is +//! incremented once. +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu new file mode 100644 index 000000000000..507520ec0de6 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp4 +{ + +namespace +{ + +constexpr int kThreads = 256; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves): the +// instruction sequence of the Triton kernel at limit 7, alpha 1.702, beta 1 (Indexer kdev17, bitwise on SM107): the +// gate clamp as min.bf16x2 before the widening, the up clamp as min.xorsign.abs, (u + 1), (g * -1.702) * log2(e), +// ex2.approx, + 1, div.full 1 / that, s * g, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t swigluOaiBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, %1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +template +__global__ void __launch_bounds__(kThreads) + swigluOaiKernel(uint32_t const* __restrict__ x, uint32_t* __restrict__ out, int64_t rows, int64_t halfD) +{ + if constexpr (kPdl) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int64_t const i = blockIdx.x * static_cast(kThreads) + threadIdx.x; + if (i >= rows * halfD) + { + return; + } + int64_t const row = i / halfD; + int64_t const col = i - row * halfD; + uint32_t const* xr = x + row * 2 * halfD; + out[i] = swigluOaiBf16x2(xr[col], xr[halfD + col]); +} + +} // namespace + +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream) +{ + int64_t const n = rows * halfD; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast((n + kThreads - 1) / kThreads)); + cfg.blockDim = dim3(kThreads); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return pdl ? cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD) + : cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD); +} + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h new file mode 100644 index 000000000000..f2898e100006 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h @@ -0,0 +1,43 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp4 +{ + +//! The dense MLP's swigluoai activation: out[m, j] = bf16(min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + +//! 1)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both contiguous, d even): bit for bit +//! the Triton silu_and_mul kernel at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With `pdl` the launch allows +//! programmatic serialization and the grid waits on its dependency, then releases its dependents before it computes +//! (the down GEMV stages its weights during the activation and reads it only after its own grid wait). +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp4 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu new file mode 100644 index 000000000000..fb28ce95250b --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu @@ -0,0 +1,711 @@ +/* + * 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/minimaxM3Sm107Tp8/draftStep.h" + +#include +#include +#include + +// One pass of the fused greedy Eagle3 draft loop of MiniMax-M3 (GB300, TP4, <= 8 requests per pass). +// +// One kernel replaces the draft lm_head GEMM, the logits all-gather / float cast / argmax of the in-tree greedy draft +// sampling and the draft loop's per-pass bookkeeping, at the in-tree precision: +// +// * The vocab-sharded BF16 lm_head weight is stored in mma.m16n8k16 A-fragment order (the same values): per 16-row +// tile and k16 step, 512 bytes; lane (g, j) holds 16 bytes = {row g k 2j..2j+1, row g+8 k 2j.., row g k 2j+8.., +// row g+8 k 2j+8..}. +// * Or, by default, in a lossless 12-bit coding of those fragments (75% of the bytes; the target head's coding, see +// targetHead.cu): per 16-row tile and unit of 4 k16 steps a lane reads three 16-byte words, the sign/mantissa bytes +// of steps 0-1, of steps 2-3, then 4 code words (exponent - base, 4 bits; 15 = an exception). The decoded values +// are the checkpoint's BF16 bits; an exception decodes to +0 and the segment's flush adds the exact product x * w +// from the tile's list (products of BF16 values are exact in FP32). +// * A persistent grid (one 16-warp CTA per SM) streams whole tiles over the full K with BF16 MMAs and FP32 +// accumulation; the BF16 head input rows sit in shared memory in B-fragment order, unchanged. +// * Each logit is rounded to BF16 (the in-tree logits are the BF16 GEMM output) and becomes a 64-bit key +// ordered(value) << 32 | ~global_index, with -0 taken as +0: the larger value wins and ties pick the lowest index, as +// torch.max / torch.argmax over the gathered logits do. Keys are max-reduced per row: warp shuffles, then shared +// memory, then atomicMax. +// * The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping +// (Lamport: 0 = empty, a key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the +// slots, and runs the pass's bookkeeping. Call n uses slot half n % 2. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ +namespace draftStep +{ +namespace +{ +constexpr int kUnroll = 8; // k16 chunks in flight per warp +constexpr int kChunkBytes = 512; // one tile's k16 step: 32 lanes x 16 bytes +constexpr long long kMaxSpins = 1LL << 31; // cross-rank key poll bound + +// Coded units (4 k16 steps, 3 x 16 bytes per lane) in flight per warp: 3 up to 4 rows, 2 at 8 rows. +template +constexpr int codedUnroll() +{ + return MP >= 8 ? 2 : 3; +} + +__device__ __forceinline__ void mma_bf16( + float (&d)[4], uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t b0, uint32_t b1) +{ + 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};\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)); +} + +// BF16 pair p (values 2p, 2p + 1) of a lane's 8 coded values: sign/mantissa bytes sm (4 per word), codes word cw. +__device__ __forceinline__ uint32_t coded_decode_pair(uint32_t sm_word, int byte0, uint32_t cw, int p, uint32_t base) +{ + uint32_t const b = __byte_perm(sm_word, 0u, byte0 == 0 ? 0x4140 : 0x4342); // {b0, 0, b1, 0} + uint32_t const q = (cw >> (4 * p)) & 0x000f000fu; + uint32_t const e = q * 128u + base; // exponent field of both halves + // Code 15 (an exception, sign/mantissa byte 0) decodes to +0. + return ((b & 0x007f007fu) | ((b << 8) & 0x80008000u) | e) & __vcmpne2(q, 0x000f000fu); +} + +// One coded unit (4 k16 steps) against the BF16 B fragments: steps alternate between the two accumulators, as the +// k16 steps of the plain BF16 head do. +template +__device__ __forceinline__ void chunk_mma_coded(uint4 const& sm01, uint4 const& sm23, uint4 const& cw, uint32_t base, + uint2 const* xs, int unit, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint32_t const smw[8] = {sm01.x, sm01.y, sm01.z, sm01.w, sm23.x, sm23.y, sm23.z, sm23.w}; + uint32_t const cws[4] = {cw.x, cw.y, cw.z, cw.w}; +#pragma unroll + for (int ci = 0; ci < 4; ++ci) + { + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[((unit * 4 + ci) * 4 + j) * MP + g]; + uint32_t const w0 = coded_decode_pair(smw[2 * ci], 0, cws[ci], 0, base); + uint32_t const w1 = coded_decode_pair(smw[2 * ci], 2, cws[ci], 1, base); + uint32_t const w2 = coded_decode_pair(smw[2 * ci + 1], 0, cws[ci], 2, base); + uint32_t const w3 = coded_decode_pair(smw[2 * ci + 1], 2, cws[ci], 3, base); + if (ci & 1) + mma_bf16(db, w0, w1, w2, w3, b.x, b.y); + else + mma_bf16(da, w0, w1, w2, w3, b.x, b.y); + } +} + +// Exceptions of `tile` with k in [kb, ke): v += x * w (exact) for the owner lanes (row g or g + 8; tokens 2j, +// 2j + 1 < M), x from the staged B fragments. Warp-uniform (shuffles). +template +__device__ __forceinline__ void coded_exceptions( + CodedArgs const& cd, int tile, int kb, int ke, uint2 const* xs, int M, int lane, float (&v)[4]) +{ + int const g = lane >> 2, j = lane & 3; + int const e0 = __ldg(cd.exc_off + tile), e1 = __ldg(cd.exc_off + tile + 1); + for (int b0 = e0; b0 < e1; b0 += 32) + { + uint32_t pos = 0xffffffffu; + float2 wv = make_float2(0.f, 0.f); + if (b0 + lane < e1) + { + pos = __ldg(cd.exc_pos + b0 + lane); + wv = __ldg(cd.exc_w + b0 + lane); + } + int const cnt = min(32, e1 - b0); + for (int q = 0; q < cnt; ++q) + { + uint32_t const pq = __shfl_sync(0xffffffffu, pos, q); + float const w0 = __shfl_sync(0xffffffffu, wv.x, q); + int const r = static_cast(pq >> 16), k = static_cast(pq & 0xffffu); + if (k >= kb && k < ke && (r & 7) == g) + { + int const st = k >> 4, rr = k & 15; + bool const hi = r >= 8; +#pragma unroll + for (int q2 = 0; q2 < 2; ++q2) + { + int const n = 2 * j + q2; + if (n < M) + { + uint2 const bw = xs[(st * 4 + ((rr & 7) >> 1)) * MP + n]; + uint32_t const word = rr >= 8 ? bw.y : bw.x; + float const xv = __uint_as_float((rr & 1) ? (word & 0xffff0000u) : (word << 16)); + float const d = __fmul_rn(xv, w0); + if (hi) + v[2 + q2] += d; + else + v[q2] += d; + } + } + } + __syncwarp(); + } + } +} + +__device__ __forceinline__ uint4 ldg_stream(uint4 const* p) +{ + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.L2::256B.v4.u32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +// The argmax value of a logit: its BF16 rounding (the in-tree logits' dtype), -0 as +0. +__device__ __forceinline__ float key_value(float v) +{ + float const r = __bfloat162float(__float2bfloat16(v)); + return r == 0.f ? 0.f : r; +} + +__device__ __forceinline__ unsigned long long make_key(float v, uint32_t index) +{ + uint32_t u = __float_as_uint(v); + u = (u & 0x80000000u) ? ~u : (u | 0x80000000u); + return (static_cast(u) << 32) | static_cast(0xFFFFFFFFu - index); +} + +__device__ __forceinline__ unsigned long long shfl_max(unsigned long long k, int o) +{ + unsigned long long const ok = __shfl_xor_sync(0xffffffffu, k, o); + return ok > k ? ok : k; +} + +// One k16 step of a tile: the lane's A fragment w against the activation B fragment of step `chunk` (rows g < MP). +// Even and odd steps accumulate separately (two independent MMA chains); the tile's sum is da + db. +template +__device__ __forceinline__ void chunk_mma( + uint4 const& w, uint2 const* xs, int chunk, int j, int g, bool b_active, float (&da)[4], float (&db)[4]) +{ + uint2 b = make_uint2(0u, 0u); + if (b_active) + b = xs[(chunk * 4 + j) * MP + g]; + if (chunk & 1) + mma_bf16(db, w.x, w.y, w.z, w.w, b.x, b.y); + else + mma_bf16(da, w.x, w.y, w.z, w.w, b.x, b.y); +} + +// A warp's pieces of the tiles it shares with other warps (its first and last local tile), kept for the warp-order sum. +struct Pieces +{ + float4 head = make_float4(0.f, 0.f, 0.f, 0.f); + float4 tail = make_float4(0.f, 0.f, 0.f, 0.f); +}; + +// Whether the warp's chunk range [q_begin, q_end) covers the CTA's whole range of local tile tl (the CTA streams +// chunks [lo, lo + span) of its local tiles). +__device__ __forceinline__ bool covers_tile(int tl, int kchunks, int lo, int span, int q_begin, int q_end) +{ + return q_begin <= max(tl * kchunks, lo) && q_end >= min((tl + 1) * kchunks, lo + span); +} + +// Local tile t's sum (D-FIX): the accumulator when one warp streamed the CTA's whole range of it, else its pieces +// from the warps that streamed it, added in warp order from 0 (each warp's first tile is its head piece, any later +// one its tail piece). Deterministic: the order is fixed by the launch shape, and with two pieces it is the same +// sum as the atomics' (two addends from 0 commute). +__device__ __forceinline__ float4 ordered_tile( + float4 const* acc_t, float4 const* stage, int t, int kchunks, int lo, int span, int lane) +{ + int const r_lo = max(t * kchunks, lo); + int const r_last = min((t + 1) * kchunks, lo + span) - 1; + int w_lo = 0, w_hi = 0; + for (int w = 1; w < kWarps; ++w) + { + int const qb = lo + w * span / kWarps; + if (qb <= r_lo) + w_lo = w; + if (qb <= r_last) + w_hi = w; + } + if (w_lo == w_hi) + return acc_t[lane]; + float4 s = make_float4(0.f, 0.f, 0.f, 0.f); + for (int w = w_lo; w <= w_hi; ++w) + { + int const qb = lo + w * span / kWarps; + float4 const x = stage[(w * 2 + (qb / kchunks == t ? 0 : 1)) * 32 + lane]; + s.x += x.x; + s.y += x.y; + s.z += x.z; + s.w += x.w; + } + return s; +} + +// Add a warp's segment of a tile into the CTA's accumulator (with a coded head, the segment's exceptions first); +// with `fixed` (D-FIX) a shared tile's piece goes to `pc` instead (ordered_tile adds it). +template +__device__ __forceinline__ void flush_tile(float4* acc, float (&da)[4], float (&db)[4], Params const& p, int tile, + int kb, int ke, uint2 const* xs, int lane, Pieces& pc, bool fixed, bool whole, bool first) +{ + float v[4] = {da[0] + db[0], da[1] + db[1], da[2] + db[2], da[3] + db[3]}; + if constexpr (CODED) + coded_exceptions(p.cd, tile, kb, ke, xs, p.M, lane, v); + if (fixed) + { + // A piece covering the CTA's whole range of the tile is its only addend (0 + v, as the atomic gave); a + // warp's other pieces (at most its first and last tile) wait in registers for the warp-order sum. + if (whole) + *acc = make_float4(0.f + v[0], 0.f + v[1], 0.f + v[2], 0.f + v[3]); + else if (first) + pc.head = make_float4(v[0], v[1], v[2], v[3]); + else + pc.tail = make_float4(v[0], v[1], v[2], v[3]); + } + else + { + float* a = reinterpret_cast(acc); + atomicAdd(a + 0, v[0]); + atomicAdd(a + 1, v[1]); + atomicAdd(a + 2, v[2]); + atomicAdd(a + 3, v[3]); + } +#pragma unroll + for (int e = 0; e < 4; ++e) + { + da[e] = 0.f; + db[e] = 0.f; + } +} + +// Balanced streaming GEMV over the fragment-ordered weight, persistent grid of 16-warp CTAs. CTA c owns tiles +// [t_begin, t_end); its (tile, chunk) space is split into 16 equal contiguous ranges, one per warp, so every warp +// streams the same number of bytes. A tile straddles at most two warps; their partial sums are added into shared +// memory (two addends: the order does not change the sum). MP = rows padded to 1 / 2 / 4 / 8. A chunk is one k16 +// step of the BF16 head, or one coded unit (4 k16 steps) of the coded head (CODED: no shared-memory prologue). +template +__global__ void __launch_bounds__(kWarps * 32, 1) + draft_step_kernel(Params const p, int kchunks, int pf_chunks, int pf_off) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(16) uint2 xs[]; // [K/16 steps][4 (j)][MP (g)], then the weight prologue + __shared__ float4 acc_s[kMaxTilesPerCta + 1][32]; // + 1: an equal-chunk span can touch one more tile + __shared__ unsigned long long red[kWarps][kMaxTokens]; + __shared__ __align__(8) unsigned long long pf_bar[kWarps]; + __shared__ bool s_last; + __shared__ int s_tok[kMaxTokens]; + + int const warp = threadIdx.x >> 5; + int const lane = threadIdx.x & 31; + int const g = lane >> 2; + int const j = lane & 3; + constexpr int kChunkK = CODED ? 64 : 16; // k per chunk + constexpr int kWords = CODED ? 3 : 1; // 16-byte words per lane and chunk + constexpr int KU = CODED ? codedUnroll() : kUnroll; + int const K = kchunks * kChunkK; + int const M = p.M; + int const r_lo = p.tile_lo; + int const r_n = (p.tile_hi > 0 ? p.tile_hi : p.n_tiles) - r_lo; + bool const eq = p.bnd != nullptr; + // Whole tiles: CTA c owns tiles [t_begin, t_end). Equal chunks: CTA c streams chunks [c0, c1) of the range, which + // start off0 chunks into tile t_begin (its first and last tiles may be shared with the neighbouring CTAs). + long long const r_chunks = static_cast(r_n) * kchunks; + long long const c0 = eq ? (static_cast(blockIdx.x) * r_chunks) / gridDim.x : 0; + long long const c1 = eq ? (static_cast(blockIdx.x + 1) * r_chunks) / gridDim.x : 0; + int const t_begin = r_lo + + (eq ? static_cast(c0 / kchunks) + : static_cast((static_cast(blockIdx.x) * r_n) / gridDim.x)); + int const t_end = r_lo + + (eq ? static_cast((c1 + kchunks - 1) / kchunks) + : static_cast((static_cast(blockIdx.x + 1) * r_n) / gridDim.x)); + int const n_local = t_end - t_begin; + int const off0 = eq ? static_cast(c0 - static_cast(t_begin - r_lo) * kchunks) : 0; + int const span = eq ? static_cast(c1 - c0) : n_local * kchunks; + int const q_begin = off0 + warp * span / kWarps; + int const q_end = off0 + (warp + 1) * span / kWarps; + // Equal chunks: whether the first / last local tile is shared with the previous / next CTA of the launch. + bool const shared_first = eq && off0 != 0; + bool const shared_last = eq && ((c1 % kchunks) != 0); + long long const base = static_cast(t_begin) * kchunks; // first chunk of this CTA + uint4 const* const wf = p.wf; + + // Weights are static: start this warp's stream before waiting on the producer of x. The warp's first pf_chunks + // chunks come in by one bulk copy into shared memory (dynamic offset pf_off), which fills the grid-dependency + // wait and the staging; the register ring starts after them. + int const n_pf = CODED ? 0 : min(pf_chunks, q_end - q_begin); + uint4* const wpf = reinterpret_cast(reinterpret_cast(xs) + pf_off) + warp * pf_chunks * 32; + if (n_pf > 0 && lane == 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + uint32_t const bytes = static_cast(n_pf) * kChunkBytes; + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(bar) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(bar), "r"(bytes) : "memory"); + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + static_cast(__cvta_generic_to_shared(wpf))), + "l"(wf + (base + q_begin) * 32), "r"(bytes), "r"(bar) + : "memory"); + } + __syncwarp(); + int const q_reg = q_begin + n_pf; + uint4 buf[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q_reg + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = ldg_stream(wf + ((base + q_reg + u) * kWords + l) * 32 + lane); + } + } + for (int i = threadIdx.x; i < n_local * 32; i += blockDim.x) + acc_s[i / 32][i % 32] = make_float4(0.f, 0.f, 0.f, 0.f); + + asm volatile("griddepcontrol.wait;" ::: "memory"); + + // Stage the head input rows in B-fragment order: item (k16 step S, row gg) = two 16-byte loads of x, four + // fragment words {x[2jj], x[2jj + 1]}, {x[2jj + 8], x[2jj + 9]} (BF16 pairs, moved unchanged). Rows >= M are 0. + constexpr int kSU = MP >= 4 ? 4 : 2; + int const n_items = (K / 16) * MP; + for (int i0 = threadIdx.x; i0 < n_items; i0 += kSU * blockDim.x) + { + uint4 xv[kSU][2]; +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + int const gg = i % MP; + int const S = i / MP; + xv[u][0] = make_uint4(0u, 0u, 0u, 0u); + xv[u][1] = make_uint4(0u, 0u, 0u, 0u); + if (i < n_items && gg < M) + { + long long const xr = p.x_rows != nullptr ? p.x_rows[gg] : gg; + uint4 const* xp = reinterpret_cast(p.x + xr * K + 16 * S); + xv[u][0] = xp[0]; + xv[u][1] = xp[1]; + } + } +#pragma unroll + for (int u = 0; u < kSU; ++u) + { + int const i = i0 + u * blockDim.x; + if (i >= n_items) + continue; + int const gg = i % MP; + int const S = i / MP; + uint32_t const lo[4] = {xv[u][0].x, xv[u][0].y, xv[u][0].z, xv[u][0].w}; + uint32_t const hi[4] = {xv[u][1].x, xv[u][1].y, xv[u][1].z, xv[u][1].w}; +#pragma unroll + for (int jj = 0; jj < 4; ++jj) + xs[(S * 4 + jj) * MP + gg] = make_uint2(lo[jj], hi[jj]); + } + } + __syncthreads(); + + bool const b_active = g < MP; + float da[4] = {0.f, 0.f, 0.f, 0.f}; + float db[4] = {0.f, 0.f, 0.f, 0.f}; + int tl = q_begin / kchunks; // local tile of the current accumulator + int c = q_begin - tl * kchunks; // chunk within that tile + int cs = c; // first chunk of the warp's segment of that tile + int const tl0 = tl; // the warp's first local tile (its head piece) + Pieces pc; + if (n_pf > 0) + { + uint32_t const bar = static_cast(__cvta_generic_to_shared(&pf_bar[warp])); + asm volatile( + "{\n .reg .pred p;\n PF_WAIT%=:\n mbarrier.try_wait.parity.shared::cta.b64 p, [%0], 0;\n" + " @!p bra PF_WAIT%=;\n}\n" ::"r"(bar) + : "memory"); + for (int q = q_begin; q < q_reg; ++q) + { + chunk_mma(wpf[(q - q_begin) * 32 + lane], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, off0, span, q_begin, q_end), tl == tl0); + c = 0; + cs = 0; + ++tl; + } + } + } + for (int q = q_reg; q < q_end; q += KU) + { + uint4 nxt[KU][kWords]; +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + KU + u < q_end) + { +#pragma unroll + for (int l = 0; l < kWords; ++l) + nxt[u][l] = ldg_stream(wf + ((base + q + KU + u) * kWords + l) * 32 + lane); + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) + { + if (q + u < q_end) + { + if constexpr (CODED) + chunk_mma_coded(buf[u][0], buf[u][kWords > 1 ? 1 : 0], buf[u][kWords > 2 ? 2 : 0], p.cd.base, + xs, c, j, g, b_active, da, db); + else + chunk_mma(buf[u][0], xs, c, j, g, b_active, da, db); + if (++c == kchunks) + { + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, K, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, off0, span, q_begin, q_end), tl == tl0); + c = 0; + cs = 0; + ++tl; + } + } + } +#pragma unroll + for (int u = 0; u < KU; ++u) +#pragma unroll + for (int l = 0; l < kWords; ++l) + buf[u][l] = nxt[u][l]; + } + if (c != 0) + flush_tile(&acc_s[tl][lane], da, db, p, t_begin + tl, cs * kChunkK, c * kChunkK, xs, lane, pc, + p.fixedOrder, covers_tile(tl, kchunks, off0, span, q_begin, q_end), tl == tl0); + __syncthreads(); + // D-FIX: stage every warp's head / tail piece. xs and the weight prologue are dead after the barrier above (the + // launcher sizes them to hold the 16 x 2 x 32 pieces). + float4* const stage = reinterpret_cast(xs); + if (p.fixedOrder) + { + stage[(warp * 2) * 32 + lane] = pc.head; + stage[(warp * 2 + 1) * 32 + lane] = pc.tail; + __syncthreads(); + } + + // Accumulator lane (g, j) of a tile: rows g / g + 8 (x, y / z, w) of tokens 2j / 2j + 1 (x, z / y, w). + unsigned long long best0 = 0ull, best1 = 0ull; // tokens 2j, 2j + 1 + for (int t = warp; t < n_local; t += kWarps) + { + float4 v = p.fixedOrder ? ordered_tile(acc_s[t], stage, t, kchunks, off0, span, lane) : acc_s[t][lane]; + if ((t == 0 && shared_first) || (t == n_local - 1 && shared_last)) + { + // A tile shared with a neighbouring CTA: add this CTA's partial into bnd; the second CTA to arrive reads + // the sum (two addends), keys the tile and restores bnd and its counter to zero for the next call. + int const tile = t_begin + t; + float* const bp = reinterpret_cast(p.bnd + static_cast(tile) * 32 + lane); + atomicAdd(bp + 0, v.x); + atomicAdd(bp + 1, v.y); + atomicAdd(bp + 2, v.z); + atomicAdd(bp + 3, v.w); + __threadfence(); + __syncwarp(); + unsigned int prev = 0; + if (lane == 0) + prev = atomicAdd(p.bnd_cnt + tile, 1u); + prev = __shfl_sync(0xffffffffu, prev, 0); + if (prev == 0) + continue; // the neighbour keys it + __threadfence(); + v = __ldcg(p.bnd + static_cast(tile) * 32 + lane); + p.bnd[static_cast(tile) * 32 + lane] = make_float4(0.f, 0.f, 0.f, 0.f); + if (lane == 0) + p.bnd_cnt[tile] = 0u; + } + int const r0 = (t_begin + t) * 16 + g; + uint32_t const i0 = static_cast(p.vocab_offset + r0); + unsigned long long k0 = make_key(key_value(v.x), i0); + unsigned long long const k0b = make_key(key_value(v.z), i0 + 8); + unsigned long long k1 = make_key(key_value(v.y), i0); + unsigned long long const k1b = make_key(key_value(v.w), i0 + 8); + k0 = k0b > k0 ? k0b : k0; + k1 = k1b > k1 ? k1b : k1; + best0 = k0 > best0 ? k0 : best0; + best1 = k1 > best1 ? k1 : best1; + } +#pragma unroll + for (int o = 4; o < 32; o <<= 1) + { + best0 = shfl_max(best0, o); + best1 = shfl_max(best1, o); + } + asm volatile("griddepcontrol.launch_dependents;"); + if (lane < 4) + { + red[warp][2 * lane] = best0; + red[warp][2 * lane + 1] = best1; + } + __syncthreads(); + if (threadIdx.x < M) + { + unsigned long long b = 0ull; +#pragma unroll + for (int w = 0; w < kWarps; ++w) + b = red[w][threadIdx.x] > b ? red[w][threadIdx.x] : b; + if (b != 0ull) + atomicMax(p.keys + threadIdx.x, b); + } + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) + { + unsigned int const prev = atomicAdd(p.counters, 1u); + s_last = prev == (p.total_ctas != 0u ? p.total_ctas : gridDim.x) - 1u; + } + __syncthreads(); + if (!s_last) + return; + + __threadfence(); + unsigned int const call = *reinterpret_cast(p.counters + 1); + int const half = static_cast(call & 1u); + if (threadIdx.x < M) + { + int const t = threadIdx.x; + unsigned long long const mine = atomicExch(p.keys + t, 0ull); + unsigned long long* slot_mc = p.ws_mc + (half * kMaxTokens + t) * p.world; + unsigned long long* slot_uc = p.ws_uc + (half * kMaxTokens + t) * p.world; + slot_mc[p.rank] = mine; + unsigned long long best = 0ull; + for (int r = 0; r < p.world; ++r) + { + unsigned long long v; + long long spins = 0; + do + { + v = *reinterpret_cast(slot_uc + r); + if (++spins > kMaxSpins) + __trap(); // a missing peer key becomes a CUDA error, never a hang (the mc all-reduce's bound) + } while (v == 0ull); + best = v > best ? v : best; + slot_uc[r] = 0ull; + } + s_tok[t] = static_cast(0xFFFFFFFFu - static_cast(best & 0xFFFFFFFFull)); + if (p.out_keys != nullptr) + p.out_keys[t] = best; + } + __syncthreads(); + if (threadIdx.x == 0) + { + p.counters[0] = 0u; + p.counters[1] = call + 1u; + } + if (threadIdx.x < M) + { + StepArgs const& st = p.st; + int const b = threadIdx.x; + int const tok = s_tok[b]; + st.draft_out[b * st.draft_stride + st.step] = tok; + st.ids_next[b] = tok; + long long const r = p.x_rows != nullptr ? p.x_rows[b] : b; + st.pos_next[b] = st.pos_src[r] + 1; + if (st.step == 0 && st.num_acc != nullptr) + { + st.kv_lens[b] -= st.draft_len - st.num_acc[b]; + st.seq_lens[b] = 1; + } + else if (st.step > 0) + { + st.kv_lens[b] += 1; + } + } + if (p.hasFold) + { + // The fold tail (H-R2-3) after this pass's tokens and bookkeeping: the barrier makes this CTA's writes above + // visible to warp 0; every earlier operand was written by grids this one waited on. All CTAs triggered their + // dependents before the exchange, which is safe: a dependent's griddepcontrol.wait waits for this whole grid + // (this tail included) to complete, so only a dependent that read the tail's outputs before its own wait + // could race, and none does (the next step's readers run in the next graph replay). + __syncthreads(); + if (threadIdx.x < 32) + { + foldTailWarp(p.fold, static_cast(threadIdx.x)); + } + } +#endif +} + +template +cudaError_t launch(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + auto kern = draft_step_kernel; + size_t const xs_bytes = static_cast(K / 16) * 4 * MP * sizeof(uint2); + // Shared-memory weight prologue: chunks per warp that fit next to the kernel's static + dynamic shared memory. + static int const sStaticSmem = [] + { + cudaFuncAttributes fa{}; + return cudaFuncGetAttributes(&fa, draft_step_kernel) == cudaSuccess + ? static_cast(fa.sharedSizeBytes) + : -1; + }(); + static int const sMaxSmemPerBlock = [] + { + int dev = 0, maxSmem = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&maxSmem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + return maxSmem; + }(); + int pf = 0; + if (!CODED && sStaticSmem >= 0) + { + long long const room + = static_cast(sMaxSmemPerBlock) - sStaticSmem - static_cast(xs_bytes) - 1024; + long long const want = static_cast(kPrefetchKb) * 1024; + pf = static_cast(std::max(0ll, std::min(room, want)) / (kWarps * kChunkBytes)); + } + // At least the D-FIX piece staging (16 warps x head / tail x 32 lanes of float4), which reuses this space. + size_t const smem = std::max(xs_bytes + static_cast(pf) * kWarps * kChunkBytes, + static_cast(kWarps) * 2 * 32 * sizeof(float4)); + cudaError_t err = cudaFuncSetAttribute(kern, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem)); + if (err != cudaSuccess) + return err; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(static_cast(grid)); + cfg.blockDim = dim3(kWarps * 32); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, kern, p, K / (CODED ? 64 : 16), pf, static_cast(xs_bytes)); +} + +template +cudaError_t dispatchRows(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + if (p.M <= 1) + return launch<1, CODED>(p, K, grid, pdl, stream); + if (p.M <= 2) + return launch<2, CODED>(p, K, grid, pdl, stream); + if (p.M <= 4) + return launch<4, CODED>(p, K, grid, pdl, stream); + return launch<8, CODED>(p, K, grid, pdl, stream); +} + +} // namespace + +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream) +{ + bool const coded = p.cd.exc_off != nullptr; + if (p.M < 1 || p.M > kMaxTokens || K % (coded ? 64 : 16) != 0 || grid < 1) + return cudaErrorInvalidValue; + return coded ? dispatchRows(p, K, grid, pdl, stream) : dispatchRows(p, K, grid, pdl, stream); +} + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h new file mode 100644 index 000000000000..485f383732fe --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h @@ -0,0 +1,119 @@ +/* + * 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 "tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ +//! One pass of the fused greedy Eagle3 draft loop of the MiniMax-M3 NVFP4 TP4 decode path (sm_103a): the BF16 +//! vocab-sharded draft head, the cross-rank argmax and the draft loop's per-pass bookkeeping in one kernel. See +//! draftStep.cu. +namespace draftStep +{ + +constexpr int kMaxTokens = 8; // head rows (requests) of one launch +constexpr int kWarps = 16; // warps per CTA +constexpr int kMaxTilesPerCta = 32; // 16-row weight tiles one CTA can own +//! Shared-memory weight prologue per CTA, in KB (capped by the free shared memory). +constexpr int kPrefetchKb = 112; + +// Draft-pass epilogue (last CTA, one thread per request row b): stores the token as draft token `step`, stages the +// next pass's input id and position, and advances the draft KV lengths the way the in-tree draft loop does between +// passes (pass 0 with num_acc: drop the rejected verify rows and switch to one token per request; later passes: +1). +struct StepArgs +{ + int* draft_out; // [nreq, draft_stride]: column `step` receives the token + int draft_stride; + int step; + int* ids_next; // [nreq] + int const* pos_src; // positions of the rows the head read (indexed like x) + int* pos_next; // [nreq] = pos_src[x_rows[b] or b] + 1 + int* kv_lens; // [nreq] draft KV lengths, updated in place + int const* num_acc; // pass 0: accepted counts (nullptr: no rollback) + int* seq_lens; // pass 0: set to 1 + int draft_len; +}; + +// The 12-bit coded head (see draftStep.cu): the per-tile exception lists and the exponent base. exc_off == nullptr: +// the head is the plain BF16 fragment-ordered copy. +struct CodedArgs +{ + int const* exc_off; // [n_tiles + 1] exception list offsets + uint32_t const* exc_pos; // row in tile << 16 | k, per exception + float2 const* exc_w; // {w, 0} per exception: the exact weight + uint32_t base; // exponent base << 7 in both BF16 halves (base << 7 | base << 23) +}; + +// Everything a draft-step launch takes (see invokeDraftStep). +struct Params +{ + __nv_bfloat16 const* x; // [*, K] head input rows + long long const* x_rows; // optional [M]: head input row of request b (pass 0: the last accepted verify row) + int M; // requests, <= kMaxTokens + uint4 const* wf; // [n_tiles][K / 16][32 lanes] 16-byte BF16 A fragments, or the 12-bit coded stream + // [n_tiles][K / 64][3][32 lanes] when cd.exc_off is set (see draftStep.cu) + int n_tiles; + int vocab_offset; // global id of this rank's first head row + unsigned long long* keys; + unsigned int* counters; + unsigned long long* ws_uc; // this rank's exchange slots [2][kMaxTokens][world] + unsigned long long* ws_mc; // their multicast address + int rank; + int world; + unsigned long long* out_keys; // optional [M]: the winning global key of each row + StepArgs st; + CodedArgs cd; + // The step's forward-tail store + next-step staging (foldTail.h), run by the last CTA after this pass's + // tokens and bookkeeping: set on the drafter's last pass only (dispatch.DRAFT_STEP_FOLD_TAIL). + FoldTailParams fold{}; + bool hasFold = false; + // D-FIX: sum the pieces of a tile that several warps stream in warp order (deterministic) instead of by + // shared-memory atomics in arrival order (dispatch.DRAFT_STEP_FIXED_ORDER). + bool fixedOrder = false; + // A pass split over several launches (one per locality domain, each on that domain's SMs): this launch streams + // tiles [tile_lo, tile_hi) (tile_hi 0 = n_tiles), and the last of total_ctas CTAs over all of the pass's launches + // runs the exchange and the bookkeeping (0 = this launch's grid: one launch per pass). + int tile_lo = 0; + int tile_hi = 0; + unsigned int total_ctas = 0; + // Equal-chunk split (set both, zero-filled [n_tiles][32] float4 / [n_tiles], restored to zero by the kernel): each + // CTA of a launch streams an equal share of its range's (tile, chunk) space instead of whole tiles, so a tile at a + // CTA boundary gets one partial sum from each of its two CTAs, added in bnd (two addends: the order does not change + // the sum) and keyed by the second to arrive (bnd_cnt). nullptr: whole tiles per CTA. + float4* bnd = nullptr; + unsigned int* bnd_cnt = nullptr; +}; + +//! Launch the draft step over `grid` 512-thread CTAs (K = the hidden size, a multiple of 16; of 64 for a coded head). +cudaError_t invokeDraftStep(Params const& p, int K, int grid, bool pdl, cudaStream_t stream); + +} // namespace draftStep +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu new file mode 100644 index 000000000000..68279dc3bfaa --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu @@ -0,0 +1,189 @@ +/* + * 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. + */ + +// The TP8 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's GPT-NeoX +// path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin negated, +// dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale E4M3 +// store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query +// saturates the same way (448, not NaN). + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ +namespace drafterRopeFp8Kv +{ +namespace +{ + +constexpr int kElemsPerThread = kHeadDim / 32; +constexpr int kBlockSize = 256; + +template +__device__ __forceinline__ float rotate(float c, float x, float s, float pair) +{ + if constexpr (kRopeMode == 0) + { + return __fadd_rn(__fmul_rn(c, x), __fmul_rn(s, pair)); + } + else if constexpr (kRopeMode == 1) + { + return __fmaf_rn(c, x, __fmul_rn(s, pair)); + } + else + { + return __fmaf_rn(s, pair, __fmul_rn(c, x)); + } +} + +__device__ __forceinline__ void storeE4m3Pairs(__nv_fp8_e4m3* dst, __nv_bfloat16 const (&v)[kElemsPerThread]) +{ +#pragma unroll + for (int i = 0; i < kElemsPerThread; i += 2) + { + __nv_bfloat162 pair; + pair.x = v[i]; + pair.y = v[i + 1]; + __nv_fp8x2_e4m3 const packed(pair); + reinterpret_cast<__nv_fp8x2_storage_t*>(dst)[i / 2] = packed.__x; + } +} + +template +__global__ void __launch_bounds__(kBlockSize) drafterRopeFp8KvKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const lane = threadIdx.x % 32; + int const warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; + int const slots = p.numHeadsQ + p.numHeadsKv; + int const token = warp / slots; + if (token >= p.numTokens) + { + return; + } + int const slot = warp % slots; + bool const isQ = slot < p.numHeadsQ; + int const head = isQ ? slot : slot - p.numHeadsQ; + int const request = token / p.qLen; + int const row = token % p.qLen; + int const cacheLen = p.seqLens[request]; + int const kvPos = cacheLen - p.qLen + row; + int const ropePos + = p.specPositionOffsets != nullptr ? p.specPositionOffsets[request * p.qLen + row] + cacheLen - p.qLen : kvPos; + int const rowWidth = (p.numHeadsQ + 2 * p.numHeadsKv) * kHeadDim; + __nv_bfloat16 const* src + = p.qkv + static_cast(token) * rowWidth + (isQ ? head * kHeadDim : (p.numHeadsQ + head) * kHeadDim); + int const half = p.rotaryDim / 2; + float2 const* coef = p.cosSin + static_cast(ropePos) * half; + + __nv_bfloat16 rotated[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + int const d = lane * kElemsPerThread + i; + bool const firstHalf = d < half; + float const x = __bfloat162float(src[d]); + float const pair = __bfloat162float(src[firstHalf ? d + half : d - half]); + float2 const cs = d < p.rotaryDim ? coef[firstHalf ? d : d - half] : make_float2(1.0f, 0.0f); + float const s = firstHalf ? -cs.y : cs.y; + rotated[i] = __float2bfloat16(rotate(cs.x, x, s, pair)); + } + + if (isQ) + { + storeE4m3Pairs( + p.qOut + (static_cast(token) * p.numHeadsQ + head) * kHeadDim + lane * kElemsPerThread, rotated); + return; + } + + int const block = kvPos / p.tokensPerBlock; + long long const inPage = (static_cast(head) * p.tokensPerBlock + kvPos % p.tokensPerBlock) * kHeadDim + + lane * kElemsPerThread; + long long const pageElems = static_cast(p.numHeadsKv) * p.tokensPerBlock * kHeadDim; + int const* table = p.blockTable + request * p.blockTableRequestStride; + long long const kPage = table[block]; + long long const vPage = table[p.blockTableKvStride + block]; + storeE4m3Pairs(p.pool + kPage * pageElems + inPage, rotated); + + __nv_bfloat16 const* vSrc = p.qkv + static_cast(token) * rowWidth + + (p.numHeadsQ + p.numHeadsKv + head) * kHeadDim + lane * kElemsPerThread; + __nv_bfloat16 v[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) + { + v[i] = vSrc[i]; + } + storeE4m3Pairs(p.pool + vPage * pageElems + inPage, v); +} + +//! With PDL and maxSharedCarveout, the maximum shared-memory carveout lets the full-context decode's large-smem CTAs +//! (the dependent this kernel releases right after its wait) co-reside on the SMs this small kernel occupies. Set once +//! per instantiation. +template +cudaError_t launchWithMaxSharedCarveout( + cudaLaunchConfig_t const& config, Params const& params, bool pdl, bool maxSharedCarveout) +{ + auto const kernel = drafterRopeFp8KvKernel; + if (pdl && maxSharedCarveout) + { + static cudaError_t const carveout = cudaFuncSetAttribute( + kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); + if (carveout != cudaSuccess) + { + return carveout; + } + } + return cudaLaunchKernelEx(&config, kernel, params); +} + +} // namespace + +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream) +{ + int const warps = params.numTokens * (params.numHeadsQ + params.numHeadsKv); + int const warpsPerBlock = kBlockSize / 32; + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((warps + warpsPerBlock - 1) / warpsPerBlock); + config.blockDim = dim3(kBlockSize); + config.dynamicSmemBytes = 0; + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + switch (params.ropeMode) + { + case 0: return launchWithMaxSharedCarveout<0>(config, params, pdl, maxSharedCarveout); + case 1: return launchWithMaxSharedCarveout<1>(config, params, pdl, maxSharedCarveout); + case 2: return launchWithMaxSharedCarveout<2>(config, params, pdl, maxSharedCarveout); + default: return cudaErrorInvalidValue; + } +} + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h new file mode 100644 index 000000000000..f5316504784a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h @@ -0,0 +1,79 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ +namespace drafterRopeFp8Kv +{ + +constexpr int kHeadDim = 128; + +struct Params +{ + //! [tokens, (numHeadsQ + 2 * numHeadsKv) * 128] BF16: the drafter layer's packed Q, K, V rows. + __nv_bfloat16 const* qkv; + //! [positions, rotaryDim / 2] (cos, sin) pairs: the layer's rotary_cos_sin table. + float2 const* cosSin; + //! [requests] cache lengths including this step's rows. + int const* seqLens; + //! [requests, qLen] spec-decoding position offsets of a multi-token pass, or null. + int const* specPositionOffsets; + //! [tokens, numHeadsQ, 128] E4M3 query out. + __nv_fp8_e4m3* qOut; + //! Base of the layer's K / V sub-pages, [*, numHeadsKv, tokensPerBlock, 128] E4M3. + __nv_fp8_e4m3* pool; + //! [requests, 2, maxBlocks] sub-page indices: K in row 0, V in row 1. + int const* blockTable; + long long blockTableRequestStride; + long long blockTableKvStride; + int numTokens; + int qLen; + int numHeadsQ; + int numHeadsKv; + int rotaryDim; + int tokensPerBlock; + //! How the rotation's two products are summed: 0 two products then the add, 1 fma(cos, x, sin * pair), + //! 2 fma(sin, pair, cos * x). + int ropeMode; +}; + +//! The MiniMax-M3 TP8 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, +//! rounded to BF16 (the in-tree preprocess's q output and cache input), then Q, K and V to saturating E4M3: K / V +//! as the in-tree preprocess stores them at unit KV scale, Q as torch's CUDA cast of the BF16 query. Rows map to +//! (request token / qLen, row token % qLen) at cache position seqLen - qLen + row. +//! +//! PDL (pdl): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction and the kernel +//! releases its dependent right after it. The dependent may read only step-invariant data before its own wait. +cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream); + +} // namespace drafterRopeFp8Kv +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu new file mode 100644 index 000000000000..be98efebeb5d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu @@ -0,0 +1,64 @@ +/* + * 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/minimaxM3Sm107Tp8/foldTail.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +// Row i of the step (lane c = column), as the Triton fold tail's program i: the store's padded rows, then the next +// step's inputs read from them (zero past each source width). +__global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) +{ + int const i = blockIdx.x; + if (i >= *p.ctrl) + { + return; + } + foldTailRow(p, i, static_cast(threadIdx.x)); +} + +__global__ void __launch_bounds__(32) foldTailOneWarpKernel(FoldTailParams const p) +{ + foldTailWarp(p, static_cast(threadIdx.x)); +} + +} // namespace + +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream) +{ + foldTailKernel<<>>(p); + return cudaGetLastError(); +} + +cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream) +{ + foldTailOneWarpKernel<<<1, 32, 0, stream>>>(p); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h new file mode 100644 index 000000000000..86572b786a81 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h @@ -0,0 +1,129 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ + +//! Operands of the speculative step's forward-tail store and next-step staging (kernels/m3_fold_tail.py). Every tensor +//! is int32 except `slots` (int64); strides are in elements. The store views are [slots, width], the sources +//! [rows, width]. +struct FoldTailParams +{ + int32_t const* ctrl; // ctrl[0]: the rows the step armed (0: write nothing) + int64_t const* slots; + int32_t const* srcNew; + int64_t sn0, sn1; + int32_t snW; + int32_t* dstNew; + int64_t dn0, dn1; + int32_t dnW; + int32_t const* srcNext; + int64_t sx0, sx1; + int32_t sxW; + int32_t* dstNext; + int64_t dx0, dx1; + int32_t dxW; + int32_t const* srcDraft; + int64_t sd0, sd1; + int32_t sdW; + int32_t* dstDraft; + int64_t dd0, dd1; + int32_t ddW; + int32_t const* srcLens; + int32_t* dstLens; + int32_t* ids; // [rows * T] + int32_t* draftIds; // [rows * (T - 1)] + int32_t* posOff; // [rows * T] + int32_t* kvOff; // [rows] + int32_t tokensPerStep; +}; + +//! One CTA of 32 lanes per row, `rows` CTAs; the first ctrl[0] rows write. Widths and tokensPerStep must be <= 32. +cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream); + +//! The same writes from ONE warp (one CTA of 32 lanes looping over the rows, foldTailWarp): the in-kernel form a caller +//! such as the drafter's last-pass draft_step runs in its last CTA; launched alone it is that form's oracle. +cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream); + +#ifdef __CUDACC__ +//! Row i of the tail at lane c (0..31), the work of foldTailKernel's CTA i, for a caller that runs the tail inside its +//! own kernel (e.g. the drafter's last-pass draft_step in its last CTA, after its final tokens are written and fenced). +//! The caller runs rows i < *p.ctrl (one warp per row, or one warp looping over the rows); same code as the kernel, so +//! the writes are bit for bit the standalone tail's. +__device__ __forceinline__ void foldTailRow(FoldTailParams const& p, int i, int c) +{ + int64_t const s = p.slots[i]; + int32_t const nw = c < p.snW ? p.srcNew[i * p.sn0 + c * p.sn1] : 0; + if (c < p.dnW) + { + p.dstNew[s * p.dn0 + c * p.dn1] = nw; + } + int32_t const nxt = c < p.sxW ? p.srcNext[i * p.sx0 + c * p.sx1] : 0; + if (c < p.dxW) + { + p.dstNext[s * p.dx0 + c * p.dx1] = nxt; + } + int32_t const drf = c < p.sdW ? p.srcDraft[i * p.sd0 + c * p.sd1] : 0; + if (c < p.ddW) + { + p.dstDraft[s * p.dd0 + c * p.dd1] = drf; + } + int32_t const ln = p.srcLens[i]; + int32_t const t = p.tokensPerStep; + int32_t const d = t - 1; + if (c == 0) + { + p.dstLens[s] = ln; + p.kvOff[i] = ln - t; + } + if (c < t) + { + p.ids[i * t + c] = c < p.dxW ? nxt : 0; + p.posOff[i * t + c] = ln; + } + if (c < d) + { + p.draftIds[i * d + c] = c < p.ddW ? drf : 0; + } +} + +//! The whole tail from one warp of a running kernel (lane = threadIdx.x % 32): rows i < *p.ctrl in order. +__device__ __forceinline__ void foldTailWarp(FoldTailParams const& p, int lane) +{ + int const rows = *p.ctrl; + for (int i = 0; i < rows; ++i) + { + foldTailRow(p, i, lane); + } +} +#endif + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu new file mode 100644 index 000000000000..e7505fa48e13 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu @@ -0,0 +1,268 @@ +/* + * 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. + */ + +// The split-KV merge for the MiniMax-M3 NVFP4 TP2 (Rubin) full-context decode. The producer is the PrimTS split-KV +// decode with its separate-reduction scratch: per (request, KV head, split, row) one normalized 16-bit partial O and +// one fp32 log2-LSE. Active splits per row follow the producer's split ranges (the causal union of the row's grouped +// Q CTA, instruction-group aligned): splits past them hold nothing and are not read. +// +// One CTA per (request, KV head, kRowsPerCta rows); per row kGroups split groups x kThreadsPerRow threads x 8 dims. +// A thread issues the loads of every split it owns (strided over the groups) before using any, folds them under its +// group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to +// BF16 once. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ +namespace fullCtxReduce +{ +namespace +{ + +constexpr int kGroups = 8; +constexpr int kThreadsPerRow = kHeadDim / 8; +constexpr int kRowsPerCta = 1; +constexpr int kSplitsPerGroup = kMaxSplits / kGroups; +constexpr int kThreads = kGroups * kThreadsPerRow * kRowsPerCta; + +__device__ int activeSplits(Params const& p, int b, int row) +{ + int const seqKv = p.seqLensKv[b]; + int visible = seqKv; + if (p.causal) + { + int const token = row / p.headsPerKv; + int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; + int const end = min(base + p.qTokensPerCta, p.seqLenQ); + visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); + } + int const tiles = (visible + p.tileKv - 1) / p.tileKv; + int const tilesPerGroup = p.splits * p.instsKv; + int const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; + int const local = max(groups * p.instsKv, p.instsKv); + return (tiles + local - 1) / local; +} + +template +__device__ float2 unpack(uint32_t word) +{ + if constexpr (kBf16Partial) + { + return __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&word)); + } + else + { + return __half22float2(*reinterpret_cast<__half2 const*>(&word)); + } +} + +template +__global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("griddepcontrol.wait;" ::: "memory"); + // The dependent (the o_proj activation quantization / GEMV) reads `out` only after its own griddepcontrol.wait, + // which waits for this whole grid, so it is released here and runs its weight-only prologue under the merge. + asm volatile("griddepcontrol.launch_dependents;"); +#endif + int const b = blockIdx.z; + int const h = blockIdx.y; + int const localRow = threadIdx.x / (kGroups * kThreadsPerRow); + int const row = blockIdx.x * kRowsPerCta + localRow; + int const g = (threadIdx.x / kThreadsPerRow) % kGroups; + int const t = threadIdx.x % kThreadsPerRow; + bool const validRow = row < p.rows; + int const active = validRow ? activeSplits(p, b, row) : 0; + + __shared__ float sMax[kRowsPerCta][kGroups]; + __shared__ float sSum[kRowsPerCta][kGroups]; + __shared__ float sAcc[kRowsPerCta][kGroups][kHeadDim]; + + uint4 const* partialO = static_cast(p.partialO); + size_t const kvIdx = static_cast(b) * p.numKvHeads + h; + float lse[kSplitsPerGroup]; + uint4 vec[kSplitsPerGroup]; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + int const s = g + i * kGroups; + if (s < active) + { + size_t const r = (kvIdx * p.splits + s) * p.rows + row; + lse[i] = __ldcg(p.partialStats + r); + vec[i] = __ldcg(partialO + r * kThreadsPerRow + t); + } + } + float groupMax = -CUDART_INF_F; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active) + { + groupMax = fmaxf(groupMax, lse[i]); + } + } + // A split of the CTA's causal union can hold no key of this row (log2-LSE -inf). Such a split, and a group of + // only such splits, is skipped rather than weighted 0: its partial O carries no defined value. + float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + float sum = 0.f; +#pragma unroll + for (int i = 0; i < kSplitsPerGroup; ++i) + { + if (g + i * kGroups < active && lse[i] != -CUDART_INF_F) + { + float const w = exp2f(lse[i] - groupMax); + sum += w; + uint32_t const* words = reinterpret_cast(&vec[i]); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = unpack(words[j]); + acc[2 * j] += w * f.x; + acc[2 * j + 1] += w * f.y; + } + } + } + if (t == 0) + { + sMax[localRow][g] = groupMax; + sSum[localRow][g] = sum; + } +#pragma unroll + for (int j = 0; j < 8; ++j) + { + sAcc[localRow][g][t * 8 + j] = acc[j]; + } + __syncthreads(); + if (g != 0 || !validRow) + { + return; + } + float rowMax = -CUDART_INF_F; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + rowMax = fmaxf(rowMax, sMax[localRow][k]); + } + float total = 0.f; + float o[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; +#pragma unroll + for (int k = 0; k < kGroups; ++k) + { + if (sMax[localRow][k] == -CUDART_INF_F) + { + continue; + } + float const scale = exp2f(sMax[localRow][k] - rowMax); + total += scale * sSum[localRow][k]; +#pragma unroll + for (int j = 0; j < 8; ++j) + { + o[j] += scale * sAcc[localRow][k][t * 8 + j]; + } + } + float const inv = total > 0.f ? 1.f / total : 0.f; + int const token = row / p.headsPerKv; + int const head = row % p.headsPerKv; + __nv_bfloat162 packed[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + packed[j] = __floats2bfloat162_rn(o[2 * j] * inv, o[2 * j + 1] * inv); + } + size_t const outRow + = (static_cast(b) * p.seqLenQ + token) * (p.numKvHeads * p.headsPerKv) + h * p.headsPerKv + head; + reinterpret_cast(p.out + outRow * kHeadDim)[t] = *reinterpret_cast(packed); + if (p.q8 == nullptr) + { + return; + } + // MXFP8 copy of the BF16 row: a 32-column block is 4 neighbouring threads x 8 columns (lanes 0..15 of warp 0 + // hold the row). Block scale = the UE8M0 exponent of ceil(amax / 448), values = satfinite E4M3 of x * 2^-scale. + float x[8]; + float amax = 0.f; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + float2 const f = __bfloat1622float2(packed[j]); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + amax = fmaxf(amax, fmaxf(fabsf(f.x), fabsf(f.y))); + } + amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 2)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + uint16_t q2[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2[j]) : "f"(x[2 * j + 1] * qs), "f"(x[2 * j] * qs)); + } + int const tok = b * p.seqLenQ + token; + int const col = (h * p.headsPerKv + head) * kHeadDim + t * 8; + *reinterpret_cast(p.q8 + static_cast(tok) * (p.numSfKTiles * 128) + col) + = *reinterpret_cast(q2); + if (t % 4 == 0) + { + int const kSf = col / 32; + int64_t const off = static_cast(tok / 128) * p.numSfKTiles * 512 + (kSf / 4) * 512 + (tok % 32) * 16 + + ((tok % 128) / 32) * 4 + kSf % 4; + p.sf[off] = static_cast(e); + } +} + +} // namespace + +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) +{ + if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1) + { + return cudaErrorInvalidValue; + } + cudaLaunchConfig_t config = {}; + config.gridDim = dim3((params.rows + kRowsPerCta - 1) / kRowsPerCta, params.numKvHeads, params.batch); + config.blockDim = dim3(kThreads); + config.stream = stream; + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + return params.bf16Partial ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) + : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); +} + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h new file mode 100644 index 000000000000..1095ec4e020d --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h @@ -0,0 +1,77 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ +//! The split-KV merge of the MiniMax-M3 TP2 full-context decode (PrimTS split-KV decode, separate-reduction scratch) +//! straight to BF16. See fullCtxReduce.cu. +namespace fullCtxReduce +{ + +constexpr int kHeadDim = 128; +constexpr int kMaxSplits = 128; + +struct Params +{ + void const* partialO; // [B, Hkv, S, R, 128] normalized 16-bit partial O per split + float const* partialStats; // [B, Hkv, S, R] log2-LSE per split + int const* seqLensKv; // [B] + __nv_bfloat16* out; // [B, SQ, Hkv * hR, 128] + int batch; + int numKvHeads; + int splits; // S, the planned split capacity (<= kMaxSplits) + int rows; // R = hR * SQ, rows [token, local head] + int seqLenQ; // SQ + int headsPerKv; // hR + int qTokensPerCta; // query tokens one producer CTA owns (grouped Q) + int tileKv; // producer K/V tile + int instsKv; // producer K/V instructions per loop step + bool causal; + bool bf16Partial; + uint8_t* q8 = nullptr; // optional E4M3 copy of out [B * SQ, Hkv * hR * 128] (the o_proj MXFP8 hand-off) + uint8_t* sf = nullptr; // its UE8M0 block scales, 128x4-swizzled as trtllm::mxfp8_quantize(x, True) lays them out + int numSfKTiles = 0; // Hkv * hR * 128 / 128 +}; + +//! Merges the active splits of every row in fp32 with one max rescale and writes BF16; with q8 set, also the MXFP8 +//! copy of the BF16 rows (bit for bit trtllm::mxfp8_quantize of out). With pdl the launch is +//! programmatic: the grid waits on the producer (griddepcontrol.wait) before reading its scratch. +//! +//! PDL contract: griddepcontrol.wait is the kernel's first instruction, so nothing (scratch, lengths, out) is read or +//! written before every prerequisite grid has completed and flushed; a producer's early trigger +//! (griddepcontrol.launch_dependents) only lets this grid launch and wait sooner. The kernel triggers its own +//! dependents right after its wait: the next kernel's griddepcontrol.wait still waits for this whole grid, so its +//! post-wait code sees all of out, and its pre-wait prologue must not read out. +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream); + +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu new file mode 100644 index 000000000000..a5a157e33fe4 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu @@ -0,0 +1,1096 @@ +/* + * 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/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/tllmException.h" + +#include +#include +#include +#include + +// The sm_107 TP8 target's own copy of minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu (the version with the Eagle3 capture +// epilogue, cap_out), verbatim but for the names: the TP8 ext carries it, so the TP8 target no longer needs an image +// whose TP2 op has cap_out (SpecTail S8). The TP2 file's notes follow unchanged. +// +// SM107 (Rubin) copy of minimaxM3/mcAllreduceNorm.cu for the sm_107 TP8 target. Two changes in the arithmetic: the +// rank sum rounds to BF16 before the residual add, the rounding points of the stock all-reduce + fused residual +// RMSNorm on this target (bf16 all-reduce output, then bf16(residual + sum)); and every call but the MoE finalize runs +// the row kernel below, whose RMS statistic is the stock fused norm's reduction tree (bitwise the stock outputs). +// +// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction +// counts, block-scaled conversions): they are compiled for every sm_10x target of the build, the sm_100f family target +// included (which runs on SM107). A device pass for any other architecture compiles them to a trap, never to a silent +// empty body, and the host entry point checks that the kernel binary it launches holds sm_10x code (see +// checkBinary). +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 +#define MINIMAX_M3_KERNELS_ENABLED 1 +#else +#define MINIMAX_M3_KERNELS_ENABLED 0 +#endif + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +namespace cg = cooperative_groups; + +namespace +{ + +constexpr int kEmpty = static_cast(0x80000000u); +constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k + +// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], +// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c +// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp +// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its +// arithmetic is the same source. +constexpr int kCapThreads = 256; +constexpr int kCapIters = 3; +constexpr int kCapWarps = kCapThreads / 32; +constexpr int kCapHidden = kCapThreads * kCapIters * 8; + +// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). +__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) +{ + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int q = 0; q < 8; ++q) + { + float const v = __bfloat162float(e[q]); + ss += v * v; + } + return ss; +} + +__device__ __forceinline__ float cap_warp_sum(float v) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. +__device__ __forceinline__ void cap_store( + float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) +{ + uint4 const wr = *reinterpret_cast(w); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int q = 0; q < 8; ++q) + oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); + *reinterpret_cast(cap) = o; +} + +__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) +{ + if ((w & 0xFFFFu) == 0x8000u) + w &= 0xFFFF0000u; + if ((w >> 16) == 0x8000u) + w &= 0x0000FFFFu; + return w; +} + +__device__ __forceinline__ float bf16_lo(uint32_t w) +{ + return __uint_as_float(w << 16); +} + +// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's +// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals +// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the +// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. +__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) +{ + float value; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); + uint32_t result; + asm("{\n" + " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" + " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" + " setp.le.f32 p_zero, %1, 0f00000000;\n" + " mov.b32 bits, %1;\n" + " shr.b32 exp_biased, bits, 23;\n" + " and.b32 exp_biased, exp_biased, 255;\n" + " and.b32 mantissa, bits, 0x7FFFFF;\n" + " setp.ne.u32 p_has_mant, mantissa, 0;\n" + " selp.u32 bump, 1, 0, p_has_mant;\n" + " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" + " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" + " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" + " @p_tiny_sub mov.u32 bump, 0;\n" + " add.u32 r, exp_biased, bump;\n" + " setp.gt.u32 p_ovf, r, 254;\n" + " selp.u32 r, 254, r, p_ovf;\n" + " selp.u32 %0, 0, r, p_zero;\n" + "}\n" + : "=r"(result) + : "f"(value)); + return result; +} + +__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) +{ + float inv; + asm("{\n" + " .reg .s32 new_exp;\n" + " .reg .b32 float_bits;\n" + " .reg .pred p_zero;\n" + " setp.eq.u32 p_zero, %1, 0;\n" + " sub.s32 new_exp, 254, %1;\n" + " max.s32 new_exp, new_exp, 0;\n" + " shl.b32 float_bits, new_exp, 23;\n" + " mov.b32 %0, float_bits;\n" + " @p_zero mov.b32 %0, 0;\n" + "}\n" + : "=f"(inv) + : "r"(ue8m0)); + return inv; +} + +// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. +__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) +{ + uint16_t pair; + asm("{\n" + " .reg .f32 f0, f1;\n" + " mul.f32 f0, %1, %3;\n" + " mul.f32 f1, %2, %3;\n" + " min.f32 f0, f0, 0f43E00000;\n" + " max.f32 f0, f0, 0fC3E00000;\n" + " min.f32 f1, f1, 0f43E00000;\n" + " max.f32 f1, f1, 0fC3E00000;\n" + " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" + "}\n" + : "=h"(pair) + : "f"(lo), "f"(hi), "f"(inv)); + return pair; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ float bf16_hi(uint32_t w) +{ + return __uint_as_float(w & 0xFFFF0000u); +} + +__device__ __forceinline__ int4 ld_volatile_v4(int const* p) +{ + int4 v; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +struct FinalizeArgs +{ + __nv_bfloat16 const* fc2; // [permuted rows, H] routed-expert FC2 output (null: plain input x) + int const* perm; // [M, topk] expanded index -> permuted row, -1 = dropped + void const* scales; // [M, topk] routing weights (fp32 or bf16) + int scale_fp32; + int topk; +}; + +// grid = min(M, MAXT) * CLUSTER CTAs, cluster (CLUSTER, 1, 1), THREADS threads. +// OPTS (each bit-identical): 1 residual + norm weight loaded at entry (before the poll); 2 cluster RMS exchange by +// st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated right after the poll; 8 max-smem +// carveout so the next kernel's CTAs co-reside early. +// OPTS 16 (NOT bit-identical with the others; a different boundary): the Eagle3 drafter's boundary, whose stock is a +// plain BF16 all-reduce, then FlashInfer fused_add_rmsnorm: the RMS statistic and the normalization take the +// unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded to BF16). +// CAP: the Eagle3 capture epilogue (cap_out) is compiled in only where a call passes cap_out, so every other call +// runs the kernel without the capture branches, shared arrays or parameters in its body. +template +__global__ void __launch_bounds__(THREADS) mc_allreduce_norm_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, FinalizeArgs fin, int M, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int world, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ alignas(16) float cluster_sq[CLUSTER]; + extern __shared__ __align__(16) uint4 cap_slice[]; // [cols / 8] this CTA's updated chunks (cap_out only) + __shared__ alignas(8) unsigned long long xbar; + cg::cluster_group cluster = cg::this_cluster(); + if constexpr ((OPTS & 2) != 0) + { + if (threadIdx.x == 0) + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(b) : "memory"); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + } + int const crank = static_cast(cluster.block_rank()); + int const t = blockIdx.x / CLUSTER; + int const cols = H / CLUSTER; + int const c0 = crank * cols + threadIdx.x * 8; + bool const active = t < M && threadIdx.x * 8 < cols; + int64_t const half_words = static_cast(MAXT) * world * (H / 2); + + // prepushed 1: the producer (m3_mx_gemv) already pushed this rank's rows into the exchange slots; 2: and it + // launched this grid only after its pushes (its grid wait passed), so the slots are polled without a grid wait; + // 3: as 2, with x in set A and x2 in set B. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + int const parity = *reinterpret_cast(flags + blockIdx.x) & 1; + uint4 rraw = make_uint4(0u, 0u, 0u, 0u); + uint4 wraw = make_uint4(0u, 0u, 0u, 0u); + if constexpr ((OPTS & 1) != 0) + { + // the residual row was written grids earlier (prepushed >= 2: the producer launched this grid only after its + // own grid wait; otherwise the grid wait above) + if (active) + { + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + wraw = *reinterpret_cast(norm_w + c0); + } + } + // Push this rank's slice into every rank's buffer. + if (active && !prepushed) + { + uint4 raw; + if (fin.fc2 != nullptr) + { + // MoE finalize: fp32 sum of the routing-weighted top-k expert rows plus the shared expert, rounded once. + float f[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + if (fin.topk <= kFinMaxTopk) + { + // All top-k indices and scales, then every expert-row load in flight at once, then the same in-order + // FP32 accumulation as the loop below (one memory round trip instead of topk; bit-identical). + int prow[kFinMaxTopk]; + float sc[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + prow[k] = -1; + sc[k] = 0.f; + if (k < fin.topk) + { + int const e = t * fin.topk + k; + prow[k] = fin.perm[e]; + sc[k] = fin.scale_fp32 ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + } + } + uint4 rv[kFinMaxTopk]; +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] >= 0) + rv[k] = *reinterpret_cast(fin.fc2 + static_cast(prow[k]) * H + c0); + } +#pragma unroll + for (int k = 0; k < kFinMaxTopk; ++k) + { + if (prow[k] < 0) + continue; + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv[k]); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc[k]; + f[2 * q + 1] += v.y * sc[k]; + } + } + } + else + for (int k = 0; k < fin.topk; ++k) + { + int const e = t * fin.topk + k; + int const prow = fin.perm[e]; + if (prow < 0) + continue; + float const sc = fin.scale_fp32 + ? static_cast(fin.scales)[e] + : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); + uint4 const rv = *reinterpret_cast(fin.fc2 + static_cast(prow) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] += v.x * sc; + f[2 * q + 1] += v.y * sc; + } + } + if (x2 != nullptr) + { + // As stock: the routed output is rounded to BF16 first, then the shared expert is added (one more + // BF16 rounding of the FP32 sum). + uint4 const rv = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const routed = __bfloat1622float2(__floats2bfloat162_rn(f[2 * q], f[2 * q + 1])); + float2 const v = __bfloat1622float2(h[q]); + f[2 * q] = routed.x + v.x; + f[2 * q + 1] = routed.y + v.y; + } + } + __nv_bfloat162* o = reinterpret_cast<__nv_bfloat162*>(&raw); +#pragma unroll + for (int q = 0; q < 4; ++q) + o[q] = __floats2bfloat162_rn(f[2 * q], f[2 * q + 1]); + } + else + { + raw = *reinterpret_cast(x + static_cast(t) * H + c0); + } + if (fin.fc2 == nullptr && x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * world + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float s[8] = {}; + float sq = 0.f; + if (active) + { + // The residual row (written grids earlier) and every rank's slot are loaded together; slots still empty are + // re-polled. Ranks are then summed in rank order, as before. + if constexpr ((OPTS & 1) == 0) + rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + constexpr int kMaxWorld = 8; + int* const p0 = ws_uc + parity * half_words + static_cast(t) * world * (H / 2) + c0 / 2; + int4 v[kMaxWorld]; + int4 v2[kMaxWorld]; + bool const two = prepushed == 3; + int* const p1 = p0 + 2 * half_words; // set B + uint32_t pending = (1u << world) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < kMaxWorld; ++r) + { + if (r < world) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr ((OPTS & 16) != 0) + { + float const x0 = s[2 * q] + rf.x, x1 = s[2 * q + 1] + rf.y; + s[2 * q] = x0; + s[2 * q + 1] = x1; + sq += x0 * x0 + x1 * x1; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + sq += uf.x * uf.x + uf.y * uf.y; + } + } + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr ((OPTS & 16) == 0) + { + if constexpr (CAP) + cap_slice[threadIdx.x] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + if constexpr ((OPTS & 4) != 0) + { + // every CTA (active or not) rotates its call parity once its own poll is done (read only by later grids) + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + sq += __shfl_xor_sync(0xffffffffu, sq, off); + if ((threadIdx.x & 31) == 0) + warp_sq[threadIdx.x / 32] = sq; + __syncthreads(); + if constexpr ((OPTS & 2) != 0) + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); // peers' xbar initialized (long since) + if (threadIdx.x == 0) + { + float total = 0.f; +#pragma unroll + for (int w = 0; w < THREADS / 32; ++w) + total += warp_sq[w]; + if constexpr ((OPTS & 2) == 0) + for (int c = 0; c < CLUSTER; ++c) + *cluster.map_shared_rank(&cluster_sq[crank], c) = total; + else + { + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t const d = static_cast(__cvta_generic_to_shared(&cluster_sq[crank])); + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * 4) + : "memory"); + for (int c = 0; c < CLUSTER; ++c) + { + uint32_t rd, rb; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); + asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), + "r"(__float_as_uint(total)), "r"(rb) + : "memory"); + } + } + } + if constexpr ((OPTS & 2) == 0) + cluster.sync(); + else + { + // every peer's sum landed (st.async complete_tx on this CTA's barrier); no cluster-wide release-arrive + uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); + uint32_t done = 0; + while (!done) + asm volatile( + "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" + " selp.u32 %0, 1, 0, p;\n}" + : "=r"(done) + : "r"(b) + : "memory"); + } + if constexpr ((OPTS & 16) == 0) + { + if constexpr (CAP) + { + // The row's chunks are spread over the cluster (CTA c holds chunks [c * per, (c + 1) * per)); every CTA + // forms capture_norm's statistic itself, its 128 threads standing in for capture threads ct and ct + 128. + __shared__ float cap_red[kCapWarps]; + int const per = cols / 8; + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + float ca = 0.f, cb = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + { + int const ka = it * kCapThreads + static_cast(threadIdx.x); + int const kb = ka + kCapThreads / 2; + ca = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[ka % per], ka / per), ca); + cb = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[kb % per], kb / per), cb); + } + // Every peer read is done before any CTA moves on (and later exits). + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + ca = cap_warp_sum(ca); + cb = cap_warp_sum(cb); + if ((threadIdx.x & 31) == 0) + { + cap_red[threadIdx.x / 32] = ca; + cap_red[threadIdx.x / 32 + kCapWarps / 2] = cb; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + if (active) + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + if (active) + { + float total = 0.f; +#pragma unroll + for (int c = 0; c < CLUSTER; ++c) + total += cluster_sq[c]; + float const inv = rsqrtf(total / static_cast(H) + eps); + if constexpr ((OPTS & 1) == 0) + wraw = *reinterpret_cast(norm_w + c0); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + } + if (out_q != nullptr) + { + // MXFP8 copy of the normed row for the next GEMM, as the CuTe DSL mxfp8 quantize of `normed` (the tree's + // decode quantize, MX-TAC; see mxfp8_ue8m0): four threads hold one 32-column scale block. + float amax = 0.f; + if (active) + { +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + if (active) + { + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((threadIdx.x & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } + } + if constexpr ((OPTS & 4) == 0) + { + __syncthreads(); + if (threadIdx.x == 0) + flags[blockIdx.x] = parity ^ 1; + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +// Refuses (raises) to launch a kernel whose loaded binary holds no sm_10x code (PTX and SASS version >= 100): the +// device passes for other architectures compile the kernels to a trap, and a build that left out every sm_10x target +// must not reach the GPU at all. +template +bool checkBinary(Kernel kernel, char const* name) +{ + cudaFuncAttributes attr{}; + TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, reinterpret_cast(kernel))); + TLLM_CHECK_WITH_INFO(attr.ptxVersion >= 100 && attr.binaryVersion >= 100, + "%s: the loaded binary holds no sm_10x code (PTX %d, SASS %d): rebuild with an sm_100f / sm_107a target", name, + attr.ptxVersion, attr.binaryVersion); + return true; +} + +template +void launch(MinimaxM3Sm107Tp8McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) +{ + int const M = p.numTokens, H = p.hidden; + TLLM_CHECK_WITH_INFO(H % (CLUSTER * 8) == 0 && H / CLUSTER / 8 <= THREADS, "unsupported hidden size %d", H); + TLLM_CHECK_WITH_INFO(M >= 1 && M <= MAXT, "%d tokens do not fit the %d-token bucket", M, MAXT); + auto kernel = p.capOut != nullptr ? mc_allreduce_norm_kernel + : mc_allreduce_norm_kernel; + static bool const binaryOk + = checkBinary(mc_allreduce_norm_kernel, "mc_allreduce_norm_kernel") + && checkBinary( + mc_allreduce_norm_kernel, "mc_allreduce_norm_kernel (capture)"); + (void) binaryOk; + // Both instantiations (with and without the capture epilogue) get the same attributes, whichever runs first. + static bool const configured = []() + { + auto setup = [](auto k) + { + bool ok = CLUSTER <= 8 + || cudaFuncSetAttribute(k, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess; + if ((OPTS & 8) != 0) + ok = ok + && cudaFuncSetAttribute( + k, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared) + == cudaSuccess; + return ok; + }; + bool const ok = setup(mc_allreduce_norm_kernel) + && setup(mc_allreduce_norm_kernel); + if (!ok) + cudaGetLastError(); + return ok; + }(); + TLLM_CHECK_WITH_INFO(configured, "cluster size not supported"); + FinalizeArgs const fin{p.fc2, p.perm, p.scales, p.scaleFp32 ? 1 : 0, p.topk}; + cudaLaunchConfig_t cfg = {}; + // The grid is sized to the call's rows: min(M, MAXT) clusters. Each (row, slice) slot keeps its own parity flag, + // read by the producers before they push and rotated only by the slot's CTA after it polled: a slot that a call + // does not use keeps its parity, and every rank runs the same calls, so the parities stay consistent when M + // changes between calls. + cfg.gridDim = dim3((M < MAXT ? M : MAXT) * CLUSTER); + cfg.blockDim = dim3(THREADS); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO((OPTS & 16) == 0 && H == kCapHidden && 2 * THREADS == kCapThreads, + "the capture epilogue takes H %d on the 128-thread cluster kernel, not the drafter boundary", kCapHidden); + cfg.dynamicSmemBytes = static_cast(H / CLUSTER / 8) * sizeof(uint4); + } + cudaLaunchAttribute attr[2]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + attr[1].id = cudaLaunchAttributeClusterDimension; + attr[1].val.clusterDim.x = CLUSTER; + attr[1].val.clusterDim.y = 1; + attr[1].val.clusterDim.z = 1; + cfg.attrs = attr; + cfg.numAttrs = 2; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, fin, M, H, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, p.world, p.outNormed, p.outResidual, p.outQuant, p.outSf, (H / 32 + 3) / 4, + p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +// ws_uc / ws_mc are this rank's unicast and multicast views of ONE exchange buffer (they alias): no __restrict__. +// prepushed >= 2 contract: the producer (m3_mx_gemv's o_proj push, the m3_moe engine's finalize push) passes its own +// grid dependency, pushes its rows into the exchange, and only then releases this grid (griddepcontrol +// .launch_dependents); this kernel therefore skips griddepcontrol.wait and reads only the pushed (Lamport-polled) +// words, its parity flags and the residual / norm weight, which were written two or more grids earlier. +// +// The TP8 row kernel: one CTA of H / 8 = 768 threads per token row, used for every call but the MoE finalize. Its +// residual add + RMSNorm is TRT-LLM's rms_norm_kernel (customAllReduceKernels.cu), the stock fused norm after the +// all-reduce on this target, operation for operation at H 6144: thread j owns columns [8j, 8j + 8), its sum of squares +// is acc += v * v over them in order, then the warp butterfly 16..1, the 24 warp sums through shared memory into a +// second 32-lane butterfly (lanes >= 24 zero), rsqrtf(acc / H + eps) and (v * denom) * w. With the BF16 rank sum +// before the residual add, the normed rows, the updated residual and the MXFP8 copy are bitwise the stock ones. The +// exchange slots and the per-(row, slice) parity flags keep the cluster kernel's layout (slice = 96 threads), so +// producers that push (the o_proj GEMV) are unchanged. +constexpr int kSlices = 8; +constexpr int kSliceThreads = 96; +constexpr int kRowThreads = 768; + +// grid = M CTAs (one per row), THREADS = H / 8 threads. WORLD ranks (compile time). DRAFTER (opts bit 16): the Eagle3 +// drafter boundary, whose stock is a plain BF16 all-reduce then FlashInfer fused_add_rmsnorm: the RMS statistic and +// the normalization take the unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded). +// ws_uc / ws_mc alias (one exchange buffer): no __restrict__. +template +__global__ void __launch_bounds__(THREADS, 1) mc_allreduce_norm_row_kernel(__nv_bfloat16 const* __restrict__ x, + __nv_bfloat16 const* __restrict__ x2, int H, __nv_bfloat16 const* __restrict__ residual, + __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, + int maxt, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, + uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, + __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) +{ +#if MINIMAX_M3_KERNELS_ENABLED + __shared__ float warp_sq[THREADS / 32]; + __shared__ float total_sh; + extern __shared__ __align__(16) uint4 cap_row[]; // [THREADS] the updated row's chunks (cap_out only) + int const t = blockIdx.x; + int const j = threadIdx.x; + int const c0 = j * 8; + int const slice = j / kSliceThreads; + int64_t const half_words = static_cast(maxt) * WORLD * (H / 2); + if (prepushed < 2) + asm volatile("griddepcontrol.wait;" ::: "memory"); + int const parity = *reinterpret_cast(flags + t * kSlices + slice) & 1; + uint4 const rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); + uint4 const wraw = *reinterpret_cast(norm_w + c0); + if (!prepushed) + { + uint4 raw = *reinterpret_cast(x + static_cast(t) * H + c0); + if (x2 != nullptr) + { + uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); + __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); + __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const fa = __bfloat1622float2(a[q]); + float2 const fb = __bfloat1622float2(b[q]); + a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); + } + } + int64_t const slot = parity * half_words + (static_cast(t) * WORLD + rank) * (H / 2) + c0 / 2; + *reinterpret_cast(ws_mc + slot) + = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), + static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); + } + asm volatile("griddepcontrol.launch_dependents;"); + float s[8] = {}; + float sq = 0.f; + { + int* const p0 = ws_uc + parity * half_words + static_cast(t) * WORLD * (H / 2) + c0 / 2; + int* const p1 = p0 + 2 * half_words; // set B + bool const two = prepushed == 3; + int4 v[WORLD]; + int4 v2[WORLD]; + uint32_t pending = (1u << WORLD) - 1u; + uint32_t pending2 = two ? pending : 0u; + long long spins = 0; + while ((pending | pending2) != 0u) + { +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if ((pending >> r) & 1u) + v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); + if ((pending2 >> r) & 1u) + v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty + && v[r].w != kEmpty) + pending &= ~(1u << r); + if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty + && v2[r].w != kEmpty) + pending2 &= ~(1u << r); + } + if (++spins > (1LL << 31)) + __trap(); + } +#pragma unroll + for (int r = 0; r < WORLD; ++r) + { + *reinterpret_cast(p0 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), + static_cast(v[r].z), static_cast(v[r].w)}; + if (two) + { + *reinterpret_cast(p1 + static_cast(r) * (H / 2)) + = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); + uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), + static_cast(v2[r].z), static_cast(v2[r].w)}; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + __nv_bfloat162 const u + = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); + ww[q] = *reinterpret_cast(&u); + } + } +#pragma unroll + for (int q = 0; q < 4; ++q) + { + s[2 * q] += bf16_lo(ww[q]); + s[2 * q + 1] += bf16_hi(ww[q]); + } + } + // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. +#pragma unroll + for (int i = 0; i < 8; ++i) + s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); + __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); + uint32_t upd[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const rf = __bfloat1622float2(rr[q]); + __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); + upd[q] = *reinterpret_cast(&u); + if constexpr (DRAFTER) + { + s[2 * q] += rf.x; + s[2 * q + 1] += rf.y; + } + else + { + float2 const uf = __bfloat1622float2(u); + s[2 * q] = uf.x; + s[2 * q + 1] = uf.y; + } + } +#pragma unroll + for (int e = 0; e < 8; ++e) + sq += s[e] * s[e]; + *reinterpret_cast(out_residual + static_cast(t) * H + c0) + = make_uint4(upd[0], upd[1], upd[2], upd[3]); + if constexpr (!DRAFTER) + { + if constexpr (CAP) + cap_row[j] = make_uint4(upd[0], upd[1], upd[2], upd[3]); + } + } + // TRT-LLM reduce_fusion::block_reduce_sum. + sq += __shfl_xor_sync(~0, sq, 16); + sq += __shfl_xor_sync(~0, sq, 8); + sq += __shfl_xor_sync(~0, sq, 4); + sq += __shfl_xor_sync(~0, sq, 2); + sq += __shfl_xor_sync(~0, sq, 1); + if ((j & 31) == 0) + warp_sq[j / 32] = sq; + __syncthreads(); + { + float val = (j & 31) < THREADS / 32 ? warp_sq[j & 31] : 0.f; + val += __shfl_xor_sync(~0, val, 16); + val += __shfl_xor_sync(~0, val, 8); + val += __shfl_xor_sync(~0, val, 4); + val += __shfl_xor_sync(~0, val, 2); + val += __shfl_xor_sync(~0, val, 1); + if (j == 0) + total_sh = val; + } + // The parity of this row's slices rotates once every thread has read the exchange (later grids only). + if (j % kSliceThreads == 0) + flags[t * kSlices + slice] = parity ^ 1; + __syncthreads(); + if constexpr (!DRAFTER) + { + if constexpr (CAP) + { + // The whole row is in this CTA: threads 0..255 run capture_norm's chains over the staged chunks. + __shared__ float cap_red[kCapWarps]; + if (j < kCapThreads) + { + float css = 0.f; +#pragma unroll + for (int it = 0; it < kCapIters; ++it) + css = cap_chunk_sq(cap_row[it * kCapThreads + j], css); + css = cap_warp_sum(css); + if ((j & 31) == 0) + cap_red[j / 32] = css; + } + __syncthreads(); + float tot = 0.f; +#pragma unroll + for (int w = 0; w < kCapWarps; ++w) + tot += cap_red[w]; + float const rs = rsqrtf(tot / H + cap_eps); + cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); + } + } + float const inv = rsqrtf(total_sh / H + eps); + __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); + uint32_t o[4]; +#pragma unroll + for (int q = 0; q < 4; ++q) + { + float2 const wf = __bfloat1622float2(wv[q]); + __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); + o[q] = *reinterpret_cast(&n); + float2 const nf = __bfloat1622float2(n); + s[2 * q] = nf.x; + s[2 * q + 1] = nf.y; + } + *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); + if (out_q != nullptr) + { + float amax = 0.f; +#pragma unroll + for (int e = 0; e < 8; ++e) + amax = fmaxf(amax, fabsf(s[e])); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); + uint32_t const code = mxfp8_ue8m0(amax); + float const inv_scale = mxfp8_inv_scale(code); + uint32_t packed[2]; +#pragma unroll + for (int h = 0; h < 2; ++h) + { + packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) + | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); + } + *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); + if ((j & 3) == 0) + { + int const kb = c0 / 32; + out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 + + ((t % 128) / 32) * 4 + kb % 4] + = static_cast(code); + } + } +#elif defined(__CUDA_ARCH__) + __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op +#endif +} + +void launchRow( + MinimaxM3Sm107Tp8McAllreduceNormParams const& p, int32_t maxTokens, bool drafter, bool pdl, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(p.hidden == kRowThreads * 8 && p.world == 2, "the row kernel takes H 6144 at 2 ranks"); + static bool const binaryOk + = checkBinary(mc_allreduce_norm_row_kernel<2, kRowThreads, false, false>, "mc_allreduce_norm_row_kernel") + && checkBinary( + mc_allreduce_norm_row_kernel<2, kRowThreads, true, false>, "mc_allreduce_norm_row_kernel (drafter)") + && checkBinary( + mc_allreduce_norm_row_kernel<2, kRowThreads, false, true>, "mc_allreduce_norm_row_kernel (capture)"); + (void) binaryOk; + TLLM_CHECK_WITH_INFO(p.numTokens >= 1 && p.numTokens <= maxTokens, "%d tokens do not fit the %d-token bucket", + p.numTokens, maxTokens); + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(p.numTokens); + cfg.blockDim = dim3(kRowThreads); + cfg.stream = stream; + if (p.capOut != nullptr) + { + TLLM_CHECK_WITH_INFO( + !drafter && p.hidden == kCapHidden, "the capture epilogue takes H %d, not the drafter", kCapHidden); + cfg.dynamicSmemBytes = static_cast(kRowThreads) * sizeof(uint4); + } + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + // The capture epilogue never runs on the drafter boundary (checked above), so there are three instantiations. + auto kernel = drafter ? mc_allreduce_norm_row_kernel<2, kRowThreads, true, false> + : p.capOut != nullptr ? mc_allreduce_norm_row_kernel<2, kRowThreads, false, true> + : mc_allreduce_norm_row_kernel<2, kRowThreads, false, false>; + cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, p.hidden, p.residual, p.normWeight, p.eps, + p.wsUc, p.wsMc, p.flags, p.rank, maxTokens, p.outNormed, p.outResidual, p.outQuant, p.outSf, + (p.hidden / 32 + 3) / 4, p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); + if (err != cudaSuccess) + { + cudaGetLastError(); + TLLM_CUDA_CHECK(err); + } +} + +constexpr int kCluster = 8; +constexpr int kThreads = 128; + +} // namespace + +bool isMinimaxM3Sm107Tp8McAllreduceNormSupported(int32_t maxTokens, int32_t opts) +{ + return (maxTokens == 8 && (opts == 7 || opts == 15 || opts == 23 || opts == 31)) || (maxTokens == 16 && opts == 15) + || (maxTokens == 48 && opts == 15); +} + +void invokeMinimaxM3Sm107Tp8McAllreduceNorm(MinimaxM3Sm107Tp8McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream) +{ + if (params.fc2 == nullptr && params.world == 2 && params.hidden == kRowThreads * 8) + { + launchRow(params, maxTokens, (opts & 16) != 0, pdl, stream); + return; + } + if (maxTokens == 8 && opts == 7) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 23) + { + launch(params, pdl, stream); + } + else if (maxTokens == 8 && opts == 31) + { + launch(params, pdl, stream); + } + else if (maxTokens == 16 && opts == 15) + { + launch(params, pdl, stream); + } + else if (maxTokens == 48 && opts == 15) + { + launch(params, pdl, stream); + } + else + { + TLLM_THROW("unsupported (bucket %d, opts %d)", maxTokens, opts); + } +} + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h new file mode 100644 index 000000000000..4b36daf66813 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h @@ -0,0 +1,83 @@ +/* + * 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 +{ + +//! Operands of the MiniMax-M3 multicast all-reduce + residual + RMSNorm kernel (decode, at most 48 tokens). +//! +//! Per rank the kernel forms the partial `bf16(x + x2)` (or, in finalize mode, +//! `bf16(bf16(sum_k scale * fc2[perm]) + x2)`, the routed sum rounded before the shared expert is added, as stock), +//! pushes it through the multicast mapping `wsMc` into every rank's exchange, polls its own unicast view `wsUc`, +//! sums the ranks in FP32 in rank order, adds the residual (BF16 round) and applies RMSNorm. With `outQuant` it +//! also writes the MXFP8 copy of the normed rows (E4M3 [numTokens, hidden] and swizzled UE8M0 scales). +struct MinimaxM3Sm107Tp8McAllreduceNormParams +{ + __nv_bfloat16 const* x{nullptr}; //!< [numTokens, hidden] partial; in finalize mode the FC2 rows [rows, hidden]. + __nv_bfloat16 const* x2{nullptr}; //!< Optional second addend [numTokens, hidden]. + __nv_bfloat16 const* fc2{nullptr}; //!< Finalize mode: the permuted FC2 output (same buffer as x); null otherwise. + int32_t const* perm{nullptr}; //!< Finalize mode: [numTokens, topk] expanded index -> permuted row, -1 dropped. + void const* scales{nullptr}; //!< Finalize mode: [numTokens, topk] routing weights (fp32 or bf16). + bool scaleFp32{false}; + int32_t topk{0}; + int32_t numTokens{0}; + int32_t hidden{0}; + __nv_bfloat16 const* residual{nullptr}; //!< [numTokens, hidden] + __nv_bfloat16 const* normWeight{nullptr}; //!< [hidden] + float eps{0.F}; + int32_t* wsUc{nullptr}; //!< This rank's unicast view of the exchange (int32 words, re-armed after reading). + int32_t* wsMc{nullptr}; //!< The multicast mapping of the exchange. + int32_t* flags{nullptr}; //!< Per-CTA call-parity counters, [maxTokens * 8]. + int32_t rank{0}; + int32_t world{0}; + __nv_bfloat16* outNormed{nullptr}; + __nv_bfloat16* outResidual{nullptr}; + uint8_t* outQuant{nullptr}; //!< Optional E4M3 [numTokens, hidden]. + uint8_t* outSf{nullptr}; //!< Optional UE8M0 scales, 128x4-swizzled. + //! 0: push x (+ x2) here; 2: a producer already pushed x and launched this grid after its pushes; 3: as 2, with x + //! in exchange set A and x2 in set B. + int32_t prepushed{0}; + //! Optional Eagle3 capture of a capture layer's boundary: capOut[t * capStride + c] = the m3_spec_tail + //! capture_norm of the updated residual row with weight capWeight [hidden] and capEps, bit for bit (hidden 6144, + //! not the drafter boundary). + __nv_bfloat16* capOut{nullptr}; + int64_t capStride{0}; + __nv_bfloat16 const* capWeight{nullptr}; + float capEps{0.F}; +}; + +//! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). +bool isMinimaxM3Sm107Tp8McAllreduceNormSupported(int32_t maxTokens, int32_t opts); + +//! Launch with 8-CTA clusters of 128 threads, one cluster per token. `maxTokens` selects the exchange bucket; +//! `opts` the epilogue variant (every variant is bit-identical). +void invokeMinimaxM3Sm107Tp8McAllreduceNorm(MinimaxM3Sm107Tp8McAllreduceNormParams const& params, int32_t maxTokens, + int32_t opts, bool pdl, cudaStream_t stream); + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu new file mode 100644 index 000000000000..3062080fac2c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu @@ -0,0 +1,340 @@ +/* + * 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/minimaxM3Sm107Tp8/producerPw8.h" + +#include +#include +#include + +// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP8 heads. +// +// The SM107 TP8 target's copy of the SM107 TP2 producerPw2.cu (itself a copy of minimaxM3/producerPw.cu, which the +// sm_103 targets keep), for its pre-wait step inputs. Same math and +// outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with --use_fast_math like +// the rest of the library), specialized to 8 query, 1 KV and 1 index head per rank (one index-K head), head_dim 128, +// rotary_dim 64, page 128, NVFP4 main KV and FP8 index-K. One warp per (token, packed head); K and V rows go to their +// head's block of the HND pools (the native kernel's head strides). +// +// The decode producer is latency bound: after griddepcontrol.wait the native kernel runs a chain of dependent global +// loads (packed row -> norm weight -> position id -> RoPE row -> cache slot -> quant scale). This kernel reads the +// constant norm weights and KV quant scale before griddepcontrol.wait and issues the packed row, position id and cache +// slot loads together after it, followed by the RoPE row; with `prewaitStepInputs` the position id, the slot and the +// RoPE row are read before the wait too, leaving only the packed row after it. With `early` it also calls +// griddepcontrol.launch_dependents on entry: consumers read its outputs only after their own griddepcontrol.wait, which +// waits for this grid to complete, so the early trigger only lets the next kernel launch and run its pre-wait prologue +// while this one runs. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ +namespace producerPw8 +{ +namespace +{ + +constexpr int kHeadDim = 128; +constexpr int kRotaryDim = 64; +constexpr int kPageSize = 128; +constexpr int kElems = kHeadDim / 32; +constexpr int kBlock = 256; + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ uint32_t sf_group_mask4() +{ + uint32_t laneId; + asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); + return 0xFu << (laneId & ~3u); +} + +__device__ __forceinline__ uint32_t fp32x8_to_e2m1(float (&a)[8]) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t val; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + ".reg .b8 byte2;\n" + ".reg .b8 byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "}" + : "=r"(val) + : "f"(a[0]), "f"(a[1]), "f"(a[2]), "f"(a[3]), "f"(a[4]), "f"(a[5]), "f"(a[6]), "f"(a[7])); + return val; +#else + return 0; +#endif +} + +__device__ __forceinline__ void store_fp8x4(__nv_fp8_e4m3* out, int64_t off, float const (&e)[kElems]) +{ + __nv_fp8x2_e4m3 const low(make_float2(e[0], e[1])); + __nv_fp8x2_e4m3 const high(make_float2(e[2], e[3])); + *reinterpret_cast(out + off) = static_cast(low.__x) | (static_cast(high.__x) << 16); +} + +__device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float secondLevelScale, uint8_t* scaleOut) +{ + __nv_bfloat16 rounded[4]; + float localMax = 0.0F; +#pragma unroll + for (int i = 0; i < 4; ++i) + { + rounded[i] = __float2bfloat16_rn(e[i]); + localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); + } + uint32_t const groupMask = sf_group_mask4(); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); + localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); + float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); + __nv_fp8_e4m3 const encodedScale(scaleValue); + if (scaleOut != nullptr) + *scaleOut = encodedScale.__x; + float const narrowedScale = static_cast(encodedScale); + float const outputScale + = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; + float q[8] = {}; +#pragma unroll + for (int i = 0; i < 4; ++i) + q[i] = __bfloat162float(rounded[i]) * outputScale; + return static_cast(fp32x8_to_e2m1(q)); +} + +template +__global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if constexpr (kEarly) + { + asm volatile("griddepcontrol.launch_dependents;"); + } +#endif + int const warpId = threadIdx.x / 32; + int const laneId = threadIdx.x % 32; + int const globalWarp = blockIdx.x * (kBlock / 32) + warpId; + int const tokenIdx = globalWarp / kTotalHeads; + int const localHead = globalWarp % kTotalHeads; + if (tokenIdx >= p.numTokens) + return; + + constexpr int kBegin = kHeadsQ; + constexpr int vBegin = kBegin + kHeadsKV; + constexpr int indexQBegin = vBegin + kHeadsKV; + constexpr int indexKHead = indexQBegin + kHeadsIndex; + bool const isQ = localHead < kBegin; + bool const isK = localHead >= kBegin && localHead < vBegin; + bool const isV = localHead >= vBegin && localHead < indexQBegin; + bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; + bool const isIndexK = localHead == indexKHead; + + // Constant data read before the wait: norm weights and the KV quant scale. + float w[kElems] = {}; + float quantScale = 0.0F; + __nv_bfloat16 const* normWeight + = isQ ? p.qWeight : (isK ? p.kWeight : (isIndexQ ? p.indexQWeight : p.indexKWeight)); + if (!isV) + { + uint2 const wv = *reinterpret_cast(normWeight + laneId * kElems); + float2 const w01 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.x)); + float2 const w23 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.y)); + w[0] = w01.x; + w[1] = w01.y; + w[2] = w23.x; + w[3] = w23.y; + } + if (isK || isV) + { + quantScale = p.kvQuantScale[isV ? 2 : 1]; + } + + int64_t const inputOffset = (static_cast(tokenIdx) * kTotalHeads + localHead) * kHeadDim + laneId * kElems; + bool const rotates = !isV && laneId * kElems < kRotaryDim; + bool const slotted = isK || isV || isIndexK; + int const coefficient = (laneId * kElems) % (kRotaryDim / 2); + // kPrewait: the position id, the cache slot and the RoPE row are read before the wait. They are per-step metadata + // written before the replay (host-to-device) or by plain-launched kernels earlier in the step (position_ids by the + // step's input preprocessing, the slots by msa_kv_update or its torch fallback), so they are final before this + // grid can start whatever its predecessors trigger; the RoPE table is constant. + int positionId = 0; + int slot = 0; + float4 c = make_float4(0.0F, 0.0F, 0.0F, 0.0F); + float4 s = c; + auto const loadStepInputs = [&]() + { + positionId = rotates ? p.positionIds[tokenIdx] : 0; + slot = slotted ? p.outCacheLoc[tokenIdx] : 0; + if (rotates) + { + int64_t const ropeRow = static_cast(positionId) * kRotaryDim; + c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); + s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); + } + }; + if constexpr (kPrewait) + { + loadStepInputs(); + } +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + // The packed row comes from the predecessor grid; without kPrewait the step inputs are issued with it, so the + // native kernel's serial chain of dependent loads becomes two round trips. + uint2 const packed = *reinterpret_cast(p.packed + inputOffset); + if constexpr (!kPrewait) + { + loadStepInputs(); + } + float cosine[kElems] = {c.x, c.y, c.z, c.w}; + float sine[kElems] = {s.x, s.y, s.z, s.w}; + + float e[kElems]; + float sumSquares = 0.0F; +#pragma unroll + for (int pair = 0; pair < 2; ++pair) + { + float2 const v = __bfloat1622float2( + *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); + e[2 * pair] = v.x; + e[2 * pair + 1] = v.y; + if (!isV) + sumSquares += v.x * v.x + v.y * v.y; + } + + if (!isV) + { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) + sumSquares += __shfl_xor_sync(0xffffffff, sumSquares, mask, 32); + float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kHeadDim) + p.eps); +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + e[i] *= rmsReciprocal * (1.0F + w[i]); + } + constexpr int kPairOffset = (kRotaryDim / 2) / kElems; +#pragma unroll + for (int i = 0; i < kElems; ++i) + { + int const dim = laneId * kElems + i; + float paired = __shfl_xor_sync(0xffffffff, e[i], kPairOffset); + if (dim < kRotaryDim) + { + if (dim < kRotaryDim / 2) + paired = -paired; + // The native kernel's contraction of e * cos + paired * sin; spelled out so this build matches it. + e[i] = __fmaf_rn(paired, sine[i], __fmul_rn(e[i], cosine[i])); + } + } + } + + if (isQ) + { + store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); + return; + } + if (isIndexQ) + { + store_fp8x4(p.indexQOut, + (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); + return; + } + + if (slot < 0) + return; + int const page = slot / kPageSize; + if (page >= p.numPages) + return; + int const withinPage = slot & (kPageSize - 1); + if (isIndexK) + { + store_fp8x4(p.indexK, + static_cast(page) * p.indexPageStride + static_cast(withinPage) * p.indexTokenStride + + laneId * kElems, + e); + return; + } + + int const head = isK ? localHead - kBegin : localHead - vBegin; + int const plane = isV ? 1 : 0; + int const scaleColumn = laneId / 4; + int64_t const scaleWithinPage = isV + ? static_cast(withinPage / 4) * 32 + static_cast(scaleColumn) * 4 + withinPage % 4 + : static_cast(withinPage) * 8 + scaleColumn; + int64_t const scaleOffset = static_cast(page) * p.scalePageStride + + static_cast(plane) * p.scalePlaneStride + static_cast(head) * p.scaleHeadStride + + scaleWithinPage; + uint8_t* scaleOut = laneId % 4 == 0 ? p.kvScale + scaleOffset : nullptr; + uint16_t const packedFp4 = quantize_fp4x4(e, quantScale, scaleOut); + int64_t const dataOffset = static_cast(page) * p.dataPageStride + + static_cast(plane) * p.dataPlaneStride + static_cast(head) * p.dataHeadStride + + static_cast(withinPage) * p.dataTokenStride + laneId * 2; + *reinterpret_cast(p.kvData + dataOffset) = packedFp4; +#endif +} + +template +cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) +{ + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; + int const warps = numTokens * kTotalHeads; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); + cfg.blockDim = dim3(kBlock); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); +} + +} // namespace + +cudaError_t invokeProducerPw8( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream) +{ + if (numTokens <= 0) + return cudaSuccess; + if (prewaitStepInputs) + { + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); + } + return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); +} + +} // namespace producerPw8 +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h new file mode 100644 index 000000000000..5b269b94587c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h @@ -0,0 +1,76 @@ +/* + * 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 +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ +//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP8 heads (8 query, 1 KV and +//! 1 index head per rank) with its constant loads before the PDL grid-dependency wait and an optional early dependents +//! trigger. Same outputs, bit for bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw8.cu. +namespace producerPw8 +{ + +struct Params +{ + __nv_bfloat16 const* packed; + __nv_fp8_e4m3* qOut; + __nv_fp8_e4m3* indexQOut; + uint8_t* kvData; + uint8_t* kvScale; + __nv_fp8_e4m3* indexK; + int const* outCacheLoc; + float const* kvQuantScale; + int64_t dataPageStride, dataPlaneStride, dataHeadStride, dataTokenStride; + int64_t scalePageStride, scalePlaneStride, scaleHeadStride; + int64_t indexPageStride, indexTokenStride, numPages; + int numTokens; + float eps; + __nv_bfloat16 const* qWeight; + __nv_bfloat16 const* kWeight; + __nv_bfloat16 const* indexQWeight; + __nv_bfloat16 const* indexKWeight; + float const* rotaryCosSin; + int const* positionIds; +}; + +//! The heads per rank this producer is built for. +constexpr int kHeadsQ = 8; +constexpr int kHeadsKV = 1; +constexpr int kHeadsIndex = 1; + +//! Launch over `numTokens` tokens, one warp per (token, packed head). `prewaitStepInputs` reads the position ids, the +//! cache slots and the RoPE rows before the grid-dependency wait (they must not be written by a PDL predecessor). +cudaError_t invokeProducerPw8( + Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream); + +} // namespace producerPw8 +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu new file mode 100644 index 000000000000..6e5f943b2438 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu @@ -0,0 +1,946 @@ +/* + * 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/minimaxM3Sm107Tp8/sampledAccept.h" + +#include "tensorrt_llm/common/assert.h" + +#include +#include + +// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). +// +// sm_107 (Rubin) TP8-owned copy of kernels/minimaxM3Sm107Tp2/sampledAccept.cu (bit for bit the same kernel: the +// seam for the TP8 SA-1R round-0 decision), itself the TP2 copy of kernels/minimaxM3/sampledAccept.cu: identical +// algorithm and arithmetic; the vocab shard limit is raised to the TP2 shard (200064 / 2 = 100032 -> kMaxVs 100352), so +// a CTA stages up to kMaxVs / kSplits BF16 logits (25 KB at 8 splits, 50 KB at 4) and the launch opts in to that +// dynamic shared memory. +// +// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +// active must stay on the generic path; the caller decides. +// +// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens +// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the +// nucleus with probability >= top_p, so most rows decide on their second exchange. +// +// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based +// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one +// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and +// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the +// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing +// and two call halves keep a rank one call ahead from overwriting a slot still being read. +// +// SA-1R (Args.one_round_k 9 / 82; MoE prototype moe/sa1r/6c753632ec, exact on 39,968 rows, eager + graph): each CTA +// also multicasts its largest bf16 logits with the round-0 record (3 per extra word). After round 0 every CTA knows M, +// the candidate logit cv and wc = __expf(cv - M); round 1 would accept iff S <= target, S = the fp32 sum (its tree) of +// w = __expf(__fmul_rn(x, tInv) - M) over every token with w > wc. x -> w is monotone (__fmul_rn by tInv > 0, fsub, +// __expf: checked exhaustively over all non-positive fp32 inputs), so if no CTA's unlisted value can beat wc (82: the +// bound tau; 9: the 9th listed), the listed w > wc are exactly round 1's summands, bitwise. Two fp32 sums of the same +// m non-negative terms differ by < m * 2^-23 * S, so accepting iff S_1r + m * 2^-21 * S_1r (round up) <= target +// implies round 1 accepts; m <= 1 is exact. Otherwise round 1 runs as before. An accepted row outputs the round-0 +// candidate, as the two-round kernel does, and consumes no Philox draw it would not. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ +namespace sampledAccept +{ +namespace +{ +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kMaxRec = kMaxWorld * kMaxSplits; +constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) +constexpr int kTagBits = 64 - 32 - kIdBits; +constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard +// Record polls trap after this many spins (the multicast all-reduce's bound): a missing writer becomes a CUDA error, +// never a hang. +constexpr long long kMaxSpins = 1LL << 31; + +// SA-1R variant kTopK: 0 none; K exact top-K list; 10 * K + P pooled (each thread keeps its top P plus its (P + 1)-th +// as a bound; the CTA lists the pool's top K plus tau = max(K-th listed, max thread bound)). +__host__ __device__ constexpr int mode_k(int m) +{ + return m >= 10 ? m / 10 : m; +} + +__host__ __device__ constexpr int mode_pool(int m) +{ + return m >= 10 ? m % 10 : 0; +} + +__host__ __device__ constexpr int mode_vals(int m) +{ + return mode_k(m) + (mode_pool(m) ? 1 : 0); +} + +// Words per record: key, logit of key, f2, f3, then the round-0 list (3 bf16 + a 16-bit tag per word). +template +struct Rec +{ + static constexpr int kXw = (mode_vals(kTopK) + 2) / 3; + static constexpr int kEx = 4 + kXw; +}; + +// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. +template +__host__ __device__ inline long long half_words(int world) +{ + return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * Rec::kEx; +} + +__device__ __forceinline__ uint32_t ordered(float v) +{ + uint32_t const b = __float_as_uint(v); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +// bf16 bits <-> an order-preserving 16-bit key. +__device__ __forceinline__ uint32_t ord16(uint32_t b) +{ + return (b & 0x8000u) ? (~b & 0xFFFFu) : (b | 0x8000u); +} + +__device__ __forceinline__ uint32_t unord16(uint32_t k) +{ + return (k & 0x8000u) ? (k & 0x7FFFu) : (~k & 0xFFFFu); +} + +// Philox4x32-10. +__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) +{ +#pragma unroll + for (int i = 0; i < 10; ++i) + { + uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; + uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; + c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); + k.x += 0x9E3779B9u; + k.y += 0xBB67AE85u; + } + return c; +} + +// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. +__device__ __forceinline__ float uniform(uint32_t bits) +{ + return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); +} + +// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. +__device__ __forceinline__ unsigned long long gkey(float lg, int id) +{ + return (static_cast(ordered(lg)) << kIdBits) + | static_cast((1u << kIdBits) - 1u - static_cast(id)); +} + +__device__ __forceinline__ int key_token(unsigned long long k) +{ + return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); +} + +__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) +{ + asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); +} + +__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) +{ + unsigned long long v; + asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); + return v; +} + +template +struct Shared +{ + __nv_bfloat16 row[kMaxVs / kSplits]; + float redf[kWarps]; + unsigned long long redk[kWarps]; + unsigned long long rk[kMaxRec]; + float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; + float wsum[kWarps]; + int owner; + int pick; + float pickx; + bool last; + float dS[kWarps]; // SA-1R: per-warp partial mass above the candidate + int dM[kWarps]; // SA-1R: per-warp count | incomplete << 16 + uint32_t wtop[kWarps][kTopK > 0 ? mode_k(kTopK) : 1]; // SA-1R: per-warp top-K keys + uint32_t wbound[kWarps]; // SA-1R pooled: per-warp max thread bound + uint32_t ctop[kTopK > 0 ? mode_vals(kTopK) : 1]; // SA-1R: the CTA's top-K keys (descending), by rank + unsigned long long rx[kMaxRec][Rec::kXw > 0 ? Rec::kXw : 1]; +}; + +template +__device__ __forceinline__ float block_max(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = s.redf[0]; + for (int i = 1; i < kWarps; ++i) + r = fmaxf(r, s.redf[i]); + __syncthreads(); + return r; +} + +// Fixed-order block sum: identical result for identical inputs. +template +__device__ __forceinline__ float block_sum(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = 0.f; + for (int i = 0; i < kWarps; ++i) + r += s.redf[i]; + __syncthreads(); + return r; +} + +template +__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + if ((threadIdx.x & 31) == 0) + s.redk[threadIdx.x >> 5] = v; + __syncthreads(); + unsigned long long r = s.redk[0]; + for (int i = 1; i < kWarps; ++i) + r = s.redk[i] > r ? s.redk[i] : r; + __syncthreads(); + return r; +} + +// Exclusive prefix of v over the threads in thread order, and the total. +template +__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + float inc = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + float const t = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) + inc += t; + } + if (lane == 31) + s.wsum[warp] = inc; + __syncthreads(); + float base = 0.f; + total = 0.f; + for (int i = 0; i < kWarps; ++i) + { + base += i < warp ? s.wsum[i] : 0.f; + total += s.wsum[i]; + } + __syncthreads(); + return base + inc - v; +} + +// Warp-wide integer max / sum / or (the redux instructions from sm_80; a butterfly otherwise, same results). +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_max_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + uint32_t const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_add_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_add_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_or_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_or_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v |= __shfl_xor_sync(0xffffffffu, v, o); + return v; +#endif +} + +// Insert key into a descending register list of K (fully unrolled: no local memory). +template +__device__ __forceinline__ void topk_insert(uint32_t (&a)[K], uint32_t key) +{ +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const hi = key > a[j] ? key : a[j]; + key = key > a[j] ? a[j] : key; + a[j] = hi; + } +} + +template +__device__ __forceinline__ void topk_pop(uint32_t (&a)[K]) +{ +#pragma unroll + for (int j = 0; j + 1 < K; ++j) + a[j] = a[j + 1]; + a[K - 1] = 0u; +} + +// Per-thread insertion with an early reject: after the first few elements almost every key fails the K-th test. +template +__device__ __forceinline__ void thread_insert(uint32_t (&a)[K], uint32_t key) +{ + if (key > a[K - 1]) + topk_insert(a, key); +} + +// The warp's K largest keys (descending) into s.wtop[warp]; published by the caller's next block barrier. +template +__device__ __forceinline__ void warp_topk(uint32_t (&a)[K], S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const m = warp_max_u32(a[0]); + if (a[0] == m && m != 0u) + topk_pop(a); + if (lane == 0) + s.wtop[warp][j] = m; + } +} + +// Pooled: K rounds over each lane's P pool keys (descending) and the warp's max thread bound. +template +__device__ __forceinline__ void warp_pool(uint32_t (&a)[P + 1], S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + uint32_t const bound = warp_max_u32(a[P]); + uint32_t pl[P]; +#pragma unroll + for (int j = 0; j < P; ++j) + pl[j] = a[j]; +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const m = warp_max_u32(pl[0]); + if (pl[0] == m && m != 0u) + topk_pop

(pl); + if (lane == 0) + s.wtop[warp][j] = m; + } + if (lane == 0) + s.wbound[warp] = bound; +} + +// A row's kSplits CTAs wait for each other's records, so they launch as one thread-block cluster: the hardware +// schedules a cluster's CTAs together, whatever other streams hold. Rows never wait on other rows. +template +__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) +{ + using R = Rec; + extern __shared__ __align__(16) unsigned char smem[]; + Shared& s = *reinterpret_cast*>(smem); + float constexpr kNegInf = -__builtin_huge_valf(); + int const tid = threadIdx.x; + int const row = blockIdx.x / kSplits; + int const split = blockIdx.x % kSplits; + int const world = a.world; + int const nrec = world * kSplits; + // The row's sampling parameters are step inputs (the spec metadata, written before the target forward), not the + // predecessor head GEMM's output, so with prewait_params they load before the dependency wait and overlap the + // head's tail; the same values either way. + float temperature = 0.f, topP = 0.f; + long long seed_in = 0, offset_in = 0; + auto load_params = [&]() + { + temperature = a.temperatures[row]; + topP = a.top_ps[row]; + seed_in = a.seeds[row]; + offset_in = a.offsets[row]; + }; + if (a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (!a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // early_trigger: the dependent grid (the first draft pass's input) may launch now and park at its own + // griddepcontrol.wait, which still returns only once this grid has completed and flushed; it skips the launch + // latency after this kernel ends. Otherwise dependents are released at the acceptance epilogue (the CTAs of + // splits 1.. exit without releasing). + if (a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const tempInv = 1.0f / (temperature + kTempEpsilon); + unsigned long long const seed = static_cast(seed_in); + unsigned long long const offset = static_cast(offset_in); + unsigned int const call = *reinterpret_cast(a.counters + 1); + // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of + // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round + // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten + // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. + auto ktag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; + }; + auto ftag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; + }; + // SA-1R round-0 list words: a 16-bit tag, never 0. + unsigned long long const xtag = (((call >> 1) << 1) & 0x7FFFu) | 0x8000u; + long long const rec_stride = static_cast(world) * kMaxSplits * R::kEx; + unsigned long long* const xu = a.ws_uc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride; + unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * R::kEx; + long long const round_stride = static_cast(kMaxRows) * rec_stride; + + // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- + __nv_bfloat16 const* src = a.logits + row * a.ld; + int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; + int const s0 = min(split * sl_len, a.Vs); + int const n = min(s0 + sl_len, a.Vs) - s0; + int const id0 = a.vocab_offset + s0; + float tmax = kNegInf; + // SA-1R: each thread's largest (ord16(bf16) << 16 | ~index) keys while staging (only rows with a nucleus cut). + bool const want_x = kTopK > 0 && topP < 1.0f; + constexpr int kK = mode_k(kTopK), kPool = mode_pool(kTopK), kVals = mode_vals(kTopK); + constexpr int kTk = kPool ? kPool + 1 : (kK > 0 ? kK : 1); + uint32_t tk[kTk]; +#pragma unroll + for (int j = 0; j < kTk; ++j) + tk[j] = 0u; + bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; + int const nvec = vec ? n / 8 : 0; + for (int v = tid; v < nvec; v += kThreads) + { + uint4 const w = *reinterpret_cast(src + s0 + 8 * v); + *reinterpret_cast(&s.row[8 * v]) = w; + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); +#pragma unroll + for (int j = 0; j < 8; ++j) + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); + if constexpr (kTopK > 0) + { + if (want_x) + { +#pragma unroll + for (int j = 0; j < 8; ++j) + { + uint32_t const key + = (ord16(__bfloat16_as_ushort(e[j])) << 16) | (0xFFFFu - static_cast(8 * v + j)); + if constexpr (kPool > 0) + topk_insert(tk, key); // branchless: P + 1 compare-swaps + else + thread_insert(tk, key); + } + } + } + } + for (int i = nvec * 8 + tid; i < n; i += kThreads) + { + s.row[i] = src[s0 + i]; + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); + if constexpr (kTopK > 0) + { + if (want_x) + topk_insert( + tk, (ord16(__bfloat16_as_ushort(src[s0 + i])) << 16) | (0xFFFFu - static_cast(i))); + } + } + if constexpr (kTopK > 0) + { + if (tid < kVals) + s.ctop[tid] = 0u; // published with s.wtop by block_max's barrier + if (want_x) + { + if constexpr (kPool > 0) + warp_pool(tk, s); // s.wtop / s.wbound published by block_max's barrier + else + warp_topk(tk, s); // s.wtop published by block_max's barrier + } + } + float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier + if constexpr (kTopK > 0) + { + // SA-1R: the CTA's top K of the warp lists by rank (keys are unique): a lane quad per key, no serial merge; + // s.ctop is read after the round-0 draw's first barrier + constexpr int kN = kWarps * kK; + constexpr int kSp = kN * 4 <= kThreads ? 4 : 1; // threads per key + if (want_x && tid < kN * kSp) + { + uint32_t const* const wl = &s.wtop[0][0]; + int const part = tid % kSp; + uint32_t const key = wl[tid / kSp]; + int rank = 0; +#pragma unroll + for (int j = part; j < kN; j += kSp) + rank += wl[j] > key ? 1 : 0; + if constexpr (kSp == 4) + { + rank += __shfl_xor_sync(0xffffffffu, rank, 1); + rank += __shfl_xor_sync(0xffffffffu, rank, 2); + } + if (part == 0 && key != 0u && rank < kK) + s.ctop[rank] = key; + } + } + float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z + unsigned long long ckey = 0ULL; + int const gid = a.rank * kMaxSplits + split; + for (int r = 0;; ++r) + { + // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) + float const ref = r == 0 ? smax : M; + float const wmin = r == 0 ? 0.f : wc; + float ts = 0.f; + for (int i = tid; i < n; i += kThreads) + { + float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); + ts += w > wmin ? w : 0.f; + } + float total; + float const excl = block_excl_scan(ts, total, s); + if constexpr (kTopK > 0) + { + // SA-1R round 0: the list words depend only on the staged slice, so they go out now (s.ctop / s.wbound are + // published by the barrier above), one lane per word, ahead of the main record: off the readers' path + if (r == 0 && want_x && tid < R::kXw) + { + uint32_t tau = 0u; + if constexpr (kPool > 0) + { + tau = s.ctop[kK - 1]; // the largest unlisted pool key is <= the K-th listed +#pragma unroll + for (int wq = 0; wq < kWarps; ++wq) + tau = s.wbound[wq] > tau ? s.wbound[wq] : tau; + } + unsigned long long w = xtag; +#pragma unroll + for (int t = 0; t < 3; ++t) + { + int const e = 3 * tid + t; + uint32_t const kv = e < kK ? s.ctop[e] : (kPool > 0 && e == kK ? tau : 0u); + // missing entries: -inf (bf16 0xFF80), w = 0 + uint32_t const b = kv != 0u ? unord16(kv >> 16) : 0xFF80u; + w |= static_cast(b) << (16 * (3 - t)); + } + st_sys(xm + 4 + tid, w); // round 0 = slot 0 + } + } + uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), + static_cast(offset), static_cast(offset >> 32)), + make_uint2(static_cast(seed), static_cast(seed >> 32))); + float const tgt = uniform(rnd.x) * total; + if (tid == 0) + s.owner = -1; + __syncthreads(); + if (ts > 0.f && excl <= tgt) + atomicMax(&s.owner, tid); + __syncthreads(); + if (tid == s.owner) + { + float acc = excl; + int pick = -1; + float px = 0.f; + if constexpr (kIlpWalk) + { + // SA-PW: the owner's strided tokens 8 at a time. The 8 x and w (the same expressions) are + // independent, so their MUFU work overlaps; the serial sum then adds them in the same order with + // the same early exit: the same acc sequence, pick and px as the one-token walk, bit for bit. + constexpr int kW = 8; + bool done = false; + for (int i0 = tid; i0 < n && !done; i0 += kW * kThreads) + { + float xs[kW], wv[kW]; +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + xs[k] = i < n ? __fmul_rn(__bfloat162float(s.row[i]), tempInv) : 0.f; + wv[k] = __expf(xs[k] - ref); + } +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + if (!done && i < n && wv[k] > wmin) + { + acc += wv[k]; + pick = i; + px = xs[k]; + if (acc > tgt) + done = true; + } + } + } + } + else + { + for (int i = tid; i < n; i += kThreads) + { + float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); + float const w = __expf(x - ref); + if (w > wmin) + { + acc += w; + pick = i; + px = x; + if (acc > tgt) + break; + } + } + } + s.pick = pick; + s.pickx = px; + } + __syncthreads(); + unsigned long long sb = 0ULL; + if (total > 0.f && s.owner >= 0) + sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); + float const smab = total; + unsigned long long const ktag = ktag_of(r); + unsigned long long const ftag = ftag_of(r); + unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; + bool const with_x = want_x && r == 0; + if (tid == 0) + { + unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; + st_sys(xw + 0, (sb << kTagBits) | ktag); + st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); + st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); + st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); + } + // ---- every record of the row, every rank --------------------------------------------------------------- + int const nw = with_x ? R::kEx : 4; + for (int t = tid; t < nrec * nw; t += kThreads) + { + int const q = t / nw, f = t % nw; + unsigned long long const* const p = xr + (q / kSplits * kMaxSplits + q % kSplits) * R::kEx + f; + unsigned long long v; + long long spins = 0; + if (f == 0) + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); + s.rk[q] = v >> kTagBits; + } + else if (f < 4) + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFFFFF00000000ULL) != ftag); + float const fv = __uint_as_float(static_cast(v)); + (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; + } + else + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFULL) != xtag); + s.rx[q][f - 4] = v; + } + } + __syncthreads(); + // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- + unsigned long long gk = 0ULL; + float gv = 0.f; + for (int q = 0; q < nrec; ++q) + { + if (s.rk[q] > gk) + { + gk = s.rk[q]; + gv = s.rv[q]; + } + } + bool accept; + if (r == 0) + { + M = kNegInf; + for (int q = 0; q < nrec; ++q) + M = fmaxf(M, s.r2[q]); + float Z = 0.f; + for (int q = 0; q < nrec; ++q) + Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); + target = topP * Z; + accept = !(topP < 1.0f); // no nucleus cut: the first draw stands + ckey = gk; + cv = gv; + if constexpr (kTopK > 0) + { + // SA-1R: every thread takes (record, entry) items; the per-warp partials are published by the round's + // barrier below and every thread combines them in warp order after it (no new barrier) + if (with_x) + { + float const wcand = __expf(cv - M); // = round 1's wmin + float part = 0.f; + int cnt = 0, inc = 0; + for (int t = tid; t < nrec * kVals; t += kThreads) + { + int const q = t / kVals, e = t % kVals; + uint32_t const b = static_cast(s.rx[q][e / 3] >> (16 * (3 - e % 3))) & 0xFFFFu; + float const w = __expf( + __fmul_rn(__bfloat162float(__ushort_as_bfloat16(static_cast(b))), tempInv) + - M); // w bitwise as round 1 computes it + if (w > wcand) + { + if (kPool > 0 && e == kK) + inc = 1; // tau (the largest unlisted value) is heavier than the candidate + else + { + part += w; + ++cnt; + if (kPool == 0 && e == kK - 1) + inc = 1; // the K-th is heavier: unlisted heavier tokens may exist + } + } + } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + part += __shfl_xor_sync(0xffffffffu, part, o); // any order: the margin covers it + cnt = static_cast(warp_add_u32(static_cast(cnt))); + inc = static_cast(warp_or_u32(static_cast(inc))); + if ((tid & 31) == 0) + { + s.dS[tid >> 5] = part; + s.dM[tid >> 5] = cnt | (inc << 16); + } + } + } + } + else + { + float S = 0.f; // mass strictly above the current candidate + for (int q = 0; q < nrec; ++q) + S += s.r2[q]; + accept = S <= target || gk == 0ULL; + if (!accept) + { + ckey = gk; + cv = gv; + } + } + wc = __expf(cv - M); + __syncthreads(); // records consumed before the next round's writes to s + bool decided = false; + if constexpr (kTopK > 0) + { + if (r == 0 && with_x) + { + float S1 = 0.f; + int m = 0, inc = 0; + for (int wq = 0; wq < kWarps; ++wq) + { + S1 += s.dS[wq]; + m += s.dM[wq] & 0xFFFF; + inc |= s.dM[wq] >> 16; + } + if (!inc) + { + float const margin = m <= 1 ? 0.f : S1 * static_cast(m) * (1.0f / 2097152.0f); + decided = __fadd_ru(S1, margin) <= target; + accept = accept || decided; + } + } + } + if (accept) + { + if (split != 0) + return; + if (tid == 0) + { + a.tokens[row] = key_token(ckey); + if (a.extra_attempts != nullptr && r > 1) + atomicAdd(a.extra_attempts, r - 1); + if (a.decided != nullptr && decided) + atomicAdd(a.decided, 1); + } + break; + } + } + + // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (!a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (tid == 0) + s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; + __syncthreads(); + if (!s.last) + return; + __threadfence(); + if (tid == 0) + { + a.counters[0] = 0u; + a.counters[1] = call + 1u; + } + if (tid < a.nreq) + { + int const* tok = a.tokens; + int const b = tid; + int const k1 = a.draft_len + 1; + int nacc = 1; + bool run = true; + for (int k = 0; k < a.draft_len; ++k) + { + int const tt = *reinterpret_cast(tok + b * k1 + k); + a.acc_tokens[b * k1 + k] = tt; + run = run && a.draft[b * a.draft_len + k] == tt; + nacc += run ? 1 : 0; + } + a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); + if (a.force_vec != nullptr) + nacc = a.force_vec[b]; + else if (a.force_total > 0) + nacc = a.force_total; + a.num_acc[b] = nacc; + a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; + if (a.next_new != nullptr) + a.next_new[b * a.next_stride] + = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); + } +} + +template +cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) +{ + static int max_clusters = 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(a.rows * kSplits); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Shared); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + if (max_clusters == 0) + { + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } + } + TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); + return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); +} + +template +cudaError_t launchRows(Args const& a, bool use_pdl, cudaStream_t stream) +{ + // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); + // more rows: 4 CTAs per row. + if (a.ilp_walk) + { + return a.rows <= 32 ? launch<8, kTopK, true>(a, use_pdl, stream) : launch<4, kTopK, true>(a, use_pdl, stream); + } + return a.rows <= 32 ? launch<8, kTopK, false>(a, use_pdl, stream) : launch<4, kTopK, false>(a, use_pdl, stream); +} + +} // namespace + +long long sampledAcceptWsWords(int world, int oneRoundK) +{ + TLLM_CHECK_WITH_INFO( + oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "SA-1R variant %d is not built", oneRoundK); + return 2 * (oneRoundK == 9 ? half_words<9>(world) : oneRoundK == 82 ? half_words<82>(world) : half_words<0>(world)); +} + +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) +{ + switch (a.one_round_k) + { + case 0: return launchRows<0>(a, usePdl, stream); + case 9: return launchRows<9>(a, usePdl, stream); + case 82: return launchRows<82>(a, usePdl, stream); + default: TLLM_CHECK_WITH_INFO(false, "SA-1R variant %d is not built", a.one_round_k); + } + return cudaErrorInvalidValue; +} + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h new file mode 100644 index 000000000000..aa52cf4a0559 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h @@ -0,0 +1,87 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ +//! Fused sampled (temperature + top-p) verification of the MiniMax-M3 Eagle3 target head over the TP-sharded +//! vocabulary, decode rows <= kMaxRows; see sampledAccept.cu. +namespace sampledAccept +{ + +constexpr int kMaxSplits = 8; // CTAs per verify row (at most) +constexpr int kMaxRows = 48; // verify rows per call +constexpr int kMaxWorld = 8; // TP ranks +constexpr int kMaxVs = 100352; // vocab shard width (TP2: 200064 / 2 = 100032, padded to a multiple of 512) +constexpr int kIdBits = 18; // token ids < 2^18 + +struct Args +{ + __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard + long long ld; + int rows, Vs, vocab_offset; + float const* temperatures; + float const* top_ps; + long long const* seeds; + long long const* offsets; + unsigned long long* ws_uc; + unsigned long long* ws_mc; + unsigned int* counters; // [0] finished rows, [1] call + int* tokens; // [rows] sampled token per row + int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows + int rank, world; + int const* draft; + int const* batch_idx; + int* acc_tokens; + int* num_acc; + long long* gather_ids; + int draft_len, nreq, force_total; + int* next_new; + int next_stride; + int const* force_vec; + int prewait_params; // nonzero: the row's sampling parameters load before the PDL dependency wait + int early_trigger; // nonzero: dependents are released right after the dependency wait, not at the epilogue + // SA-1R (one-round decision): 0 the two-round kernel; 9 each CTA's top-9 logits ride its round-0 record; 82 pooled + // (each thread's top 2, the CTA lists the pool's top 8 plus the bound of everything unlisted). A row whose first + // candidate is provably in the nucleus then decides in round 0, with the same token as the two-round kernel. + int one_round_k; + int* decided; // optional [1]: rows decided in round 0 (SA-1R), summed + int ilp_walk; // nonzero: the owner walk takes its tokens 8 at a time (same weights, same summation order) +}; + +//! Words (uint64) of the multicast exchange workspace for `world` ranks and SA-1R variant `oneRoundK`: two call halves. +long long sampledAcceptWsWords(int world, int oneRoundK = 0); + +//! Launch the sampled verification over a.rows * (8 or 4) co-resident CTAs (8 per row up to 32 rows). +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream); + +} // namespace sampledAccept +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu new file mode 100644 index 000000000000..8cbf5d90175c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu @@ -0,0 +1,89 @@ +/* + * 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/minimaxM3Sm107Tp8/siluAndMul.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kThreads = 128; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves). The +// instruction sequence of the Triton kernel's PTX at its defaults (host/silu/silu_and_mul_plain.ptx): g * -log2(e) as +// one product, ex2.approx, + 1, div.full 1 / that, g * s, (u + 0.0) as an IEEE add, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t siluMulBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 gl, gh, ul, uh, t0, t1, e0, e1, s0, s1, p0, p1, q0, q1, r0, r1, one, zero;\n\t" + ".reg .b16 hl, hh;\n\t" + "mov.b32 {hl, hh}, %1;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mul.rn.f32 t0, gl, 0fBFB8AA3B;\n\tmul.rn.f32 t1, gh, 0fBFB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 p0, gl, s0;\n\tmul.rn.f32 p1, gh, s1;\n\t" + "mov.b32 zero, 0f00000000;\n\tadd.rn.f32 q0, ul, zero;\n\tadd.rn.f32 q1, uh, zero;\n\t" + "mul.rn.f32 r0, q0, p0;\n\tmul.rn.f32 r1, q1, p1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +// Eight outputs per thread: 16-byte loads of the gate and the up half, one 16-byte store. +__global__ void __launch_bounds__(kThreads) + siluAndMulKernel(uint4 const* __restrict__ x, uint4* __restrict__ out, int64_t rows, int64_t vecs) +{ + int64_t const row = blockIdx.y; + int64_t const v = static_cast(blockIdx.x) * kThreads + threadIdx.x; + if (row >= rows || v >= vecs) + { + return; + } + uint4 const g = x[row * 2 * vecs + v]; + uint4 const u = x[row * 2 * vecs + vecs + v]; + uint4 o; + o.x = siluMulBf16x2(g.x, u.x); + o.y = siluMulBf16x2(g.y, u.y); + o.z = siluMulBf16x2(g.z, u.z); + o.w = siluMulBf16x2(g.w, u.w); + out[row * vecs + v] = o; +} + +} // namespace + +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream) +{ + int64_t const vecs = d / 8; + dim3 const grid(static_cast((vecs + kThreads - 1) / kThreads), static_cast(rows)); + siluAndMulKernel<<>>( + reinterpret_cast(x), reinterpret_cast(out), rows, vecs); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h new file mode 100644 index 000000000000..669e33eb5cf5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h @@ -0,0 +1,40 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ + +//! out[m, j] = bf16(silu(g) * (u + 0)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both +//! contiguous, d a multiple of 8, 16-byte aligned rows): bit for bit the Triton silu_and_mul kernel at its defaults +//! (no limit, alpha 1, beta 0), the plain SwiGLU of the Eagle3 drafter MLP. +cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu new file mode 100644 index 000000000000..d51138cc15d2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu @@ -0,0 +1,746 @@ +/* + * 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. + */ + +// MiniMax-M3 TP8 long-context sparse decode attention in one launch (8 query heads on one KV head). +// +// Grid: one cluster of 16 CTAs per query token; CTA `r` attends the token's `r`-th selected block. The QK and PV +// MMAs are transposed (S^T = K q^T, O^T = V^T P^T) so the 8 heads fill the MMA n dimension. Each CTA pushes its +// normalized partial over DSMEM so that CTA `r` holds (head r / 2, dims (r % 2) * 64 + [0, 64)) of all 16 partials, +// merges them, and writes the BF16 row and, optionally, its MXFP8 copy and swizzled UE8M0 scales. +// +// Per-block math: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and PV on FP16 tensor cores with +// FP32 accumulation, log2-domain softmax with either the linear denominator or the log-domain recurrence. +// +// PREWAIT stages everything that does not depend on the top-16 table before the grid dependency wait: the request's +// seq_lens entry, the q rows and the request's block-table row (in dynamic shared memory after Smem8, when it has at +// most kBtCap entries). The arithmetic is the same in both variants. +// +// ONE_DIV merge (MSA_ONE_DIV_MERGE): each CTA pushes its unnormalized partial with its row max m and sum l, and the +// merge forms o = sum_c w_c acc_c / sum_c w_c l_c with w_c = 2^(m_c - max m), one division per output, instead of +// normalizing every partial and weighting by its LSE. Deterministic (fixed chunk order, no atomics); a different +// rounding sequence from the Triton merge, inside the same band of the float64 reference. + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kDim = 128; +constexpr int kPage = 128; +constexpr int kTopk = 16; +constexpr int kThreads = 128; +constexpr int kPitch = kDim + 8; // halves per smem row: 272 B, conflict-free ldmatrix +constexpr float kNegInf = -__builtin_huge_valf(); +constexpr int kHeads8 = 8; +constexpr int kHalfDim = kDim / 2; +constexpr int kBtCap = 8192; // block-table entries staged in shared memory (8192 x 128 = 1M tokens) + +using Params = MinimaxM3Tp8SparseDecodeLc8Params; + +struct Smem8 +{ + __half qs[kHeads8 * kPitch]; + __half vs[kPage * kPitch]; + __half ps[kHeads8 * kPitch]; // P^T: [head][token] + float recv[kTopk][kHalfDim]; // (head rank / 2, dims (rank % 2) * 64 + [0, 64)) partials of the 16 CTAs + float recv_lse[kTopk]; // the chunk LSEs (log2); with ONE_DIV, the chunk row maxima + float recv_l[kTopk]; // ONE_DIV: the chunk row sums + float part[kHalfDim]; + float ex[kTopk]; // the 16 chunk exponentials (ONE_DIV: weighted row sums), formed once per CTA + float wgt[kTopk]; // the 16 merge weights, formed once per CTA + float red[4][kHeads8]; +}; + +template +constexpr size_t smemBytes() +{ + return PREWAIT ? sizeof(Smem8) + kBtCap * sizeof(int) : sizeof(Smem8); +} + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__device__ __forceinline__ uint32_t smem_u32(void const* p) +{ + return static_cast(__cvta_generic_to_shared(p)); +} + +__device__ __forceinline__ uint32_t e2m1x2_to_f16x2(uint32_t byte) +{ + uint32_t out; + asm("{\n .reg .b8 s;\n cvt.u8.u32 s, %1;\n cvt.rn.f16x2.e2m1x2 %0, s;\n}" : "=r"(out) : "r"(byte)); + return out; +} + +__device__ __forceinline__ uint32_t e4m3x2_to_f16x2(uint16_t pair) +{ + uint32_t out; + asm("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(out) : "h"(pair)); + return out; +} + +__device__ __forceinline__ float ex2(float x) +{ + float y; + asm("ex2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ float lg2(float x) +{ + float y; + asm("lg2.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +__device__ __forceinline__ void ldsm_x4(uint32_t (&r)[4], uint32_t addr) +{ + 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"(addr)); +} + +__device__ __forceinline__ void ldsm_x4_t(uint32_t (&r)[4], uint32_t addr) +{ + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) + : "r"(addr)); +} + +__device__ __forceinline__ void mma(float (&c)[4], uint32_t const (&a)[4], uint32_t b0, uint32_t b1) +{ + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.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"(b0), "r"(b1)); +} + +__device__ __forceinline__ void cluster_sync() +{ + asm volatile("barrier.cluster.arrive.release.aligned;\nbarrier.cluster.wait.acquire.aligned;" ::: "memory"); +} + +__device__ __forceinline__ uint32_t map_cluster(void const* local, uint32_t rank) +{ + uint32_t remote; + asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(remote) : "r"(smem_u32(local)), "r"(rank)); + return remote; +} + +__device__ __forceinline__ void st_cluster_f32(uint32_t addr, float a) +{ + asm volatile("st.shared::cluster.f32 [%0], %1;" ::"r"(addr), "f"(a) : "memory"); +} + +// One packed NVFP4 row (64 bytes, 128 values) with its 8 E4M3 scales -> 128 fp16 values in smem. +__device__ __forceinline__ void dequant_row(__half* dst, uint8_t const* packed, uint8_t const (&sc)[8], bool valid) +{ + uint4 const* src = reinterpret_cast(packed); + uint32_t* d = reinterpret_cast(dst); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + uint32_t words[4] = {w.x, w.y, w.z, w.w}; + // 16 bytes = 32 values = scale columns 2c and 2c + 1. +#pragma unroll + for (int h = 0; h < 2; ++h) + { + uint8_t const s = valid ? sc[2 * c + h] : 0; + uint32_t const sh = e4m3x2_to_f16x2(static_cast(s | (s << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); +#pragma unroll + for (int j = 0; j < 2; ++j) + { + uint32_t const word = words[2 * h + j]; +#pragma unroll + for (int b = 0; b < 4; ++b) + { + uint32_t const v2 = e2m1x2_to_f16x2((word >> (8 * b)) & 0xFF); + __half2 const x = __hmul2(*reinterpret_cast<__half2 const*>(&v2), scale); + d[c * 16 + h * 8 + j * 4 + b] = *reinterpret_cast(&x); + } + } + } + } +} + +// q: 8 heads x 128 e4m3 -> fp16; threads 0..63 = (head, 16-value chunk). +__device__ __forceinline__ void load_q(Smem8& s, Params const& p, int tok, int tid) +{ + if (tid < kHeads8 * 8) + { + int const h = tid >> 3; + int const c = (tid & 7) * 16; + uint4 const w = *reinterpret_cast(p.q + tok * p.q_tok + h * p.q_head + c); + uint32_t const words[4] = {w.x, w.y, w.z, w.w}; + uint32_t* d = reinterpret_cast(s.qs + h * kPitch + c); +#pragma unroll + for (int i = 0; i < 4; ++i) + { + d[2 * i] = e4m3x2_to_f16x2(static_cast(words[i] & 0xFFFF)); + d[2 * i + 1] = e4m3x2_to_f16x2(static_cast(words[i] >> 16)); + } + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +// This CTA's 16-byte-aligned slice of [base, base + bytes) into L2, at most 16 KiB per bulk prefetch. +__device__ __forceinline__ void prefetch_l2_slice(uint8_t const* base, int64_t bytes) +{ + if (base == nullptr || bytes <= 0) + return; + int64_t const n = gridDim.x; + int64_t const slice = ((bytes + n - 1) / n + 15) & ~int64_t{15}; + int64_t const off = static_cast(blockIdx.x) * slice; + if (off >= bytes) + return; + int64_t const len = min(slice, bytes - off) & ~int64_t{15}; + for (int64_t o = 0; o < len; o += 16384) + { + uint32_t const sz = static_cast(min(int64_t{16384}, len - o)); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(base + off + o), "r"(sz) : "memory"); + } +} +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +template +__global__ void __launch_bounds__(kThreads) m3_tp8_lc8_attn_kernel(__grid_constant__ Params const p) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + extern __shared__ __align__(128) uint8_t smem_raw[]; + Smem8& s = *reinterpret_cast(smem_raw); + int const tid = threadIdx.x; + int const lane = tid & 31; + int const warp = tid >> 5; + int const g = lane >> 2; + int const qd = lane & 3; + int const tok = blockIdx.x / kTopk; + uint32_t rank; + asm volatile("mov.u32 %0, %%cluster_ctarank;" : "=r"(rank)); + + asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); + // The o_proj weights are read-only for the whole step: start pulling them into L2 now, while HBM is idle in the + // latency-bound attention phase (one lane of warps 0 and 1; nothing waits on these). + if (tid == 0) + prefetch_l2_slice(p.pf0, p.pf0_bytes); + else if (tid == 32) + prefetch_l2_slice(p.pf1, p.pf1_bytes); + if constexpr (!PREWAIT) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + + int const req = tok / p.dql; + int const qoff = tok - req * p.dql; + int const seq = p.seq_lens[req]; + int const kv_len = max(seq - p.dql + qoff + 1, 0); + int const nblocks = (kv_len + kPage - 1) / kPage; + int const real_topk = min(kTopk, nblocks); + // The next sparse layer's block that its CTA (tok, rank) selected at the previous step (a decode step selects + // mostly the blocks it selected one step earlier): its K/V pages and scales into L2 for that layer, from one lane + // of warp 2. The id buffer was written by a kernel of the previous step, long complete; a recorded id at or past + // this request's block count (another request's row) or a page outside the pool is skipped. + if (p.n1_next != nullptr && tid == 64) + { + int const nb = p.n1_next[tok * kTopk + rank]; + if (nb >= 0 && nb < nblocks) + { + int64_t const pg = p.block_table[req * p.bt_stride + nb]; + if (pg >= 0 && pg < p.n1_pages) + { + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_k + pg * p.k_blk), "r"(kPage * 64) + : "memory"); + asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_v + pg * p.v_blk), "r"(kPage * 64) + : "memory"); + asm volatile( + "cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_ksc + pg * p.ksc_blk), "r"(kPage * 8) + : "memory"); + asm volatile( + "cp.async.bulk.prefetch.L2.global [%0], %1;" ::"l"(p.n1_vsc + pg * p.vsc_blk), "r"(kPage * 8) + : "memory"); + } + } + } + + // The request's block-table row, staged before the grid dependency (written before the graph runs). + int* const s_bt = reinterpret_cast(smem_raw + sizeof(Smem8)); + bool bt_smem = false; + if constexpr (PREWAIT) + { + int const nbt = (max(seq, 0) + kPage - 1) / kPage; + bt_smem = nbt <= kBtCap; + if (bt_smem) + { + for (int i = tid; i < nbt; i += kThreads) + { + s_bt[i] = __ldg(p.block_table + req * p.bt_stride + i); + } + } + // Before the wait only behind the select: an identity step has no scorer / select, so this grid's + // predecessor is the producer itself, which releases its dependents at entry, before it writes q. There q is + // read after the wait, whether the blocks come from the lengths (identity) or from the staged table. + if (!p.identity && !p.identity_step) + load_q(s, p, tok, tid); + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int blk = static_cast(rank) >= real_topk ? -1 + : p.identity ? static_cast(rank) + : p.topk[tok * p.t_tok + rank * p.t_k]; + // Fail closed on an id past the row (a stale or corrupt table), as the stock Triton producer does: the block is + // skipped. Read through the block table it would be another request's page, every token masked, and its softmax + // ex2(-inf - (-inf)) = NaN. + if (blk >= nblocks) + blk = -1; + bool const active = blk >= 0; + // Record this step's selection for this layer's next-step prefetch (read by the previous sparse layer's CTAs). + if (p.n1_self != nullptr && tid == 0) + p.n1_self[tok * kTopk + rank] = blk; + if constexpr (PREWAIT) + { + if (p.identity || p.identity_step) + load_q(s, p, tok, tid); + __syncthreads(); // s_bt and the q rows are complete + } + else + { + load_q(s, p, tok, tid); + } + + // Lane holds heads 2 * qd + {0, 1}; acc[mt]: dims warp * 32 + mt * 16 + g + {0, 8}. + float m_row[2] = {kNegInf, kNegInf}; + float l_row[2] = {0.F, 0.F}; + float acc[2][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + acc[j][i] = 0.F; + + // K rows n = warp * 32 + j * 8 + g stay packed in registers. As the MMA A operand, m-tile mt takes rows + // j = 2 mt (row g) and 2 mt + 1 (row g + 8); k-step ks takes packed bytes 8 ks + qd and 8 ks + 4 + qd. + uint32_t kw[4][16]; + uint2 ksc[4]; + if (active) + { + int64_t const page = bt_smem ? s_bt[blk] : p.block_table[req * p.bt_stride + blk]; +#pragma unroll + for (int j = 0; j < 4; ++j) + { + int const n = warp * 32 + j * 8 + g; + bool const valid = blk * kPage + n < kv_len; + uint4 const* src = reinterpret_cast(p.k + page * p.k_blk + n * 64); +#pragma unroll + for (int c = 0; c < 4; ++c) + { + uint4 const w = valid ? __ldg(src + c) : make_uint4(0, 0, 0, 0); + kw[j][4 * c] = w.x; + kw[j][4 * c + 1] = w.y; + kw[j][4 * c + 2] = w.z; + kw[j][4 * c + 3] = w.w; + } + ksc[j] = valid ? __ldg(reinterpret_cast(p.ksc + page * p.ksc_blk + n * 8)) : make_uint2(0, 0); + } + bool const valid = blk * kPage + tid < kv_len; + uint8_t vsc[8]; + if (valid) + { + uint8_t const* vb = p.vsc + page * p.vsc_blk + (tid >> 2) * 32 + (tid & 3); +#pragma unroll + for (int c = 0; c < 8; ++c) + vsc[c] = __ldg(vb + 4 * c); + } + dequant_row(s.vs + tid * kPitch, p.v + page * p.v_blk + tid * 64, vsc, valid); + } + __syncthreads(); + + if (active) + { + float const scale_log2 = p.sm_scale * 1.4426950409F * __ldg(p.kgs); + int const q4 = qd * 8; + // S^T = K q^T for this warp's 32 tokens: rows (tokens) g / g + 8 of each m-tile, columns (heads) 2 qd + {0, 1}. + float sacc[2][4]; +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int i = 0; i < 4; ++i) + sacc[j][i] = 0.F; +#pragma unroll + for (int ks2 = 0; ks2 < kDim / 16; ks2 += 2) + { + uint32_t bq[4]; + ldsm_x4(bq, smem_u32(s.qs + (lane & 7) * kPitch + ks2 * 16 + (lane >> 3) * 8)); +#pragma unroll + for (int u = 0; u < 2; ++u) + { + int const ks = ks2 + u; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) + { + uint32_t a[4]; +#pragma unroll + for (int hi = 0; hi < 2; ++hi) + { + int const j = 2 * mt + hi; + uint32_t const scw = ks < 4 ? ksc[j].x : ksc[j].y; + uint8_t const sc = static_cast(scw >> (8 * (ks & 3))); + uint32_t const sh = e4m3x2_to_f16x2(static_cast(sc | (sc << 8))); + __half2 const scale = *reinterpret_cast<__half2 const*>(&sh); + uint32_t const r0 = e2m1x2_to_f16x2((kw[j][2 * ks] >> q4) & 0xFF); + uint32_t const r1 = e2m1x2_to_f16x2((kw[j][2 * ks + 1] >> q4) & 0xFF); + __half2 const h0 = __hmul2(*reinterpret_cast<__half2 const*>(&r0), scale); + __half2 const h1 = __hmul2(*reinterpret_cast<__half2 const*>(&r1), scale); + a[hi] = *reinterpret_cast(&h0); + a[2 + hi] = *reinterpret_cast(&h1); + } + mma(sacc[mt], a, bq[2 * u], bq[2 * u + 1]); + } + } + } + float mloc[2] = {kNegInf, kNegInf}; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const t = warp * 32 + mt * 16 + g + 8 * (i >> 1); + float const v = blk * kPage + t < kv_len ? 0.F + sacc[mt][i] * scale_log2 : kNegInf; + sacc[mt][i] = v; + mloc[i & 1] = fmaxf(mloc[i & 1], v); + } +#pragma unroll + for (int e = 0; e < 2; ++e) + { + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 4)); + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 8)); + mloc[e] = fmaxf(mloc[e], __shfl_xor_sync(0xffffffff, mloc[e], 16)); + } + if (g == 0) + { + s.red[warp][2 * qd] = mloc[0]; + s.red[warp][2 * qd + 1] = mloc[1]; + } + __syncthreads(); + // Every thread reduces its two heads' four warp maxima itself (the same expression, so the same bits): no + // shared row-max stage and its second barrier. +#pragma unroll + for (int e = 0; e < 2; ++e) + { + int const row = 2 * qd + e; + m_row[e] = fmaxf(fmaxf(s.red[0][row], s.red[1][row]), fmaxf(s.red[2][row], s.red[3][row])); + } + // A row with no valid token has m = -inf: its probabilities are ex2(-inf - 0) = 0, never NaN. + float const m_use[2] = {m_row[0] == kNegInf ? 0.F : m_row[0], m_row[1] == kNegInf ? 0.F : m_row[1]}; + float lloc[2] = {0.F, 0.F}; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + float const pv = ex2(sacc[mt][i] - m_use[i & 1]); + lloc[i & 1] += pv; + int const t = warp * 32 + mt * 16 + g + 8 * (i >> 1); + s.ps[(2 * qd + (i & 1)) * kPitch + t] = __float2half_rn(pv); + } +#pragma unroll + for (int e = 0; e < 2; ++e) + { + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 4); + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 8); + lloc[e] += __shfl_xor_sync(0xffffffff, lloc[e], 16); + } + __syncthreads(); // row-max reads of red done before red is reused; ps complete + if (g == 0) + { + s.red[warp][2 * qd] = lloc[0]; + s.red[warp][2 * qd + 1] = lloc[1]; + } + // O^T[warp's 32 dims, 8 heads] = V^T P^T over all 128 tokens. +#pragma unroll + for (int kk = 0; kk < kPage; kk += 32) + { + uint32_t b[4]; + ldsm_x4(b, smem_u32(s.ps + (lane & 7) * kPitch + kk + (lane >> 3) * 8)); +#pragma unroll + for (int u = 0; u < 2; ++u) + { +#pragma unroll + for (int mt = 0; mt < 2; ++mt) + { + uint32_t a[4]; + int const t = kk + u * 16 + (lane & 7) + ((lane >> 4) & 1) * 8; + int const d = warp * 32 + mt * 16 + ((lane >> 3) & 1) * 8; + ldsm_x4_t(a, smem_u32(s.vs + t * kPitch + d)); + mma(acc[mt], a, b[2 * u], b[2 * u + 1]); + } + } + } + __syncthreads(); + // Likewise the row sums, in the same order (red is not written again: the merge uses ex / wgt). +#pragma unroll + for (int e = 0; e < 2; ++e) + { + int const row = 2 * qd + e; + l_row[e] = (s.red[0][row] + s.red[1][row]) + (s.red[2][row] + s.red[3][row]); + } + } + + // Normalize; push (head, dim) to CTA 2 * head + dim / 64, slot recv[rank][dim % 64], plus each head's LSE (log2) + // to both of its CTAs. Posted DSMEM stores, made visible by the cluster barrier. + asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); + { + float const vgs = __ldg(p.vgs); + float scale[2], lse[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) + { + if constexpr (ONE_DIV) + { + // Unnormalized: the merge divides once. An inactive CTA pushes m = -inf, l = 0 (weight 0). + scale[e] = 1.F; + lse[e] = m_row[e]; + } + else if constexpr (LINEAR) + { + scale[e] = l_row[e] > 0.F ? 1.0F / l_row[e] : 0.F; + lse[e] = l_row[e] > 0.F ? m_row[e] + lg2(l_row[e]) : kNegInf; + } + else + { + float const lse_i + = active && m_row[e] > kNegInf ? m_row[e] + lg2(ex2(kNegInf - m_row[e]) + l_row[e]) : kNegInf; + scale[e] = lse_i > kNegInf ? ex2(m_row[e] - lse_i) : 0.F; + lse[e] = lse_i; + } + scale[e] = scale[e] * vgs; + } + uint32_t dst[2]; +#pragma unroll + for (int e = 0; e < 2; ++e) + dst[e] = map_cluster(&s.recv[rank][0], (2 * qd + e) * 2 + (warp >> 1)); +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int i = 0; i < 4; ++i) + { + int const dl = (warp & 1) * 32 + mt * 16 + g + 8 * (i >> 1); + st_cluster_f32(dst[i & 1] + dl * 4, acc[mt][i] * scale[i & 1]); + } + if (warp == 0 && g == 0) + { +#pragma unroll + for (int e = 0; e < 2; ++e) +#pragma unroll + for (int hf = 0; hf < 2; ++hf) + { + st_cluster_f32(map_cluster(&s.recv_lse[rank], (2 * qd + e) * 2 + hf), lse[e]); + if constexpr (ONE_DIV) + { + st_cluster_f32(map_cluster(&s.recv_l[rank], (2 * qd + e) * 2 + hf), l_row[e]); + } + } + } + } + cluster_sync(); + + // CTA `rank` merges head rank / 2, dims (rank % 2) * 64 + dl; thread = (chunk half, dl). + int const h = static_cast(rank) >> 1; + int const dh = static_cast(rank) & 1; + int const dl = tid & (kHalfDim - 1); + int const hc = tid >> 6; + if constexpr (ONE_DIV) + { + // Thread c < 16 takes chunk c's weight 2^(m_c - max m) (wgt) and its weighted row sum (ex); the output is + // the weighted partials over the weighted row sums, both summed in chunk order, one division. + if (tid < kTopk) + { + float m_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + m_max = fmaxf(m_max, s.recv_lse[c]); + m_max = m_max == kNegInf ? 0.F : m_max; + float const w = ex2(s.recv_lse[tid] - m_max); + s.wgt[tid] = w; + s.ex[tid] = w * s.recv_l[tid]; + } + __syncthreads(); + } + else + { + // The merge weights (exp2 of each chunk's LSE minus the largest, divided by their sum in chunk order) are + // formed once per CTA: thread c < 16 takes chunk c's exponential, then divides it by the sum each of those + // threads forms in the same order. Same values in the same order as every thread forming all 16: bitwise the + // same rows. + if (tid < kTopk) + { + float lse_max = kNegInf; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + lse_max = fmaxf(lse_max, s.recv_lse[c]); + lse_max = lse_max == kNegInf ? 0.F : lse_max; + s.ex[tid] = ex2(s.recv_lse[tid] - lse_max); + } + __syncthreads(); + if (tid < kTopk) + { + float denom = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + denom += s.ex[c]; + denom = denom > 0.F ? denom : 1.F; + s.wgt[tid] = __fdiv_rn(s.ex[tid], denom); + } + __syncthreads(); + } + float o = 0.F; +#pragma unroll + for (int c = 0; c < kTopk / 2; ++c) + { + int const cc = hc * (kTopk / 2) + c; + o += s.recv[cc][dl] * s.wgt[cc]; + } + if (hc) + s.part[dl] = o; + __syncthreads(); + if (!hc) + { + o += s.part[dl]; + if constexpr (ONE_DIV) + { + float den = 0.F; +#pragma unroll + for (int c = 0; c < kTopk; ++c) + den += s.ex[c]; + o = den > 0.F ? __fdiv_rn(o, den) : 0.F; + } + int const col = h * kDim + dh * kHalfDim + dl; + __nv_bfloat16 const ob = __float2bfloat16_rn(o); + p.out[tok * p.out_stride + col] = ob; + if constexpr (MXFP8) + { + float const x = __bfloat162float(ob); + float amax = fabsf(x); +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + float const xq = x * qs; + uint16_t q2; + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2) : "f"(0.F), "f"(xq)); + p.q8[static_cast(tok) * (p.num_sf_k_tiles * 128) + col] = static_cast(q2 & 0xFF); + if (lane == 0) + { + int const k_sf = col / 32; + int64_t const off = static_cast(tok / 128) * p.num_sf_k_tiles * 512 + (k_sf / 4) * 512 + + (tok % 32) * 16 + ((tok % 128) / 32) * 4 + k_sf % 4; + p.sf[off] = static_cast(e); + } + } + } +#endif // defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +cudaError_t configure() +{ + auto kernel = m3_tp8_lc8_attn_kernel; + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smemBytes())); + if (err != cudaSuccess) + { + return err; + } + return cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1); +} + +template +cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) +{ + // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs + // every variant eagerly first). + static cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + auto kernel = m3_tp8_lc8_attn_kernel; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(num_tokens * kTopk); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smemBytes(); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kTopk; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + return cudaLaunchKernelEx(&cfg, kernel, prm); +} + +template +cudaError_t launchPrewait(Params const& prm, int num_tokens, bool prewait, bool use_pdl, cudaStream_t stream) +{ + if (prm.one_div) + { + return prewait ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); + } + return prewait ? launch(prm, num_tokens, use_pdl, stream) + : launch(prm, num_tokens, use_pdl, stream); +} + +} // namespace + +cudaError_t invokeSparseDecodeLc8m(MinimaxM3Tp8SparseDecodeLc8Params const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream) +{ + if (linearSoftmax) + { + return mxfp8 ? launchPrewait(params, numTokens, prewait, usePdl, stream) + : launchPrewait(params, numTokens, prewait, usePdl, stream); + } + return mxfp8 ? launchPrewait(params, numTokens, prewait, usePdl, stream) + : launchPrewait(params, numTokens, prewait, usePdl, stream); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h new file mode 100644 index 000000000000..7f59a0f7ff59 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h @@ -0,0 +1,105 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ + +//! Operands of the MiniMax-M3 TP8 one-launch long-context sparse decode. Strides are in elements of the pointed-to +//! type. Geometry is fixed: 8 query heads on one KV head, head_dim 128, page 128, top-k 16. +struct MinimaxM3Tp8SparseDecodeLc8Params +{ + uint8_t const* q; // e4m3 [T, 8, 128] + int64_t q_tok, q_head; + uint8_t const* k; // packed e2m1 [pages, 1, 128, 64] + uint8_t const* v; + int64_t k_blk, v_blk; + uint8_t const* ksc; // e4m3 [pages, 1, 128, 8], token-major + uint8_t const* vsc; // e4m3, 4x4 token-quad order + int64_t ksc_blk, vsc_blk; + float const* kgs; + float const* vgs; + int const* topk; // [T, 16] view (kv head 0) + int64_t t_tok, t_k; + int const* block_table; + int64_t bt_stride; + int const* seq_lens; + int dql; + float sm_scale; + __nv_bfloat16* out; // [T, 1024] + int64_t out_stride; + uint8_t* q8; // e4m3 [rows, num_sf_k_tiles * 128] + uint8_t* sf; + int num_sf_k_tiles; + // Optional L2 prefetch of two read-only ranges (the layer's o_proj weight and scales), split across the grid's + // CTAs and issued before the grid dependency: a hint, so the output does not depend on it. nullptr = none. + uint8_t const* pf0; + int64_t pf0_bytes; + uint8_t const* pf1; + int64_t pf1_bytes; + // Optional next-layer KV prefetch: CTA (token, r) records its selected block id in this layer's [rows, 16] id + // buffer (n1_self) and pulls into L2 the K/V pages and scales of the block the next sparse layer's CTA (token, r) + // selected at the previous step (n1_next, from that layer's pools n1_k / n1_v / n1_ksc / n1_vsc, this layer's + // strides). Hints only. nullptr = off. + int* n1_self; + int const* n1_next; + uint8_t const* n1_k; + uint8_t const* n1_v; + uint8_t const* n1_ksc; + uint8_t const* n1_vsc; + int64_t n1_pages; // pages in the next layer's pool; a recorded block outside it is skipped + // Identity selection (short-sequence steps, every request within 16 blocks): CTA r attends block r of its + // token's request when r < min(16, blocks), from the device sequence length; topk is not read. + int identity; + // The step is an identity step (short sequences: no scorer / select ran this step), whether or not this kernel + // derives its blocks from the lengths (identity) or reads the staged table: q is then read after the grid wait, + // because the predecessor is the producer, which releases its dependents at entry, before it writes q. + int identity_step; + // One-division merge (MSA_ONE_DIV_MERGE): the CTAs push unnormalized partials with their row max and sum, and the + // merge divides once per output. Deterministic, a different rounding sequence from the per-chunk normalization. + int one_div; +}; + +//! One cluster of 16 CTAs per query token: CTA r attends the token's r-th selected block with transposed MMAs (the 8 +//! heads fill the MMA n dimension), then CTA r merges (head r / 2, dims (r % 2) * 64 + [0, 64)) of the 16 partials +//! over DSMEM into the BF16 [T, 1024] row (and, when mxfp8, its MXFP8 copy with swizzled UE8M0 scales). Dependents +//! are released right after the grid dependency wait. +//! +//! With prewait, the request's seq_lens entry, its q rows (widened to FP16) and its block-table row (up to 8192 +//! entries, in dynamic shared memory; longer rows are read from global memory) are staged before the grid dependency +//! wait; outputs are bit-identical to prewait=false. +//! +//! Instantiated variants: linearSoftmax x mxfp8 x prewait, all in {false, true}. The device code is SM100-family +//! only; on other architectures the kernel compiles to an empty body, so callers check the device first. +cudaError_t invokeSparseDecodeLc8m(MinimaxM3Tp8SparseDecodeLc8Params const& params, int32_t numTokens, + bool linearSoftmax, bool mxfp8, bool prewait, bool usePdl, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu new file mode 100644 index 000000000000..fd101e7f521c --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu @@ -0,0 +1,58 @@ +/* + * 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/minimaxM3Sm107Tp8/specOffsets.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kThreads = 256; + +__global__ void __launch_bounds__(kThreads) + specOffsetsKernel(int32_t* __restrict__ pos, int32_t const* __restrict__ posOff, int64_t nPos, + int32_t* __restrict__ kv, int32_t const* __restrict__ kvOff, int64_t nKv) +{ + for (int64_t i = threadIdx.x; i < nPos; i += kThreads) + { + pos[i] = static_cast(static_cast(pos[i]) + static_cast(posOff[i])); + } + for (int64_t j = threadIdx.x; j < nKv; j += kThreads) + { + kv[j] = static_cast(static_cast(kv[j]) + static_cast(kvOff[j])); + } +} + +} // namespace + +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream) +{ + specOffsetsKernel<<<1, kThreads, 0, stream>>>(pos, posOff, nPos, kv, kvOff, nKv); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h new file mode 100644 index 000000000000..3cfd05b73369 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h @@ -0,0 +1,41 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ + +//! The overlap scheduler's speculative step correction in one launch: pos[i] += posOff[i] for i < nPos and +//! kv[j] += kvOff[j] for j < nKv (int32, wrapping as the ATen adds do). A plain launch: it is a step-input writer that +//! pre-wait readers rely on (rubin-20261015/PREWAIT_AUDIT.md item b). +cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, + int64_t nKv, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu new file mode 100644 index 000000000000..b1ab46a36088 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu @@ -0,0 +1,111 @@ +/* + * 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/minimaxM3Sm107Tp8/stepPrologue.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kThreads = 512; +constexpr int kWarps = kThreads / 32; + +// One warp per copy. Plain (coherent) loads: a device copy may read bytes another warp of this CTA just stored, which +// the CTA barrier between the two kinds makes visible; a non-coherent load could return the stale line. +__device__ __forceinline__ void copyWarp(uint8_t* dst, uint8_t const* src, int64_t bytes, int lane) +{ + auto const d = reinterpret_cast(dst); + auto const s = reinterpret_cast(src); + if (((d | s | static_cast(bytes)) & 15) == 0) + { + auto* d16 = reinterpret_cast(dst); + auto const* s16 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 16; i += 32) + { + d16[i] = s16[i]; + } + } + else if (((d | s | static_cast(bytes)) & 3) == 0) + { + auto* d4 = reinterpret_cast(dst); + auto const* s4 = reinterpret_cast(src); + for (int64_t i = lane; i < bytes / 4; i += 32) + { + d4[i] = s4[i]; + } + } + else + { + for (int64_t i = lane; i < bytes; i += 32) + { + dst[i] = src[i]; + } + } +} + +// One CTA: a step's window is a few dozen copies of a few hundred bytes, so the launch is the cost; a CTA barrier +// orders the host copies before the device copies where the parallel scatter needed a second launch. +__global__ void __launch_bounds__(kThreads) stepPrologueKernel(uint8_t* window, int64_t* epoch, int64_t maxDesc) +{ + auto* header = reinterpret_cast(window); + int64_t const numHost = min(header[0], maxDesc); + int64_t const numDevice = min(header[1], maxDesc); + int64_t const* desc = header + 2; + int const warp = threadIdx.x / 32; + int const lane = threadIdx.x % 32; + if (epoch != nullptr && threadIdx.x == 0) + { + *epoch += 1; + } + for (int64_t i = warp; i < numHost; i += kWarps) + { + int64_t const* c = desc + 3 * i; + copyWarp(reinterpret_cast(c[0]), window + c[1], c[2], lane); + } + __syncthreads(); + for (int64_t i = warp; i < numDevice; i += kWarps) + { + // The device copies' slots follow all header[0] host copies, as the parallel scatter indexes them. + int64_t const* c = desc + 3 * (header[0] + i); + copyWarp(reinterpret_cast(c[0]), reinterpret_cast(c[1]), c[2], lane); + } + // Every thread has read the header and the descriptors: consume the window. + __syncthreads(); + if (threadIdx.x < 2) + { + header[threadIdx.x] = 0; + } +} + +} // namespace + +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream) +{ + stepPrologueKernel<<<1, kThreads, 0, stream>>>(window, epoch, maxDesc); + return cudaGetLastError(); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h new file mode 100644 index 000000000000..c125dede4361 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h @@ -0,0 +1,47 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ + +//! The decode graph's step-input prologue in one launch (step_inputs.StepInputPack with early_copy): the work of the +//! two parallel scatter launches and the header clear that follow them, in that order. +//! +//! The window is the pack's staging layout (kernels/stepInputsScatter.h): an int64 header `[numHostCopies, +//! numDeviceCopies]`, then one `(dst, src, bytes)` int64 triple per copy, host copies first; a host copy's `src` is a +//! byte offset into the window, a device copy's a device address. Up to `maxDesc` copies of each kind are written (the +//! slots the parallel scatter's grid covers), every host copy before any device copy, so a device copy may read a +//! destination a host copy filled; the copies of one kind must not overlap (the pack issues the window only then). +//! Afterwards the header is zeroed, so a replay with no new upload writes nothing. With `epoch` non-null, it is +//! incremented once. +cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu new file mode 100644 index 000000000000..4fad864f2d64 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu @@ -0,0 +1,100 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107Tp8 +{ + +namespace +{ + +constexpr int kThreads = 256; + +// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves): the +// instruction sequence of the Triton kernel at limit 7, alpha 1.702, beta 1 (Indexer kdev17, bitwise on SM107): the +// gate clamp as min.bf16x2 before the widening, the up clamp as min.xorsign.abs, (u + 1), (g * -1.702) * log2(e), +// ex2.approx, + 1, div.full 1 / that, s * g, the product, one rounding to bf16. +__device__ __forceinline__ uint32_t swigluOaiBf16x2(uint32_t gate, uint32_t up) +{ + uint32_t r; + asm("{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, %1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" + : "=r"(r) + : "r"(gate), "r"(up)); + return r; +} + +template +__global__ void __launch_bounds__(kThreads) + swigluOaiKernel(uint32_t const* __restrict__ x, uint32_t* __restrict__ out, int64_t rows, int64_t halfD) +{ + if constexpr (kPdl) + { + asm volatile("griddepcontrol.wait;" ::: "memory"); + asm volatile("griddepcontrol.launch_dependents;"); + } + int64_t const i = blockIdx.x * static_cast(kThreads) + threadIdx.x; + if (i >= rows * halfD) + { + return; + } + int64_t const row = i / halfD; + int64_t const col = i - row * halfD; + uint32_t const* xr = x + row * 2 * halfD; + out[i] = swigluOaiBf16x2(xr[col], xr[halfD + col]); +} + +} // namespace + +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream) +{ + int64_t const n = rows * halfD; + cudaLaunchConfig_t cfg{}; + cfg.gridDim = dim3(static_cast((n + kThreads - 1) / kThreads)); + cfg.blockDim = dim3(kThreads); + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + cfg.attrs = attr; + cfg.numAttrs = 1; + return pdl ? cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD) + : cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD); +} + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h new file mode 100644 index 000000000000..78033e018997 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h @@ -0,0 +1,43 @@ +/* + * 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 +{ +namespace minimaxM3Sm107Tp8 +{ + +//! The dense MLP's swigluoai activation: out[m, j] = bf16(min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + +//! 1)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both contiguous, d even): bit for bit +//! the Triton silu_and_mul kernel at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With `pdl` the launch allows +//! programmatic serialization and the grid waits on its dependency, then releases its dependents before it computes +//! (the down GEMV stages its weights during the activation and reads it only after its own grid wait). +cudaError_t invokeSwigluOai( + uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream); + +} // namespace minimaxM3Sm107Tp8 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/quantization.cu b/cpp/tensorrt_llm/kernels/quantization.cu index f6cf68053094..33a3801a5232 100644 --- a/cpp/tensorrt_llm/kernels/quantization.cu +++ b/cpp/tensorrt_llm/kernels/quantization.cu @@ -194,6 +194,9 @@ void invokeFP4Quantization(int b, int m, int n, T const* input, float const* SFS //////////////////////////////////////////////////////////////////////////////////////////////////// // MXFP8 Quantization +// Kill switch for the SM107 quantize_mxfp8_rows path of invokeMxFP8Quantization. +constexpr bool kMxfp8QuantSm107Rows = true; + template void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, int64_t* output, int32_t* SFOuput, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream) @@ -207,6 +210,33 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, return; } + // SM107: one unpadded matrix with swizzled K32 scales runs quantize_mxfp8_rows (same math, one column chunk per + // thread, 8 rows in flight); bitwise the same output. + if constexpr (SF_VEC_SIZE == 32 && SF_OUTPUT_VEC_SIZE == 32) + { + int const colThreads = n / CVT_ELTS_PER_THREAD; + if (kMxfp8QuantSm107Rows && tensorrt_llm::common::getSMVersion() == 107 && b == 1 && n == padded_n + && layout == QuantizationSFLayout::SWIZZLED && n % 128 == 0 && colThreads <= 1024 && colThreads % 32 == 0) + { + constexpr int kRows = 8; + int const ctasPerSm = colThreads > 512 ? 1 : 2; + dim3 grid(std::min(int(PadUpFn(m, 128) + kRows - 1) / kRows, multiProcessorCount * ctasPerSm)); + cudaLaunchConfig_t config; + config.gridDim = grid; + config.blockDim = dim3(colThreads); + 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, quantize_mxfp8_rows, m, n, input, + reinterpret_cast(output), reinterpret_cast(SFOuput))); + return; + } + } + // Grid, Block size. // Each thread converts 8 values. dim3 block(std::min(int(padded_n / CVT_ELTS_PER_THREAD), 512)); diff --git a/cpp/tensorrt_llm/kernels/quantization.cuh b/cpp/tensorrt_llm/kernels/quantization.cuh index 926c2439e3ed..3b6b813dd10f 100644 --- a/cpp/tensorrt_llm/kernels/quantization.cuh +++ b/cpp/tensorrt_llm/kernels/quantization.cuh @@ -980,6 +980,64 @@ quantize_with_block_size( #endif } +//! MXFP8 (K32 UE8M0 scales, swizzled 128x4 layout) quantization of one unpadded [numRows, numCols] matrix with the +//! per-element math of quantize_with_block_size (cvt_warp_fp16_to_mxfp8): one 8-element column chunk per thread +//! (blockDim.x == numCols / 8) and ROWS rows loaded before any is converted, so more bytes are in flight per thread. +//! Requires numCols % 128 == 0 and numCols / 8 <= 1024 a multiple of 32. +template +__global__ void __launch_bounds__(1024) + quantize_mxfp8_rows(int32_t numRows, int32_t numCols, Type const* in, uint64_t* out, uint32_t* SFout) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using PackedVec = PackedVec; + static constexpr int CVT_NUM_THREADS_PER_SF = 32 / CVT_ELTS_PER_THREAD; + int const colIdx = threadIdx.x; + int const numColThreads = numCols / CVT_ELTS_PER_THREAD; + int const numPaddedRowsForSf = PadUpFn(numRows, 128); + std::optional const optionalBatchIdx = 0; + std::optional const optionalNumRows = numRows; + + cudaGridDependencySynchronize(); + for (int rowBase = blockIdx.x * ROWS; rowBase < numPaddedRowsForSf; rowBase += gridDim.x * ROWS) + { + PackedVec inVecs[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; r++) + { + if (rowBase + r < numRows) + { + inVecs[r] = reinterpret_cast( + in)[static_cast(rowBase + r) * numColThreads + colIdx]; + } + } +#pragma unroll + for (int r = 0; r < ROWS; r++) + { + int const rowIdx = rowBase + r; + if (rowIdx >= numPaddedRowsForSf) + { + break; + } + auto sf_out = cvt_quant_get_sf_out_offset( + optionalBatchIdx, rowIdx, colIdx, optionalNumRows, numCols / 32, SFout, QuantizationSFLayout::SWIZZLED); + if (rowIdx < numRows) + { + out[static_cast(rowIdx) * numColThreads + colIdx] + = cvt_warp_fp16_to_mxfp8(inVecs[r], sf_out); + } + else + { + cvt_store_replicated_sf<1>(sf_out, 0x00); + } + } + } + // See quantize_with_block_size: every thread's stores are visible before any thread triggers completion. + __threadfence(); + __syncthreads(); + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + __global__ void block_scale_interleave_kernel( int numbatches, int numRows, int numCols, uint8_t const* SFIn, uint8_t* SFOutput); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index deb8884ad99d..ef5aa4a6ad5b 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -255,7 +255,7 @@ class TllmGenFmhaKernel } } - static bool shouldUseNvrtc(FmhaOptions const& options) + bool shouldUseNvrtc(FmhaOptions const& options) const { #ifdef TLLM_RUBIN_FEATURES // Spcompress kernels must use precompiled cubins. @@ -264,7 +264,20 @@ class TllmGenFmhaKernel return false; } #endif // TLLM_RUBIN_FEATURES - // Sparse MQA/GQA uses NVRTC path for now because no model really uses it. + // The cubin table has no dense paged-KV kernels for 128-token pages, and the page size is + // part of the kernel hash. On SM107, generate those BF16/E4M3 context and KeepsMmaAb + // generation kernels for the actual cache layout instead. + constexpr int kPageSize128 = 128; + if (mSM == kSM_107 + && (isContextKernel(options.mFmhaKernelType) || isKeepsMmaAbForGenerationKernel(options.mFmhaKernelType)) + && !options.mIsMlaGen && !isTokenSparse(options.mSparseType) && isPagedKv(options.mQkvLayout) + && options.mNumTokensPerPage == kPageSize128 && options.mDtypeQ == options.mDtypeKv + && (options.mDtypeQ == tg::Dtype::Bfloat16 || options.mDtypeQ == tg::Dtype::E4m3) + && options.mDtypeOut == tg::Dtype::Bfloat16) + { + return true; + } + // Sparse MQA/GQA uses NVRTC path for now because no model really uses it. if (isStaticTokenSparse(options.mSparseType) && !options.mIsMlaGen) { return true; @@ -376,6 +389,20 @@ class TllmGenFmhaKernel 1152, 1280, 1408, 1536, 1664, 1792, 1920, 2048, 2176, 2304, 2432, 2560, 2688, 2816, 2944, 3072, 3328, 3584, 3840, 4096, 4352, 4608, 4864, 5120, 5376, 5632, 5888, 6144, 6656, 7168, 7680, 8192, 8704, 9216, 9728, 10240, 10752, 11264, 11776, 12288, 13312, 14336, 16384, 20480, 24576, 32768, 40960, 49152, 65536, 98304, 131072}; + // The same grid continued past 131072 in 32768-token steps to 1M tokens: kernel selection still changes there + // (e.g. MultiCtasKv vs MultiCtasKvCga), so a long-context step whose + // KV length falls between two grid points otherwise compiles its kernel inside a served request (~4-5 s). + // Selection is a cheap heuristic and compiled kernels are cached by config, so the finer grid only adds the + // compiles of configs the coarser grid missed. Kill: TRTLLM_FMHA_LONG_WARMUP=0. + inline static std::vector const kLongWarmupSeqLenKvCandidates = [] + { + std::vector sizes = kDefaultWarmupSeqLenKvCandidates; + for (int size = 131072 + 32768; size <= 1048576; size += 32768) + { + sizes.push_back(size); + } + return sizes; + }(); // Selection of seqLenQ: typically our prefill kernels would not use JIT and seqLenQ is fixed during decode. // However seqLenQ can change when we use generation kernel to do prefill. SeqLenQ is sensitive in the range 1-16. // Below sequence are chosen to cover most sparse attention cases. @@ -412,6 +439,11 @@ class TllmGenFmhaKernel std::tie(options, optionsFromArgs, ctaDim) = selectKernelAutotuned(options, optionsFromArgs, params.mMultiProcessorCount); + // Same override as run(), so the warmed config matches the dispatched one. + if (mNumEltsPerSageAttnBlkP + mNumEltsPerSageAttnBlkV > 0) + { + options.mTileScheduler = TileScheduler::Static; + } checkFmhaOptions(options, optionsFromArgs); updateFmhaOptions(options, optionsFromArgs); @@ -422,10 +454,12 @@ class TllmGenFmhaKernel if (shouldUseNvrtc(options)) { + // Compile exactly what run() compiles (same target arch), so the warmed kernel is the one run() hits. + FmhaOptions const nvrtcOptions = nvrtcCompileOptions(options); FmhaConfig fmhaConfig; - fmhaConfig.mOptions = options; + fmhaConfig.mOptions = nvrtcOptions; std::ostringstream sstream; - populateJsonConfig(options, sstream); + populateJsonConfig(nvrtcOptions, sstream); fmhaConfig.mGenCfgJsonStr = sstream.str(); fmhaConfig.mExecPath = getExecPath().c_str(); @@ -443,9 +477,73 @@ class TllmGenFmhaKernel } } + // Context-phase grid: the kernel variant (e.g. static vs persistent tile scheduler, tile shape) follows the CTA + // count, i.e. batch x ceil(seqLenQ / tileQ), so cover each power of two and the tile boundary just past it + // (chunked prefill runs chunk + 1 query tokens), with every KV length >= seqLenQ (cached prefix) from the KV grid. + inline static std::vector const kContextWarmupBatchSizeCandidates = {1, 2, 4, 8, 16, 32, 64, 128, 256}; + inline static std::vector const kContextWarmupSeqLenQCandidates = {1, 128, 129, 256, 257, 512, 513, 1024, 1025, + 2048, 2049, 4096, 4097, 8192, 8193, 16384, 16385, 32768, 32769, 65536, 65537, 131072}; + + void runContextJITWarmupGrid(RunnerParams const& runnerParams) + { + static bool const contextWarmup + = tensorrt_llm::common::getIntEnv("TRTLLM_FMHA_CONTEXT_WARMUP").value_or(1) != 0; + int const maxBatchSize = runnerParams.mJITWarmupMaxNumRequests; + int const maxSeqLenQ = runnerParams.mJITWarmupMaxSeqLenQ; + int const maxSeqLenKv = runnerParams.mJITWarmupMaxSeqLenKv; + if (!contextWarmup || maxBatchSize <= 0 || maxSeqLenQ <= 0 || maxSeqLenKv <= 0) + { + return; + } + static bool const longWarmup = tensorrt_llm::common::getIntEnv("TRTLLM_FMHA_LONG_WARMUP").value_or(1) != 0; + std::vector const seqLenKvCandidates = makeWarmupCandidateSizes( + longWarmup ? kLongWarmupSeqLenKvCandidates : kDefaultWarmupSeqLenKvCandidates, maxSeqLenKv); + auto warmupParams = runnerParams; + for (int batchSize : makeWarmupCandidateSizes(kContextWarmupBatchSizeCandidates, maxBatchSize)) + { + warmupParams.mBatchSize = batchSize; + for (int seqLenQ : makeWarmupCandidateSizes(kContextWarmupSeqLenQCandidates, maxSeqLenQ)) + { + warmupParams.mMaxSeqLenQ = seqLenQ; + for (int seqLenKv : seqLenKvCandidates) + { + if (seqLenKv < seqLenQ) + { + continue; + } + warmupParams.mMaxSeqLenKv = seqLenKv; + warmupParams.mSumOfSeqLensQ = static_cast( + std::min(static_cast(batchSize) * seqLenQ, std::numeric_limits::max())); + warmupParams.mSumOfSeqLensKv = static_cast( + std::min(static_cast(batchSize) * seqLenKv, std::numeric_limits::max())); + warmupOneKernel(warmupParams); + } + } + } + } + void runJITWarmupGridIfRequested(RunnerParams const& runnerParams) { - if (!runnerParams.mJITWarmup || runnerParams.mKernelType != FmhaKernelType::Generation) + if (!runnerParams.mJITWarmup) + { + return; + } + if (isContextKernel(runnerParams.mKernelType)) + { + // Only SM107 needs the dense paged-context NVRTC warmup grid. + // Other architectures retain their precompiled context path. + if (mSM != kSM_107) + { + return; + } + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + TLLM_CUDA_CHECK(cudaStreamIsCapturing(runnerParams.stream, &captureStatus)); + TLLM_CHECK_WITH_INFO(captureStatus == cudaStreamCaptureStatusNone, + "TRTLLM-Gen FMHA JIT warmup must not run during CUDA graph capture."); + runContextJITWarmupGrid(runnerParams); + return; + } + if (runnerParams.mKernelType != FmhaKernelType::Generation) { return; } @@ -472,7 +570,9 @@ class TllmGenFmhaKernel std::vector seqLenQCandidates = useGenKernelForPrefill ? makeWarmupCandidateSizes(kDefaultWarmupSeqLenQCandidates, maxSeqLenQ) : std::vector{runnerParams.mMaxSeqLenQ}; - std::vector seqLenKvCandidates = makeWarmupCandidateSizes(kDefaultWarmupSeqLenKvCandidates, maxSeqLenKv); + static bool const longWarmup = tensorrt_llm::common::getIntEnv("TRTLLM_FMHA_LONG_WARMUP").value_or(1) != 0; + std::vector seqLenKvCandidates = makeWarmupCandidateSizes( + longWarmup ? kLongWarmupSeqLenKvCandidates : kDefaultWarmupSeqLenKvCandidates, maxSeqLenKv); auto warmupParams = runnerParams; @@ -582,14 +682,7 @@ class TllmGenFmhaKernel if (shouldUseNvrtc(options)) { // nvrtc path - uses mFmhaInterface member for kernel caching - FmhaOptions nvrtcOptions = options; -#ifdef TLLM_RUBIN_FEATURES - // On Rubin, compile non-spcompress NVRTC kernels with Sm100f target for compatibility. - if (tg::isArchRubin(options.mCudaArch) && !options.mUsesSpcompress) - { - nvrtcOptions.mCudaArch = tg::CudaArch::Sm100f; - } -#endif // TLLM_RUBIN_FEATURES + FmhaOptions const nvrtcOptions = nvrtcCompileOptions(options); FmhaConfig fmhaConfig; fmhaConfig.mOptions = nvrtcOptions; @@ -682,6 +775,20 @@ class TllmGenFmhaKernel } private: + // The options NVRTC compiles with; shared by run() and warmupOneKernel() so both produce the same kernel. + FmhaOptions nvrtcCompileOptions(FmhaOptions const& options) const + { + FmhaOptions nvrtcOptions = options; +#ifdef TLLM_RUBIN_FEATURES + // On Rubin, compile non-spcompress NVRTC kernels with Sm100f target for compatibility. + if (tg::isArchRubin(options.mCudaArch) && !options.mUsesSpcompress) + { + nvrtcOptions.mCudaArch = tg::CudaArch::Sm100f; + } +#endif // TLLM_RUBIN_FEATURES + return nvrtcOptions; + } + // Autotuner entry for checkIfKernelExist, warmupOneKernel, and run. SM107 may select // CgaSmemReduction configs where clusterDimX * mMaxNumCtasPerSeqKv must stay <= 16; // the pre/post clamps below enforce that. Other architectures keep the legacy path. diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp index 059d29e43d85..6e7a6f1fbef9 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3RouterGemvOp.cpp @@ -43,6 +43,14 @@ void checkSmVersion(char const* op) TORCH_CHECK(sm == 100 || sm == 103, op, " requires SM100 or SM103, got SM", sm); } +// The TF32 tensor-core router also runs on SM107 (Rubin; its kernels are sm_100f-family cubins), there only with +// 16-CTA clusters: one reduction order on every rank (minimaxM3RouterTcCluster16 lets callers fall back otherwise). +void checkSmVersionTc(char const* op) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 100 || sm == 103 || sm == 107, op, " requires SM100, SM103 or SM107, got SM", sm); +} + void checkActivation(torch::Tensor const& x, torch::Tensor const& w, int64_t maxM) { TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.stride(1) == 1, @@ -96,7 +104,7 @@ std::vector minimaxM3RouterGemvFp4( std::vector minimaxM3RouterTc( torch::Tensor const& x, torch::Tensor const& w, std::optional const& sfScale, bool pdl) { - checkSmVersion("minimax_m3_router_tc"); + checkSmVersionTc("minimax_m3_router_tc"); TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.stride(1) == 1, "x must be a 2-D bf16 CUDA tensor with unit column stride"); int64_t const M = x.size(0), K = x.size(1); @@ -105,6 +113,10 @@ std::vector minimaxM3RouterTc( "minimax_m3_router_tc takes 1..48 rows and a contiguous fp32 [128, K] weight"); TORCH_CHECK(K % 1024 == 0 && (x.stride(0) * 2) % 16 == 0, "K must be a multiple of 1024 with 16-byte rows"); c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(tensorrt_llm::common::getSMVersion() != 107 + || tensorrt_llm::kernels::minimaxM3RouterTcCluster(static_cast(K)) == 16, + "minimax_m3_router_tc on SM107 runs only with 16-CTA clusters; this device cannot co-schedule them (use the " + "stock router)"); auto out = torch::empty({M, w.size(0)}, x.options().dtype(at::kFloat)); torch::Tensor q, sf; uint8_t* qp = nullptr; @@ -130,6 +142,14 @@ std::vector minimaxM3RouterTc( return {out}; } +//! Whether `minimax_m3_router_tc` runs on this device with 16-CTA clusters for hidden size `k` (device-static; the +//! SM107 caller falls back to the stock router otherwise). +bool minimaxM3RouterTcCluster16(int64_t k) +{ + checkSmVersionTc("minimax_m3_router_tc_cluster16"); + return tensorrt_llm::kernels::minimaxM3RouterTcCluster(static_cast(k)) == 16; +} + } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -139,6 +159,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def("minimax_m3_router_gemv(Tensor x, Tensor w, int rows, bool pdl) -> Tensor"); m.def("minimax_m3_router_gemv_fp4(Tensor x, Tensor w, Tensor sf_scale, bool pdl) -> Tensor[]"); m.def("minimax_m3_router_tc(Tensor x, Tensor w, Tensor? sf_scale, bool pdl) -> Tensor[]"); + m.def("minimax_m3_router_tc_cluster16(int k) -> bool", &tensorrt_llm::torch_ext::minimaxM3RouterTcCluster16); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp new file mode 100644 index 000000000000..99c182799cab --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp @@ -0,0 +1,125 @@ +/* + * 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/minimaxM3Sm107Tp2/denseQkNormRopeFp8.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace dqk = tensorrt_llm::kernels::minimaxM3Sm107Tp2::denseQkNormRopeFp8; + +//! QK RMSNorm + partial RoPE of a packed BF16 [tokens, (q + k + v) * 128] row of the MiniMax-M3 TP2 dense layers, +//! returned as a new FP8 E4M3 tensor of the same shape (V copy-cast; Q/K rounded to BF16 before E4M3, as stock). +torch::Tensor minimaxM3Sm107Tp2DenseQkNormRopeFp8(torch::Tensor const& qkv, int64_t numHeadsQ, int64_t numHeadsK, + int64_t numHeadsV, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, + double base, bool isNeox, torch::Tensor const& positionIds, bool useGemma) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); + TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); + TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); + TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + at::cuda::CUDAGuard guard(qkv.device()); + auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); + dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(rotaryDim), static_cast(eps), + qWeight.data_ptr(), kWeight.data_ptr(), static_cast(base), !isNeox, positionIds.data_ptr(), + useGemma, nullptr, at::cuda::getCurrentCUDAStream(qkv.get_device())); + return out; +} + +//! As minimaxM3Sm107Tp2DenseQkNormRopeFp8, and the first numLiveTokens rows' K and V are also stored into the layer's +//! FP8 pool kvBuffers (the HND view [pages, 2, kv heads, tokens per block, 128]) at their int32 cache slots; with +//! qCompact, every row's Q is also stored there compact ([tokens, q heads, 128]). +torch::Tensor minimaxM3Sm107Tp2DenseQkNormRopeFp8Kv(torch::Tensor const& qkv, int64_t numHeadsQ, int64_t numHeadsK, + int64_t numHeadsV, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, + double base, bool isNeox, torch::Tensor const& positionIds, bool useGemma, torch::Tensor const& kvBuffers, + torch::Tensor const& outCacheLoc, int64_t numLiveTokens, std::optional const& qCompact) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); + TORCH_CHECK(numHeadsK == numHeadsV, "the pool keeps one head count for K and V"); + TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); + TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); + TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + TORCH_CHECK(kvBuffers.is_cuda() && kvBuffers.scalar_type() == at::kFloat8_e4m3fn && kvBuffers.dim() == 5); + TORCH_CHECK(kvBuffers.size(1) == 2 && kvBuffers.size(2) == numHeadsK && kvBuffers.size(4) == dqk::kHeadDim + && kvBuffers.stride(4) == 1, + "kv_buffers must be the HND pool view [pages, 2, kv heads, tokens per block, 128]"); + TORCH_CHECK(outCacheLoc.is_cuda() && outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous()); + TORCH_CHECK(numLiveTokens >= 0 && numLiveTokens <= qkv.size(0) && numLiveTokens <= outCacheLoc.numel()); + at::cuda::CUDAGuard guard(qkv.device()); + auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); + dqk::KvCacheView view; + view.base = static_cast<__nv_fp8_e4m3*>(kvBuffers.data_ptr()); + view.outCacheLoc = outCacheLoc.data_ptr(); + view.pageStride = kvBuffers.stride(0); + view.kvStride = kvBuffers.stride(1); + view.headStride = kvBuffers.stride(2); + view.tokenStride = kvBuffers.stride(3); + view.tokensPerBlock = static_cast(kvBuffers.size(3)); + view.numLiveTokens = static_cast(numLiveTokens); + if (qCompact.has_value()) + { + TORCH_CHECK(qCompact->is_cuda() && qCompact->is_contiguous() && qCompact->scalar_type() == at::kFloat8_e4m3fn + && qCompact->numel() == qkv.size(0) * numHeadsQ * dqk::kHeadDim, + "q_compact must be a contiguous FP8 [tokens, q heads, 128] tensor"); + view.qCompact = static_cast<__nv_fp8_e4m3*>(qCompact->data_ptr()); + } + dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(rotaryDim), static_cast(eps), + qWeight.data_ptr(), kWeight.data_ptr(), static_cast(base), !isNeox, positionIds.data_ptr(), + useGemma, &view, at::cuda::getCurrentCUDAStream(qkv.get_device())); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, " + "bool use_gemma) -> Tensor"); + m.def( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, " + "bool use_gemma, Tensor(a!) kv_buffers, Tensor out_cache_loc, int num_live_tokens, Tensor(b!)? q_compact=None) " + "-> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2DenseQkNormRopeFp8); + m.impl("minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2DenseQkNormRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DraftStepOp.cpp new file mode 100644 index 000000000000..cc452ef4c776 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DraftStepOp.cpp @@ -0,0 +1,220 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/draftStep.h" + +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace ds = tensorrt_llm::kernels::minimaxM3Sm107Tp2::draftStep; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_sm107tp2_draft_step expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107tp2_draft_step requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +void checkInt(torch::Tensor const& t, int64_t n, char const* name) +{ + TORCH_CHECK(t.scalar_type() == at::kInt && t.is_contiguous() && t.numel() >= n, "draft_step: ", name, + " must be contiguous int32 with at least ", n, " elements"); +} + +} // namespace + +//! One fused greedy draft pass over `rows` <= 8 requests: the BF16 head (this rank's fragment-ordered shard wf +//! [N / 16, K / 16, 256], or with excOff / excPos / excW / expBase its 12-bit coding, wf uint8 [N * K * 3 / 2]; +//! global ids vocabOffset..) over x[xRows[b]] (or x[b]), the global argmax across the TP ranks through the multicast +//! slots wsMcPtr / wsUc, then draftOut[b, step] = idsNext[b] = token, posNext[b] = posSrc[xRows[b] or b] + 1 and the +//! KV-length bookkeeping (see draftStep.h). +void minimaxM3Sm107Tp2DraftStep(torch::Tensor const& x, OptTensor const& xRows, int64_t rows, torch::Tensor const& wf, + int64_t vocabOffset, torch::Tensor& keys, torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, + int64_t rank, int64_t world, torch::Tensor& draftOut, int64_t step, torch::Tensor& idsNext, + torch::Tensor const& posSrc, torch::Tensor& posNext, torch::Tensor& kvLens, OptTensor const& numAcc, + OptTensor const& seqLens, int64_t draftLen, int64_t numCtas, bool pdl, OptTensor const& outKeys, + OptTensor const& excOff, OptTensor const& excPos, OptTensor const& excW, int64_t expBase, int64_t tileLo, + int64_t tileHi, int64_t totalCtas, OptTensor const& bnd, OptTensor const& bndCnt) +{ + checkArch(x); + for (torch::Tensor const* t : std::initializer_list{ + &wf, &keys, &counters, &wsUc, &draftOut, &idsNext, &posSrc, &posNext, &kvLens}) + { + TORCH_CHECK(t->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + for (OptTensor const* t : std::initializer_list{ + &xRows, &numAcc, &seqLens, &outKeys, &excOff, &excPos, &excW, &bnd, &bndCnt}) + { + TORCH_CHECK(!t->has_value() || (*t)->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "draft_step: x must be contiguous BF16 [rows, K]"); + int const K = static_cast(x.size(1)); + TORCH_CHECK(K % 16 == 0, "draft_step: K must be a multiple of 16"); + TORCH_CHECK(rows >= 1 && rows <= ds::kMaxTokens, "draft_step: 1..", ds::kMaxTokens, " rows"); + bool const coded = excOff.has_value(); + TORCH_CHECK(coded == excPos.has_value() && coded == excW.has_value(), + "draft_step: exc_off, exc_pos and exc_w come together"); + int nTiles = 0; + if (coded) + { + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "draft_step: a coded head needs K a multiple of 64"); + TORCH_CHECK(excOff->scalar_type() == at::kInt && excOff->is_contiguous() && excOff->numel() >= 2, + "draft_step: exc_off must be int32 [tiles + 1]"); + nTiles = static_cast(excOff->numel() - 1); + TORCH_CHECK(wf.scalar_type() == at::kByte && wf.is_contiguous() + && wf.numel() == static_cast(nTiles) * 16 * K * 3 / 2, + "draft_step: a coded wf must be the contiguous uint8 stream of [tiles * 16, K]"); + TORCH_CHECK(excPos->scalar_type() == at::kInt && excPos->is_contiguous(), "draft_step: exc_pos must be int32"); + TORCH_CHECK(excW->scalar_type() == at::kFloat && excW->is_contiguous() && excW->numel() == 2 * excPos->numel(), + "draft_step: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "draft_step: exponent base out of range"); + } + else + { + TORCH_CHECK(wf.scalar_type() == at::kBFloat16 && wf.dim() == 3 && wf.is_contiguous() && wf.size(1) == K / 16 + && wf.size(2) == 256, + "draft_step: wf must be the fragment-ordered BF16 weight [N / 16, K / 16, 256]"); + nTiles = static_cast(wf.size(0)); + } + int64_t const rangeHi = tileHi > 0 ? tileHi : nTiles; + TORCH_CHECK(tileLo >= 0 && tileLo < rangeHi && rangeHi <= nTiles, "draft_step: tile range"); + TORCH_CHECK(nTiles >= 1 && numCtas >= 1 && (rangeHi - tileLo + numCtas - 1) / numCtas <= ds::kMaxTilesPerCta, + "draft_step: too many weight tiles per CTA"); + TORCH_CHECK(totalCtas == 0 || totalCtas >= numCtas, "draft_step: total_ctas"); + TORCH_CHECK(bnd.has_value() == bndCnt.has_value(), "draft_step: bnd and bnd_cnt come together"); + if (bnd.has_value()) + { + TORCH_CHECK(bnd->scalar_type() == at::kFloat && bnd->is_contiguous() && bnd->numel() >= int64_t{nTiles} * 128, + "draft_step: bnd must be zero-filled fp32 [tiles, 32, 4]"); + TORCH_CHECK(bndCnt->scalar_type() == at::kInt && bndCnt->is_contiguous() && bndCnt->numel() >= nTiles, + "draft_step: bnd_cnt must be zero-filled int32 [tiles]"); + // At least one tile per CTA: a tile is then shared by at most two CTAs (two addends in bnd). + TORCH_CHECK(rangeHi - tileLo >= numCtas, "draft_step: an equal-chunk split needs a tile per CTA"); + } + TORCH_CHECK(vocabOffset >= 0 && vocabOffset + 16LL * nTiles <= 0xFFFFFFFFLL, "draft_step: vocab offset"); + if (xRows.has_value()) + { + TORCH_CHECK(xRows->scalar_type() == at::kLong && xRows->is_contiguous() && xRows->numel() >= rows, + "draft_step: x_rows must be contiguous int64 [rows]"); + } + else + { + TORCH_CHECK(x.size(0) >= rows, "draft_step: x must hold a row per request"); + } + TORCH_CHECK(world >= 1 && rank >= 0 && rank < world, "draft_step: rank / world"); + TORCH_CHECK(keys.scalar_type() == at::kLong && keys.numel() >= rows, "draft_step: keys must be int64 [>= rows]"); + TORCH_CHECK(counters.scalar_type() == at::kInt && counters.numel() >= 2, "draft_step: counters must be int32 [2]"); + TORCH_CHECK(wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, + "draft_step: ws_uc must hold [2][8][world] slots"); + TORCH_CHECK(draftOut.scalar_type() == at::kInt && draftOut.dim() == 2 && draftOut.stride(1) == 1 + && draftOut.size(0) >= rows && step >= 0 && step < draftOut.size(1), + "draft_step: draft_out must be int32 [rows, > step] with unit column stride"); + checkInt(idsNext, rows, "ids_next"); + checkInt(posNext, rows, "pos_next"); + checkInt(kvLens, rows, "kv_lens"); + TORCH_CHECK(posSrc.scalar_type() == at::kInt && posSrc.is_contiguous(), "draft_step: pos_src must be int32"); + TORCH_CHECK(numAcc.has_value() == seqLens.has_value(), "draft_step: num_acc and seq_lens come together"); + TORCH_CHECK(!numAcc.has_value() || step == 0, "draft_step: the KV rollback (num_acc) is pass 0's"); + if (numAcc.has_value()) + { + checkInt(*numAcc, rows, "num_acc"); + checkInt(*seqLens, rows, "seq_lens"); + } + if (outKeys.has_value()) + { + TORCH_CHECK(outKeys->scalar_type() == at::kLong && outKeys->is_contiguous() && outKeys->numel() >= rows, + "draft_step: out_keys must be contiguous int64 [rows]"); + } + ds::Params p{}; + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_rows = xRows.has_value() ? reinterpret_cast(xRows->data_ptr()) : nullptr; + p.M = static_cast(rows); + p.wf = reinterpret_cast(wf.data_ptr()); + p.n_tiles = nTiles; + p.vocab_offset = static_cast(vocabOffset); + p.keys = reinterpret_cast(keys.data_ptr()); + p.counters = reinterpret_cast(counters.data_ptr()); + p.ws_uc = reinterpret_cast(wsUc.data_ptr()); + p.ws_mc = reinterpret_cast(wsMcPtr); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.out_keys = outKeys.has_value() ? reinterpret_cast(outKeys->data_ptr()) : nullptr; + p.st = ds::StepArgs{draftOut.data_ptr(), static_cast(draftOut.stride(0)), static_cast(step), + idsNext.data_ptr(), posSrc.data_ptr(), posNext.data_ptr(), kvLens.data_ptr(), + numAcc.has_value() ? numAcc->data_ptr() : nullptr, + seqLens.has_value() ? seqLens->data_ptr() : nullptr, static_cast(draftLen)}; + p.cd = ds::CodedArgs{coded ? excOff->data_ptr() : nullptr, + coded ? reinterpret_cast(excPos->data_ptr()) : nullptr, + coded ? reinterpret_cast(excW->data_ptr()) : nullptr, + static_cast((expBase << 7) | (expBase << 23))}; + p.tile_lo = static_cast(tileLo); + p.tile_hi = static_cast(tileHi); + p.total_ctas = static_cast(totalCtas); + p.bnd = bnd.has_value() ? reinterpret_cast(bnd->data_ptr()) : nullptr; + p.bnd_cnt = bnd.has_value() ? reinterpret_cast(bndCnt->data_ptr()) : nullptr; + check(ds::invokeDraftStep(p, K, static_cast(numCtas), pdl, at::cuda::getCurrentCUDAStream(x.get_device()))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// Mutations through the raw multicast address ws_mc_ptr are not visible to the schema: that buffer is owned by the +// Python workspace and never aliased by the op's tensor arguments. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp2_draft_step(Tensor x, Tensor? x_rows, int rows, Tensor wf, int vocab_offset, Tensor(a!) " + "keys, " + "Tensor(b!) counters, Tensor(c!) ws_uc, int ws_mc_ptr, int rank, int world, Tensor(d!) draft_out, int step, " + "Tensor(e!) ids_next, Tensor pos_src, Tensor(f!) pos_next, Tensor(g!) kv_lens, Tensor? num_acc, " + "Tensor(h!)? seq_lens, int draft_len, int num_ctas, bool pdl, Tensor(i!)? out_keys, Tensor? exc_off=None, " + "Tensor? exc_pos=None, Tensor? exc_w=None, int exp_base=0, int tile_lo=0, int tile_hi=0, int total_ctas=0, " + "Tensor(j!)? bnd=None, Tensor(k!)? bnd_cnt=None) " + "-> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp2_draft_step", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2DraftStep); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp new file mode 100644 index 000000000000..6b154d20d3f2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp @@ -0,0 +1,103 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp2::drafterRopeFp8Kv; + +//! The MiniMax-M3 TP2 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h): RoPE on +//! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages +//! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. +void minimaxM3Sm107Tp2DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, + torch::Tensor const& seqLens, std::optional const& specPositionOffsets, torch::Tensor& kvPages, + torch::Tensor const& blockTables, torch::Tensor& qOut, int64_t qLen, int64_t numHeadsQ, int64_t numHeadsKv, + int64_t rotaryDim, int64_t tokensPerBlock, int64_t ropeMode, bool pdl, bool maxSharedCarveout) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + 2 * numHeadsKv) * drk::kHeadDim, "qkv is not [tokens, (Hq + 2 Hkv) * 128]"); + TORCH_CHECK(qLen >= 1 && qkv.size(0) % qLen == 0, "the rows are not whole passes of q_len rows"); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= drk::kHeadDim && rotaryDim % 16 == 0, "unsupported rotary_dim"); + TORCH_CHECK(cosSin.is_cuda() && cosSin.is_contiguous() && cosSin.scalar_type() == at::kFloat + && cosSin.numel() % rotaryDim == 0, + "cos_sin must be the fp32 [positions, rotary_dim / 2] (cos, sin) table"); + int64_t const requests = qkv.size(0) / qLen; + TORCH_CHECK(seqLens.is_cuda() && seqLens.scalar_type() == at::kInt && seqLens.numel() >= requests); + if (specPositionOffsets.has_value()) + { + TORCH_CHECK(specPositionOffsets->is_cuda() && specPositionOffsets->is_contiguous() + && specPositionOffsets->scalar_type() == at::kInt && specPositionOffsets->numel() >= requests * qLen); + } + TORCH_CHECK(kvPages.is_cuda() && kvPages.is_contiguous() && kvPages.dim() == 4 + && kvPages.scalar_type() == at::kFloat8_e4m3fn && kvPages.size(1) == numHeadsKv + && kvPages.size(2) == tokensPerBlock && kvPages.size(3) == drk::kHeadDim, + "kv_pages must be the E4M3 [sub-pages, Hkv, tokens_per_block, 128] cache"); + TORCH_CHECK(blockTables.is_cuda() && blockTables.scalar_type() == at::kInt && blockTables.dim() == 3 + && blockTables.size(0) >= requests && blockTables.size(1) == 2 && blockTables.stride(2) == 1, + "block_tables must be the int32 [requests, 2, max blocks] K / V sub-page table"); + TORCH_CHECK(qOut.is_cuda() && qOut.is_contiguous() && qOut.scalar_type() == at::kFloat8_e4m3fn + && qOut.numel() == qkv.size(0) * numHeadsQ * drk::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{qkv.device()}; + drk::Params params{}; + params.qkv = reinterpret_cast<__nv_bfloat16 const*>(qkv.data_ptr()); + params.cosSin = reinterpret_cast(cosSin.data_ptr()); + params.seqLens = seqLens.data_ptr(); + params.specPositionOffsets = specPositionOffsets.has_value() ? specPositionOffsets->data_ptr() : nullptr; + params.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + params.pool = reinterpret_cast<__nv_fp8_e4m3*>(kvPages.data_ptr()); + params.blockTable = blockTables.data_ptr(); + params.blockTableRequestStride = blockTables.stride(0); + params.blockTableKvStride = blockTables.stride(1); + params.numTokens = static_cast(qkv.size(0)); + params.qLen = static_cast(qLen); + params.numHeadsQ = static_cast(numHeadsQ); + params.numHeadsKv = static_cast(numHeadsKv); + params.rotaryDim = static_cast(rotaryDim); + params.tokensPerBlock = static_cast(tokensPerBlock); + params.ropeMode = static_cast(ropeMode); + cudaError_t const err + = drk::launch(params, pdl, maxSharedCarveout, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp2_drafter_rope_fp8_kv: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp2_drafter_rope_fp8_kv(Tensor qkv, Tensor cos_sin, Tensor seq_lens, " + "Tensor? spec_position_offsets, Tensor(a!) kv_pages, Tensor block_tables, Tensor(b!) q_out, int q_len, " + "int num_heads_q, int num_heads_kv, int rotary_dim, int tokens_per_block, int rope_mode, bool pdl, bool " + "max_shared_carveout=True) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp2_drafter_rope_fp8_kv", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2DrafterRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp new file mode 100644 index 000000000000..de42edcc82c0 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FoldTailOp.cpp @@ -0,0 +1,130 @@ +/* + * 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/minimaxM3Sm107Tp2/foldTail.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxWidth = 32; + +void check2d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 2 && t.size(1) <= kMaxWidth, + "minimax_m3_sm107_tp2_fold_tail: ", name, " must be a 2-D int32 CUDA view at most 32 wide"); +} + +void check1d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp2_fold_tail: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +} // namespace + +//! The speculative step's forward-tail store and next-step staging in one launch (the C++ form of the Triton +//! kernels/m3_fold_tail.py, the same writes). The caller checks the row counts against the buffers (fold_tail). +void minimaxM3Sm107Tp2FoldTail(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + TORCH_CHECK(ctrl.is_cuda() && ctrl.scalar_type() == torch::kInt32 && ctrl.numel() >= 1, + "minimax_m3_sm107_tp2_fold_tail: ctrl must be an int32 CUDA tensor"); + TORCH_CHECK(slots.is_cuda() && slots.scalar_type() == torch::kInt64 && slots.dim() == 1 && slots.stride(0) == 1, + "minimax_m3_sm107_tp2_fold_tail: slots must be a contiguous 1-D int64 CUDA tensor"); + TORCH_CHECK(rows > 0 && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, + "minimax_m3_sm107_tp2_fold_tail: rows must be positive and tokens_per_step in [1, 32]"); + check2d(srcNew, "src_new"); + check2d(dstNew, "dst_new"); + check2d(srcNext, "src_next"); + check2d(dstNext, "dst_next"); + check2d(srcDraft, "src_draft"); + check2d(dstDraft, "dst_draft"); + check1d(srcLens, "src_lens"); + check1d(dstLens, "dst_lens"); + check1d(ids, "ids"); + check1d(draftIds, "draft_ids"); + check1d(posOff, "pos_off"); + check1d(kvOff, "kv_off"); + c10::cuda::CUDAGuard const deviceGuard{ctrl.device()}; + + tensorrt_llm::kernels::minimaxM3Sm107Tp2::FoldTailParams p{}; + p.ctrl = ctrl.data_ptr(); + p.slots = slots.data_ptr(); + p.srcNew = srcNew.data_ptr(); + p.sn0 = srcNew.stride(0); + p.sn1 = srcNew.stride(1); + p.snW = static_cast(srcNew.size(1)); + p.dstNew = dstNew.data_ptr(); + p.dn0 = dstNew.stride(0); + p.dn1 = dstNew.stride(1); + p.dnW = static_cast(dstNew.size(1)); + p.srcNext = srcNext.data_ptr(); + p.sx0 = srcNext.stride(0); + p.sx1 = srcNext.stride(1); + p.sxW = static_cast(srcNext.size(1)); + p.dstNext = dstNext.data_ptr(); + p.dx0 = dstNext.stride(0); + p.dx1 = dstNext.stride(1); + p.dxW = static_cast(dstNext.size(1)); + p.srcDraft = srcDraft.data_ptr(); + p.sd0 = srcDraft.stride(0); + p.sd1 = srcDraft.stride(1); + p.sdW = static_cast(srcDraft.size(1)); + p.dstDraft = dstDraft.data_ptr(); + p.dd0 = dstDraft.stride(0); + p.dd1 = dstDraft.stride(1); + p.ddW = static_cast(dstDraft.size(1)); + p.srcLens = srcLens.data_ptr(); + p.dstLens = dstLens.data_ptr(); + p.ids = ids.data_ptr(); + p.draftIds = draftIds.data_ptr(); + p.posOff = posOff.data_ptr(); + p.kvOff = kvOff.data_ptr(); + p.tokensPerStep = static_cast(tokensPerStep); + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeFoldTail( + p, static_cast(rows), at::cuda::getCurrentCUDAStream(ctrl.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_fold_tail launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_fold_tail(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) dst_new, " + "Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_fold_tail", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2FoldTail); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp new file mode 100644 index 000000000000..8db408008dc2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp @@ -0,0 +1,74 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107Tp2::fullCtxReduce; + +//! Merges the PrimTS split-KV decode scratch (partial_o [B, Hkv, S, R, 128] 16-bit, partial_stats [B, Hkv, S, R] +//! fp32 log2-LSE) of the MiniMax-M3 TP2 full-context decode into out (BF16 [B * SQ, Hkv * hR, 128]). +void minimaxM3Sm107Tp2FullCtxReduce(torch::Tensor const& partialO, torch::Tensor const& partialStats, + torch::Tensor const& seqLensKv, torch::Tensor& out, int64_t seqLenQ, int64_t qTokensPerCta, int64_t tileKv, + int64_t instsKv, bool causal, bool pdl) +{ + TORCH_CHECK(partialO.is_cuda() && partialO.is_contiguous() && partialO.dim() == 5); + TORCH_CHECK(partialO.scalar_type() == at::kHalf || partialO.scalar_type() == at::kBFloat16); + TORCH_CHECK(partialO.size(4) == fcr::kHeadDim && partialO.size(2) <= fcr::kMaxSplits); + TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() + && partialStats.sizes() == partialO.sizes().slice(0, 4)); + TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + int const rows = static_cast(partialO.size(3)); + TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); + int const headsPerKv = rows / static_cast(seqLenQ); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.is_contiguous() + && out.numel() == partialO.size(0) * seqLenQ * partialO.size(1) * headsPerKv * fcr::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{partialO.device()}; + fcr::Params params{partialO.data_ptr(), partialStats.data_ptr(), seqLensKv.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), static_cast(partialO.size(0)), + static_cast(partialO.size(1)), static_cast(partialO.size(2)), rows, static_cast(seqLenQ), + headsPerKv, static_cast(qTokensPerCta), static_cast(tileKv), static_cast(instsKv), causal, + partialO.scalar_type() == at::kBFloat16}; + cudaError_t const err + = fcr::invokeFullCtxReduce(params, pdl, at::cuda::getCurrentCUDAStream(partialO.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_fullctx_reduce: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_fullctx_reduce(Tensor partial_o, Tensor partial_stats, Tensor seq_lens_kv, " + "Tensor(a!) out, int seq_len_q, int q_tokens_per_cta, int tile_kv, int insts_kv, bool causal, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_fullctx_reduce", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2FullCtxReduce); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp new file mode 100644 index 000000000000..a0f27263a152 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp @@ -0,0 +1,181 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h" + +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxTokens = 48; +constexpr int64_t kCluster = 8; +constexpr int64_t kMaxWorld = 8; + +} // namespace + +//! MiniMax-M3 SM107 TP2 multicast all-reduce (stock rounding points: BF16 rank sum, then the residual add) + residual + +//! RMSNorm (+ optional MoE finalize, + optional MXFP8 output). +//! +//! @param x BF16 [M, H] partial, or with `perm` the permuted FC2 output [rows, H]. +//! @param x2 Optional BF16 [M, H] second addend (the shared expert in finalize mode). +//! @param residual BF16 [M, H], 1 <= M <= 48. +//! @param normWeight BF16 [H]. +//! @param wsUc This rank's unicast view of the exchange workspace (int32 words); re-armed in place. +//! @param wsMcPtr Device address of the workspace's multicast mapping. +//! @param flags Int32 per-CTA call-parity counters, at least maxt * 8; updated in place. +//! @param cluster Cluster size; 8 is the only one built. +//! @param perm, scales Finalize mode: int32 [M, topk] permuted-row index and fp32 / bf16 [M, topk] routing weights. +//! @param maxt Token capacity of the workspace bucket (8, 16 or 48). +//! @param quant Also return the MXFP8 copy of the normed rows (E4M3 [M, H], swizzled UE8M0 scales). +//! @param prepushed 0, or 2 / 3 when producers already pushed this rank's rows (see the kernel header). +//! @param opts Epilogue variant; every variant is bit-identical. +//! @param capOut Optional BF16 [>= M, H] view (row stride a multiple of 8, 16-byte aligned) that receives the Eagle3 +//! capture of the updated residual: m3_spec_tail capture_norm with capW [H] and capEps, bit for bit (H 6144; not the +//! drafter boundary, opts bit 16). +//! @return [normed, updated residual] or, with `quant`, [normed, updated residual, E4M3 rows, scale bytes]. +std::vector minimaxM3Sm107Tp2McAllreduceNorm(torch::Tensor const& x, + std::optional const& x2, torch::Tensor const& residual, torch::Tensor const& normWeight, double eps, + torch::Tensor& wsUc, int64_t wsMcPtr, torch::Tensor& flags, int64_t rank, int64_t world, int64_t cluster, bool pdl, + std::optional const& perm, std::optional const& scales, int64_t maxt, bool quant, + int64_t prepushed, int64_t opts, std::optional const& capOut, + std::optional const& capW, double capEps) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 107, "minimax_m3_sm107_tp2_mc_allreduce_norm requires SM107, got SM", sm); + TORCH_CHECK(cluster == kCluster, "minimax_m3_sm107_tp2_mc_allreduce_norm: only cluster 8 is built, got ", cluster); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107Tp2McAllreduceNormSupported( + static_cast(maxt), static_cast(opts)), + "minimax_m3_sm107_tp2_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); + TORCH_CHECK( + world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); + + tensorrt_llm::kernels::MinimaxM3Sm107Tp2McAllreduceNormParams p; + if (x2.has_value()) + { + TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), + "x2 must be a contiguous bf16 tensor shaped like residual"); + p.x2 = reinterpret_cast<__nv_bfloat16 const*>(x2->data_ptr()); + } + if (perm.has_value()) + { + // Finalize mode: x is the permuted FC2 output [rows, H]. + TORCH_CHECK(scales.has_value() && perm->scalar_type() == at::kInt && perm->is_contiguous() && perm->dim() == 2 + && perm->size(0) == residual.size(0), + "perm must be a contiguous int32 [M, topk] tensor, with scales"); + TORCH_CHECK(scales->is_contiguous() && scales->numel() == perm->numel() + && (scales->scalar_type() == at::kFloat || scales->scalar_type() == at::kBFloat16), + "scales must be a contiguous fp32 or bf16 tensor with one weight per perm entry"); + TORCH_CHECK(x.dim() == 2 && x.size(1) == residual.size(1), "fc2 must be [rows, H]"); + p.fc2 = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.perm = perm->data_ptr(); + p.scales = scales->data_ptr(); + p.scaleFp32 = scales->scalar_type() == at::kFloat; + p.topk = static_cast(perm->size(1)); + } + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2, + "x must be a contiguous 2-D bf16 CUDA tensor"); + TORCH_CHECK(residual.dim() == 2 && residual.size(0) >= 1 && residual.size(0) <= std::min(kMaxTokens, maxt), + "residual must hold 1..min(48, bucket) rows"); + TORCH_CHECK(residual.scalar_type() == at::kBFloat16 && residual.is_contiguous() + && (perm.has_value() || residual.sizes() == x.sizes()), + "residual must be a contiguous bf16 tensor shaped like x"); + TORCH_CHECK( + normWeight.scalar_type() == at::kBFloat16 && normWeight.numel() == residual.size(1), "norm_w must be bf16 [H]"); + TORCH_CHECK(wsUc.scalar_type() == at::kInt && wsUc.is_cuda(), "ws_uc must be an int32 CUDA tensor"); + TORCH_CHECK(flags.scalar_type() == at::kInt && flags.is_contiguous() && flags.numel() >= maxt * kCluster, + "flags must be a contiguous int32 tensor of at least bucket * 8 counters"); + + int64_t const m = residual.size(0), h = residual.size(1); + TORCH_CHECK(!quant || h % 128 == 0, "MXFP8 output needs H % 128 == 0"); + c10::cuda::CUDAGuard const deviceGuard{residual.device()}; + auto outNormed = torch::empty_like(residual); + auto outResidual = torch::empty_like(residual); + auto outQ = quant ? torch::empty({m, h}, residual.options().dtype(at::kFloat8_e4m3fn)) : torch::Tensor(); + auto outSf = quant ? torch::empty({(m + 127) / 128 * 128 * (h / 32)}, residual.options().dtype(at::kByte)) + : torch::Tensor(); + + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.numTokens = static_cast(m); + p.hidden = static_cast(h); + p.residual = reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()); + p.normWeight = reinterpret_cast<__nv_bfloat16 const*>(normWeight.data_ptr()); + p.eps = static_cast(eps); + p.wsUc = wsUc.data_ptr(); + p.wsMc = reinterpret_cast(wsMcPtr); + p.flags = flags.data_ptr(); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.outNormed = reinterpret_cast<__nv_bfloat16*>(outNormed.data_ptr()); + p.outResidual = reinterpret_cast<__nv_bfloat16*>(outResidual.data_ptr()); + p.outQuant = quant ? reinterpret_cast(outQ.data_ptr()) : nullptr; + p.outSf = quant ? outSf.data_ptr() : nullptr; + p.prepushed = static_cast(prepushed); + if (capOut.has_value()) + { + TORCH_CHECK((opts & 16) == 0, "the capture epilogue is not built for the drafter boundary (opts bit 16)"); + TORCH_CHECK(capOut->scalar_type() == at::kBFloat16 && capOut->dim() == 2 && capOut->size(0) >= m + && capOut->size(1) == h && capOut->stride(1) == 1 && capOut->stride(0) % 8 == 0 + && reinterpret_cast(capOut->data_ptr()) % 16 == 0 && capOut->device() == residual.device(), + "cap_out must be a bf16 [>= M, H] view with unit column stride, a row stride that is a multiple of 8 and " + "16-byte alignment"); + TORCH_CHECK( + capW.has_value() && capW->scalar_type() == at::kBFloat16 && capW->is_contiguous() && capW->numel() == h, + "cap_w must be a contiguous bf16 [H] tensor"); + p.capOut = reinterpret_cast<__nv_bfloat16*>(capOut->data_ptr()); + p.capStride = capOut->stride(0); + p.capWeight = reinterpret_cast<__nv_bfloat16 const*>(capW->data_ptr()); + p.capEps = static_cast(capEps); + } + tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp2McAllreduceNorm(p, static_cast(maxt), + static_cast(opts), pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); + if (quant) + { + return {outNormed, outResidual, outQ, outSf}; + } + return {outNormed, outResidual}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_mc_allreduce_norm(Tensor x, Tensor? x2, Tensor residual, Tensor norm_w, float eps, " + "Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) flags, int rank, int world, int cluster, bool pdl, " + "Tensor? perm, Tensor? scales, int maxt, bool quant, int prepushed, int opts, Tensor(c!)? cap_out=None, " + "Tensor? cap_w=None, float cap_eps=0.) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_mc_allreduce_norm", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2McAllreduceNorm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp new file mode 100644 index 000000000000..ed50e6d10ba2 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp @@ -0,0 +1,185 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp2::producerPw2; + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPageSize = 128; + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "sm107_tp2_producer_pw: ", name, " must be a CUDA tensor"); +} + +} // namespace + +//! The SM107 TP2 NVFP4 QKV + indexer norm/RoPE/KV-insert producer (same arguments and outputs as +//! minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert at 32 query, 2 KV and 2 index heads, head_dim 128, rotary_dim 64): +//! returns (q FP8 [T, 32, 128], index_q FP8 [T, 2, 128]) and writes the NVFP4 K/V of both KV heads and the FP8 index-K +//! of each token's cache slot. `early` triggers the dependents on entry; `prewait_step_inputs` reads the position ids, +//! the slots and the RoPE rows before the grid-dependency wait (none of them may be written by a PDL predecessor). +std::vector minimaxM3Sm107Tp2ProducerPw(torch::Tensor const& packed, torch::Tensor& kvData, + torch::Tensor& kvScale, torch::Tensor& indexK, torch::Tensor const& outCacheLoc, torch::Tensor const& kvQuantScale, + double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds, bool pdl, + bool early, bool prewaitStepInputs) +{ + checkCuda(packed, "packed"); + for (torch::Tensor const* t : std::initializer_list{&kvData, &kvScale, &indexK, &outCacheLoc, + &kvQuantScale, &qWeight, &kWeight, &indexQWeight, &indexKWeight, &rotaryCosSin, &positionIds}) + { + TORCH_CHECK(t->is_cuda() && t->get_device() == packed.get_device(), + "sm107_tp2_producer_pw: every operand must be on packed's CUDA device"); + } + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0 + && reinterpret_cast(kvData.data_ptr()) % 2 == 0 + && reinterpret_cast(indexK.data_ptr()) % 4 == 0, + "sm107_tp2_producer_pw: packed and the caches must start at the vector-load / store alignment"); + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp2_producer_pw requires SM100 or newer"); + int64_t const headsQ = pw::kHeadsQ; + int64_t const headsKV = pw::kHeadsKV; + int64_t const headsIndex = pw::kHeadsIndex; + int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; + TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim + && packed.is_contiguous(), + "sm107_tp2_producer_pw: packed must be contiguous BF16 [T, 39 * 128]"); + // The kernel's layout (the native producer's checkMinimaxM3Nvfp4HndKVPools at two KV heads): HND pools + // [pages, 2 planes, 2 heads, 128 tokens, 64 data / 8 scale bytes], dense token rows and head blocks, planes and + // pages that do not overlap. + TORCH_CHECK(kvData.scalar_type() == at::kByte && kvData.dim() == 5 && kvData.size(0) > 0 && kvData.size(1) == 2 + && kvData.size(2) == headsKV && kvData.size(3) == kPageSize && kvData.size(4) == kHeadDim / 2, + "sm107_tp2_producer_pw: kv_data must be the NVFP4 pool [pages, 2, 2, 128, 64] uint8"); + TORCH_CHECK(kvScale.scalar_type() == at::kByte && kvScale.dim() == 5 && kvScale.size(0) == kvData.size(0) + && kvScale.size(1) == 2 && kvScale.size(2) == headsKV && kvScale.size(3) == kPageSize + && kvScale.size(4) == kHeadDim / 16, + "sm107_tp2_producer_pw: kv_scale must be the E4M3 scale pool [pages, 2, 2, 128, 8] uint8"); + TORCH_CHECK(kvData.stride(4) == 1 && kvData.stride(3) == kvData.size(4) + && kvData.stride(2) == kvData.size(3) * kvData.stride(3), + "sm107_tp2_producer_pw: kv_data must have dense token rows and head blocks"); + TORCH_CHECK(kvScale.stride(4) == 1 && kvScale.stride(3) == kvScale.size(4) + && kvScale.stride(2) == kvScale.size(3) * kvScale.stride(3), + "sm107_tp2_producer_pw: kv_scale must have dense token rows and head blocks"); + TORCH_CHECK(kvData.stride(1) >= kvData.size(2) * kvData.stride(2) && kvData.stride(0) >= 2 * kvData.stride(1) + && kvScale.stride(1) >= kvScale.size(2) * kvScale.stride(2) && kvScale.stride(0) >= 2 * kvScale.stride(1), + "sm107_tp2_producer_pw: pool planes or pages overlap"); + TORCH_CHECK(kvData.stride(0) % 2 == 0 && kvData.stride(1) % 2 == 0, + "sm107_tp2_producer_pw: kv_data page and plane strides must keep the uint16 stores aligned"); + TORCH_CHECK(indexK.scalar_type() == at::kFloat8_e4m3fn && indexK.dim() == 4 && indexK.size(0) == kvData.size(0) + && indexK.size(1) == 1 && indexK.size(2) == kPageSize && indexK.size(3) == kHeadDim && indexK.stride(3) == 1 + && indexK.stride(2) == kHeadDim && indexK.stride(0) >= kPageSize * kHeadDim && indexK.stride(0) % 4 == 0, + "sm107_tp2_producer_pw: index_k must be FP8 [pages, 1, 128, 128] with a page stride keeping the 32-bit stores " + "aligned"); + int64_t const numTokens = packed.size(0); + TORCH_CHECK( + outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous() && outCacheLoc.numel() >= numTokens, + "sm107_tp2_producer_pw: out_cache_loc must be int32 [>= T]"); + TORCH_CHECK( + positionIds.scalar_type() == at::kInt && positionIds.is_contiguous() && positionIds.numel() >= numTokens, + "sm107_tp2_producer_pw: position_ids must be int32 [>= T]"); + TORCH_CHECK(kvQuantScale.scalar_type() == at::kFloat && kvQuantScale.is_contiguous() && kvQuantScale.numel() >= 3, + "sm107_tp2_producer_pw: kv_quant_scale must be fp32 [>= 3]"); + TORCH_CHECK(rotaryCosSin.scalar_type() == at::kFloat && rotaryCosSin.is_contiguous() && rotaryCosSin.dim() == 3 + && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == 32 + && reinterpret_cast(rotaryCosSin.data_ptr()) % 16 == 0, + "sm107_tp2_producer_pw: rotary_cos_sin must be fp32 [positions, 2, 32], 16-byte aligned (read as float4)"); + for (auto const* w : {&qWeight, &kWeight, &indexQWeight, &indexKWeight}) + { + TORCH_CHECK(w->scalar_type() == at::kBFloat16 && w->is_contiguous() && w->numel() == kHeadDim + && reinterpret_cast(w->data_ptr()) % 8 == 0, + "sm107_tp2_producer_pw: norm weights must be BF16 [128], 8-byte aligned (read as 4-value words)"); + } + auto const fp8 = packed.options().dtype(at::kFloat8_e4m3fn); + auto qOut = torch::empty({numTokens, headsQ, kHeadDim}, fp8); + auto indexQOut = torch::empty({numTokens, headsIndex, kHeadDim}, fp8); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + pw::Params p{}; + p.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + p.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + p.indexQOut = reinterpret_cast<__nv_fp8_e4m3*>(indexQOut.data_ptr()); + p.kvData = kvData.data_ptr(); + p.kvScale = kvScale.data_ptr(); + p.indexK = reinterpret_cast<__nv_fp8_e4m3*>(indexK.data_ptr()); + p.outCacheLoc = outCacheLoc.data_ptr(); + p.kvQuantScale = kvQuantScale.data_ptr(); + p.dataPageStride = kvData.stride(0); + p.dataPlaneStride = kvData.stride(1); + p.dataHeadStride = kvData.stride(2); + p.dataTokenStride = kvData.stride(3); + p.scalePageStride = kvScale.stride(0); + p.scalePlaneStride = kvScale.stride(1); + p.scaleHeadStride = kvScale.stride(2); + p.indexPageStride = indexK.stride(0); + p.indexTokenStride = indexK.stride(2); + p.numPages = kvData.size(0); + p.numTokens = static_cast(numTokens); + p.eps = static_cast(eps); + p.qWeight = reinterpret_cast<__nv_bfloat16 const*>(qWeight.data_ptr()); + p.kWeight = reinterpret_cast<__nv_bfloat16 const*>(kWeight.data_ptr()); + p.indexQWeight = reinterpret_cast<__nv_bfloat16 const*>(indexQWeight.data_ptr()); + p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); + p.rotaryCosSin = rotaryCosSin.data_ptr(); + p.positionIds = positionIds.data_ptr(); + cudaError_t const err = pw::invokeProducerPw2(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "sm107_tp2_producer_pw launch: ", cudaGetErrorString(err)); + } + return {qOut, indexQOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_producer_pw(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) index_k, " + "Tensor out_cache_loc, Tensor kv_quant_scale, float eps, Tensor q_weight, Tensor k_weight, " + "Tensor index_q_weight, Tensor index_k_weight, Tensor rotary_cos_sin, Tensor position_ids, bool pdl, " + "bool early, bool prewait_step_inputs=False) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_producer_pw", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2ProducerPw); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp new file mode 100644 index 000000000000..415042d5c082 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp @@ -0,0 +1,110 @@ +/* + * 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/minimaxM3Sm107Tp2/sampledAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3Sm107Tp2::sampledAccept; + +//! uint64 words of the sampled-accept exchange workspace for `world` ranks. +int64_t minimaxM3Sm107Tp2SampledAcceptWsWords(int64_t world) +{ + return sa::sampledAcceptWsWords(static_cast(world)); +} + +//! Sampled (temperature + top-p) verification from this rank's BF16 logits shard [rows, Vs], rows = nreq * +//! (draft_len + 1) <= 48: the sampled token of every row (tokens), then the strict acceptance against draft (accTokens, +//! numAcc, gatherIds, optional nextNew column 0). Per-row Philox (seeds, offsets); the exchange goes through the +//! multicast workspace wsMcPtr / wsUc with the (finished rows, call) counters. +void minimaxM3Sm107Tp2SampledAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor const& temperatures, + torch::Tensor const& topPs, torch::Tensor const& seeds, torch::Tensor const& offsets, torch::Tensor& wsUc, + int64_t wsMcPtr, torch::Tensor& counters, torch::Tensor& tokens, int64_t rank, int64_t world, + torch::Tensor const& draft, torch::Tensor const& batchIdx, torch::Tensor& accTokens, torch::Tensor& numAcc, + torch::Tensor& gatherIds, int64_t forceTotal, std::optional const& nextNew, + std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, + bool prewaitParams, bool earlyTrigger, bool ilpWalk) +{ + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + sa::Args a; + a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + a.ld = logits.stride(0); + a.rows = static_cast(logits.size(0)); + a.Vs = static_cast(logits.size(1)); + TORCH_CHECK(a.Vs <= sa::kMaxVs && a.rows <= sa::kMaxRows && world <= sa::kMaxWorld + && static_cast(vocabOffset) + a.Vs <= (1LL << sa::kIdBits)); + a.vocab_offset = static_cast(vocabOffset); + a.temperatures = temperatures.data_ptr(); + a.top_ps = topPs.data_ptr(); + a.seeds = reinterpret_cast(seeds.data_ptr()); + a.offsets = reinterpret_cast(offsets.data_ptr()); + a.ws_uc = reinterpret_cast(wsUc.data_ptr()); + a.ws_mc = reinterpret_cast(wsMcPtr); + a.counters = reinterpret_cast(counters.data_ptr()); + a.tokens = tokens.data_ptr(); + a.extra_attempts = extraAttempts.has_value() ? extraAttempts->data_ptr() : nullptr; + a.rank = static_cast(rank); + a.world = static_cast(world); + a.draft = draft.data_ptr(); + a.batch_idx = batchIdx.data_ptr(); + a.acc_tokens = accTokens.data_ptr(); + a.num_acc = numAcc.data_ptr(); + a.gather_ids = reinterpret_cast(gatherIds.data_ptr()); + a.draft_len = static_cast(draft.size(1)); + a.nreq = static_cast(draft.size(0)); + a.force_total = static_cast(forceTotal); + a.next_new = nextNew.has_value() ? nextNew->data_ptr() : nullptr; + a.next_stride = nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0; + a.force_vec = forceVec.has_value() ? forceVec->data_ptr() : nullptr; + a.prewait_params = prewaitParams ? 1 : 0; + a.early_trigger = earlyTrigger ? 1 : 0; + a.ilp_walk = ilpWalk ? 1 : 0; + cudaError_t const err = sa::invokeSampledAccept(a, usePdl, at::cuda::getCurrentCUDAStream(logits.get_device())); + TORCH_CHECK(err == cudaSuccess, "m3_sampled_accept launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// ws_mc_ptr is the multicast address of ws_uc's mapping: the op writes through it, which the schema cannot show. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + // No tensor arguments, so no dispatch key: registered with the schema as a catch-all kernel. + m.def("minimax_m3_sm107tp2_sampled_accept_ws_words(int world) -> int", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SampledAcceptWsWords); + m.def( + "minimax_m3_sm107tp2_sampled_accept(Tensor logits, int vocab_offset, Tensor temperatures, Tensor top_ps, " + "Tensor seeds, Tensor offsets, Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) counters, Tensor(c!) tokens, " + "int rank, int world, Tensor draft, Tensor batch_idx, Tensor(d!) acc_tokens, Tensor(e!) num_acc, " + "Tensor(f!) gather_ids, int force_total, Tensor(g!)? next_new, Tensor? force_vec, " + "Tensor(h!)? extra_attempts, bool use_pdl, bool prewait_params=True, bool early_trigger=False, " + "bool ilp_walk=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp2_sampled_accept", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SampledAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SelectLongOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SelectLongOp.cpp new file mode 100644 index 000000000000..b54249840c93 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SelectLongOp.cpp @@ -0,0 +1,85 @@ +/* + * 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/minimaxM3Sm107Tp2/selectLong2.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, into a caller-owned table. +//! +//! @param scores Float32 CUDA tensor [num_kv_heads, num_blocks, total_q], any strides. +//! @param nValidBlocks Contiguous int32 CUDA tensor with at least total_q elements. +//! @param out Int32 CUDA tensor of logical shape [total_q, num_kv_heads, 16] with unit stride in the last dim (the +//! head-major backing the target hands over is a permuted view). Every entry is written. +//! @param triggerPhase Where the dependent grid is released without earlyTrigger: 0 at this grid's exit, 1 right +//! after every score of the row is loaded (both after this grid's dependency wait). +void minimaxM3Sm107Tp2SelectLong(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, torch::Tensor& out, + int64_t initBlocks, int64_t localBlocks, bool usePdl, bool earlyTrigger, int64_t triggerPhase) +{ + TORCH_CHECK(scores.is_cuda() && scores.scalar_type() == torch::kFloat32 && scores.dim() == 3, + "minimax_m3_sm107_tp2_select_long: scores must be a 3-D float32 CUDA tensor [num_kv_heads, num_blocks, " + "total_q]"); + TORCH_CHECK(nValidBlocks.is_cuda() && nValidBlocks.scalar_type() == torch::kInt32 && nValidBlocks.is_contiguous() + && nValidBlocks.numel() >= scores.size(2), + "minimax_m3_sm107_tp2_select_long: n_valid_blocks must be contiguous int32 CUDA with at least total_q " + "elements"); + int64_t const numKvHeads = scores.size(0); + int64_t const numBlocks = scores.size(1); + int64_t const totalQ = scores.size(2); + TORCH_CHECK(numBlocks <= tensorrt_llm::kernels::minimaxM3Sm107Tp2::kSelectLong2MaxBlocks, + "minimax_m3_sm107_tp2_select_long supports at most ", + tensorrt_llm::kernels::minimaxM3Sm107Tp2::kSelectLong2MaxBlocks, " blocks, got ", numBlocks); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kInt32 && out.dim() == 3 && out.size(0) == totalQ + && out.size(1) == numKvHeads && out.size(2) == 16 && out.stride(2) == 1, + "minimax_m3_sm107_tp2_select_long: out must be int32 [total_q, num_kv_heads, 16] with unit last-dim stride"); + + c10::cuda::CUDAGuard const deviceGuard{scores.device()}; + auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device()); + cudaError_t const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSelectLong2(scores.data_ptr(), + scores.stride(0), scores.stride(1), scores.stride(2), nValidBlocks.data_ptr(), out.data_ptr(), + out.stride(0), out.stride(1), static_cast(numKvHeads), static_cast(numBlocks), + static_cast(totalQ), static_cast(initBlocks), static_cast(localBlocks), usePdl, + earlyTrigger, static_cast(triggerPhase), stream); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_sm107_tp2_select_long launch failed: ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_select_long(Tensor scores, Tensor n_valid_blocks, Tensor(a!) out, int init_blocks, " + "int local_blocks, bool use_pdl=False, bool early_trigger=False, int trigger_phase=0) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_select_long", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SelectLong); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp new file mode 100644 index 000000000000..52f60de1b7ca --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp @@ -0,0 +1,66 @@ +/* + * 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/minimaxM3Sm107Tp2/siluAndMul.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp2/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, +//! d]. +void minimaxM3Sm107Tp2SiluAndMul(torch::Tensor const& x, torch::Tensor& out) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp2_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp2_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp2_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); + TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp2_silu_and_mul: at most 65535 rows"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err + = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), + reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_silu_and_mul launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SiluAndMul); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp new file mode 100644 index 000000000000..fb8fc9eaa8ee --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp @@ -0,0 +1,169 @@ +/* + * 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/minimaxM3Sm107Tp2/sparseDecodeLc2.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeadsPerKv = 16; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kTopK = 16; + +void checkBytes(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, "minimax_m3_sm107_tp2_sparse_decode_lc: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3 long-context sparse decode over the selected NVFP4 blocks, SM107 TP2 geometry (16 query +//! heads per KV head, 2 KV heads per rank, head_dim 128, page 128, top-k 16). +//! +//! @param q E4M3 bytes [T, 16 * H, 128], unit stride in the last dim. +//! @param k, v Packed E2M1 bytes [pages, H, 128, 64] (the per-layer HND views). +//! @param kBlockScale, vBlockScale E4M3 bytes [pages, H, 128, 8]: K token-major, V in 4x4 token-quad order. +//! @param kGlobalScale, vGlobalScale One-element float32 per-tensor dequant scales. +//! @param topk Int32 [H, T, 16] ascending selected block ids with -1 tail (the head-major selector table). +//! @param out BF16 [T, H * 2048], written. q8 / sf, when given, receive the MXFP8 copy and its swizzled UE8M0 scales. +//! @param n1Ids Optional int32 [H, rows, 16] (unit last-dim stride): this layer's selected block ids at its previous +//! decode step, read to prefetch those blocks into L2 before the grid dependency and overwritten with this step's; +//! query rows at or past `rows` skip it. Initialize to -1. A hint: `out` does not depend on it. +void minimaxM3Sm107Tp2SparseDecodeLc(torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& kBlockScale, torch::Tensor const& vBlockScale, torch::Tensor const& kGlobalScale, + torch::Tensor const& vGlobalScale, torch::Tensor const& topk, torch::Tensor const& blockTable, + torch::Tensor const& seqLens, torch::Tensor& out, std::optional q8, std::optional sf, + int64_t numSfKTiles, int64_t decodeQueryLen, double smScale, bool linearSoftmax, bool usePdl, bool prewait, + std::optional n1Ids) +{ + TORCH_CHECK(q.is_cuda(), "minimax_m3_sm107_tp2_sparse_decode_lc: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m3_sm107_tp2_sparse_decode_lc requires an SM100-family GPU"); + checkBytes(q, "q"); + checkBytes(k, "k"); + checkBytes(v, "v"); + checkBytes(kBlockScale, "k_block_scale"); + checkBytes(vBlockScale, "v_block_scale"); + TORCH_CHECK(k.dim() == 4 && k.size(2) == 128 && k.size(3) == kHeadDim / 2 && k.stride(3) == 1 + && v.sizes() == k.sizes() && v.strides() == k.strides(), + "minimax_m3_sm107_tp2_sparse_decode_lc: k and v must be [pages, H, 128, 64] with matching strides"); + int64_t const kvHeads = k.size(1); + TORCH_CHECK(kBlockScale.dim() == 4 && kBlockScale.size(1) == kvHeads && kBlockScale.size(3) == 8 + && kBlockScale.stride(3) == 1 && kBlockScale.stride(2) == 8 && vBlockScale.sizes() == kBlockScale.sizes() + && vBlockScale.strides() == kBlockScale.strides(), + "minimax_m3_sm107_tp2_sparse_decode_lc: block scales must be contiguous [pages, H, 128, 8] per page"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeadsPerKv * kvHeads && q.size(2) == kHeadDim && q.stride(2) == 1, + "minimax_m3_sm107_tp2_sparse_decode_lc: q must be [T, 16 * H, 128] with unit last-dim stride"); + TORCH_CHECK(kGlobalScale.scalar_type() == torch::kFloat32 && vGlobalScale.scalar_type() == torch::kFloat32, + "minimax_m3_sm107_tp2_sparse_decode_lc: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 3 && topk.size(0) == kvHeads + && topk.size(1) == q.size(0) && topk.size(2) == kTopK, + "minimax_m3_sm107_tp2_sparse_decode_lc: topk must be int32 [H, T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_sm107_tp2_sparse_decode_lc: block_table and seq_lens must be int32"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.size(0) == q.size(0) + && out.size(1) == kvHeads * kHeadsPerKv * kHeadDim && out.stride(1) == 1, + "minimax_m3_sm107_tp2_sparse_decode_lc: out must be bfloat16 [T, H * 2048] with unit last-dim stride"); + TORCH_CHECK( + q8.has_value() == sf.has_value(), "minimax_m3_sm107_tp2_sparse_decode_lc: pass both q8 and sf, or neither"); + bool const mxfp8 = q8.has_value(); + if (mxfp8) + { + checkBytes(*q8, "q8"); + checkBytes(*sf, "sf"); + } + + tensorrt_llm::kernels::MinimaxM3Sm107Tp2SparseDecodeLcParams prm{}; + prm.q = reinterpret_cast(q.data_ptr()); + prm.q_tok = q.stride(0); + prm.q_head = q.stride(1); + prm.k = reinterpret_cast(k.data_ptr()); + prm.v = reinterpret_cast(v.data_ptr()); + prm.k_blk = k.stride(0); + prm.v_blk = v.stride(0); + prm.k_h = k.stride(1); + prm.v_h = v.stride(1); + prm.ksc = reinterpret_cast(kBlockScale.data_ptr()); + prm.vsc = reinterpret_cast(vBlockScale.data_ptr()); + prm.ksc_blk = kBlockScale.stride(0); + prm.vsc_blk = vBlockScale.stride(0); + prm.ksc_h = kBlockScale.stride(1); + prm.vsc_h = vBlockScale.stride(1); + prm.kgs = kGlobalScale.data_ptr(); + prm.vgs = vGlobalScale.data_ptr(); + prm.topk = topk.data_ptr(); + prm.t_h = topk.stride(0); + prm.t_tok = topk.stride(1); + prm.t_k = topk.stride(2); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.seq_lens = seqLens.data_ptr(); + prm.dql = static_cast(decodeQueryLen); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + prm.out_stride = out.stride(0); + prm.q8 = mxfp8 ? reinterpret_cast(q8->data_ptr()) : nullptr; + prm.sf = mxfp8 ? reinterpret_cast(sf->data_ptr()) : nullptr; + prm.num_sf_k_tiles = static_cast(numSfKTiles); + prm.num_kv_heads = static_cast(kvHeads); + if (n1Ids.has_value()) + { + TORCH_CHECK(n1Ids->is_cuda() && n1Ids->scalar_type() == torch::kInt32 && n1Ids->dim() == 3 + && n1Ids->size(0) == kvHeads && n1Ids->size(2) == kTopK && n1Ids->stride(2) == 1 + && n1Ids->stride(1) == kTopK, + "minimax_m3_sm107_tp2_sparse_decode_lc: n1_ids must be int32 [H, rows, 16] with dense rows"); + prm.n1_ids = n1Ids->data_ptr(); + prm.n1_h = n1Ids->stride(0); + prm.n1_rows = static_cast(n1Ids->size(1)); + prm.n1_pages = k.size(0); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp2SparseDecodeLc( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, prewait, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_sparse_decode_lc launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp2_sparse_decode_lc(Tensor q, Tensor k, Tensor v, Tensor k_block_scale, " + "Tensor v_block_scale, Tensor k_global_scale, Tensor v_global_scale, Tensor topk, Tensor block_table, " + "Tensor seq_lens, Tensor(a!) out, Tensor(b!)? q8, Tensor(c!)? sf, int num_sf_k_tiles, int decode_query_len, " + "float sm_scale, bool linear_softmax, bool use_pdl, bool prewait, Tensor(d!)? n1_ids=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_sparse_decode_lc", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SparseDecodeLc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp new file mode 100644 index 000000000000..87b6ae599742 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp @@ -0,0 +1,74 @@ +/* + * 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/minimaxM3Sm107Tp2/specOffsets.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void check(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp2_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +} // namespace + +//! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch +//! (kernels/minimaxM3Sm107Tp2/specOffsets.h). +void minimaxM3Sm107Tp2SpecOffsets( + torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) +{ + check(pos, "pos"); + check(posOff, "pos_off"); + check(kv, "kv"); + check(kvOff, "kv_off"); + TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), + "minimax_m3_sm107_tp2_spec_offsets: each tensor and its offsets must have the same length"); + if (pos.numel() == 0 && kv.numel() == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{pos.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSpecOffsets(pos.data_ptr(), + posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), + at::cuda::getCurrentCUDAStream(pos.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_spec_offsets launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SpecOffsets); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp new file mode 100644 index 000000000000..9e8229ba9e16 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp @@ -0,0 +1,68 @@ +/* + * 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/minimaxM3Sm107Tp2/stepPrologue.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp2/stepPrologue.h). +//! +//! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the +//! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses +//! the caller guarantees are live. +//! @param epoch Optional one-element int64 step counter, incremented once. +void minimaxM3Sm107Tp2StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) +{ + TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), + "minimax_m3_sm107_tp2_step_prologue: window must be a contiguous uint8 CUDA tensor"); + TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, + "minimax_m3_sm107_tp2_step_prologue: window is shorter than its header and max_desc descriptors"); + c10::cuda::CUDAGuard const deviceGuard{window.device()}; + int64_t* epochPtr = nullptr; + if (epoch.has_value()) + { + TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 + && epoch->device() == window.device(), + "minimax_m3_sm107_tp2_step_prologue: epoch must be a one-element int64 tensor on the window's device"); + epochPtr = epoch->data_ptr(); + } + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeStepPrologue( + window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_step_prologue launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2StepPrologue); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp new file mode 100644 index 000000000000..2839f6036ffc --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp @@ -0,0 +1,65 @@ +/* + * 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/minimaxM3Sm107Tp2/swigluOai.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp2/swigluOai.h): x bf16 [rows, 2d] -> out bf16 +//! [rows, d], PDL-launched with an early release when `pdl`. +void minimaxM3Sm107Tp2SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp2_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp2_swiglu_oai: x rows must be 4-byte aligned with an even half width"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp2_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSwigluOai( + reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, + at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_swiglu_oai launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SwigluOai); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp new file mode 100644 index 000000000000..d92f2c1cd29b --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp @@ -0,0 +1,128 @@ +/* + * 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/minimaxM3Sm107Tp4/denseQkNormRopeFp8.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace dqk = tensorrt_llm::kernels::minimaxM3Sm107Tp4::denseQkNormRopeFp8; + +//! QK RMSNorm + partial RoPE of a packed BF16 [tokens, (q + k + v) * 128] row of the MiniMax-M3 TP4 dense layers, +//! returned as a new FP8 E4M3 tensor of the same shape (V copy-cast; Q/K rounded to BF16 before E4M3, as stock). +torch::Tensor minimaxM3Sm107Tp4DenseQkNormRopeFp8(torch::Tensor const& qkv, int64_t numHeadsQ, int64_t numHeadsK, + int64_t numHeadsV, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, + double base, bool isNeox, torch::Tensor const& positionIds, bool useGemma) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); + TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); + TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); + TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + at::cuda::CUDAGuard guard(qkv.device()); + auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); + cudaError_t const err + = dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(rotaryDim), + static_cast(eps), qWeight.data_ptr(), kWeight.data_ptr(), static_cast(base), !isNeox, + positionIds.data_ptr(), useGemma, nullptr, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "dense_qk_norm_rope_fp8 launch: ", cudaGetErrorString(err)); + return out; +} + +//! As minimaxM3Sm107Tp4DenseQkNormRopeFp8, and the first numLiveTokens rows' K and V are also stored into the layer's +//! FP8 pool kvBuffers (the HND view [pages, 2, kv heads, tokens per block, 128]) at their int32 cache slots; with +//! qCompact, every row's Q is also stored there compact ([tokens, q heads, 128]). +torch::Tensor minimaxM3Sm107Tp4DenseQkNormRopeFp8Kv(torch::Tensor const& qkv, int64_t numHeadsQ, int64_t numHeadsK, + int64_t numHeadsV, int64_t rotaryDim, double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, + double base, bool isNeox, torch::Tensor const& positionIds, bool useGemma, torch::Tensor const& kvBuffers, + torch::Tensor const& outCacheLoc, int64_t numLiveTokens, std::optional const& qCompact) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); + TORCH_CHECK(numHeadsK == numHeadsV, "the pool keeps one head count for K and V"); + TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); + TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); + TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + TORCH_CHECK(kvBuffers.is_cuda() && kvBuffers.scalar_type() == at::kFloat8_e4m3fn && kvBuffers.dim() == 5); + TORCH_CHECK(kvBuffers.size(1) == 2 && kvBuffers.size(2) == numHeadsK && kvBuffers.size(4) == dqk::kHeadDim + && kvBuffers.stride(4) == 1, + "kv_buffers must be the HND pool view [pages, 2, kv heads, tokens per block, 128]"); + TORCH_CHECK(outCacheLoc.is_cuda() && outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous()); + TORCH_CHECK(numLiveTokens >= 0 && numLiveTokens <= qkv.size(0) && numLiveTokens <= outCacheLoc.numel()); + at::cuda::CUDAGuard guard(qkv.device()); + auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); + dqk::KvCacheView view; + view.base = static_cast<__nv_fp8_e4m3*>(kvBuffers.data_ptr()); + view.outCacheLoc = outCacheLoc.data_ptr(); + view.pageStride = kvBuffers.stride(0); + view.kvStride = kvBuffers.stride(1); + view.headStride = kvBuffers.stride(2); + view.tokenStride = kvBuffers.stride(3); + view.tokensPerBlock = static_cast(kvBuffers.size(3)); + view.numLiveTokens = static_cast(numLiveTokens); + if (qCompact.has_value()) + { + TORCH_CHECK(qCompact->is_cuda() && qCompact->is_contiguous() && qCompact->scalar_type() == at::kFloat8_e4m3fn + && qCompact->numel() == qkv.size(0) * numHeadsQ * dqk::kHeadDim, + "q_compact must be a contiguous FP8 [tokens, q heads, 128] tensor"); + view.qCompact = static_cast<__nv_fp8_e4m3*>(qCompact->data_ptr()); + } + cudaError_t const err + = dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), + static_cast(numHeadsK), static_cast(numHeadsV), static_cast(rotaryDim), + static_cast(eps), qWeight.data_ptr(), kWeight.data_ptr(), static_cast(base), !isNeox, + positionIds.data_ptr(), useGemma, &view, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "dense_qk_norm_rope_fp8_kv launch: ", cudaGetErrorString(err)); + return out; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, " + "bool use_gemma) -> Tensor"); + m.def( + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv(Tensor qkv, int num_heads_q, int num_heads_k, int num_heads_v, " + "int rotary_dim, float eps, Tensor q_weight, Tensor k_weight, float base, bool is_neox, Tensor position_ids, " + "bool use_gemma, Tensor(a!) kv_buffers, Tensor out_cache_loc, int num_live_tokens, Tensor(b!)? q_compact=None) " + "-> Tensor"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_dense_qk_norm_rope_fp8", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4DenseQkNormRopeFp8); + m.impl("minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4DenseQkNormRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DraftStepOp.cpp new file mode 100644 index 000000000000..b0a49e8a29fb --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DraftStepOp.cpp @@ -0,0 +1,215 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h" +#include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h" + +#include +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace ds = tensorrt_llm::kernels::minimaxM3Sm107Tp4::draftStep; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_sm107tp4_draft_step expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107tp4_draft_step requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +void checkInt(torch::Tensor const& t, int64_t n, char const* name) +{ + TORCH_CHECK(t.scalar_type() == at::kInt && t.is_contiguous() && t.numel() >= n, "draft_step: ", name, + " must be contiguous int32 with at least ", n, " elements"); +} + +} // namespace + +//! One fused greedy draft pass over `rows` <= 8 requests: the BF16 head (this rank's fragment-ordered shard wf +//! [N / 16, K / 16, 256], or with excOff / excPos / excW / expBase its 12-bit coding, wf uint8 [N * K * 3 / 2]; +//! global ids vocabOffset..) over x[xRows[b]] (or x[b]), the global argmax across the TP ranks through the multicast +//! slots wsMcPtr / wsUc, then draftOut[b, step] = idsNext[b] = token, posNext[b] = posSrc[xRows[b] or b] + 1 and the +//! KV-length bookkeeping (see draftStep.h). +void minimaxM3Sm107Tp4DraftStep(torch::Tensor const& x, OptTensor const& xRows, int64_t rows, torch::Tensor const& wf, + int64_t vocabOffset, torch::Tensor& keys, torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, + int64_t rank, int64_t world, torch::Tensor& draftOut, int64_t step, torch::Tensor& idsNext, + torch::Tensor const& posSrc, torch::Tensor& posNext, torch::Tensor& kvLens, OptTensor const& numAcc, + OptTensor const& seqLens, int64_t draftLen, int64_t numCtas, bool pdl, OptTensor const& outKeys, + OptTensor const& excOff, OptTensor const& excPos, OptTensor const& excW, int64_t expBase, + std::optional> fold, int64_t foldRows, int64_t foldTokensPerStep, bool fixedOrder) +{ + checkArch(x); + for (torch::Tensor const* t : std::initializer_list{ + &wf, &keys, &counters, &wsUc, &draftOut, &idsNext, &posSrc, &posNext, &kvLens}) + { + TORCH_CHECK(t->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + for (OptTensor const* t : + std::initializer_list{&xRows, &numAcc, &seqLens, &outKeys, &excOff, &excPos, &excW}) + { + TORCH_CHECK(!t->has_value() || (*t)->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "draft_step: x must be contiguous BF16 [rows, K]"); + int const K = static_cast(x.size(1)); + TORCH_CHECK(K % 16 == 0, "draft_step: K must be a multiple of 16"); + TORCH_CHECK(rows >= 1 && rows <= ds::kMaxTokens, "draft_step: 1..", ds::kMaxTokens, " rows"); + bool const coded = excOff.has_value(); + TORCH_CHECK(coded == excPos.has_value() && coded == excW.has_value(), + "draft_step: exc_off, exc_pos and exc_w come together"); + int nTiles = 0; + if (coded) + { + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "draft_step: a coded head needs K a multiple of 64"); + TORCH_CHECK(excOff->scalar_type() == at::kInt && excOff->is_contiguous() && excOff->numel() >= 2, + "draft_step: exc_off must be int32 [tiles + 1]"); + nTiles = static_cast(excOff->numel() - 1); + TORCH_CHECK(wf.scalar_type() == at::kByte && wf.is_contiguous() + && wf.numel() == static_cast(nTiles) * 16 * K * 3 / 2, + "draft_step: a coded wf must be the contiguous uint8 stream of [tiles * 16, K]"); + TORCH_CHECK(excPos->scalar_type() == at::kInt && excPos->is_contiguous(), "draft_step: exc_pos must be int32"); + TORCH_CHECK(excW->scalar_type() == at::kFloat && excW->is_contiguous() && excW->numel() == 2 * excPos->numel(), + "draft_step: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "draft_step: exponent base out of range"); + } + else + { + TORCH_CHECK(wf.scalar_type() == at::kBFloat16 && wf.dim() == 3 && wf.is_contiguous() && wf.size(1) == K / 16 + && wf.size(2) == 256, + "draft_step: wf must be the fragment-ordered BF16 weight [N / 16, K / 16, 256]"); + nTiles = static_cast(wf.size(0)); + } + TORCH_CHECK(nTiles >= 1 && numCtas >= 1 && (nTiles + numCtas - 1) / numCtas <= ds::kMaxTilesPerCta, + "draft_step: too many weight tiles per CTA"); + TORCH_CHECK(vocabOffset >= 0 && vocabOffset + 16LL * nTiles <= 0xFFFFFFFFLL, "draft_step: vocab offset"); + if (xRows.has_value()) + { + TORCH_CHECK(xRows->scalar_type() == at::kLong && xRows->is_contiguous() && xRows->numel() >= rows, + "draft_step: x_rows must be contiguous int64 [rows]"); + } + else + { + TORCH_CHECK(x.size(0) >= rows, "draft_step: x must hold a row per request"); + } + TORCH_CHECK(world >= 1 && rank >= 0 && rank < world, "draft_step: rank / world"); + TORCH_CHECK(keys.scalar_type() == at::kLong && keys.numel() >= rows, "draft_step: keys must be int64 [>= rows]"); + TORCH_CHECK(counters.scalar_type() == at::kInt && counters.numel() >= 2, "draft_step: counters must be int32 [2]"); + TORCH_CHECK(wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, + "draft_step: ws_uc must hold [2][8][world] slots"); + TORCH_CHECK(draftOut.scalar_type() == at::kInt && draftOut.dim() == 2 && draftOut.stride(1) == 1 + && draftOut.size(0) >= rows && step >= 0 && step < draftOut.size(1), + "draft_step: draft_out must be int32 [rows, > step] with unit column stride"); + checkInt(idsNext, rows, "ids_next"); + checkInt(posNext, rows, "pos_next"); + checkInt(kvLens, rows, "kv_lens"); + TORCH_CHECK(posSrc.scalar_type() == at::kInt && posSrc.is_contiguous(), "draft_step: pos_src must be int32"); + TORCH_CHECK(numAcc.has_value() == seqLens.has_value(), "draft_step: num_acc and seq_lens come together"); + TORCH_CHECK(!numAcc.has_value() || step == 0, "draft_step: the KV rollback (num_acc) is pass 0's"); + if (numAcc.has_value()) + { + checkInt(*numAcc, rows, "num_acc"); + checkInt(*seqLens, rows, "seq_lens"); + } + if (outKeys.has_value()) + { + TORCH_CHECK(outKeys->scalar_type() == at::kLong && outKeys->is_contiguous() && outKeys->numel() >= rows, + "draft_step: out_keys must be contiguous int64 [rows]"); + } + ds::Params p{}; + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_rows = xRows.has_value() ? reinterpret_cast(xRows->data_ptr()) : nullptr; + p.M = static_cast(rows); + p.wf = reinterpret_cast(wf.data_ptr()); + p.n_tiles = nTiles; + p.vocab_offset = static_cast(vocabOffset); + p.keys = reinterpret_cast(keys.data_ptr()); + p.counters = reinterpret_cast(counters.data_ptr()); + p.ws_uc = reinterpret_cast(wsUc.data_ptr()); + p.ws_mc = reinterpret_cast(wsMcPtr); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.out_keys = outKeys.has_value() ? reinterpret_cast(outKeys->data_ptr()) : nullptr; + p.st = ds::StepArgs{draftOut.data_ptr(), static_cast(draftOut.stride(0)), static_cast(step), + idsNext.data_ptr(), posSrc.data_ptr(), posNext.data_ptr(), kvLens.data_ptr(), + numAcc.has_value() ? numAcc->data_ptr() : nullptr, + seqLens.has_value() ? seqLens->data_ptr() : nullptr, static_cast(draftLen)}; + p.cd = ds::CodedArgs{coded ? excOff->data_ptr() : nullptr, + coded ? reinterpret_cast(excPos->data_ptr()) : nullptr, + coded ? reinterpret_cast(excW->data_ptr()) : nullptr, + static_cast((expBase << 7) | (expBase << 23))}; + if (fold.has_value()) + { + // The 14 fold-tail tensors in buildParams' order (ctrl, slots, src_new, dst_new, src_next, dst_next, + // src_draft, dst_draft, src_lens, dst_lens, ids, draft_ids, pos_off, kv_off): the in-kernel tail's operands, + // packed exactly as the standalone fold tail packs them. + std::vector f = *fold; + TORCH_CHECK(f.size() == 14, "minimax_m3_sm107tp4_draft_step: fold takes 14 tensors, got ", f.size()); + p.fold = minimaxM3Sm107Tp4FoldTailParams::buildParams(f[0], f[1], foldRows, f[2], f[3], f[4], f[5], f[6], f[7], + f[8], f[9], f[10], f[11], f[12], f[13], foldTokensPerStep); + p.hasFold = true; + } + p.fixedOrder = fixedOrder; + check(ds::invokeDraftStep(p, K, static_cast(numCtas), pdl, at::cuda::getCurrentCUDAStream(x.get_device()))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// Mutations through the raw multicast address ws_mc_ptr are not visible to the schema: that buffer is owned by the +// Python workspace and never aliased by the op's tensor arguments. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp4_draft_step(Tensor x, Tensor? x_rows, int rows, Tensor wf, int vocab_offset, " + "Tensor(a!) keys, " + "Tensor(b!) counters, Tensor(c!) ws_uc, int ws_mc_ptr, int rank, int world, Tensor(d!) draft_out, int step, " + "Tensor(e!) ids_next, Tensor pos_src, Tensor(f!) pos_next, Tensor(g!) kv_lens, Tensor? num_acc, " + "Tensor(h!)? seq_lens, int draft_len, int num_ctas, bool pdl, Tensor(i!)? out_keys, Tensor? exc_off=None, " + "Tensor? exc_pos=None, Tensor? exc_w=None, int exp_base=0" + ", Tensor(j!)[]? fold=None, int fold_rows=0, int fold_tokens_per_step=0, bool fixed_order=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_draft_step", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4DraftStep); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp new file mode 100644 index 000000000000..034521d29796 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp @@ -0,0 +1,103 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp4::drafterRopeFp8Kv; + +//! The MiniMax-M3 TP4 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h): RoPE on +//! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages +//! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. +void minimaxM3Sm107Tp4DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, + torch::Tensor const& seqLens, std::optional const& specPositionOffsets, torch::Tensor& kvPages, + torch::Tensor const& blockTables, torch::Tensor& qOut, int64_t qLen, int64_t numHeadsQ, int64_t numHeadsKv, + int64_t rotaryDim, int64_t tokensPerBlock, int64_t ropeMode, bool pdl, bool maxSharedCarveout) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + 2 * numHeadsKv) * drk::kHeadDim, "qkv is not [tokens, (Hq + 2 Hkv) * 128]"); + TORCH_CHECK(qLen >= 1 && qkv.size(0) % qLen == 0, "the rows are not whole passes of q_len rows"); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= drk::kHeadDim && rotaryDim % 16 == 0, "unsupported rotary_dim"); + TORCH_CHECK(cosSin.is_cuda() && cosSin.is_contiguous() && cosSin.scalar_type() == at::kFloat + && cosSin.numel() % rotaryDim == 0, + "cos_sin must be the fp32 [positions, rotary_dim / 2] (cos, sin) table"); + int64_t const requests = qkv.size(0) / qLen; + TORCH_CHECK(seqLens.is_cuda() && seqLens.scalar_type() == at::kInt && seqLens.numel() >= requests); + if (specPositionOffsets.has_value()) + { + TORCH_CHECK(specPositionOffsets->is_cuda() && specPositionOffsets->is_contiguous() + && specPositionOffsets->scalar_type() == at::kInt && specPositionOffsets->numel() >= requests * qLen); + } + TORCH_CHECK(kvPages.is_cuda() && kvPages.is_contiguous() && kvPages.dim() == 4 + && kvPages.scalar_type() == at::kFloat8_e4m3fn && kvPages.size(1) == numHeadsKv + && kvPages.size(2) == tokensPerBlock && kvPages.size(3) == drk::kHeadDim, + "kv_pages must be the E4M3 [sub-pages, Hkv, tokens_per_block, 128] cache"); + TORCH_CHECK(blockTables.is_cuda() && blockTables.scalar_type() == at::kInt && blockTables.dim() == 3 + && blockTables.size(0) >= requests && blockTables.size(1) == 2 && blockTables.stride(2) == 1, + "block_tables must be the int32 [requests, 2, max blocks] K / V sub-page table"); + TORCH_CHECK(qOut.is_cuda() && qOut.is_contiguous() && qOut.scalar_type() == at::kFloat8_e4m3fn + && qOut.numel() == qkv.size(0) * numHeadsQ * drk::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{qkv.device()}; + drk::Params params{}; + params.qkv = reinterpret_cast<__nv_bfloat16 const*>(qkv.data_ptr()); + params.cosSin = reinterpret_cast(cosSin.data_ptr()); + params.seqLens = seqLens.data_ptr(); + params.specPositionOffsets = specPositionOffsets.has_value() ? specPositionOffsets->data_ptr() : nullptr; + params.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + params.pool = reinterpret_cast<__nv_fp8_e4m3*>(kvPages.data_ptr()); + params.blockTable = blockTables.data_ptr(); + params.blockTableRequestStride = blockTables.stride(0); + params.blockTableKvStride = blockTables.stride(1); + params.numTokens = static_cast(qkv.size(0)); + params.qLen = static_cast(qLen); + params.numHeadsQ = static_cast(numHeadsQ); + params.numHeadsKv = static_cast(numHeadsKv); + params.rotaryDim = static_cast(rotaryDim); + params.tokensPerBlock = static_cast(tokensPerBlock); + params.ropeMode = static_cast(ropeMode); + cudaError_t const err + = drk::launch(params, pdl, maxSharedCarveout, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp4_drafter_rope_fp8_kv: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp4_drafter_rope_fp8_kv(Tensor qkv, Tensor cos_sin, Tensor seq_lens, " + "Tensor? spec_position_offsets, Tensor(a!) kv_pages, Tensor block_tables, Tensor(b!) q_out, int q_len, " + "int num_heads_q, int num_heads_kv, int rotary_dim, int tokens_per_block, int rope_mode, bool pdl, bool " + "max_shared_carveout=True) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_drafter_rope_fp8_kv", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4DrafterRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp new file mode 100644 index 000000000000..3df36817814d --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp @@ -0,0 +1,89 @@ +/* + * 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/minimaxM3Sm107Tp4/foldTail.h" +#include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The speculative step's forward-tail store and next-step staging in one launch (the C++ form of the Triton +//! kernels/m3_fold_tail.py, the same writes). The caller checks the row counts against the buffers (fold_tail). +static void foldTailImpl(bool oneWarp, torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + c10::cuda::CUDAGuard const deviceGuard{ctrl.device()}; + auto const p = minimaxM3Sm107Tp4FoldTailParams::buildParams(ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, + srcDraft, dstDraft, srcLens, dstLens, ids, draftIds, posOff, kvOff, tokensPerStep); + auto const stream = at::cuda::getCurrentCUDAStream(ctrl.get_device()); + auto const err = oneWarp + ? tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeFoldTailOneWarp(p, stream) + : tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeFoldTail(p, static_cast(rows), stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_fold_tail launch failed: ", cudaGetErrorString(err)); +} + +//! The tail as its standalone kernel (one CTA per row). +void minimaxM3Sm107Tp4FoldTail(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + foldTailImpl(false, ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, srcDraft, dstDraft, srcLens, dstLens, ids, + draftIds, posOff, kvOff, tokensPerStep); +} + +//! The tail's in-kernel form (one warp over the rows, foldTailWarp), the oracle of a caller that folds it in. +void minimaxM3Sm107Tp4FoldTailOneWarp(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + foldTailImpl(true, ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, srcDraft, dstDraft, srcLens, dstLens, ids, + draftIds, posOff, kvOff, tokensPerStep); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_fold_tail(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) dst_new, " + "Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); + m.def( + "minimax_m3_sm107_tp4_fold_tail_one_warp(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) " + "dst_new, Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_fold_tail", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4FoldTail); + m.impl("minimax_m3_sm107_tp4_fold_tail_one_warp", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4FoldTailOneWarp); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h new file mode 100644 index 000000000000..b73820b6fc54 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h @@ -0,0 +1,111 @@ +/* + * 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/kernels/minimaxM3Sm107Tp4/foldTail.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace minimaxM3Sm107Tp4FoldTailParams +{ + +constexpr int64_t kMaxWidth = 32; + +inline void check2d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 2 && t.size(1) <= kMaxWidth, + "minimax_m3_sm107_tp4_fold_tail: ", name, " must be a 2-D int32 CUDA view at most 32 wide"); +} + +inline void check1d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp4_fold_tail: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +//! The fold tail's operands checked and packed (the standalone fold tail op and any op that runs the tail in-kernel, +//! e.g. the drafter's last-pass draft_step, build them identically). +inline tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams buildParams(torch::Tensor const& ctrl, + torch::Tensor const& slots, int64_t rows, torch::Tensor const& srcNew, torch::Tensor& dstNew, + torch::Tensor const& srcNext, torch::Tensor& dstNext, torch::Tensor const& srcDraft, torch::Tensor& dstDraft, + torch::Tensor const& srcLens, torch::Tensor& dstLens, torch::Tensor& ids, torch::Tensor& draftIds, + torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + TORCH_CHECK(ctrl.is_cuda() && ctrl.scalar_type() == torch::kInt32 && ctrl.numel() >= 1, + "minimax_m3_sm107_tp4_fold_tail: ctrl must be an int32 CUDA tensor"); + TORCH_CHECK(slots.is_cuda() && slots.scalar_type() == torch::kInt64 && slots.dim() == 1 && slots.stride(0) == 1, + "minimax_m3_sm107_tp4_fold_tail: slots must be a contiguous 1-D int64 CUDA tensor"); + TORCH_CHECK(rows > 0 && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, + "minimax_m3_sm107_tp4_fold_tail: rows must be positive and tokens_per_step in [1, 32]"); + check2d(srcNew, "src_new"); + check2d(dstNew, "dst_new"); + check2d(srcNext, "src_next"); + check2d(dstNext, "dst_next"); + check2d(srcDraft, "src_draft"); + check2d(dstDraft, "dst_draft"); + check1d(srcLens, "src_lens"); + check1d(dstLens, "dst_lens"); + check1d(ids, "ids"); + check1d(draftIds, "draft_ids"); + check1d(posOff, "pos_off"); + check1d(kvOff, "kv_off"); + tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams p{}; + p.ctrl = ctrl.data_ptr(); + p.slots = slots.data_ptr(); + p.srcNew = srcNew.data_ptr(); + p.sn0 = srcNew.stride(0); + p.sn1 = srcNew.stride(1); + p.snW = static_cast(srcNew.size(1)); + p.dstNew = dstNew.data_ptr(); + p.dn0 = dstNew.stride(0); + p.dn1 = dstNew.stride(1); + p.dnW = static_cast(dstNew.size(1)); + p.srcNext = srcNext.data_ptr(); + p.sx0 = srcNext.stride(0); + p.sx1 = srcNext.stride(1); + p.sxW = static_cast(srcNext.size(1)); + p.dstNext = dstNext.data_ptr(); + p.dx0 = dstNext.stride(0); + p.dx1 = dstNext.stride(1); + p.dxW = static_cast(dstNext.size(1)); + p.srcDraft = srcDraft.data_ptr(); + p.sd0 = srcDraft.stride(0); + p.sd1 = srcDraft.stride(1); + p.sdW = static_cast(srcDraft.size(1)); + p.dstDraft = dstDraft.data_ptr(); + p.dd0 = dstDraft.stride(0); + p.dd1 = dstDraft.stride(1); + p.ddW = static_cast(dstDraft.size(1)); + p.srcLens = srcLens.data_ptr(); + p.dstLens = dstLens.data_ptr(); + p.ids = ids.data_ptr(); + p.draftIds = draftIds.data_ptr(); + p.posOff = posOff.data_ptr(); + p.kvOff = kvOff.data_ptr(); + p.tokensPerStep = static_cast(tokensPerStep); + return p; +} + +} // namespace minimaxM3Sm107Tp4FoldTailParams +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp new file mode 100644 index 000000000000..e2f65f63ae33 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp @@ -0,0 +1,92 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107Tp4::fullCtxReduce; + +//! Merges the PrimTS split-KV decode scratch (partial_o [B, Hkv, S, R, 128] 16-bit, partial_stats [B, Hkv, S, R] +//! fp32 log2-LSE) of the MiniMax-M3 TP4 full-context decode into out (BF16 [B * SQ, Hkv * hR, 128]); with q8 / sf +//! also its MXFP8 copy (E4M3 [>= B * SQ, Hkv * hR * 128] + 128x4-swizzled UE8M0 scales, as mxfp8_quantize(out, True)). +void minimaxM3Sm107Tp4FullCtxReduce(torch::Tensor const& partialO, torch::Tensor const& partialStats, + torch::Tensor const& seqLensKv, torch::Tensor& out, int64_t seqLenQ, int64_t qTokensPerCta, int64_t tileKv, + int64_t instsKv, bool causal, bool pdl, std::optional const& q8, + std::optional const& sf) +{ + TORCH_CHECK(partialO.is_cuda() && partialO.is_contiguous() && partialO.dim() == 5); + TORCH_CHECK(partialO.scalar_type() == at::kHalf || partialO.scalar_type() == at::kBFloat16); + TORCH_CHECK(partialO.size(4) == fcr::kHeadDim && partialO.size(2) <= fcr::kMaxSplits); + TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() + && partialStats.sizes() == partialO.sizes().slice(0, 4)); + TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + int const rows = static_cast(partialO.size(3)); + TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); + int const headsPerKv = rows / static_cast(seqLenQ); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.is_contiguous() + && out.numel() == partialO.size(0) * seqLenQ * partialO.size(1) * headsPerKv * fcr::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{partialO.device()}; + fcr::Params params{partialO.data_ptr(), partialStats.data_ptr(), seqLensKv.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), static_cast(partialO.size(0)), + static_cast(partialO.size(1)), static_cast(partialO.size(2)), rows, static_cast(seqLenQ), + headsPerKv, static_cast(qTokensPerCta), static_cast(tileKv), static_cast(instsKv), causal, + partialO.scalar_type() == at::kBFloat16}; + TORCH_CHECK(q8.has_value() == sf.has_value()); + if (q8.has_value()) + { + int64_t const width = partialO.size(1) * headsPerKv * fcr::kHeadDim; + int64_t const tokens = partialO.size(0) * seqLenQ; + int64_t const tileRows = (tokens + 127) / 128 * 128; + TORCH_CHECK(q8->is_cuda() && q8->is_contiguous() && q8->scalar_type() == at::kFloat8_e4m3fn && q8->dim() == 2 + && q8->size(1) == width && q8->size(0) >= tokens); + TORCH_CHECK(sf->is_cuda() && sf->is_contiguous() && sf->scalar_type() == at::kByte + && sf->numel() >= tileRows * (width / 32)); + params.q8 = reinterpret_cast(q8->data_ptr()); + params.sf = sf->data_ptr(); + params.numSfKTiles = static_cast(width / 128); + } + cudaError_t const err + = fcr::invokeFullCtxReduce(params, pdl, at::cuda::getCurrentCUDAStream(partialO.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp4_fullctx_reduce: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp4_fullctx_reduce(Tensor partial_o, Tensor partial_stats, Tensor seq_lens_kv, " + "Tensor(a!) out, int seq_len_q, int q_tokens_per_cta, int tile_kv, int insts_kv, bool causal, bool pdl, " + "Tensor(b!)? q8=None, Tensor(c!)? sf=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_fullctx_reduce", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4FullCtxReduce); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4McAllreduceNormOp.cpp new file mode 100644 index 000000000000..27442555f3b4 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4McAllreduceNormOp.cpp @@ -0,0 +1,181 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h" + +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxTokens = 48; +constexpr int64_t kCluster = 8; +constexpr int64_t kMaxWorld = 8; + +} // namespace + +//! MiniMax-M3 SM107 TP4 multicast all-reduce (stock rounding points: BF16 rank sum, then the residual add) + residual + +//! RMSNorm (+ optional MoE finalize, + optional MXFP8 output). +//! +//! @param x BF16 [M, H] partial, or with `perm` the permuted FC2 output [rows, H]. +//! @param x2 Optional BF16 [M, H] second addend (the shared expert in finalize mode). +//! @param residual BF16 [M, H], 1 <= M <= 48. +//! @param normWeight BF16 [H]. +//! @param wsUc This rank's unicast view of the exchange workspace (int32 words); re-armed in place. +//! @param wsMcPtr Device address of the workspace's multicast mapping. +//! @param flags Int32 per-CTA call-parity counters, at least maxt * 8; updated in place. +//! @param cluster Cluster size; 8 is the only one built. +//! @param perm, scales Finalize mode: int32 [M, topk] permuted-row index and fp32 / bf16 [M, topk] routing weights. +//! @param maxt Token capacity of the workspace bucket (8, 16 or 48). +//! @param quant Also return the MXFP8 copy of the normed rows (E4M3 [M, H], swizzled UE8M0 scales). +//! @param prepushed 0, or 2 / 3 when producers already pushed this rank's rows (see the kernel header). +//! @param opts Epilogue variant; every variant is bit-identical. +//! @param capOut Optional BF16 [>= M, H] view (row stride a multiple of 8, 16-byte aligned) that receives the Eagle3 +//! capture of the updated residual: m3_spec_tail capture_norm with capW [H] and capEps, bit for bit (H 6144; not the +//! drafter boundary, opts bit 16). +//! @return [normed, updated residual] or, with `quant`, [normed, updated residual, E4M3 rows, scale bytes]. +std::vector minimaxM3Sm107Tp4McAllreduceNorm(torch::Tensor const& x, + std::optional const& x2, torch::Tensor const& residual, torch::Tensor const& normWeight, double eps, + torch::Tensor& wsUc, int64_t wsMcPtr, torch::Tensor& flags, int64_t rank, int64_t world, int64_t cluster, bool pdl, + std::optional const& perm, std::optional const& scales, int64_t maxt, bool quant, + int64_t prepushed, int64_t opts, std::optional const& capOut, + std::optional const& capW, double capEps) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 107, "minimax_m3_sm107_tp4_mc_allreduce_norm requires SM107, got SM", sm); + TORCH_CHECK(cluster == kCluster, "minimax_m3_sm107_tp4_mc_allreduce_norm: only cluster 8 is built, got ", cluster); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107Tp4McAllreduceNormSupported( + static_cast(maxt), static_cast(opts)), + "minimax_m3_sm107_tp4_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); + TORCH_CHECK( + world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); + + tensorrt_llm::kernels::MinimaxM3Sm107Tp4McAllreduceNormParams p; + if (x2.has_value()) + { + TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), + "x2 must be a contiguous bf16 tensor shaped like residual"); + p.x2 = reinterpret_cast<__nv_bfloat16 const*>(x2->data_ptr()); + } + if (perm.has_value()) + { + // Finalize mode: x is the permuted FC2 output [rows, H]. + TORCH_CHECK(scales.has_value() && perm->scalar_type() == at::kInt && perm->is_contiguous() && perm->dim() == 2 + && perm->size(0) == residual.size(0), + "perm must be a contiguous int32 [M, topk] tensor, with scales"); + TORCH_CHECK(scales->is_contiguous() && scales->numel() == perm->numel() + && (scales->scalar_type() == at::kFloat || scales->scalar_type() == at::kBFloat16), + "scales must be a contiguous fp32 or bf16 tensor with one weight per perm entry"); + TORCH_CHECK(x.dim() == 2 && x.size(1) == residual.size(1), "fc2 must be [rows, H]"); + p.fc2 = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.perm = perm->data_ptr(); + p.scales = scales->data_ptr(); + p.scaleFp32 = scales->scalar_type() == at::kFloat; + p.topk = static_cast(perm->size(1)); + } + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2, + "x must be a contiguous 2-D bf16 CUDA tensor"); + TORCH_CHECK(residual.dim() == 2 && residual.size(0) >= 1 && residual.size(0) <= std::min(kMaxTokens, maxt), + "residual must hold 1..min(48, bucket) rows"); + TORCH_CHECK(residual.scalar_type() == at::kBFloat16 && residual.is_contiguous() + && (perm.has_value() || residual.sizes() == x.sizes()), + "residual must be a contiguous bf16 tensor shaped like x"); + TORCH_CHECK( + normWeight.scalar_type() == at::kBFloat16 && normWeight.numel() == residual.size(1), "norm_w must be bf16 [H]"); + TORCH_CHECK(wsUc.scalar_type() == at::kInt && wsUc.is_cuda(), "ws_uc must be an int32 CUDA tensor"); + TORCH_CHECK(flags.scalar_type() == at::kInt && flags.is_contiguous() && flags.numel() >= maxt * kCluster, + "flags must be a contiguous int32 tensor of at least bucket * 8 counters"); + + int64_t const m = residual.size(0), h = residual.size(1); + TORCH_CHECK(!quant || h % 128 == 0, "MXFP8 output needs H % 128 == 0"); + c10::cuda::CUDAGuard const deviceGuard{residual.device()}; + auto outNormed = torch::empty_like(residual); + auto outResidual = torch::empty_like(residual); + auto outQ = quant ? torch::empty({m, h}, residual.options().dtype(at::kFloat8_e4m3fn)) : torch::Tensor(); + auto outSf = quant ? torch::empty({(m + 127) / 128 * 128 * (h / 32)}, residual.options().dtype(at::kByte)) + : torch::Tensor(); + + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.numTokens = static_cast(m); + p.hidden = static_cast(h); + p.residual = reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()); + p.normWeight = reinterpret_cast<__nv_bfloat16 const*>(normWeight.data_ptr()); + p.eps = static_cast(eps); + p.wsUc = wsUc.data_ptr(); + p.wsMc = reinterpret_cast(wsMcPtr); + p.flags = flags.data_ptr(); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.outNormed = reinterpret_cast<__nv_bfloat16*>(outNormed.data_ptr()); + p.outResidual = reinterpret_cast<__nv_bfloat16*>(outResidual.data_ptr()); + p.outQuant = quant ? reinterpret_cast(outQ.data_ptr()) : nullptr; + p.outSf = quant ? outSf.data_ptr() : nullptr; + p.prepushed = static_cast(prepushed); + if (capOut.has_value()) + { + TORCH_CHECK((opts & 16) == 0, "the capture epilogue is not built for the drafter boundary (opts bit 16)"); + TORCH_CHECK(capOut->scalar_type() == at::kBFloat16 && capOut->dim() == 2 && capOut->size(0) >= m + && capOut->size(1) == h && capOut->stride(1) == 1 && capOut->stride(0) % 8 == 0 + && reinterpret_cast(capOut->data_ptr()) % 16 == 0 && capOut->device() == residual.device(), + "cap_out must be a bf16 [>= M, H] view with unit column stride, a row stride that is a multiple of 8 and " + "16-byte alignment"); + TORCH_CHECK( + capW.has_value() && capW->scalar_type() == at::kBFloat16 && capW->is_contiguous() && capW->numel() == h, + "cap_w must be a contiguous bf16 [H] tensor"); + p.capOut = reinterpret_cast<__nv_bfloat16*>(capOut->data_ptr()); + p.capStride = capOut->stride(0); + p.capWeight = reinterpret_cast<__nv_bfloat16 const*>(capW->data_ptr()); + p.capEps = static_cast(capEps); + } + tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp4McAllreduceNorm(p, static_cast(maxt), + static_cast(opts), pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); + if (quant) + { + return {outNormed, outResidual, outQ, outSf}; + } + return {outNormed, outResidual}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_mc_allreduce_norm(Tensor x, Tensor? x2, Tensor residual, Tensor norm_w, float eps, " + "Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) flags, int rank, int world, int cluster, bool pdl, " + "Tensor? perm, Tensor? scales, int maxt, bool quant, int prepushed, int opts, Tensor(c!)? cap_out=None, " + "Tensor? cap_w=None, float cap_eps=0.) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_mc_allreduce_norm", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4McAllreduceNorm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp new file mode 100644 index 000000000000..3ae68e7b4681 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp @@ -0,0 +1,185 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp4::producerPw4; + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPageSize = 128; + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "sm107_tp4_producer_pw: ", name, " must be a CUDA tensor"); +} + +} // namespace + +//! The SM107 TP4 NVFP4 QKV + indexer norm/RoPE/KV-insert producer (same arguments and outputs as +//! minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert at 16 query, 1 KV and 1 index head, head_dim 128, rotary_dim 64): +//! returns (q FP8 [T, 8, 128], index_q FP8 [T, 1, 128]) and writes the NVFP4 K/V of the KV head and the FP8 index-K +//! of each token's cache slot. `early` triggers the dependents on entry; `prewait_step_inputs` reads the position ids, +//! the slots and the RoPE rows before the grid-dependency wait (none of them may be written by a PDL predecessor). +std::vector minimaxM3Sm107Tp4ProducerPw(torch::Tensor const& packed, torch::Tensor& kvData, + torch::Tensor& kvScale, torch::Tensor& indexK, torch::Tensor const& outCacheLoc, torch::Tensor const& kvQuantScale, + double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds, bool pdl, + bool early, bool prewaitStepInputs) +{ + checkCuda(packed, "packed"); + for (torch::Tensor const* t : std::initializer_list{&kvData, &kvScale, &indexK, &outCacheLoc, + &kvQuantScale, &qWeight, &kWeight, &indexQWeight, &indexKWeight, &rotaryCosSin, &positionIds}) + { + TORCH_CHECK(t->is_cuda() && t->get_device() == packed.get_device(), + "sm107_tp4_producer_pw: every operand must be on packed's CUDA device"); + } + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0 + && reinterpret_cast(kvData.data_ptr()) % 2 == 0 + && reinterpret_cast(indexK.data_ptr()) % 4 == 0, + "sm107_tp4_producer_pw: packed and the caches must start at the vector-load / store alignment"); + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp4_producer_pw requires SM100 or newer"); + int64_t const headsQ = pw::kHeadsQ; + int64_t const headsKV = pw::kHeadsKV; + int64_t const headsIndex = pw::kHeadsIndex; + int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; + TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim + && packed.is_contiguous(), + "sm107_tp4_producer_pw: packed must be contiguous BF16 [T, 20 * 128]"); + // The kernel's layout (the native producer's checkMinimaxM3Nvfp4HndKVPools at one KV head): HND pools + // [pages, 2 planes, 1 head, 128 tokens, 64 data / 8 scale bytes], dense token rows and head blocks, planes and + // pages that do not overlap. + TORCH_CHECK(kvData.scalar_type() == at::kByte && kvData.dim() == 5 && kvData.size(0) > 0 && kvData.size(1) == 2 + && kvData.size(2) == headsKV && kvData.size(3) == kPageSize && kvData.size(4) == kHeadDim / 2, + "sm107_tp4_producer_pw: kv_data must be the NVFP4 pool [pages, 2, 1, 128, 64] uint8"); + TORCH_CHECK(kvScale.scalar_type() == at::kByte && kvScale.dim() == 5 && kvScale.size(0) == kvData.size(0) + && kvScale.size(1) == 2 && kvScale.size(2) == headsKV && kvScale.size(3) == kPageSize + && kvScale.size(4) == kHeadDim / 16, + "sm107_tp4_producer_pw: kv_scale must be the E4M3 scale pool [pages, 2, 1, 128, 8] uint8"); + TORCH_CHECK(kvData.stride(4) == 1 && kvData.stride(3) == kvData.size(4) + && kvData.stride(2) == kvData.size(3) * kvData.stride(3), + "sm107_tp4_producer_pw: kv_data must have dense token rows and head blocks"); + TORCH_CHECK(kvScale.stride(4) == 1 && kvScale.stride(3) == kvScale.size(4) + && kvScale.stride(2) == kvScale.size(3) * kvScale.stride(3), + "sm107_tp4_producer_pw: kv_scale must have dense token rows and head blocks"); + TORCH_CHECK(kvData.stride(1) >= kvData.size(2) * kvData.stride(2) && kvData.stride(0) >= 2 * kvData.stride(1) + && kvScale.stride(1) >= kvScale.size(2) * kvScale.stride(2) && kvScale.stride(0) >= 2 * kvScale.stride(1), + "sm107_tp4_producer_pw: pool planes or pages overlap"); + TORCH_CHECK(kvData.stride(0) % 2 == 0 && kvData.stride(1) % 2 == 0, + "sm107_tp4_producer_pw: kv_data page and plane strides must keep the uint16 stores aligned"); + TORCH_CHECK(indexK.scalar_type() == at::kFloat8_e4m3fn && indexK.dim() == 4 && indexK.size(0) == kvData.size(0) + && indexK.size(1) == 1 && indexK.size(2) == kPageSize && indexK.size(3) == kHeadDim && indexK.stride(3) == 1 + && indexK.stride(2) == kHeadDim && indexK.stride(0) >= kPageSize * kHeadDim && indexK.stride(0) % 4 == 0, + "sm107_tp4_producer_pw: index_k must be FP8 [pages, 1, 128, 128] with a page stride keeping the 32-bit stores " + "aligned"); + int64_t const numTokens = packed.size(0); + TORCH_CHECK( + outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous() && outCacheLoc.numel() >= numTokens, + "sm107_tp4_producer_pw: out_cache_loc must be int32 [>= T]"); + TORCH_CHECK( + positionIds.scalar_type() == at::kInt && positionIds.is_contiguous() && positionIds.numel() >= numTokens, + "sm107_tp4_producer_pw: position_ids must be int32 [>= T]"); + TORCH_CHECK(kvQuantScale.scalar_type() == at::kFloat && kvQuantScale.is_contiguous() && kvQuantScale.numel() >= 3, + "sm107_tp4_producer_pw: kv_quant_scale must be fp32 [>= 3]"); + TORCH_CHECK(rotaryCosSin.scalar_type() == at::kFloat && rotaryCosSin.is_contiguous() && rotaryCosSin.dim() == 3 + && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == 32 + && reinterpret_cast(rotaryCosSin.data_ptr()) % 16 == 0, + "sm107_tp4_producer_pw: rotary_cos_sin must be fp32 [positions, 2, 32], 16-byte aligned (read as float4)"); + for (auto const* w : {&qWeight, &kWeight, &indexQWeight, &indexKWeight}) + { + TORCH_CHECK(w->scalar_type() == at::kBFloat16 && w->is_contiguous() && w->numel() == kHeadDim + && reinterpret_cast(w->data_ptr()) % 8 == 0, + "sm107_tp4_producer_pw: norm weights must be BF16 [128], 8-byte aligned (read as 4-value words)"); + } + auto const fp8 = packed.options().dtype(at::kFloat8_e4m3fn); + auto qOut = torch::empty({numTokens, headsQ, kHeadDim}, fp8); + auto indexQOut = torch::empty({numTokens, headsIndex, kHeadDim}, fp8); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + pw::Params p{}; + p.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + p.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + p.indexQOut = reinterpret_cast<__nv_fp8_e4m3*>(indexQOut.data_ptr()); + p.kvData = kvData.data_ptr(); + p.kvScale = kvScale.data_ptr(); + p.indexK = reinterpret_cast<__nv_fp8_e4m3*>(indexK.data_ptr()); + p.outCacheLoc = outCacheLoc.data_ptr(); + p.kvQuantScale = kvQuantScale.data_ptr(); + p.dataPageStride = kvData.stride(0); + p.dataPlaneStride = kvData.stride(1); + p.dataHeadStride = kvData.stride(2); + p.dataTokenStride = kvData.stride(3); + p.scalePageStride = kvScale.stride(0); + p.scalePlaneStride = kvScale.stride(1); + p.scaleHeadStride = kvScale.stride(2); + p.indexPageStride = indexK.stride(0); + p.indexTokenStride = indexK.stride(2); + p.numPages = kvData.size(0); + p.numTokens = static_cast(numTokens); + p.eps = static_cast(eps); + p.qWeight = reinterpret_cast<__nv_bfloat16 const*>(qWeight.data_ptr()); + p.kWeight = reinterpret_cast<__nv_bfloat16 const*>(kWeight.data_ptr()); + p.indexQWeight = reinterpret_cast<__nv_bfloat16 const*>(indexQWeight.data_ptr()); + p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); + p.rotaryCosSin = rotaryCosSin.data_ptr(); + p.positionIds = positionIds.data_ptr(); + cudaError_t const err = pw::invokeProducerPw4(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "sm107_tp4_producer_pw launch: ", cudaGetErrorString(err)); + } + return {qOut, indexQOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_producer_pw(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) index_k, " + "Tensor out_cache_loc, Tensor kv_quant_scale, float eps, Tensor q_weight, Tensor k_weight, " + "Tensor index_q_weight, Tensor index_k_weight, Tensor rotary_cos_sin, Tensor position_ids, bool pdl, " + "bool early, bool prewait_step_inputs=False) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_producer_pw", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4ProducerPw); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp new file mode 100644 index 000000000000..4f7fda5cac90 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp @@ -0,0 +1,113 @@ +/* + * 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/minimaxM3Sm107Tp4/sampledAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3Sm107Tp4::sampledAccept; + +//! uint64 words of the sampled-accept exchange workspace for `world` ranks and SA-1R variant `oneRoundK`. +int64_t minimaxM3Sm107Tp4SampledAcceptWsWords(int64_t world, int64_t oneRoundK) +{ + return sa::sampledAcceptWsWords(static_cast(world), static_cast(oneRoundK)); +} + +//! Sampled (temperature + top-p) verification from this rank's BF16 logits shard [rows, Vs], rows = nreq * +//! (draft_len + 1) <= 48: the sampled token of every row (tokens), then the strict acceptance against draft (accTokens, +//! numAcc, gatherIds, optional nextNew column 0). Per-row Philox (seeds, offsets); the exchange goes through the +//! multicast workspace wsMcPtr / wsUc with the (finished rows, call) counters. +void minimaxM3Sm107Tp4SampledAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor const& temperatures, + torch::Tensor const& topPs, torch::Tensor const& seeds, torch::Tensor const& offsets, torch::Tensor& wsUc, + int64_t wsMcPtr, torch::Tensor& counters, torch::Tensor& tokens, int64_t rank, int64_t world, + torch::Tensor const& draft, torch::Tensor const& batchIdx, torch::Tensor& accTokens, torch::Tensor& numAcc, + torch::Tensor& gatherIds, int64_t forceTotal, std::optional const& nextNew, + std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, + bool prewaitParams, bool earlyTrigger, int64_t oneRoundK, std::optional const& decided, bool ilpWalk) +{ + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + sa::Args a; + a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + a.ld = logits.stride(0); + a.rows = static_cast(logits.size(0)); + a.Vs = static_cast(logits.size(1)); + TORCH_CHECK(a.Vs <= sa::kMaxVs && a.rows <= sa::kMaxRows && world <= sa::kMaxWorld + && static_cast(vocabOffset) + a.Vs <= (1LL << sa::kIdBits)); + a.vocab_offset = static_cast(vocabOffset); + a.temperatures = temperatures.data_ptr(); + a.top_ps = topPs.data_ptr(); + a.seeds = reinterpret_cast(seeds.data_ptr()); + a.offsets = reinterpret_cast(offsets.data_ptr()); + a.ws_uc = reinterpret_cast(wsUc.data_ptr()); + a.ws_mc = reinterpret_cast(wsMcPtr); + a.counters = reinterpret_cast(counters.data_ptr()); + a.tokens = tokens.data_ptr(); + a.extra_attempts = extraAttempts.has_value() ? extraAttempts->data_ptr() : nullptr; + a.rank = static_cast(rank); + a.world = static_cast(world); + a.draft = draft.data_ptr(); + a.batch_idx = batchIdx.data_ptr(); + a.acc_tokens = accTokens.data_ptr(); + a.num_acc = numAcc.data_ptr(); + a.gather_ids = reinterpret_cast(gatherIds.data_ptr()); + a.draft_len = static_cast(draft.size(1)); + a.nreq = static_cast(draft.size(0)); + a.force_total = static_cast(forceTotal); + a.next_new = nextNew.has_value() ? nextNew->data_ptr() : nullptr; + a.next_stride = nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0; + a.force_vec = forceVec.has_value() ? forceVec->data_ptr() : nullptr; + a.prewait_params = prewaitParams ? 1 : 0; + a.early_trigger = earlyTrigger ? 1 : 0; + TORCH_CHECK(oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "one_round_k must be 0, 9 or 82, got ", oneRoundK); + a.one_round_k = static_cast(oneRoundK); + a.decided = decided.has_value() ? decided->data_ptr() : nullptr; + a.ilp_walk = ilpWalk ? 1 : 0; + cudaError_t const err = sa::invokeSampledAccept(a, usePdl, at::cuda::getCurrentCUDAStream(logits.get_device())); + TORCH_CHECK(err == cudaSuccess, "m3_sampled_accept launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// ws_mc_ptr is the multicast address of ws_uc's mapping: the op writes through it, which the schema cannot show. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + // No tensor arguments, so no dispatch key: registered with the schema as a catch-all kernel. + m.def("minimax_m3_sm107tp4_sampled_accept_ws_words(int world, int one_round_k=0) -> int", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SampledAcceptWsWords); + m.def( + "minimax_m3_sm107tp4_sampled_accept(Tensor logits, int vocab_offset, Tensor temperatures, Tensor top_ps, " + "Tensor seeds, Tensor offsets, Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) counters, Tensor(c!) tokens, " + "int rank, int world, Tensor draft, Tensor batch_idx, Tensor(d!) acc_tokens, Tensor(e!) num_acc, " + "Tensor(f!) gather_ids, int force_total, Tensor(g!)? next_new, Tensor? force_vec, " + "Tensor(h!)? extra_attempts, bool use_pdl, bool prewait_params=True, bool early_trigger=False, " + "int one_round_k=0, Tensor(i!)? decided=None, bool ilp_walk=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp4_sampled_accept", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SampledAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SelectLongOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SelectLongOp.cpp new file mode 100644 index 000000000000..b1cf6be45569 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SelectLongOp.cpp @@ -0,0 +1,87 @@ +/* + * 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/minimaxM3Sm107Tp4/selectLong2.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, into a caller-owned table. +//! +//! @param scores Float32 CUDA tensor [num_kv_heads, num_blocks, total_q], any strides. +//! @param nValidBlocks Contiguous int32 CUDA tensor with at least total_q elements. +//! @param out Int32 CUDA tensor of logical shape [total_q, num_kv_heads, 16] with unit stride in the last dim (the +//! head-major backing the target hands over is a permuted view). Every entry is written. +//! @param triggerPhase Where the dependent grid is released without earlyTrigger: 0 at this grid's exit, 1 right +//! after every score of the row is loaded (both after this grid's dependency wait). +//! @param fastTail Rank rows of at most 64 candidates on warp 0 and read the valid count before the wait (same picks). +void minimaxM3Sm107Tp4SelectLong(torch::Tensor const& scores, torch::Tensor const& nValidBlocks, torch::Tensor& out, + int64_t initBlocks, int64_t localBlocks, bool usePdl, bool earlyTrigger, int64_t triggerPhase, bool fastTail) +{ + TORCH_CHECK(scores.is_cuda() && scores.scalar_type() == torch::kFloat32 && scores.dim() == 3, + "minimax_m3_sm107_tp4_select_long: scores must be a 3-D float32 CUDA tensor [num_kv_heads, num_blocks, " + "total_q]"); + TORCH_CHECK(nValidBlocks.is_cuda() && nValidBlocks.scalar_type() == torch::kInt32 && nValidBlocks.is_contiguous() + && nValidBlocks.numel() >= scores.size(2), + "minimax_m3_sm107_tp4_select_long: n_valid_blocks must be contiguous int32 CUDA with at least total_q " + "elements"); + int64_t const numKvHeads = scores.size(0); + int64_t const numBlocks = scores.size(1); + int64_t const totalQ = scores.size(2); + TORCH_CHECK(numBlocks <= tensorrt_llm::kernels::minimaxM3Sm107Tp4::kSelectLong2MaxBlocks, + "minimax_m3_sm107_tp4_select_long supports at most ", + tensorrt_llm::kernels::minimaxM3Sm107Tp4::kSelectLong2MaxBlocks, " blocks, got ", numBlocks); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kInt32 && out.dim() == 3 && out.size(0) == totalQ + && out.size(1) == numKvHeads && out.size(2) == 16 && out.stride(2) == 1, + "minimax_m3_sm107_tp4_select_long: out must be int32 [total_q, num_kv_heads, 16] with unit last-dim stride"); + + c10::cuda::CUDAGuard const deviceGuard{scores.device()}; + auto const stream = at::cuda::getCurrentCUDAStream(scores.get_device()); + cudaError_t const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSelectLong2(scores.data_ptr(), + scores.stride(0), scores.stride(1), scores.stride(2), nValidBlocks.data_ptr(), out.data_ptr(), + out.stride(0), out.stride(1), static_cast(numKvHeads), static_cast(numBlocks), + static_cast(totalQ), static_cast(initBlocks), static_cast(localBlocks), usePdl, + earlyTrigger, static_cast(triggerPhase), fastTail, stream); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "minimax_m3_sm107_tp4_select_long launch failed: ", cudaGetErrorString(err)); + } +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_select_long(Tensor scores, Tensor n_valid_blocks, Tensor(a!) out, int init_blocks, " + "int local_blocks, bool use_pdl=False, bool early_trigger=False, int trigger_phase=0, bool fast_tail=True) -> " + "()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_select_long", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SelectLong); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp new file mode 100644 index 000000000000..b312cf79a715 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp @@ -0,0 +1,66 @@ +/* + * 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/minimaxM3Sm107Tp4/siluAndMul.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp4/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, +//! d]. +void minimaxM3Sm107Tp4SiluAndMul(torch::Tensor const& x, torch::Tensor& out) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp4_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp4_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp4_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); + TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp4_silu_and_mul: at most 65535 rows"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err + = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), + reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_silu_and_mul launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SiluAndMul); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp new file mode 100644 index 000000000000..53121e3b1e7a --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp @@ -0,0 +1,191 @@ +/* + * 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/minimaxM3Sm107Tp4/sparseDecodeLc4.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeads = 16; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kTopK = 16; + +void checkBytes(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, "minimax_m3_sm107_tp4_sparse_decode_lc: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3 long-context sparse decode over the selected NVFP4 blocks (TP4 geometry: 16 query heads on +//! one KV head, head_dim 128, page 128, top-k 16). +//! +//! @param q E4M3 bytes [T, 16, 128], unit stride in the last dim. +//! @param k, v Packed E2M1 bytes [pages, 1, 128, 64] (the per-layer HND views). +//! @param kBlockScale, vBlockScale E4M3 bytes [pages, 1, 128, 8]: K token-major, V in 4x4 token-quad order. +//! @param kGlobalScale, vGlobalScale One-element float32 per-tensor dequant scales. +//! @param topk Int32 [T, 16] ascending selected block ids with -1 tail (kv head 0 of the head-major table). +//! @param out BF16 [T, 2048], written. q8 / sf, when given, receive the MXFP8 copy and its swizzled UE8M0 scales. +void minimaxM3Sm107Tp4SparseDecodeLc(torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& kBlockScale, torch::Tensor const& vBlockScale, torch::Tensor const& kGlobalScale, + torch::Tensor const& vGlobalScale, torch::Tensor const& topk, torch::Tensor const& blockTable, + torch::Tensor const& seqLens, torch::Tensor& out, std::optional q8, std::optional sf, + int64_t numSfKTiles, int64_t decodeQueryLen, double smScale, bool linearSoftmax, bool usePdl, bool earlyTrigger, + std::optional prefetchA, std::optional prefetchB, std::optional n1Self, + std::optional n1Next, std::optional n1K, std::optional n1V, + std::optional n1KScale, std::optional n1VScale, bool identity, bool oneDiv, + bool identityStep, bool btPrewait) +{ + TORCH_CHECK(earlyTrigger, "minimax_m3_sm107_tp4_sparse_decode_lc is instantiated with early_trigger=True only"); + TORCH_CHECK(q.is_cuda(), "minimax_m3_sm107_tp4_sparse_decode_lc: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m3_sm107_tp4_sparse_decode_lc requires an SM100-family GPU"); + checkBytes(q, "q"); + checkBytes(k, "k"); + checkBytes(v, "v"); + checkBytes(kBlockScale, "k_block_scale"); + checkBytes(vBlockScale, "v_block_scale"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeads && q.size(2) == kHeadDim && q.stride(2) == 1, + "minimax_m3_sm107_tp4_sparse_decode_lc: q must be [T, 16, 128] with unit last-dim stride"); + TORCH_CHECK(kGlobalScale.scalar_type() == torch::kFloat32 && vGlobalScale.scalar_type() == torch::kFloat32, + "minimax_m3_sm107_tp4_sparse_decode_lc: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m3_sm107_tp4_sparse_decode_lc: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_sm107_tp4_sparse_decode_lc: block_table and seq_lens must be int32"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.size(0) == q.size(0) + && out.size(1) == kHeads * kHeadDim && out.stride(1) == 1, + "minimax_m3_sm107_tp4_sparse_decode_lc: out must be bfloat16 [T, 2048] with unit last-dim stride"); + TORCH_CHECK( + q8.has_value() == sf.has_value(), "minimax_m3_sm107_tp4_sparse_decode_lc: pass both q8 and sf, or neither"); + bool const mxfp8 = q8.has_value(); + if (mxfp8) + { + checkBytes(*q8, "q8"); + checkBytes(*sf, "sf"); + } + + tensorrt_llm::kernels::minimaxM3Sm107Tp4::MinimaxM3SparseDecodeLcParams prm{}; + prm.q = reinterpret_cast(q.data_ptr()); + prm.q_tok = q.stride(0); + prm.q_head = q.stride(1); + prm.k = reinterpret_cast(k.data_ptr()); + prm.v = reinterpret_cast(v.data_ptr()); + prm.k_blk = k.stride(0); + prm.v_blk = v.stride(0); + prm.ksc = reinterpret_cast(kBlockScale.data_ptr()); + prm.vsc = reinterpret_cast(vBlockScale.data_ptr()); + prm.ksc_blk = kBlockScale.stride(0); + prm.vsc_blk = vBlockScale.stride(0); + prm.kgs = kGlobalScale.data_ptr(); + prm.vgs = vGlobalScale.data_ptr(); + prm.topk = topk.data_ptr(); + prm.t_tok = topk.stride(0); + prm.t_k = topk.stride(1); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.seq_lens = seqLens.data_ptr(); + prm.dql = static_cast(decodeQueryLen); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + prm.out_stride = out.stride(0); + prm.q8 = mxfp8 ? reinterpret_cast(q8->data_ptr()) : nullptr; + prm.sf = mxfp8 ? reinterpret_cast(sf->data_ptr()) : nullptr; + prm.num_sf_k_tiles = static_cast(numSfKTiles); + prm.identity = identity ? 1 : 0; + prm.identity_step = (identity || identityStep) ? 1 : 0; + prm.one_div = oneDiv ? 1 : 0; + prm.bt_prewait = btPrewait ? 1 : 0; + auto const prefetchRange = [](std::optional const& t, uint8_t const*& ptr, int64_t& bytes) + { + ptr = nullptr; + bytes = 0; + if (t.has_value()) + { + TORCH_CHECK(t->is_cuda() && t->is_contiguous(), + "minimax_m3_sm107_tp4_sparse_decode_lc: a prefetch range must be a contiguous CUDA tensor"); + ptr = reinterpret_cast(t->data_ptr()); + bytes = t->numel() * static_cast(t->element_size()); + } + }; + prefetchRange(prefetchA, prm.pf0, prm.pf0_bytes); + prefetchRange(prefetchB, prm.pf1, prm.pf1_bytes); + bool const n1 = n1Self.has_value(); + TORCH_CHECK(n1 == n1Next.has_value() && n1 == n1K.has_value() && n1 == n1V.has_value() && n1 == n1KScale.has_value() + && n1 == n1VScale.has_value(), + "minimax_m3_sm107_tp4_sparse_decode_lc: pass all six next-layer prefetch tensors, or none"); + if (n1) + { + for (auto const* t : {&*n1Self, &*n1Next}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == torch::kInt32 && t->is_contiguous() && t->dim() == 2 + && t->size(1) == kTopK && t->size(0) >= q.size(0), + "minimax_m3_sm107_tp4_sparse_decode_lc: the id buffers must be contiguous int32 [>= T, 16]"); + } + TORCH_CHECK(n1K->stride(0) == k.stride(0) && n1V->stride(0) == v.stride(0) + && n1KScale->stride(0) == kBlockScale.stride(0) && n1VScale->stride(0) == vBlockScale.stride(0) + && n1K->size(0) == k.size(0), + "minimax_m3_sm107_tp4_sparse_decode_lc: the next layer's pools must have this layer's page strides"); + prm.n1_self = n1Self->data_ptr(); + prm.n1_next = n1Next->data_ptr(); + prm.n1_k = reinterpret_cast(n1K->data_ptr()); + prm.n1_v = reinterpret_cast(n1V->data_ptr()); + prm.n1_ksc = reinterpret_cast(n1KScale->data_ptr()); + prm.n1_vsc = reinterpret_cast(n1VScale->data_ptr()); + prm.n1_pages = n1K->size(0); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSparseDecodeLc4( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, earlyTrigger, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_sparse_decode_lc launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp4_sparse_decode_lc(Tensor q, Tensor k, Tensor v, Tensor k_block_scale, Tensor " + "v_block_scale, " + "Tensor k_global_scale, Tensor v_global_scale, Tensor topk, Tensor block_table, Tensor seq_lens, " + "Tensor(a!) out, Tensor(b!)? q8, Tensor(c!)? sf, int num_sf_k_tiles, int decode_query_len, float sm_scale, " + "bool linear_softmax, bool use_pdl, bool early_trigger, Tensor? prefetch_a=None, Tensor? prefetch_b=None, " + "Tensor(d!)? n1_self=None, Tensor? n1_next=None, Tensor? n1_k=None, Tensor? n1_v=None, Tensor? " + "n1_k_scale=None, " + "Tensor? n1_v_scale=None, bool identity=False, bool one_div=False, bool identity_step=False, " + "bool bt_prewait=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_sparse_decode_lc", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SparseDecodeLc); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp new file mode 100644 index 000000000000..2142baa18599 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp @@ -0,0 +1,74 @@ +/* + * 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/minimaxM3Sm107Tp4/specOffsets.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void check(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp4_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +} // namespace + +//! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch +//! (kernels/minimaxM3Sm107Tp4/specOffsets.h). +void minimaxM3Sm107Tp4SpecOffsets( + torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) +{ + check(pos, "pos"); + check(posOff, "pos_off"); + check(kv, "kv"); + check(kvOff, "kv_off"); + TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), + "minimax_m3_sm107_tp4_spec_offsets: each tensor and its offsets must have the same length"); + if (pos.numel() == 0 && kv.numel() == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{pos.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSpecOffsets(pos.data_ptr(), + posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), + at::cuda::getCurrentCUDAStream(pos.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_spec_offsets launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SpecOffsets); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp new file mode 100644 index 000000000000..59009e461d33 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp @@ -0,0 +1,68 @@ +/* + * 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/minimaxM3Sm107Tp4/stepPrologue.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp4/stepPrologue.h). +//! +//! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the +//! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses +//! the caller guarantees are live. +//! @param epoch Optional one-element int64 step counter, incremented once. +void minimaxM3Sm107Tp4StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) +{ + TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), + "minimax_m3_sm107_tp4_step_prologue: window must be a contiguous uint8 CUDA tensor"); + TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, + "minimax_m3_sm107_tp4_step_prologue: window is shorter than its header and max_desc descriptors"); + c10::cuda::CUDAGuard const deviceGuard{window.device()}; + int64_t* epochPtr = nullptr; + if (epoch.has_value()) + { + TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 + && epoch->device() == window.device(), + "minimax_m3_sm107_tp4_step_prologue: epoch must be a one-element int64 tensor on the window's device"); + epochPtr = epoch->data_ptr(); + } + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeStepPrologue( + window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_step_prologue launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4StepPrologue); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp new file mode 100644 index 000000000000..1172d809a817 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp @@ -0,0 +1,65 @@ +/* + * 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/minimaxM3Sm107Tp4/swigluOai.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp4/swigluOai.h): x bf16 [rows, 2d] -> out bf16 +//! [rows, d], PDL-launched with an early release when `pdl`. +void minimaxM3Sm107Tp4SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp4_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp4_swiglu_oai: x rows must be 4-byte aligned with an even half width"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp4_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSwigluOai( + reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, + at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_swiglu_oai launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SwigluOai); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp new file mode 100644 index 000000000000..bee82725b41f --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp @@ -0,0 +1,236 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h" +#include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h" + +#include +#include +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace ds = tensorrt_llm::kernels::minimaxM3Sm107Tp8::draftStep; + +using OptTensor = std::optional; + +void checkArch(torch::Tensor const& t) +{ + TORCH_CHECK(t.is_cuda(), "minimax_m3_sm107tp8_draft_step expects CUDA tensors"); + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107tp8_draft_step requires SM100 or newer"); +} + +void check(cudaError_t err) +{ + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, cudaGetErrorString(err)); + } +} + +void checkInt(torch::Tensor const& t, int64_t n, char const* name) +{ + TORCH_CHECK(t.scalar_type() == at::kInt && t.is_contiguous() && t.numel() >= n, "draft_step: ", name, + " must be contiguous int32 with at least ", n, " elements"); +} + +} // namespace + +//! One fused greedy draft pass over `rows` <= 8 requests: the BF16 head (this rank's fragment-ordered shard wf +//! [N / 16, K / 16, 256], or with excOff / excPos / excW / expBase its 12-bit coding, wf uint8 [N * K * 3 / 2]; +//! global ids vocabOffset..) over x[xRows[b]] (or x[b]), the global argmax across the TP ranks through the multicast +//! slots wsMcPtr / wsUc, then draftOut[b, step] = idsNext[b] = token, posNext[b] = posSrc[xRows[b] or b] + 1 and the +//! KV-length bookkeeping (see draftStep.h). +void minimaxM3Sm107Tp8DraftStep(torch::Tensor const& x, OptTensor const& xRows, int64_t rows, torch::Tensor const& wf, + int64_t vocabOffset, torch::Tensor& keys, torch::Tensor& counters, torch::Tensor& wsUc, int64_t wsMcPtr, + int64_t rank, int64_t world, torch::Tensor& draftOut, int64_t step, torch::Tensor& idsNext, + torch::Tensor const& posSrc, torch::Tensor& posNext, torch::Tensor& kvLens, OptTensor const& numAcc, + OptTensor const& seqLens, int64_t draftLen, int64_t numCtas, bool pdl, OptTensor const& outKeys, + OptTensor const& excOff, OptTensor const& excPos, OptTensor const& excW, int64_t expBase, int64_t tileLo, + int64_t tileHi, int64_t totalCtas, OptTensor const& bnd, OptTensor const& bndCnt, + std::optional> fold, int64_t foldRows, int64_t foldTokensPerStep, bool fixedOrder) +{ + checkArch(x); + for (torch::Tensor const* t : std::initializer_list{ + &wf, &keys, &counters, &wsUc, &draftOut, &idsNext, &posSrc, &posNext, &kvLens}) + { + TORCH_CHECK(t->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + for (OptTensor const* t : std::initializer_list{ + &xRows, &numAcc, &seqLens, &outKeys, &excOff, &excPos, &excW, &bnd, &bndCnt}) + { + TORCH_CHECK(!t->has_value() || (*t)->device() == x.device(), "draft_step: every tensor must be on x's device"); + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + TORCH_CHECK(x.scalar_type() == at::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "draft_step: x must be contiguous BF16 [rows, K]"); + int const K = static_cast(x.size(1)); + TORCH_CHECK(K % 16 == 0, "draft_step: K must be a multiple of 16"); + TORCH_CHECK(rows >= 1 && rows <= ds::kMaxTokens, "draft_step: 1..", ds::kMaxTokens, " rows"); + bool const coded = excOff.has_value(); + TORCH_CHECK(coded == excPos.has_value() && coded == excW.has_value(), + "draft_step: exc_off, exc_pos and exc_w come together"); + int nTiles = 0; + if (coded) + { + TORCH_CHECK(K % 64 == 0 && K / 16 <= 0xFFFF, "draft_step: a coded head needs K a multiple of 64"); + TORCH_CHECK(excOff->scalar_type() == at::kInt && excOff->is_contiguous() && excOff->numel() >= 2, + "draft_step: exc_off must be int32 [tiles + 1]"); + nTiles = static_cast(excOff->numel() - 1); + TORCH_CHECK(wf.scalar_type() == at::kByte && wf.is_contiguous() + && wf.numel() == static_cast(nTiles) * 16 * K * 3 / 2, + "draft_step: a coded wf must be the contiguous uint8 stream of [tiles * 16, K]"); + TORCH_CHECK(excPos->scalar_type() == at::kInt && excPos->is_contiguous(), "draft_step: exc_pos must be int32"); + TORCH_CHECK(excW->scalar_type() == at::kFloat && excW->is_contiguous() && excW->numel() == 2 * excPos->numel(), + "draft_step: exc_w must be fp32 [E, 2]"); + TORCH_CHECK(expBase >= 0 && expBase <= 240, "draft_step: exponent base out of range"); + } + else + { + TORCH_CHECK(wf.scalar_type() == at::kBFloat16 && wf.dim() == 3 && wf.is_contiguous() && wf.size(1) == K / 16 + && wf.size(2) == 256, + "draft_step: wf must be the fragment-ordered BF16 weight [N / 16, K / 16, 256]"); + nTiles = static_cast(wf.size(0)); + } + int64_t const rangeHi = tileHi > 0 ? tileHi : nTiles; + TORCH_CHECK(tileLo >= 0 && tileLo < rangeHi && rangeHi <= nTiles, "draft_step: tile range"); + TORCH_CHECK(nTiles >= 1 && numCtas >= 1 && (rangeHi - tileLo + numCtas - 1) / numCtas <= ds::kMaxTilesPerCta, + "draft_step: too many weight tiles per CTA"); + TORCH_CHECK(totalCtas == 0 || totalCtas >= numCtas, "draft_step: total_ctas"); + TORCH_CHECK(bnd.has_value() == bndCnt.has_value(), "draft_step: bnd and bnd_cnt come together"); + if (bnd.has_value()) + { + TORCH_CHECK(bnd->scalar_type() == at::kFloat && bnd->is_contiguous() && bnd->numel() >= int64_t{nTiles} * 128, + "draft_step: bnd must be zero-filled fp32 [tiles, 32, 4]"); + TORCH_CHECK(bndCnt->scalar_type() == at::kInt && bndCnt->is_contiguous() && bndCnt->numel() >= nTiles, + "draft_step: bnd_cnt must be zero-filled int32 [tiles]"); + // At least one tile per CTA: a tile is then shared by at most two CTAs (two addends in bnd). + TORCH_CHECK(rangeHi - tileLo >= numCtas, "draft_step: an equal-chunk split needs a tile per CTA"); + } + TORCH_CHECK(vocabOffset >= 0 && vocabOffset + 16LL * nTiles <= 0xFFFFFFFFLL, "draft_step: vocab offset"); + if (xRows.has_value()) + { + TORCH_CHECK(xRows->scalar_type() == at::kLong && xRows->is_contiguous() && xRows->numel() >= rows, + "draft_step: x_rows must be contiguous int64 [rows]"); + } + else + { + TORCH_CHECK(x.size(0) >= rows, "draft_step: x must hold a row per request"); + } + TORCH_CHECK(world >= 1 && rank >= 0 && rank < world, "draft_step: rank / world"); + TORCH_CHECK(keys.scalar_type() == at::kLong && keys.numel() >= rows, "draft_step: keys must be int64 [>= rows]"); + TORCH_CHECK(counters.scalar_type() == at::kInt && counters.numel() >= 2, "draft_step: counters must be int32 [2]"); + TORCH_CHECK(wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, + "draft_step: ws_uc must hold [2][8][world] slots"); + TORCH_CHECK(draftOut.scalar_type() == at::kInt && draftOut.dim() == 2 && draftOut.stride(1) == 1 + && draftOut.size(0) >= rows && step >= 0 && step < draftOut.size(1), + "draft_step: draft_out must be int32 [rows, > step] with unit column stride"); + checkInt(idsNext, rows, "ids_next"); + checkInt(posNext, rows, "pos_next"); + checkInt(kvLens, rows, "kv_lens"); + TORCH_CHECK(posSrc.scalar_type() == at::kInt && posSrc.is_contiguous(), "draft_step: pos_src must be int32"); + TORCH_CHECK(numAcc.has_value() == seqLens.has_value(), "draft_step: num_acc and seq_lens come together"); + TORCH_CHECK(!numAcc.has_value() || step == 0, "draft_step: the KV rollback (num_acc) is pass 0's"); + if (numAcc.has_value()) + { + checkInt(*numAcc, rows, "num_acc"); + checkInt(*seqLens, rows, "seq_lens"); + } + if (outKeys.has_value()) + { + TORCH_CHECK(outKeys->scalar_type() == at::kLong && outKeys->is_contiguous() && outKeys->numel() >= rows, + "draft_step: out_keys must be contiguous int64 [rows]"); + } + ds::Params p{}; + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.x_rows = xRows.has_value() ? reinterpret_cast(xRows->data_ptr()) : nullptr; + p.M = static_cast(rows); + p.wf = reinterpret_cast(wf.data_ptr()); + p.n_tiles = nTiles; + p.vocab_offset = static_cast(vocabOffset); + p.keys = reinterpret_cast(keys.data_ptr()); + p.counters = reinterpret_cast(counters.data_ptr()); + p.ws_uc = reinterpret_cast(wsUc.data_ptr()); + p.ws_mc = reinterpret_cast(wsMcPtr); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.out_keys = outKeys.has_value() ? reinterpret_cast(outKeys->data_ptr()) : nullptr; + p.st = ds::StepArgs{draftOut.data_ptr(), static_cast(draftOut.stride(0)), static_cast(step), + idsNext.data_ptr(), posSrc.data_ptr(), posNext.data_ptr(), kvLens.data_ptr(), + numAcc.has_value() ? numAcc->data_ptr() : nullptr, + seqLens.has_value() ? seqLens->data_ptr() : nullptr, static_cast(draftLen)}; + p.cd = ds::CodedArgs{coded ? excOff->data_ptr() : nullptr, + coded ? reinterpret_cast(excPos->data_ptr()) : nullptr, + coded ? reinterpret_cast(excW->data_ptr()) : nullptr, + static_cast((expBase << 7) | (expBase << 23))}; + p.tile_lo = static_cast(tileLo); + p.tile_hi = static_cast(tileHi); + p.total_ctas = static_cast(totalCtas); + p.bnd = bnd.has_value() ? reinterpret_cast(bnd->data_ptr()) : nullptr; + p.bnd_cnt = bnd.has_value() ? reinterpret_cast(bndCnt->data_ptr()) : nullptr; + if (fold.has_value()) + { + // The 14 fold-tail tensors in buildParams' order (ctrl, slots, src_new, dst_new, src_next, dst_next, + // src_draft, dst_draft, src_lens, dst_lens, ids, draft_ids, pos_off, kv_off): the in-kernel tail's operands, + // packed exactly as the standalone fold tail packs them. + std::vector f = *fold; + TORCH_CHECK(f.size() == 14, "minimax_m3_sm107tp8_draft_step: fold takes 14 tensors, got ", f.size()); + p.fold = minimaxM3Sm107Tp8FoldTailParams::buildParams(f[0], f[1], foldRows, f[2], f[3], f[4], f[5], f[6], f[7], + f[8], f[9], f[10], f[11], f[12], f[13], foldTokensPerStep); + p.hasFold = true; + } + p.fixedOrder = fixedOrder; + check(ds::invokeDraftStep(p, K, static_cast(numCtas), pdl, at::cuda::getCurrentCUDAStream(x.get_device()))); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// Mutations through the raw multicast address ws_mc_ptr are not visible to the schema: that buffer is owned by the +// Python workspace and never aliased by the op's tensor arguments. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp8_draft_step(Tensor x, Tensor? x_rows, int rows, Tensor wf, int vocab_offset, Tensor(a!) " + "keys, " + "Tensor(b!) counters, Tensor(c!) ws_uc, int ws_mc_ptr, int rank, int world, Tensor(d!) draft_out, int step, " + "Tensor(e!) ids_next, Tensor pos_src, Tensor(f!) pos_next, Tensor(g!) kv_lens, Tensor? num_acc, " + "Tensor(h!)? seq_lens, int draft_len, int num_ctas, bool pdl, Tensor(i!)? out_keys, Tensor? exc_off=None, " + "Tensor? exc_pos=None, Tensor? exc_w=None, int exp_base=0, int tile_lo=0, int tile_hi=0, int total_ctas=0, " + "Tensor(j!)? bnd=None, Tensor(k!)? bnd_cnt=None" + ", Tensor(l!)[]? fold=None, int fold_rows=0, int fold_tokens_per_step=0, bool fixed_order=False) " + "-> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp8_draft_step", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8DraftStep); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp new file mode 100644 index 000000000000..ec93068fc374 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp @@ -0,0 +1,103 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h" + +#include +#include +#include + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp8::drafterRopeFp8Kv; + +//! The MiniMax-M3 TP8 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h): RoPE on +//! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages +//! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. +void minimaxM3Sm107Tp8DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, + torch::Tensor const& seqLens, std::optional const& specPositionOffsets, torch::Tensor& kvPages, + torch::Tensor const& blockTables, torch::Tensor& qOut, int64_t qLen, int64_t numHeadsQ, int64_t numHeadsKv, + int64_t rotaryDim, int64_t tokensPerBlock, int64_t ropeMode, bool pdl, bool maxSharedCarveout) +{ + TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); + TORCH_CHECK(qkv.size(1) == (numHeadsQ + 2 * numHeadsKv) * drk::kHeadDim, "qkv is not [tokens, (Hq + 2 Hkv) * 128]"); + TORCH_CHECK(qLen >= 1 && qkv.size(0) % qLen == 0, "the rows are not whole passes of q_len rows"); + TORCH_CHECK(rotaryDim > 0 && rotaryDim <= drk::kHeadDim && rotaryDim % 16 == 0, "unsupported rotary_dim"); + TORCH_CHECK(cosSin.is_cuda() && cosSin.is_contiguous() && cosSin.scalar_type() == at::kFloat + && cosSin.numel() % rotaryDim == 0, + "cos_sin must be the fp32 [positions, rotary_dim / 2] (cos, sin) table"); + int64_t const requests = qkv.size(0) / qLen; + TORCH_CHECK(seqLens.is_cuda() && seqLens.scalar_type() == at::kInt && seqLens.numel() >= requests); + if (specPositionOffsets.has_value()) + { + TORCH_CHECK(specPositionOffsets->is_cuda() && specPositionOffsets->is_contiguous() + && specPositionOffsets->scalar_type() == at::kInt && specPositionOffsets->numel() >= requests * qLen); + } + TORCH_CHECK(kvPages.is_cuda() && kvPages.is_contiguous() && kvPages.dim() == 4 + && kvPages.scalar_type() == at::kFloat8_e4m3fn && kvPages.size(1) == numHeadsKv + && kvPages.size(2) == tokensPerBlock && kvPages.size(3) == drk::kHeadDim, + "kv_pages must be the E4M3 [sub-pages, Hkv, tokens_per_block, 128] cache"); + TORCH_CHECK(blockTables.is_cuda() && blockTables.scalar_type() == at::kInt && blockTables.dim() == 3 + && blockTables.size(0) >= requests && blockTables.size(1) == 2 && blockTables.stride(2) == 1, + "block_tables must be the int32 [requests, 2, max blocks] K / V sub-page table"); + TORCH_CHECK(qOut.is_cuda() && qOut.is_contiguous() && qOut.scalar_type() == at::kFloat8_e4m3fn + && qOut.numel() == qkv.size(0) * numHeadsQ * drk::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{qkv.device()}; + drk::Params params{}; + params.qkv = reinterpret_cast<__nv_bfloat16 const*>(qkv.data_ptr()); + params.cosSin = reinterpret_cast(cosSin.data_ptr()); + params.seqLens = seqLens.data_ptr(); + params.specPositionOffsets = specPositionOffsets.has_value() ? specPositionOffsets->data_ptr() : nullptr; + params.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + params.pool = reinterpret_cast<__nv_fp8_e4m3*>(kvPages.data_ptr()); + params.blockTable = blockTables.data_ptr(); + params.blockTableRequestStride = blockTables.stride(0); + params.blockTableKvStride = blockTables.stride(1); + params.numTokens = static_cast(qkv.size(0)); + params.qLen = static_cast(qLen); + params.numHeadsQ = static_cast(numHeadsQ); + params.numHeadsKv = static_cast(numHeadsKv); + params.rotaryDim = static_cast(rotaryDim); + params.tokensPerBlock = static_cast(tokensPerBlock); + params.ropeMode = static_cast(ropeMode); + cudaError_t const err + = drk::launch(params, pdl, maxSharedCarveout, at::cuda::getCurrentCUDAStream(qkv.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp8_drafter_rope_fp8_kv: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp8_drafter_rope_fp8_kv(Tensor qkv, Tensor cos_sin, Tensor seq_lens, " + "Tensor? spec_position_offsets, Tensor(a!) kv_pages, Tensor block_tables, Tensor(b!) q_out, int q_len, " + "int num_heads_q, int num_heads_kv, int rotary_dim, int tokens_per_block, int rope_mode, bool pdl, bool " + "max_shared_carveout=True) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp8_drafter_rope_fp8_kv", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8DrafterRopeFp8Kv); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp new file mode 100644 index 000000000000..e4d5ca2f6b54 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp @@ -0,0 +1,89 @@ +/* + * 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/minimaxM3Sm107Tp8/foldTail.h" +#include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The speculative step's forward-tail store and next-step staging in one launch (the C++ form of the Triton +//! kernels/m3_fold_tail.py, the same writes). The caller checks the row counts against the buffers (fold_tail). +static void foldTailImpl(bool oneWarp, torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + c10::cuda::CUDAGuard const deviceGuard{ctrl.device()}; + auto const p = minimaxM3Sm107Tp8FoldTailParams::buildParams(ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, + srcDraft, dstDraft, srcLens, dstLens, ids, draftIds, posOff, kvOff, tokensPerStep); + auto const stream = at::cuda::getCurrentCUDAStream(ctrl.get_device()); + auto const err = oneWarp + ? tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeFoldTailOneWarp(p, stream) + : tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeFoldTail(p, static_cast(rows), stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_fold_tail launch failed: ", cudaGetErrorString(err)); +} + +//! The tail as its standalone kernel (one CTA per row). +void minimaxM3Sm107Tp8FoldTail(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + foldTailImpl(false, ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, srcDraft, dstDraft, srcLens, dstLens, ids, + draftIds, posOff, kvOff, tokensPerStep); +} + +//! The tail's in-kernel form (one warp over the rows, foldTailWarp), the oracle of a caller that folds it in. +void minimaxM3Sm107Tp8FoldTailOneWarp(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, + torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, + torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, + torch::Tensor& ids, torch::Tensor& draftIds, torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + foldTailImpl(true, ctrl, slots, rows, srcNew, dstNew, srcNext, dstNext, srcDraft, dstDraft, srcLens, dstLens, ids, + draftIds, posOff, kvOff, tokensPerStep); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp8_fold_tail(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) dst_new, " + "Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); + m.def( + "minimax_m3_sm107_tp8_fold_tail_one_warp(Tensor ctrl, Tensor slots, int rows, Tensor src_new, Tensor(a!) " + "dst_new, Tensor src_next, Tensor(b!) dst_next, Tensor src_draft, Tensor(c!) dst_draft, Tensor src_lens, " + "Tensor(d!) dst_lens, Tensor(e!) ids, Tensor(f!) draft_ids, Tensor(g!) pos_off, Tensor(h!) kv_off, " + "int tokens_per_step) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_fold_tail", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8FoldTail); + m.impl("minimax_m3_sm107_tp8_fold_tail_one_warp", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8FoldTailOneWarp); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h new file mode 100644 index 000000000000..cc5cde28c9d9 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h @@ -0,0 +1,111 @@ +/* + * 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/kernels/minimaxM3Sm107Tp8/foldTail.h" + +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ +namespace minimaxM3Sm107Tp8FoldTailParams +{ + +constexpr int64_t kMaxWidth = 32; + +inline void check2d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 2 && t.size(1) <= kMaxWidth, + "minimax_m3_sm107_tp8_fold_tail: ", name, " must be a 2-D int32 CUDA view at most 32 wide"); +} + +inline void check1d(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp8_fold_tail: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +//! The fold tail's operands checked and packed (the standalone fold tail op and any op that runs the tail in-kernel, +//! e.g. the drafter's last-pass draft_step, build them identically). +inline tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams buildParams(torch::Tensor const& ctrl, + torch::Tensor const& slots, int64_t rows, torch::Tensor const& srcNew, torch::Tensor& dstNew, + torch::Tensor const& srcNext, torch::Tensor& dstNext, torch::Tensor const& srcDraft, torch::Tensor& dstDraft, + torch::Tensor const& srcLens, torch::Tensor& dstLens, torch::Tensor& ids, torch::Tensor& draftIds, + torch::Tensor& posOff, torch::Tensor& kvOff, int64_t tokensPerStep) +{ + TORCH_CHECK(ctrl.is_cuda() && ctrl.scalar_type() == torch::kInt32 && ctrl.numel() >= 1, + "minimax_m3_sm107_tp8_fold_tail: ctrl must be an int32 CUDA tensor"); + TORCH_CHECK(slots.is_cuda() && slots.scalar_type() == torch::kInt64 && slots.dim() == 1 && slots.stride(0) == 1, + "minimax_m3_sm107_tp8_fold_tail: slots must be a contiguous 1-D int64 CUDA tensor"); + TORCH_CHECK(rows > 0 && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, + "minimax_m3_sm107_tp8_fold_tail: rows must be positive and tokens_per_step in [1, 32]"); + check2d(srcNew, "src_new"); + check2d(dstNew, "dst_new"); + check2d(srcNext, "src_next"); + check2d(dstNext, "dst_next"); + check2d(srcDraft, "src_draft"); + check2d(dstDraft, "dst_draft"); + check1d(srcLens, "src_lens"); + check1d(dstLens, "dst_lens"); + check1d(ids, "ids"); + check1d(draftIds, "draft_ids"); + check1d(posOff, "pos_off"); + check1d(kvOff, "kv_off"); + tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams p{}; + p.ctrl = ctrl.data_ptr(); + p.slots = slots.data_ptr(); + p.srcNew = srcNew.data_ptr(); + p.sn0 = srcNew.stride(0); + p.sn1 = srcNew.stride(1); + p.snW = static_cast(srcNew.size(1)); + p.dstNew = dstNew.data_ptr(); + p.dn0 = dstNew.stride(0); + p.dn1 = dstNew.stride(1); + p.dnW = static_cast(dstNew.size(1)); + p.srcNext = srcNext.data_ptr(); + p.sx0 = srcNext.stride(0); + p.sx1 = srcNext.stride(1); + p.sxW = static_cast(srcNext.size(1)); + p.dstNext = dstNext.data_ptr(); + p.dx0 = dstNext.stride(0); + p.dx1 = dstNext.stride(1); + p.dxW = static_cast(dstNext.size(1)); + p.srcDraft = srcDraft.data_ptr(); + p.sd0 = srcDraft.stride(0); + p.sd1 = srcDraft.stride(1); + p.sdW = static_cast(srcDraft.size(1)); + p.dstDraft = dstDraft.data_ptr(); + p.dd0 = dstDraft.stride(0); + p.dd1 = dstDraft.stride(1); + p.ddW = static_cast(dstDraft.size(1)); + p.srcLens = srcLens.data_ptr(); + p.dstLens = dstLens.data_ptr(); + p.ids = ids.data_ptr(); + p.draftIds = draftIds.data_ptr(); + p.posOff = posOff.data_ptr(); + p.kvOff = kvOff.data_ptr(); + p.tokensPerStep = static_cast(tokensPerStep); + return p; +} + +} // namespace minimaxM3Sm107Tp8FoldTailParams +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp new file mode 100644 index 000000000000..29b9a0ab391b --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp @@ -0,0 +1,92 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107Tp8::fullCtxReduce; + +//! Merges the PrimTS split-KV decode scratch (partial_o [B, Hkv, S, R, 128] 16-bit, partial_stats [B, Hkv, S, R] +//! fp32 log2-LSE) of the MiniMax-M3 TP8 full-context decode into out (BF16 [B * SQ, Hkv * hR, 128]); with q8 / sf +//! also its MXFP8 copy (E4M3 [>= B * SQ, Hkv * hR * 128] + 128x4-swizzled UE8M0 scales, as mxfp8_quantize(out, True)). +void minimaxM3Sm107Tp8FullCtxReduce(torch::Tensor const& partialO, torch::Tensor const& partialStats, + torch::Tensor const& seqLensKv, torch::Tensor& out, int64_t seqLenQ, int64_t qTokensPerCta, int64_t tileKv, + int64_t instsKv, bool causal, bool pdl, std::optional const& q8, + std::optional const& sf) +{ + TORCH_CHECK(partialO.is_cuda() && partialO.is_contiguous() && partialO.dim() == 5); + TORCH_CHECK(partialO.scalar_type() == at::kHalf || partialO.scalar_type() == at::kBFloat16); + TORCH_CHECK(partialO.size(4) == fcr::kHeadDim && partialO.size(2) <= fcr::kMaxSplits); + TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() + && partialStats.sizes() == partialO.sizes().slice(0, 4)); + TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + int const rows = static_cast(partialO.size(3)); + TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); + int const headsPerKv = rows / static_cast(seqLenQ); + TORCH_CHECK(out.scalar_type() == at::kBFloat16 && out.is_contiguous() + && out.numel() == partialO.size(0) * seqLenQ * partialO.size(1) * headsPerKv * fcr::kHeadDim); + c10::cuda::CUDAGuard const deviceGuard{partialO.device()}; + fcr::Params params{partialO.data_ptr(), partialStats.data_ptr(), seqLensKv.data_ptr(), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), static_cast(partialO.size(0)), + static_cast(partialO.size(1)), static_cast(partialO.size(2)), rows, static_cast(seqLenQ), + headsPerKv, static_cast(qTokensPerCta), static_cast(tileKv), static_cast(instsKv), causal, + partialO.scalar_type() == at::kBFloat16}; + TORCH_CHECK(q8.has_value() == sf.has_value()); + if (q8.has_value()) + { + int64_t const width = partialO.size(1) * headsPerKv * fcr::kHeadDim; + int64_t const tokens = partialO.size(0) * seqLenQ; + int64_t const tileRows = (tokens + 127) / 128 * 128; + TORCH_CHECK(q8->is_cuda() && q8->is_contiguous() && q8->scalar_type() == at::kFloat8_e4m3fn && q8->dim() == 2 + && q8->size(1) == width && q8->size(0) >= tokens); + TORCH_CHECK(sf->is_cuda() && sf->is_contiguous() && sf->scalar_type() == at::kByte + && sf->numel() >= tileRows * (width / 32)); + params.q8 = reinterpret_cast(q8->data_ptr()); + params.sf = sf->data_ptr(); + params.numSfKTiles = static_cast(width / 128); + } + cudaError_t const err + = fcr::invokeFullCtxReduce(params, pdl, at::cuda::getCurrentCUDAStream(partialO.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107tp8_fullctx_reduce: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107tp8_fullctx_reduce(Tensor partial_o, Tensor partial_stats, Tensor seq_lens_kv, " + "Tensor(a!) out, int seq_len_q, int q_tokens_per_cta, int tile_kv, int insts_kv, bool causal, bool pdl, " + "Tensor(b!)? q8=None, Tensor(c!)? sf=None) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp8_fullctx_reduce", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8FullCtxReduce); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp new file mode 100644 index 000000000000..eb3759b28a58 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp @@ -0,0 +1,181 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h" + +#include +#include +#include + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kMaxTokens = 48; +constexpr int64_t kCluster = 8; +constexpr int64_t kMaxWorld = 8; + +} // namespace + +//! MiniMax-M3 SM107 TP8 multicast all-reduce (stock rounding points: BF16 rank sum, then the residual add) + residual + +//! RMSNorm (+ optional MoE finalize, + optional MXFP8 output). +//! +//! @param x BF16 [M, H] partial, or with `perm` the permuted FC2 output [rows, H]. +//! @param x2 Optional BF16 [M, H] second addend (the shared expert in finalize mode). +//! @param residual BF16 [M, H], 1 <= M <= 48. +//! @param normWeight BF16 [H]. +//! @param wsUc This rank's unicast view of the exchange workspace (int32 words); re-armed in place. +//! @param wsMcPtr Device address of the workspace's multicast mapping. +//! @param flags Int32 per-CTA call-parity counters, at least maxt * 8; updated in place. +//! @param cluster Cluster size; 8 is the only one built. +//! @param perm, scales Finalize mode: int32 [M, topk] permuted-row index and fp32 / bf16 [M, topk] routing weights. +//! @param maxt Token capacity of the workspace bucket (8, 16 or 48). +//! @param quant Also return the MXFP8 copy of the normed rows (E4M3 [M, H], swizzled UE8M0 scales). +//! @param prepushed 0, or 2 / 3 when producers already pushed this rank's rows (see the kernel header). +//! @param opts Epilogue variant; every variant is bit-identical. +//! @param capOut Optional BF16 [>= M, H] view (row stride a multiple of 8, 16-byte aligned) that receives the Eagle3 +//! capture of the updated residual: m3_spec_tail capture_norm with capW [H] and capEps, bit for bit (H 6144; not the +//! drafter boundary, opts bit 16). +//! @return [normed, updated residual] or, with `quant`, [normed, updated residual, E4M3 rows, scale bytes]. +std::vector minimaxM3Sm107Tp8McAllreduceNorm(torch::Tensor const& x, + std::optional const& x2, torch::Tensor const& residual, torch::Tensor const& normWeight, double eps, + torch::Tensor& wsUc, int64_t wsMcPtr, torch::Tensor& flags, int64_t rank, int64_t world, int64_t cluster, bool pdl, + std::optional const& perm, std::optional const& scales, int64_t maxt, bool quant, + int64_t prepushed, int64_t opts, std::optional const& capOut, + std::optional const& capW, double capEps) +{ + int const sm = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sm == 107, "minimax_m3_sm107_tp8_mc_allreduce_norm requires SM107, got SM", sm); + TORCH_CHECK(cluster == kCluster, "minimax_m3_sm107_tp8_mc_allreduce_norm: only cluster 8 is built, got ", cluster); + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107Tp8McAllreduceNormSupported( + static_cast(maxt), static_cast(opts)), + "minimax_m3_sm107_tp8_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); + TORCH_CHECK( + world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); + + tensorrt_llm::kernels::MinimaxM3Sm107Tp8McAllreduceNormParams p; + if (x2.has_value()) + { + TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), + "x2 must be a contiguous bf16 tensor shaped like residual"); + p.x2 = reinterpret_cast<__nv_bfloat16 const*>(x2->data_ptr()); + } + if (perm.has_value()) + { + // Finalize mode: x is the permuted FC2 output [rows, H]. + TORCH_CHECK(scales.has_value() && perm->scalar_type() == at::kInt && perm->is_contiguous() && perm->dim() == 2 + && perm->size(0) == residual.size(0), + "perm must be a contiguous int32 [M, topk] tensor, with scales"); + TORCH_CHECK(scales->is_contiguous() && scales->numel() == perm->numel() + && (scales->scalar_type() == at::kFloat || scales->scalar_type() == at::kBFloat16), + "scales must be a contiguous fp32 or bf16 tensor with one weight per perm entry"); + TORCH_CHECK(x.dim() == 2 && x.size(1) == residual.size(1), "fc2 must be [rows, H]"); + p.fc2 = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.perm = perm->data_ptr(); + p.scales = scales->data_ptr(); + p.scaleFp32 = scales->scalar_type() == at::kFloat; + p.topk = static_cast(perm->size(1)); + } + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kBFloat16 && x.is_contiguous() && x.dim() == 2, + "x must be a contiguous 2-D bf16 CUDA tensor"); + TORCH_CHECK(residual.dim() == 2 && residual.size(0) >= 1 && residual.size(0) <= std::min(kMaxTokens, maxt), + "residual must hold 1..min(48, bucket) rows"); + TORCH_CHECK(residual.scalar_type() == at::kBFloat16 && residual.is_contiguous() + && (perm.has_value() || residual.sizes() == x.sizes()), + "residual must be a contiguous bf16 tensor shaped like x"); + TORCH_CHECK( + normWeight.scalar_type() == at::kBFloat16 && normWeight.numel() == residual.size(1), "norm_w must be bf16 [H]"); + TORCH_CHECK(wsUc.scalar_type() == at::kInt && wsUc.is_cuda(), "ws_uc must be an int32 CUDA tensor"); + TORCH_CHECK(flags.scalar_type() == at::kInt && flags.is_contiguous() && flags.numel() >= maxt * kCluster, + "flags must be a contiguous int32 tensor of at least bucket * 8 counters"); + + int64_t const m = residual.size(0), h = residual.size(1); + TORCH_CHECK(!quant || h % 128 == 0, "MXFP8 output needs H % 128 == 0"); + c10::cuda::CUDAGuard const deviceGuard{residual.device()}; + auto outNormed = torch::empty_like(residual); + auto outResidual = torch::empty_like(residual); + auto outQ = quant ? torch::empty({m, h}, residual.options().dtype(at::kFloat8_e4m3fn)) : torch::Tensor(); + auto outSf = quant ? torch::empty({(m + 127) / 128 * 128 * (h / 32)}, residual.options().dtype(at::kByte)) + : torch::Tensor(); + + p.x = reinterpret_cast<__nv_bfloat16 const*>(x.data_ptr()); + p.numTokens = static_cast(m); + p.hidden = static_cast(h); + p.residual = reinterpret_cast<__nv_bfloat16 const*>(residual.data_ptr()); + p.normWeight = reinterpret_cast<__nv_bfloat16 const*>(normWeight.data_ptr()); + p.eps = static_cast(eps); + p.wsUc = wsUc.data_ptr(); + p.wsMc = reinterpret_cast(wsMcPtr); + p.flags = flags.data_ptr(); + p.rank = static_cast(rank); + p.world = static_cast(world); + p.outNormed = reinterpret_cast<__nv_bfloat16*>(outNormed.data_ptr()); + p.outResidual = reinterpret_cast<__nv_bfloat16*>(outResidual.data_ptr()); + p.outQuant = quant ? reinterpret_cast(outQ.data_ptr()) : nullptr; + p.outSf = quant ? outSf.data_ptr() : nullptr; + p.prepushed = static_cast(prepushed); + if (capOut.has_value()) + { + TORCH_CHECK((opts & 16) == 0, "the capture epilogue is not built for the drafter boundary (opts bit 16)"); + TORCH_CHECK(capOut->scalar_type() == at::kBFloat16 && capOut->dim() == 2 && capOut->size(0) >= m + && capOut->size(1) == h && capOut->stride(1) == 1 && capOut->stride(0) % 8 == 0 + && reinterpret_cast(capOut->data_ptr()) % 16 == 0 && capOut->device() == residual.device(), + "cap_out must be a bf16 [>= M, H] view with unit column stride, a row stride that is a multiple of 8 and " + "16-byte alignment"); + TORCH_CHECK( + capW.has_value() && capW->scalar_type() == at::kBFloat16 && capW->is_contiguous() && capW->numel() == h, + "cap_w must be a contiguous bf16 [H] tensor"); + p.capOut = reinterpret_cast<__nv_bfloat16*>(capOut->data_ptr()); + p.capStride = capOut->stride(0); + p.capWeight = reinterpret_cast<__nv_bfloat16 const*>(capW->data_ptr()); + p.capEps = static_cast(capEps); + } + tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp8McAllreduceNorm(p, static_cast(maxt), + static_cast(opts), pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); + if (quant) + { + return {outNormed, outResidual, outQ, outSf}; + } + return {outNormed, outResidual}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp8_mc_allreduce_norm(Tensor x, Tensor? x2, Tensor residual, Tensor norm_w, float eps, " + "Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) flags, int rank, int world, int cluster, bool pdl, " + "Tensor? perm, Tensor? scales, int maxt, bool quant, int prepushed, int opts, Tensor(c!)? cap_out=None, " + "Tensor? cap_w=None, float cap_eps=0.) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_mc_allreduce_norm", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8McAllreduceNorm); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp new file mode 100644 index 000000000000..acd7061eadbb --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp @@ -0,0 +1,185 @@ +/* + * 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/cudaUtils.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h" + +#include +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp8::producerPw8; + +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPageSize = 128; + +void checkCuda(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda(), "sm107_tp8_producer_pw: ", name, " must be a CUDA tensor"); +} + +} // namespace + +//! The SM107 TP8 NVFP4 QKV + indexer norm/RoPE/KV-insert producer (same arguments and outputs as +//! minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert at 8 query, 1 KV and 1 index head, head_dim 128, rotary_dim 64): +//! returns (q FP8 [T, 8, 128], index_q FP8 [T, 1, 128]) and writes the NVFP4 K/V of the KV head and the FP8 index-K +//! of each token's cache slot. `early` triggers the dependents on entry; `prewait_step_inputs` reads the position ids, +//! the slots and the RoPE rows before the grid-dependency wait (none of them may be written by a PDL predecessor). +std::vector minimaxM3Sm107Tp8ProducerPw(torch::Tensor const& packed, torch::Tensor& kvData, + torch::Tensor& kvScale, torch::Tensor& indexK, torch::Tensor const& outCacheLoc, torch::Tensor const& kvQuantScale, + double eps, torch::Tensor const& qWeight, torch::Tensor const& kWeight, torch::Tensor const& indexQWeight, + torch::Tensor const& indexKWeight, torch::Tensor const& rotaryCosSin, torch::Tensor const& positionIds, bool pdl, + bool early, bool prewaitStepInputs) +{ + checkCuda(packed, "packed"); + for (torch::Tensor const* t : std::initializer_list{&kvData, &kvScale, &indexK, &outCacheLoc, + &kvQuantScale, &qWeight, &kWeight, &indexQWeight, &indexKWeight, &rotaryCosSin, &positionIds}) + { + TORCH_CHECK(t->is_cuda() && t->get_device() == packed.get_device(), + "sm107_tp8_producer_pw: every operand must be on packed's CUDA device"); + } + TORCH_CHECK(reinterpret_cast(packed.data_ptr()) % 8 == 0 + && reinterpret_cast(kvData.data_ptr()) % 2 == 0 + && reinterpret_cast(indexK.data_ptr()) % 4 == 0, + "sm107_tp8_producer_pw: packed and the caches must start at the vector-load / store alignment"); + c10::cuda::CUDAGuard const deviceGuard{packed.device()}; + static int const sSmVersion = tensorrt_llm::common::getSMVersion(); + TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp8_producer_pw requires SM100 or newer"); + int64_t const headsQ = pw::kHeadsQ; + int64_t const headsKV = pw::kHeadsKV; + int64_t const headsIndex = pw::kHeadsIndex; + int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; + TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim + && packed.is_contiguous(), + "sm107_tp8_producer_pw: packed must be contiguous BF16 [T, 12 * 128]"); + // The kernel's layout (the native producer's checkMinimaxM3Nvfp4HndKVPools at one KV head): HND pools + // [pages, 2 planes, 1 head, 128 tokens, 64 data / 8 scale bytes], dense token rows and head blocks, planes and + // pages that do not overlap. + TORCH_CHECK(kvData.scalar_type() == at::kByte && kvData.dim() == 5 && kvData.size(0) > 0 && kvData.size(1) == 2 + && kvData.size(2) == headsKV && kvData.size(3) == kPageSize && kvData.size(4) == kHeadDim / 2, + "sm107_tp8_producer_pw: kv_data must be the NVFP4 pool [pages, 2, 1, 128, 64] uint8"); + TORCH_CHECK(kvScale.scalar_type() == at::kByte && kvScale.dim() == 5 && kvScale.size(0) == kvData.size(0) + && kvScale.size(1) == 2 && kvScale.size(2) == headsKV && kvScale.size(3) == kPageSize + && kvScale.size(4) == kHeadDim / 16, + "sm107_tp8_producer_pw: kv_scale must be the E4M3 scale pool [pages, 2, 1, 128, 8] uint8"); + TORCH_CHECK(kvData.stride(4) == 1 && kvData.stride(3) == kvData.size(4) + && kvData.stride(2) == kvData.size(3) * kvData.stride(3), + "sm107_tp8_producer_pw: kv_data must have dense token rows and head blocks"); + TORCH_CHECK(kvScale.stride(4) == 1 && kvScale.stride(3) == kvScale.size(4) + && kvScale.stride(2) == kvScale.size(3) * kvScale.stride(3), + "sm107_tp8_producer_pw: kv_scale must have dense token rows and head blocks"); + TORCH_CHECK(kvData.stride(1) >= kvData.size(2) * kvData.stride(2) && kvData.stride(0) >= 2 * kvData.stride(1) + && kvScale.stride(1) >= kvScale.size(2) * kvScale.stride(2) && kvScale.stride(0) >= 2 * kvScale.stride(1), + "sm107_tp8_producer_pw: pool planes or pages overlap"); + TORCH_CHECK(kvData.stride(0) % 2 == 0 && kvData.stride(1) % 2 == 0, + "sm107_tp8_producer_pw: kv_data page and plane strides must keep the uint16 stores aligned"); + TORCH_CHECK(indexK.scalar_type() == at::kFloat8_e4m3fn && indexK.dim() == 4 && indexK.size(0) == kvData.size(0) + && indexK.size(1) == 1 && indexK.size(2) == kPageSize && indexK.size(3) == kHeadDim && indexK.stride(3) == 1 + && indexK.stride(2) == kHeadDim && indexK.stride(0) >= kPageSize * kHeadDim && indexK.stride(0) % 4 == 0, + "sm107_tp8_producer_pw: index_k must be FP8 [pages, 1, 128, 128] with a page stride keeping the 32-bit stores " + "aligned"); + int64_t const numTokens = packed.size(0); + TORCH_CHECK( + outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous() && outCacheLoc.numel() >= numTokens, + "sm107_tp8_producer_pw: out_cache_loc must be int32 [>= T]"); + TORCH_CHECK( + positionIds.scalar_type() == at::kInt && positionIds.is_contiguous() && positionIds.numel() >= numTokens, + "sm107_tp8_producer_pw: position_ids must be int32 [>= T]"); + TORCH_CHECK(kvQuantScale.scalar_type() == at::kFloat && kvQuantScale.is_contiguous() && kvQuantScale.numel() >= 3, + "sm107_tp8_producer_pw: kv_quant_scale must be fp32 [>= 3]"); + TORCH_CHECK(rotaryCosSin.scalar_type() == at::kFloat && rotaryCosSin.is_contiguous() && rotaryCosSin.dim() == 3 + && rotaryCosSin.size(1) == 2 && rotaryCosSin.size(2) == 32 + && reinterpret_cast(rotaryCosSin.data_ptr()) % 16 == 0, + "sm107_tp8_producer_pw: rotary_cos_sin must be fp32 [positions, 2, 32], 16-byte aligned (read as float4)"); + for (auto const* w : {&qWeight, &kWeight, &indexQWeight, &indexKWeight}) + { + TORCH_CHECK(w->scalar_type() == at::kBFloat16 && w->is_contiguous() && w->numel() == kHeadDim + && reinterpret_cast(w->data_ptr()) % 8 == 0, + "sm107_tp8_producer_pw: norm weights must be BF16 [128], 8-byte aligned (read as 4-value words)"); + } + auto const fp8 = packed.options().dtype(at::kFloat8_e4m3fn); + auto qOut = torch::empty({numTokens, headsQ, kHeadDim}, fp8); + auto indexQOut = torch::empty({numTokens, headsIndex, kHeadDim}, fp8); + if (numTokens == 0) + { + return {qOut, indexQOut}; + } + pw::Params p{}; + p.packed = reinterpret_cast<__nv_bfloat16 const*>(packed.data_ptr()); + p.qOut = reinterpret_cast<__nv_fp8_e4m3*>(qOut.data_ptr()); + p.indexQOut = reinterpret_cast<__nv_fp8_e4m3*>(indexQOut.data_ptr()); + p.kvData = kvData.data_ptr(); + p.kvScale = kvScale.data_ptr(); + p.indexK = reinterpret_cast<__nv_fp8_e4m3*>(indexK.data_ptr()); + p.outCacheLoc = outCacheLoc.data_ptr(); + p.kvQuantScale = kvQuantScale.data_ptr(); + p.dataPageStride = kvData.stride(0); + p.dataPlaneStride = kvData.stride(1); + p.dataHeadStride = kvData.stride(2); + p.dataTokenStride = kvData.stride(3); + p.scalePageStride = kvScale.stride(0); + p.scalePlaneStride = kvScale.stride(1); + p.scaleHeadStride = kvScale.stride(2); + p.indexPageStride = indexK.stride(0); + p.indexTokenStride = indexK.stride(2); + p.numPages = kvData.size(0); + p.numTokens = static_cast(numTokens); + p.eps = static_cast(eps); + p.qWeight = reinterpret_cast<__nv_bfloat16 const*>(qWeight.data_ptr()); + p.kWeight = reinterpret_cast<__nv_bfloat16 const*>(kWeight.data_ptr()); + p.indexQWeight = reinterpret_cast<__nv_bfloat16 const*>(indexQWeight.data_ptr()); + p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); + p.rotaryCosSin = rotaryCosSin.data_ptr(); + p.positionIds = positionIds.data_ptr(); + cudaError_t const err = pw::invokeProducerPw8(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + at::cuda::getCurrentCUDAStream(packed.get_device())); + if (err != cudaSuccess) + { + cudaGetLastError(); + TORCH_CHECK(false, "sm107_tp8_producer_pw launch: ", cudaGetErrorString(err)); + } + return {qOut, indexQOut}; +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp8_producer_pw(Tensor packed, Tensor(a!) kv_data, Tensor(b!) kv_scale, Tensor(c!) index_k, " + "Tensor out_cache_loc, Tensor kv_quant_scale, float eps, Tensor q_weight, Tensor k_weight, " + "Tensor index_q_weight, Tensor index_k_weight, Tensor rotary_cos_sin, Tensor position_ids, bool pdl, " + "bool early, bool prewait_step_inputs=False) -> Tensor[]"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_producer_pw", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8ProducerPw); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp new file mode 100644 index 000000000000..f6200c9aa906 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp @@ -0,0 +1,113 @@ +/* + * 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/minimaxM3Sm107Tp8/sampledAccept.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace sa = tensorrt_llm::kernels::minimaxM3Sm107Tp8::sampledAccept; + +//! uint64 words of the sampled-accept exchange workspace for `world` ranks and SA-1R variant `oneRoundK`. +int64_t minimaxM3Sm107Tp8SampledAcceptWsWords(int64_t world, int64_t oneRoundK) +{ + return sa::sampledAcceptWsWords(static_cast(world), static_cast(oneRoundK)); +} + +//! Sampled (temperature + top-p) verification from this rank's BF16 logits shard [rows, Vs], rows = nreq * +//! (draft_len + 1) <= 48: the sampled token of every row (tokens), then the strict acceptance against draft (accTokens, +//! numAcc, gatherIds, optional nextNew column 0). Per-row Philox (seeds, offsets); the exchange goes through the +//! multicast workspace wsMcPtr / wsUc with the (finished rows, call) counters. +void minimaxM3Sm107Tp8SampledAccept(torch::Tensor const& logits, int64_t vocabOffset, torch::Tensor const& temperatures, + torch::Tensor const& topPs, torch::Tensor const& seeds, torch::Tensor const& offsets, torch::Tensor& wsUc, + int64_t wsMcPtr, torch::Tensor& counters, torch::Tensor& tokens, int64_t rank, int64_t world, + torch::Tensor const& draft, torch::Tensor const& batchIdx, torch::Tensor& accTokens, torch::Tensor& numAcc, + torch::Tensor& gatherIds, int64_t forceTotal, std::optional const& nextNew, + std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, + bool prewaitParams, bool earlyTrigger, int64_t oneRoundK, std::optional const& decided, bool ilpWalk) +{ + c10::cuda::CUDAGuard const deviceGuard{logits.device()}; + sa::Args a; + a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); + a.ld = logits.stride(0); + a.rows = static_cast(logits.size(0)); + a.Vs = static_cast(logits.size(1)); + TORCH_CHECK(a.Vs <= sa::kMaxVs && a.rows <= sa::kMaxRows && world <= sa::kMaxWorld + && static_cast(vocabOffset) + a.Vs <= (1LL << sa::kIdBits)); + a.vocab_offset = static_cast(vocabOffset); + a.temperatures = temperatures.data_ptr(); + a.top_ps = topPs.data_ptr(); + a.seeds = reinterpret_cast(seeds.data_ptr()); + a.offsets = reinterpret_cast(offsets.data_ptr()); + a.ws_uc = reinterpret_cast(wsUc.data_ptr()); + a.ws_mc = reinterpret_cast(wsMcPtr); + a.counters = reinterpret_cast(counters.data_ptr()); + a.tokens = tokens.data_ptr(); + a.extra_attempts = extraAttempts.has_value() ? extraAttempts->data_ptr() : nullptr; + a.rank = static_cast(rank); + a.world = static_cast(world); + a.draft = draft.data_ptr(); + a.batch_idx = batchIdx.data_ptr(); + a.acc_tokens = accTokens.data_ptr(); + a.num_acc = numAcc.data_ptr(); + a.gather_ids = reinterpret_cast(gatherIds.data_ptr()); + a.draft_len = static_cast(draft.size(1)); + a.nreq = static_cast(draft.size(0)); + a.force_total = static_cast(forceTotal); + a.next_new = nextNew.has_value() ? nextNew->data_ptr() : nullptr; + a.next_stride = nextNew.has_value() ? static_cast(nextNew->stride(0)) : 0; + a.force_vec = forceVec.has_value() ? forceVec->data_ptr() : nullptr; + a.prewait_params = prewaitParams ? 1 : 0; + a.early_trigger = earlyTrigger ? 1 : 0; + TORCH_CHECK(oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "one_round_k must be 0, 9 or 82, got ", oneRoundK); + a.one_round_k = static_cast(oneRoundK); + a.decided = decided.has_value() ? decided->data_ptr() : nullptr; + a.ilp_walk = ilpWalk ? 1 : 0; + cudaError_t const err = sa::invokeSampledAccept(a, usePdl, at::cuda::getCurrentCUDAStream(logits.get_device())); + TORCH_CHECK(err == cudaSuccess, "m3_sampled_accept launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +// ws_mc_ptr is the multicast address of ws_uc's mapping: the op writes through it, which the schema cannot show. +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + // No tensor arguments, so no dispatch key: registered with the schema as a catch-all kernel. + m.def("minimax_m3_sm107tp8_sampled_accept_ws_words(int world, int one_round_k=0) -> int", + &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SampledAcceptWsWords); + m.def( + "minimax_m3_sm107tp8_sampled_accept(Tensor logits, int vocab_offset, Tensor temperatures, Tensor top_ps, " + "Tensor seeds, Tensor offsets, Tensor(a!) ws_uc, int ws_mc_ptr, Tensor(b!) counters, Tensor(c!) tokens, " + "int rank, int world, Tensor draft, Tensor batch_idx, Tensor(d!) acc_tokens, Tensor(e!) num_acc, " + "Tensor(f!) gather_ids, int force_total, Tensor(g!)? next_new, Tensor? force_vec, " + "Tensor(h!)? extra_attempts, bool use_pdl, bool prewait_params=True, bool early_trigger=False, " + "int one_round_k=0, Tensor(i!)? decided=None, bool ilp_walk=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107tp8_sampled_accept", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SampledAccept); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp new file mode 100644 index 000000000000..a6cda645ba7f --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp @@ -0,0 +1,66 @@ +/* + * 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/minimaxM3Sm107Tp8/siluAndMul.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp8/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, +//! d]. +void minimaxM3Sm107Tp8SiluAndMul(torch::Tensor const& x, torch::Tensor& out) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp8_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp8_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, + "minimax_m3_sm107_tp8_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); + TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp8_silu_and_mul: at most 65535 rows"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err + = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), + reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_silu_and_mul launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SiluAndMul); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp new file mode 100644 index 000000000000..dc80380b48bb --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp @@ -0,0 +1,189 @@ +/* + * 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/minimaxM3Sm107Tp8/sparseDecodeLc8m.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +constexpr int64_t kHeads = 8; +constexpr int64_t kHeadDim = 128; +constexpr int64_t kTopK = 16; + +void checkBytes(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.element_size() == 1, "minimax_m3_sm107_tp8_sparse_decode_lc8: ", name, + " must be a one-byte CUDA view (uint8 / e4m3 / packed e2m1)"); +} + +} // namespace + +//! One-launch MiniMax-M3 TP8 long-context sparse decode over the selected NVFP4 blocks (8 query heads on one KV head, +//! head_dim 128, page 128, top-k 16). Dependents are released right after the grid dependency wait. +//! +//! @param q E4M3 bytes [T, 8, 128], unit stride in the last dim. +//! @param k, v Packed E2M1 bytes [pages, 1, 128, 64] (the per-layer HND views). +//! @param kBlockScale, vBlockScale E4M3 bytes [pages, 1, 128, 8]: K token-major, V in 4x4 token-quad order. +//! @param kGlobalScale, vGlobalScale One-element float32 per-tensor dequant scales. +//! @param topk Int32 [T, 16] ascending selected block ids with -1 tail (kv head 0 of the head-major table). +//! @param out BF16 [T, 1024], written. q8 / sf, when given, receive the MXFP8 copy and its swizzled UE8M0 scales. +//! @param prewait Stage seq_lens, q and the request's block-table row before the grid dependency wait (same output). +void minimaxM3Sm107Tp8SparseDecodeLc8(torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v, + torch::Tensor const& kBlockScale, torch::Tensor const& vBlockScale, torch::Tensor const& kGlobalScale, + torch::Tensor const& vGlobalScale, torch::Tensor const& topk, torch::Tensor const& blockTable, + torch::Tensor const& seqLens, torch::Tensor& out, std::optional q8, std::optional sf, + int64_t numSfKTiles, int64_t decodeQueryLen, double smScale, bool linearSoftmax, bool usePdl, bool prewait, + std::optional prefetchA, std::optional prefetchB, std::optional n1Self, + std::optional n1Next, std::optional n1K, std::optional n1V, + std::optional n1KScale, std::optional n1VScale, bool identity, bool oneDiv, + bool identityStep) +{ + TORCH_CHECK(q.is_cuda(), "minimax_m3_sm107_tp8_sparse_decode_lc8: q must be a CUDA tensor"); + c10::cuda::CUDAGuard const deviceGuard{q.device()}; + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 10, + "minimax_m3_sm107_tp8_sparse_decode_lc8 requires an SM100-family GPU"); + checkBytes(q, "q"); + checkBytes(k, "k"); + checkBytes(v, "v"); + checkBytes(kBlockScale, "k_block_scale"); + checkBytes(vBlockScale, "v_block_scale"); + TORCH_CHECK(q.dim() == 3 && q.size(1) == kHeads && q.size(2) == kHeadDim && q.stride(2) == 1, + "minimax_m3_sm107_tp8_sparse_decode_lc8: q must be [T, 8, 128] with unit last-dim stride"); + TORCH_CHECK(kGlobalScale.scalar_type() == torch::kFloat32 && vGlobalScale.scalar_type() == torch::kFloat32, + "minimax_m3_sm107_tp8_sparse_decode_lc8: global scales must be float32"); + TORCH_CHECK(topk.scalar_type() == torch::kInt32 && topk.dim() == 2 && topk.size(1) == kTopK, + "minimax_m3_sm107_tp8_sparse_decode_lc8: topk must be int32 [T, 16]"); + TORCH_CHECK(blockTable.scalar_type() == torch::kInt32 && seqLens.scalar_type() == torch::kInt32, + "minimax_m3_sm107_tp8_sparse_decode_lc8: block_table and seq_lens must be int32"); + TORCH_CHECK(out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.size(0) == q.size(0) + && out.size(1) == kHeads * kHeadDim && out.stride(1) == 1, + "minimax_m3_sm107_tp8_sparse_decode_lc8: out must be bfloat16 [T, 1024] with unit last-dim stride"); + TORCH_CHECK( + q8.has_value() == sf.has_value(), "minimax_m3_sm107_tp8_sparse_decode_lc8: pass both q8 and sf, or neither"); + bool const mxfp8 = q8.has_value(); + if (mxfp8) + { + checkBytes(*q8, "q8"); + checkBytes(*sf, "sf"); + } + + tensorrt_llm::kernels::minimaxM3Sm107Tp8::MinimaxM3Tp8SparseDecodeLc8Params prm{}; + prm.q = reinterpret_cast(q.data_ptr()); + prm.q_tok = q.stride(0); + prm.q_head = q.stride(1); + prm.k = reinterpret_cast(k.data_ptr()); + prm.v = reinterpret_cast(v.data_ptr()); + prm.k_blk = k.stride(0); + prm.v_blk = v.stride(0); + prm.ksc = reinterpret_cast(kBlockScale.data_ptr()); + prm.vsc = reinterpret_cast(vBlockScale.data_ptr()); + prm.ksc_blk = kBlockScale.stride(0); + prm.vsc_blk = vBlockScale.stride(0); + prm.kgs = kGlobalScale.data_ptr(); + prm.vgs = vGlobalScale.data_ptr(); + prm.topk = topk.data_ptr(); + prm.t_tok = topk.stride(0); + prm.t_k = topk.stride(1); + prm.block_table = blockTable.data_ptr(); + prm.bt_stride = blockTable.stride(0); + prm.seq_lens = seqLens.data_ptr(); + prm.dql = static_cast(decodeQueryLen); + prm.sm_scale = static_cast(smScale); + prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); + prm.out_stride = out.stride(0); + prm.q8 = mxfp8 ? reinterpret_cast(q8->data_ptr()) : nullptr; + prm.sf = mxfp8 ? reinterpret_cast(sf->data_ptr()) : nullptr; + prm.num_sf_k_tiles = static_cast(numSfKTiles); + prm.identity = identity ? 1 : 0; + prm.identity_step = (identity || identityStep) ? 1 : 0; + prm.one_div = oneDiv ? 1 : 0; + auto const prefetchRange = [](std::optional const& t, uint8_t const*& ptr, int64_t& bytes) + { + ptr = nullptr; + bytes = 0; + if (t.has_value()) + { + TORCH_CHECK(t->is_cuda() && t->is_contiguous(), + "minimax_m3_sm107_tp8_sparse_decode_lc8: a prefetch range must be a contiguous CUDA tensor"); + ptr = reinterpret_cast(t->data_ptr()); + bytes = t->numel() * static_cast(t->element_size()); + } + }; + prefetchRange(prefetchA, prm.pf0, prm.pf0_bytes); + prefetchRange(prefetchB, prm.pf1, prm.pf1_bytes); + bool const n1 = n1Self.has_value(); + TORCH_CHECK(n1 == n1Next.has_value() && n1 == n1K.has_value() && n1 == n1V.has_value() && n1 == n1KScale.has_value() + && n1 == n1VScale.has_value(), + "minimax_m3_sm107_tp8_sparse_decode_lc8: pass all six next-layer prefetch tensors, or none"); + if (n1) + { + for (auto const* t : {&*n1Self, &*n1Next}) + { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == torch::kInt32 && t->is_contiguous() && t->dim() == 2 + && t->size(1) == kTopK && t->size(0) >= q.size(0), + "minimax_m3_sm107_tp8_sparse_decode_lc8: the id buffers must be contiguous int32 [>= T, 16]"); + } + TORCH_CHECK(n1K->stride(0) == k.stride(0) && n1V->stride(0) == v.stride(0) + && n1KScale->stride(0) == kBlockScale.stride(0) && n1VScale->stride(0) == vBlockScale.stride(0) + && n1K->size(0) == k.size(0), + "minimax_m3_sm107_tp8_sparse_decode_lc8: the next layer's pools must have this layer's page strides"); + prm.n1_self = n1Self->data_ptr(); + prm.n1_next = n1Next->data_ptr(); + prm.n1_k = reinterpret_cast(n1K->data_ptr()); + prm.n1_v = reinterpret_cast(n1V->data_ptr()); + prm.n1_ksc = reinterpret_cast(n1KScale->data_ptr()); + prm.n1_vsc = reinterpret_cast(n1VScale->data_ptr()); + prm.n1_pages = n1K->size(0); + } + + auto const stream = at::cuda::getCurrentCUDAStream(q.get_device()); + cudaError_t const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSparseDecodeLc8m( + prm, static_cast(q.size(0)), linearSoftmax, mxfp8, prewait, usePdl, stream); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_sparse_decode_lc8 launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "minimax_m3_sm107_tp8_sparse_decode_lc8(Tensor q, Tensor k, Tensor v, Tensor k_block_scale, Tensor " + "v_block_scale, " + "Tensor k_global_scale, Tensor v_global_scale, Tensor topk, Tensor block_table, Tensor seq_lens, " + "Tensor(a!) out, Tensor(b!)? q8, Tensor(c!)? sf, int num_sf_k_tiles, int decode_query_len, float sm_scale, " + "bool linear_softmax, bool use_pdl, bool prewait, Tensor? prefetch_a=None, Tensor? prefetch_b=None, " + "Tensor(d!)? n1_self=None, Tensor? n1_next=None, Tensor? n1_k=None, Tensor? n1_v=None, Tensor? " + "n1_k_scale=None, Tensor? n1_v_scale=None, bool identity=False, bool one_div=False, bool identity_step=False) " + "-> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_sparse_decode_lc8", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SparseDecodeLc8); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp new file mode 100644 index 000000000000..401aa5cc61fd --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp @@ -0,0 +1,74 @@ +/* + * 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/minimaxM3Sm107Tp8/specOffsets.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +void check(torch::Tensor const& t, char const* name) +{ + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, + "minimax_m3_sm107_tp8_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); +} + +} // namespace + +//! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch +//! (kernels/minimaxM3Sm107Tp8/specOffsets.h). +void minimaxM3Sm107Tp8SpecOffsets( + torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) +{ + check(pos, "pos"); + check(posOff, "pos_off"); + check(kv, "kv"); + check(kvOff, "kv_off"); + TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), + "minimax_m3_sm107_tp8_spec_offsets: each tensor and its offsets must have the same length"); + if (pos.numel() == 0 && kv.numel() == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{pos.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSpecOffsets(pos.data_ptr(), + posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), + at::cuda::getCurrentCUDAStream(pos.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_spec_offsets launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SpecOffsets); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp new file mode 100644 index 000000000000..f340b50dc7dd --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp @@ -0,0 +1,68 @@ +/* + * 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/minimaxM3Sm107Tp8/stepPrologue.h" + +#include +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp8/stepPrologue.h). +//! +//! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the +//! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses +//! the caller guarantees are live. +//! @param epoch Optional one-element int64 step counter, incremented once. +void minimaxM3Sm107Tp8StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) +{ + TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), + "minimax_m3_sm107_tp8_step_prologue: window must be a contiguous uint8 CUDA tensor"); + TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, + "minimax_m3_sm107_tp8_step_prologue: window is shorter than its header and max_desc descriptors"); + c10::cuda::CUDAGuard const deviceGuard{window.device()}; + int64_t* epochPtr = nullptr; + if (epoch.has_value()) + { + TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 + && epoch->device() == window.device(), + "minimax_m3_sm107_tp8_step_prologue: epoch must be a one-element int64 tensor on the window's device"); + epochPtr = epoch->data_ptr(); + } + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeStepPrologue( + window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_step_prologue launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8StepPrologue); +} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp new file mode 100644 index 000000000000..8b5e59751f24 --- /dev/null +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp @@ -0,0 +1,65 @@ +/* + * 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/minimaxM3Sm107Tp8/swigluOai.h" + +#include +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp8/swigluOai.h): x bf16 [rows, 2d] -> out bf16 +//! [rows, d], PDL-launched with an early release when `pdl`. +void minimaxM3Sm107Tp8SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) +{ + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), + "minimax_m3_sm107_tp8_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); + int64_t const rows = x.size(0); + int64_t const d = x.size(1) / 2; + TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp8_swiglu_oai: x rows must be 4-byte aligned with an even half width"); + TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() + && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, + "minimax_m3_sm107_tp8_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); + if (rows == 0) + { + return; + } + c10::cuda::CUDAGuard const deviceGuard{x.device()}; + auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSwigluOai( + reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, + at::cuda::getCurrentCUDAStream(x.get_device())); + TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_swiglu_oai launch failed: ", cudaGetErrorString(err)); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SwigluOai); +} diff --git a/tensorrt_llm/_device_copy.py b/tensorrt_llm/_device_copy.py index b54667235c45..4adc9e601464 100644 --- a/tensorrt_llm/_device_copy.py +++ b/tensorrt_llm/_device_copy.py @@ -20,6 +20,14 @@ def __init__(self, tensor: torch.Tensor, *, pin_memory: bool = False) -> None: def tensor(self) -> torch.Tensor: return self._tensor + @property + def host_values(self) -> torch.Tensor | None: + """Last recorded host values, reshaped to the submitted view; callers must not mutate them.""" + if self._last_host_values is None: + return None + count = self._last_host_shape.numel() + return self._last_host_values[:count].view(self._last_host_shape) + def invalidate(self) -> None: """Forget a pending or externally overwritten copy.""" self._last_host_values = None diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py index a031ea7f9e70..8e589f83d1c0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/fake_ops.py @@ -93,3 +93,85 @@ def _( residual.new_empty(((m + 127) // 128 * 128 * (h // 32),), dtype=torch.uint8), ] return outputs + + +# The MiniMax-M3 sm_107 TP2 / TP4 / TP8 decode ops: one producer and one all-reduce + norm op per TP width (Q and +# index-Q heads per rank from the kernels' producerPw{2,4,8}.h), and the dense QK-norm + RoPE ops. +def _register_minimax_m3_sm107_fakes(tp: int, heads_q: int, heads_index: int) -> None: + @torch.library.register_fake(f"trtllm::minimax_m3_sm107_tp{tp}_producer_pw") + def _( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + prewait_step_inputs=False, + ): + # [Q, index Q] in FP8; the K / V / index-K rows go into the caches in place. + num_tokens = packed.shape[0] + return [ + packed.new_empty((num_tokens, heads_q, 128), dtype=torch.float8_e4m3fn), + packed.new_empty((num_tokens, heads_index, 128), dtype=torch.float8_e4m3fn), + ] + + @torch.library.register_fake(f"trtllm::minimax_m3_sm107_tp{tp}_mc_allreduce_norm") + def _( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + prepushed, + opts, + cap_out=None, + cap_w=None, + cap_eps=0.0, + ): + # [normed, updated residual] (+ [E4M3 rows, swizzled scale bytes] with quant); cap_out is written in place. + outputs = [torch.empty_like(residual), torch.empty_like(residual)] + if quant: + m, h = residual.shape + outputs += [ + residual.new_empty((m, h), dtype=torch.float8_e4m3fn), + residual.new_empty(((m + 127) // 128 * 128 * (h // 32),), dtype=torch.uint8), + ] + return outputs + + +for _tp, _heads_q, _heads_index in ((2, 32, 2), (4, 16, 1), (8, 8, 1)): + _register_minimax_m3_sm107_fakes(_tp, _heads_q, _heads_index) + + +def _minimax_m3_sm107_dense_qk_norm_rope_fp8_fake(qkv, *args, **kwargs): + # A new FP8 tensor shaped like qkv; the KV variant also stores K / V (and the compact Q) in place. + return torch.empty_like(qkv, dtype=torch.float8_e4m3fn) + + +for _name in ( + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8", + "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv", + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8", + "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv", +): + torch.library.register_fake(f"trtllm::{_name}")(_minimax_m3_sm107_dense_qk_norm_rope_fp8_fake) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py new file mode 100644 index 000000000000..61f0a4bad521 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/__init__.py @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 NVFP4 / sm_107 / tp2.""" + +from ..kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py new file mode 100644 index 000000000000..4418d2908b33 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py @@ -0,0 +1,6453 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_107 / tp2 -- self-contained modeling code. + +MiniMax-M3's language model: 60 layers, hidden 6144, 64 query heads and 4 KV heads of dim 128. Layers 0-2 are dense +attention with a dense SwiGLU-OAI MLP; layers 3-59 are MSA block-sparse attention (an index branch scores 128-token +blocks and the top 16 are attended) with a MoE of 128 routed experts (top 4, sigmoid scores with a correction bias, +renormalized and scaled by 2) plus one shared expert. The checkpoint is a modelopt mixed-precision export: NVFP4 routed +experts, MXFP8 projections, fp32 router, bf16 embedding and head; the KV cache is NVFP4 for the sparse layers and FP8 +for the dense ones. + +``tp2`` is two Rubin GPUs: head-split attention (32 query heads, two KV heads and two index heads per rank, one +index-K head) and tensor-parallel experts (every expert on every rank, intermediate 1536 per rank), no attention DP. +Speculative decoding, when configured, is one-model Eagle3 (draft length 3): the BF16 drafter (``M3Eagle3Draft``) and +the BF16 target head, both at checkpoint precision; decode steps of up to 2 requests verify and draft through fused +kernels at that precision (``M3Eagle3Worker``). + +**The core** (``MiniMaxM3Tp2ModelingV2``) holds the weights (``weights.MODEL_WEIGHTS``) and the runtime resources the +steps use -- one MSA attention backend per layer, the TP all-reduce communicator, the RoPE table, the MXFP8 GEMM +dispatch state the engine tunes -- and dispatches each step, on the host, from its context count alone +(``_select_target``): + +* ``DecodeTarget`` takes every pure generation step, of 1..``WIDE_MAX_TOKENS`` rows (up to 48 requests, each one token + or a draft verify of 1 + 3 rows). Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids it runs the decode-step + kernels and picks each layer boundary's collective by the step's token bucket (``DECODE_PLANS``: tiny 1..8, b16 + 9..16, b32 17..32); it runs its own generic layer for the first layer when the fused embedding does not take the + step, and for every layer of a wider step, with the wide cells (the 48-row generic boundaries and MC finalize, the + tensor-core router, the Triton MoE finalize, the 33..48-row GEMVs, the chunked greedy verify). The fast-path cells + go up to 128 rows; a 129..192-row step takes the stock-exact paths. +* ``PrefillTarget`` takes every step with context rows: prefill chunks and mixed steps, every layer in its generic + form over the MSA backend. + +The two targets share no body: each owns its layer, attention and MoE compositions and its ``step_args``, the one +projection of the step metadata its forward reads. What both call are the module-level op adapters, each picking the +kernel for one call from shapes and layouts, which read no step metadata themselves. + +Every fast path the decode step can take, the rows each covers (``FAST_PATHS``, each paired with the entry whose +``CELLS`` certify it) and what a step that cannot take one does are below: **uncovered** (outside every covered +range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise +``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). A row range is empty unless a +landed cell has certified it at the TP2 geometry on Rubin. Every fast path a GEN rank takes prints a rank-tagged +engagement line once per (path, rows) (``taken``). Kernel modules report through ``ops/m3_report``, which this +module wires to the same policy at import. + +**Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the +next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the +long-context decode for the o_proj. It publishes it in ``ops/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +MXFP8 projection on a decode graph takes it, and every take clears the slot. The decode boundaries' push targets (the +o_proj / dense-down / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's exchange) +are passed explicitly from the layer to the producer and back. + +Compiled kernels (the MoE engine, the decode GEMVs with their push variants, the index scorer, the long-row select, +the full-context decode plans) are built before the executor's warmup; a forward never compiles one, and one it finds +missing is a miss. Other state a fast path needs (multicast workspaces, push descriptors) is built on the eager decode +steps; a CUDA-graph capture only reads it, so a capture that finds it missing is a miss (``capture_missed``). + +**What this target asserts rather than adapts**: SM 10.7, the topology and checkpoint (``routing.py``); the +MiniMax-M3 MSA sparse attention runtime with the fused five-way QKV/index projection, an FP8 indexer and an FP8 or +NVFP4 KV cache, no torch.compile and no piecewise prefill graphs (``_validate_runtime``); text-only input (the vision +tower is not built and a step carrying image input raises). +""" + +from __future__ import annotations + +import atexit +import copy +import dataclasses +import os +import threading +import time +from dataclasses import dataclass +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + resolved_mode, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, +) +from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + _ensure_draft_vocab_size, + _get_requested_draft_moe_backend, + _set_draft_kv_cache_quant_algo, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + TensorParallelMode, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import ( + eager_on_graph, + is_in_breakable_cuda_graph, +) +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen, get_model_extra_attrs +from tensorrt_llm._utils import mpi_rank, prefer_pinned +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import weights as _weights +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import ( + m3_drafter_attn, + m3_fold_tail, + m3_mc_allreduce, + m3_prequant, + m3_report, + m3_spec_tail, +) +from .ops import m3_fullctx_decode as m3_fullctx_decode_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_shard_accept as m3_shard_accept_kernel +from .ops import msa_kv_update as _msa_kv_update_kernel +from .ops import triton_sparse_decode as _sparse_decode_kernel +from .ops.activation.m3_silu_and_mul import m3_silu_and_mul +from .ops.attention.m3_draft_attention import m3_draft_attention +from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope +from .ops.attention.m3_fullctx_decode import m3_fullctx_decode +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_msa_attention import m3_msa_attention +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 +from .ops.attention.m3_qkv_insert import m3_qkv_insert +from .ops.attention.m3_select_blocks import m3_select_blocks +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_select_long_dsl import m3_select_long_dsl +from .ops.attention.m3_sparse_decode import m3_sparse_decode +from .ops.attention.m3_sparse_decode_lc import m3_sparse_decode_lc +from .ops.comm.m3_allreduce import m3_allreduce +from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm +from .ops.embedding.m3_embedding import m3_embedding +from .ops.gemm import m3_ugpu_head as m3_ugpu_head_op +from .ops.gemm.m3_bf16_linear import m3_bf16_linear +from .ops.gemm.m3_dense_swiglu import ALPHA as DENSE_ALPHA +from .ops.gemm.m3_dense_swiglu import BETA as DENSE_BETA +from .ops.gemm.m3_dense_swiglu import LIMIT as DENSE_LIMIT +from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu +from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS +from .ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for +from .ops.gemm.m3_fp32_linear import m3_fp32_linear +from .ops.gemm.m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + m3_mx_gemv, + served_max_rows, +) +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.gemm.m3_ugpu_head import m3_ugpu_head +from .ops.long_select import long_select_blocks +from .ops.m3_moe_finalize import m3_moe_finalize +from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .ops.spec.m3_capture_norm import m3_capture_norm +from .ops.spec.m3_chunk_norm import m3_chunk_norm +from .ops.spec.m3_draft_input import m3_draft_input +from .ops.spec.m3_draft_silu import m3_draft_silu +from .ops.spec.m3_draft_step import m3_draft_step +from .ops.spec.m3_force_vec import m3_force_vec, m3_force_vec_req +from .ops.spec.m3_sampled_accept import m3_sampled_accept +from .ops.spec.m3_shard_accept import m3_shard_accept +from .ops.spec.m3_spec_offsets import m3_spec_offsets +from .ops.spec.m3_step_prologue import m3_step_prologue + +_SM = (10, 7) +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The 1..8-row decode kernels' step bound: 2 requests verifying 1 + 3 drafts (the disaggregated GEN worker's batch). +TINY_MAX_TOKENS = 8 # capacity: fixed: the 1..8-row decode kernels' step bound +# Batch-4 plan: steps of 9..16 rows (3-4 requests verifying 1 + 3 drafts) also take the decode layer, on the "b16" +# plan below. Every fast path keeps its own rows (each up to TINY_MAX_TOKENS unless it names its 16-row rows), so a +# path not raised for 16 rows declines to the generic kernel inside the decode layer. False: every step above +# TINY_MAX_TOKENS rows runs the generic layer. +B4_DECODE = True +B4_MAX_TOKENS = 16 if B4_DECODE else TINY_MAX_TOKENS +# Batch-8 plan: steps of 17..32 rows (5-8 requests verifying 1 + 3 drafts) also take the decode layer, on the "b32" +# plan below; every path keeps the rows it names (a path raised only to 16 rows declines a 17..32-row step). False: +# the b16 plan's 16-row cap. +B8_DECODE = B4_DECODE +DECODE_MAX_TOKENS = 32 if B8_DECODE else B4_MAX_TOKENS +# The widest decode step this target takes (aggregated serving, up to 48 requests x (1 + 3 drafts) verify rows): decode +# rows 1..WIDE_MAX_TOKENS; the fast-path cells go up to 128 rows and a 129..192-row step takes the stock-exact paths +# (the Triton sparse-decode producer + merge, the generic GEMM). Steps past DECODE_MAX_TOKENS run the generic layer +# with the wide cells below (GENERIC_MC_*, ROUTER_TC_GEN, MOE_TRITON_FINALIZE, the W-G1 GEMVs, the wide verify). +WIDE_MAX_TOKENS = 192 # capacity: fixed: 48 requests x (1 + 3 drafts) + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 tp4 measurements in the comments).""" + + name: str + max_tokens: int + # The attention boundary's multicast all-reduce + residual + RMSNorm (5.05 vs 5.27 us native at 1 token). + attn_mc_ar: bool + # The same kernel at the MoE boundary with the shared + routed add folded in: wins up to 48 tokens. + moe_mc_ar: bool + # The MoE-boundary all-reduce also emits the next QKV activation in MXFP8 for the decode GEMV / GEMM. + qkv_prequant: bool + + +# tiny: every step up to 8 tokens (c1..c8 without speculation, the c1 verify, every draft pass); small: the c2 +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (ops/m3_mc_allreduce.BUCKETS). +# The multicast all-reduce + residual + RMSNorm at both boundaries of every TP2 decode step (one 8-token bucket), the +# MoE boundary's also publishing the next QKV activation in MXFP8. SM107 TP2 hop: 3.5..4.4 us at 1..8 rows against +# 8.7..9.5 us for the stock AUTO (NCCL) all-reduce + fused residual RMSNorm, and 13.3..14.0 us at the MoE boundary +# (stock add + all-reduce + mxfp8_quantize); bitwise equal to stock (ops/comm/m3_mc_allreduce_norm). +# False puts both boundaries back on the built-in all-reduce. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +MC_AR_TP2 = True +# B3: the 9..32-row steps' attention and MoE boundaries on the same multicast kernel (the 16-token bucket at 9..16 rows, +# epilogue option 15; the 48-token bucket at 17..32: on TP2 the row kernel launches one CTA per row, so the bucket +# only sizes the exchange slots); bitwise equal to stock at every row count (the op's 9..32-row cells). False: the +# built-in all-reduce on those steps. +B4_MC_AR = MC_AR_TP2 +PLAN_MC_AR_MAX_TOKENS = DECODE_MAX_TOKENS if B4_MC_AR else TINY_MAX_TOKENS +# On those steps the MoE (past the engine's 8 rows) returns its unfinalized expert outputs and the multicast finalize +# kernel does the finalize + all-reduce + residual + RMSNorm in one launch (its cells cover 9..48 rows bitwise) +# instead of the routed finalize then the boundary. False: the routed finalize, then the boundary. +B4_MC_FINALIZE = B4_MC_AR +MC_FINALIZE_MIN_TOKENS = 9 # capacity: fixed: past the MoE engine's 8 rows +# B7 (SpecTail TR-5b): on the 9..32-row steps the boundary all-reduce (and the MoE's multicast finalize) also publishes +# the next QKV's MXFP8 activation (its quant cells at 9..32 rows are bitwise mxfp8_quantize), so QKV takes the MXFP8 +# GEMV on the producer-quantized x (795376: QKV at 16 rows 4.78 vs 2.59 us/call at 8; b8 census 794573: the generic +# MXFP8 GEMMs 1,196 + quantize 256 us/step at c8 vs 537 for the GEMVs at c2). +B4_QKV_PREQUANT = B4_MC_AR +B8_QKV_PREQUANT = B4_MC_AR +DECODE_PLANS = ( + ( + DecodePlan( + "tiny", + TINY_MAX_TOKENS, + attn_mc_ar=MC_AR_TP2, + moe_mc_ar=MC_AR_TP2, + qkv_prequant=MC_AR_TP2, + ), + ) + # b16: the 9..16-row steps; their boundaries on the multicast all-reduce's 16-token bucket. + + ( + ( + DecodePlan( + "b16", + B4_MAX_TOKENS, + attn_mc_ar=B4_MC_AR, + moe_mc_ar=B4_MC_AR, + qkv_prequant=B4_QKV_PREQUANT, + ), + ) + if B4_MAX_TOKENS > TINY_MAX_TOKENS + else () + ) + # b32: the 17..32-row steps; their boundaries on the 48-token bucket. + + ( + ( + DecodePlan( + "b32", + DECODE_MAX_TOKENS, + attn_mc_ar=B4_MC_AR, + moe_mc_ar=B4_MC_AR, + qkv_prequant=B8_QKV_PREQUANT, + ), + ) + if DECODE_MAX_TOKENS > B4_MAX_TOKENS + else () + ) +) + +# The dense layers' decode attention (trtllm-gen, msa_decode.dense_decode_max_seq_len): the max_seq_len it is handed +# per request count. It only sizes the multi-CTA KV split (every CTA still covers its share of the actual length; the +# result changes only by the split-merge order), and the model maximum over-splits AgentX-length caches at one +# request: GB300 TP4, 16 heads, FP8 KV, 64K context: 15.1 us at 16K vs 28.6 us at 1M. Two requests are neutral (the +# grid is unchanged) and keep the step's bound. + + +# TP2 skeleton: no cap (the built-in split) until a TP2 cap is measured and certified. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) + +# The dense layers' decode attention on PrimTS split-KV decode (ops/m3_fullctx_decode) instead of the trtllm-gen +# call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the +# standalone reducer. Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 +# (29.3), (1, 1) 15.7 (28.0), (4, 2) 24.5 (29.3), (1, 2) 20.2 (28.1). The stock split is capped at 16 cluster CTAs +# per (request, KV head), i.e. 32 of 212 CTAs at one request. At one request (the E2E batch) the K/V stream is +# latency-bound, so the b1 recipes keep 12 K/V stages in flight (the load pipeline only; the same bits): locked +# clock, cold L2, 165k, (4, 1) 22.6 vs 26.9 and (1, 1) 20.7 vs 23.0 µs with the default 4 stages. +FULLCTX_DECODE_RECIPES: Mapping[Tuple[int, int], Mapping[str, Any]] = MappingProxyType( + { + (4, 1): MappingProxyType( + { + "config": {"tile_size_q": 32, "groups_tokens_heads_q": True, "kv_stages": 12}, + "splits": 48, + } + ), + (1, 1): MappingProxyType({"config": {"kv_stages": 12}, "splits": 72}), + (4, 2): MappingProxyType( + {"config": {"tile_size_q": 32, "groups_tokens_heads_q": True}, "splits": 24} + ), + (1, 2): MappingProxyType({"config": {}, "splits": 40}), + } +) +# False puts the dense layers back on the trtllm-gen call at every shape. +FULLCTX_DECODE_ENABLED = True +# The dense layers' FP8 q/k/v (decode steps of the hybrid cache): the TP2 op rounds Q/K to BF16 and then to E4M3 in +# its store, the bits of the BF16 fused_qk_norm_rope output cast to FP8. False runs that BF16 kernel and a separate cast +# (the same bits, one more kernel). +DENSE_QKV_FP8_VIA_BF16 = True +# On a full-context decode step that op also stores the new K/V rows into the layer's paged pool (the bytes +# write_msa_phase_kv's indexed copies store, one launch instead of about ten). False writes them with +# write_msa_phase_kv. +DENSE_KV_INSERT_FUSED = True +# The engagement line of that fused store (dispatch.taken, one per row count and rank). +DENSE_KV_INSERT_PATH = "m3_dense_kv_insert" +# The drafter MLP's C++ SwiGLU (DRAFT_SILU_CPP, spec_worker._m3_draft_mlp): its own taken / not-taken lines. +DRAFT_SILU_PATH = "m3_draft_silu" +# The split-KV merge on the BF16 op (minimax_m3_sm107_tp2_fullctx_reduce) instead of PrimTS's FP16 reducer and a cast: +# Rubin, 165k, graph-us per call (1, 1) 14.5 vs 18.5, (4, 1) 17.3 vs 21.3, (1, 2) 18.5 vs 23.1, (4, 2) 23.3 vs 27.5. +FULLCTX_DECODE_OWN_REDUCER = True +# The K/V length bound the full-context plans are compiled for; a step whose cache length bound is longer is a miss. It +# covers the TP2 AgentX config's max_seq_len (1 << 20) plus the speculative rows the cache reserves past it. +FULLCTX_DECODE_MAX_KV_LEN = 1 << 21 # capacity: fixed: the full-context plans' compiled K/V bound +# The Eagle3 drafter layer's generation phase on the same full-context decode (ops/m3_drafter_attn: the in-tree +# RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its +# split at 8 CTAs per (request, KV head) on SM 10.7: 45-46 us per draft pass at 165k. +DRAFTER_FULLCTX_ENABLED = True +# On those passes the drafter's RoPE, FP8 query and FP8 K/V append run as one launch (the op +# minimax_m3_sm107tp2_drafter_rope_fp8_kv, a copy of the certified TP4/TP8 kernel at the TP2 heads: the in-tree +# preprocess's GPT-NeoX rotation and unit-scale E4M3 store, and the torch E4M3 cast of its BF16 query) instead of +# thop.trtllm_gen_generation_preprocess and the cast (3.6 vs 1.6 us per pass, 792311). +DRAFTER_ROPE_FUSED = True +# The engagement line of that launch (dispatch.taken, one per row count and rank). +DRAFTER_ROPE_PATH = "m3_drafter_rope" +# How the rotation's products are summed, as the in-tree preprocess build sums them (0: two products then the add, +# 1: fma(cos, x, sin * pair), 2: fma(sin, pair, cos * x)); the op's bitwise test pins it. +DRAFTER_ROPE_MODE = 1 +# thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request +# multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by +# that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, +# token-0 drafts). ops/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# calls: what the build writes (padded rows included), so bitwise the op with a fresh map. False:. +PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the index scorer launches and streams its pre-wait index-K prefetch while the producer +# runs (its outputs are read only after the scorer's own grid wait). Rubin TP2, 57-layer PDL chain behind a 21 MB +# pre-wait L2 prefetch (the 165K scorer's shape): 3.59 (native) -> 1.93 us per layer at 4 rows, 3.64 -> 2.00 at 8. +PRODUCER_EARLY_TRIGGER = True +# The producer also reads its per-step inputs (position ids, cache slots, RoPE rows) before its grid wait, leaving only +# the packed QKV row after it. Their writers are host-to-device copies before the replay and plain-launched kernels +# earlier in the step (the input preprocessing, msa_kv_update or its torch fallback), so they are final before the +# producer can start; a PDL-launched writer of any of them would have to turn this off. The call site's +# position_ids.reshape(-1).contiguous().to(torch.int32) and packed.contiguous() launch a plain torch copy right before +# the producer when the dtype or layout differs: also a plain writer, so the contract holds. Bit-identical outputs. +# PRE-WAIT WRITER CONTRACT (rubin-20261015/PREWAIT_AUDIT.md item b): the kernels that read step inputs (seq_lens, +# block tables, positions, cache slots, the sampled-accept pre-wait parameters) before their grid-dependency wait rely +# on every writer of them being a host-to-device copy before the replay or a PLAIN launch earlier in the step: +# msa_kv_update, the step-input prologue (Triton scatter or the C++ op), the spec fold tail (Triton or C++), the pack +# scatter and any step-head op (the spec step correction, STEP_SPEC_OFFSETS_ONE_LAUNCH). Never give one of them a PDL +# attribute with an early trigger. +PRODUCER_PREWAIT_STEP_INPUTS = True + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. SM107: 1.55..1.65 us against +# 3.4..6.0 us for the stock embedding + RMSNorm at 1..8 rows (the gather bitwise, the norm one ulp off on ~5e-6 of the +# elements: summation order); it also hands the first layer its residual, +# so layer 0 takes the decode layer instead of the generic one. False runs layer 0 on the generic layer. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +EMBED_FIRSTNORM_TP2 = True +# B5 (SpecTail TR-5b): the same kernel on the 9..16-row steps (one CTA per row, so each row's arithmetic is the 1..8-row +# one; the op's 9..16-row cells), which hands layer 0 its residual so it takes the decode layer. Not past 16 rows: +# there the stock embedding + RMSNorm it replaces differs from it by more than at 1..8 rows (7..21 normed elements, one +# past 1 ulp, at 19..28 rows; gate 795453 / debug 795562), so 17..32-row steps keep the stock layer 0. +# False: 1..8 rows. +B4_EMBED_FIRSTNORM = True +EMBED_FIRSTNORM_MAX_TOKENS = ( + (B4_MAX_TOKENS if B4_EMBED_FIRSTNORM else TINY_MAX_TOKENS) if EMBED_FIRSTNORM_TP2 else 0 +) +# The o_proj pushes its rows into the attention boundary's multicast all-reduce (prepushed mode 2) up to this many +# tokens; wider steps run the GEMV and a plain multicast all-reduce. SM107 TP2: -0.5..-1.2 us per layer, bitwise equal +# to the unpushed boundary. False keeps the GEMV and the boundary separate. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +ATTN_PUSH_TP2 = True +# B6 (SpecTail TR-5b): the o_proj also pushes on the 9..16-row steps, into the boundary's 16-token bucket; the GEMV's +# push epilogue covers MMA width 16 with the bucket's tokens in its compile key, and the pushed boundary is bitwise the +# unpushed one (the op's prepushed 9..16-row cells). Not past 16 rows: the o_proj GEMV keeps 16 rows +# (_GEMV_B4_ROWS_ONLY), so a 17..32-row o_proj runs the generic GEMM. +# False: 1..8 rows. +B4_ATTN_PUSH = B4_MC_AR +ATTN_PUSH_MAX_TOKENS = ( + (min(PLAN_MC_AR_MAX_TOKENS, B4_MAX_TOKENS) if B4_ATTN_PUSH else TINY_MAX_TOKENS) + if (MC_AR_TP2 and ATTN_PUSH_TP2) + else 0 +) +# The dense layers' MLP down projection pushes its rows into the post-MLP boundary's multicast all-reduce (prepushed +# mode 2), as the o_proj does into the attention boundary: the all-reduce no longer waits for the whole down GEMV grid +# before loading and pushing the rows itself. SM107 TP2 (stack-ON nsys 773230): the two non-capture dense layers' +# post-MLP hops end 5.5-5.8 us after the down GEMV vs 2.4-2.5 us for a pushed hop. Bitwise equal to the unpushed +# boundary (the same bf16 words). False keeps the down GEMV and the boundary separate. +DENSE_PUSH_TP2 = True +DENSE_PUSH_MAX_TOKENS = TINY_MAX_TOKENS if (MC_AR_TP2 and DENSE_PUSH_TP2) else 0 +# The generic layer's boundaries (prefill, mixed steps, Eagle3 capture layers, decode steps past DECODE_MAX_TOKENS): the +# most rows that take the multicast all-reduce + residual + RMSNorm, alone and with the shared + routed add folded in: +# the multicast all-reduce's widest bucket (48 tokens), as the agg canon runs it (v1.8.3; GSM T0 / T1 803743 / +# 803744), under the same constant (MC_AR_TP2; off = the built-in all-reduce). +GENERIC_MC_AR_MAX_TOKENS = 48 if MC_AR_TP2 else 0 # capacity: fixed: the widest bucket +GENERIC_MC_AR_FOLDED_MAX_TOKENS = 48 if MC_AR_TP2 else 0 # capacity: fixed: the widest bucket +# The generic layer's MoE boundary past DECODE_MAX_TOKENS rows, up to this many (at most the finalize kernel's 48, +# m3_mc_allreduce.FINALIZE_MAX_TOKENS): the routed experts stay unfinalized and their finalize (+ the shared expert) +# runs inside the boundary's multicast all-reduce + residual + next RMSNorm (m3_mc_finalize_allreduce_norm), as the +# decode layer's MC_FINALIZE does. 48 as the agg canon runs it (v1.8.3; the same gates as above). +GENERIC_MC_FINALIZE_MAX_TOKENS = 48 if MC_AR_TP2 else 0 # capacity: fixed: the finalize rows +# The MoE finalize (+ the shared expert add) of a finalized MoE forward with at least MOE_TRITON_FINALIZE_MIN_TOKENS +# rows (prefill, mixed and wide decode past the folds) runs as ops/m3_moe_finalize: the trtllm-gen runner returns +# its FC2 output unfinalized and one Triton kernel gathers the top-k rows together, bitwise to the runner's own +# finalize and the separate shared add (153 -> 98 us per layer at 16K rows on SM107). False: the runner finalizes. +MOE_TRITON_FINALIZE = True +MOE_TRITON_FINALIZE_MIN_TOKENS = 49 # capacity: fixed: past the generic MC finalize's 48 rows +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch, ops/moe/m3_moe_engine at +# the TP2 width 1536). Rubin TP2 ablation chain (57 layers behind a 58.4 MB PDL producer, row = (P + MoE) - P, both +# orders): 31.3 / 50.7 us/layer at 4 / 8 tokens against the built-in MoE's 42.4 / 55.4 (FlashInfer MXFP8 shared +# expert); bit for bit the built-in block's shared.add_(routed). KILL_MOE_ENGINE: the built-in MoE at every row count. +KILL_MOE_ENGINE = False +MOE_ENGINE_MAX_TOKENS = 0 if KILL_MOE_ENGINE else 8 +# The engine's push of its sum into the MoE boundary's multicast all-reduce (prepushed mode 2) on the rows the +# engine serves, with the multicast boundary (DECODE_PLANS' moe_mc_ar): the boundary reads the pushed words instead +# of loading and pushing the rows itself. Bit for bit the unpushed boundary (comm/_m3_moe_push_op_matrix.py). +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS if MC_AR_TP2 else 0 +# When the engine pushes its sum into the boundary all-reduce, every CTA releases the PDL dependents right after its +# grid dependency wait instead of after the push, so the all-reduce (which polls the pushed tagged words, not the grid) +# becomes resident on the SMs the engine frees. Scheduling only: the same bits (the push oracle checks on vs off). +# Rubin TP2 greedy early column, same-job A,B + B,A (778191 / 778192): c1 -15 us, about -15 at 323K (within noise). +# False: releases after the push. +MOE_EARLY_AR_RELEASE = True +# A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel +# over the target lm_head's BF16 shard logits (ops/m3_shard_accept), at the configured draft length. +SHARD_ACCEPT_DRAFT_LEN = 3 # capacity: fixed: the built Eagle3 draft length +# B4 spec tail: the fused verify (shard / sampled accept), the fused first drafter inputs (chunk / capture norm) and the +# fused draft loop (draft_step, its draft input, the drafter boundary all-reduce) also take the 9..32-row steps, up to +# DECODE_MAX_TOKENS verify rows (DECODE_MAX_TOKENS // (draft length + 1) requests; the draft step's head takes at most +# 8). Every kernel is row-local (per-row norms, per-request keys and Philox draws), so each row gets the bits it gets +# at 1 / 2 requests; the op cells at the new rows certify it. +# False: TINY_MAX_TOKENS rows. +B4_SPEC_TAIL = True +SPEC_TAIL_MAX_TOKENS = DECODE_MAX_TOKENS if B4_SPEC_TAIL else TINY_MAX_TOKENS +SPEC_TAIL_MAX_REQUESTS = SPEC_TAIL_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) +# W2 (SpecTail, agg wide decode): the greedy fused verify (ops/m3_shard_accept) also takes the +# 33..WIDE_MAX_TOKENS-row decode-only steps, as whole-request chunks of the kernel's certified 48-row geometry (12 +# requests per launch at draft length 3; m3_shard_accept.WIDE_SPEC_TAIL), instead of the gathered target logits + +# argmax; and the in-tree draft loop takes the drafter's top-1 from its local head shard (W2b). Sampled verifies keep +# the fused path only up to SAMPLED_VERIFY_MAX_REQUESTS (past it, the in-tree sampler). Every other fused piece keeps +# SPEC_TAIL_MAX_TOKENS rows (the drafter then runs its in-tree loop). Rubin TP2 W1b 797141 (verify + 3 drafter top-1 +# passes, stock -> fused). Greedy chunks and the shard top-1 are bitwise the gathered argmax (lowest index on ties). +WIDE_SPEC_TAIL = m3_shard_accept_kernel.WIDE_SPEC_TAIL +WIDE_VERIFY_MAX_TOKENS = WIDE_MAX_TOKENS if WIDE_SPEC_TAIL else SPEC_TAIL_MAX_TOKENS +WIDE_VERIFY_MAX_REQUESTS = WIDE_VERIFY_MAX_TOKENS // (SHARD_ACCEPT_DRAFT_LEN + 1) +# 0: (the gathered logits + argmax run). Rubin TP2 greedy verify (microbench): stock 120.4 us at b1; fused +# with the coded head 117.0 us; the BF16-head + shard-accept arm is not measured yet. +SHARD_ACCEPT_MAX_REQUESTS = WIDE_VERIFY_MAX_REQUESTS +# A sampled (temperature + top-p, advanced_sampling_mode no_topk) Eagle3 verify step of up to this many requests +# verifies with the fused sampled acceptance (ops/m3_sampled_accept, the TP2-shard kernel) over the target +# lm_head's BF16 shard logits, at the configured draft length (the AgentX contract). 0: (the gathered logits +# and the in-tree sampler run). Rubin TP2 b1/b2 (p5 microbench): 180.1 / 181.8 -> 115.1 / 113.9 us per verify. +SAMPLED_VERIFY_MAX_REQUESTS = SPEC_TAIL_MAX_REQUESTS +# A step verified by the in-tree path takes the fused first drafter inputs (the capture norms and the first pass's +# gather) up to this many requests. +FIRST_DRAFTER_INPUTS_MAX_REQUESTS = 0 # capacity: fixed: off at TP2 (0) +# A step verified by the fused kernels (dispatch.SHARD_ACCEPT / SAMPLED_ACCEPT) of up to this many requests takes the +# fused first drafter inputs (spec_worker._m3_first_drafter_inputs: the accepted tokens as they lie, the first pass's +# gather rows, the fc_norm chunk norm or the armed capture norm, then the in-tree fc). 0: (the in-tree first +# drafter inputs run, the captures stay in-tree and the fused draft loop steps aside). +FUSED_FIRST_INPUTS_MAX_REQUESTS = SPEC_TAIL_MAX_REQUESTS +# On a fused-verify step of up to this many verify rows the Eagle3 captures write rmsnorm(bf16(h + residual)) * +# fc_norm[i] in one kernel (ops/spec/m3_capture_norm) once the fc_norm weights are stacked; the fused first drafter +# inputs then skip their chunk norm. 0: (the in-tree capture, then the chunk norm). +CAPTURE_NORM_MAX_ROWS = SPEC_TAIL_MAX_TOKENS +# The first draft pass's per-layer fc_norm over the captured columns in one kernel (ops/spec/m3_chunk_norm) at up +# to this many verify rows; wider first passes take the in-tree first drafter inputs. 0:. +CHUNK_NORM_MAX_ROWS = SPEC_TAIL_MAX_TOKENS +# The fused draft loop's two all-reduce boundaries per pass (o_proj -> post-attention norm, down -> final norm) of up +# to this many rows run the multicast all-reduce + residual + RMSNorm (ops/m3_mc_allreduce, opts bit 16: the +# drafter's stock rounding, normalizing the unrounded FP32 sum as FlashInfer fused_add_rmsnorm does) instead of the +# built-in all-reduce + RMSNorm. 0:. Rubin TP2 (p16 microbench): 10.0..11.0 -> 4.6..5.3 us per boundary. +# (The drafter epilogue, opts bit 16, is built for the 8-token bucket only: a B4 spec-tail first pass of 16 / 32 rows +# takes the built-in all-reduce + RMSNorm; its later passes of up to 8 rows stay multicast.) +DRAFT_MC_AR_MAX_ROWS = 8 # capacity: fixed: the drafter epilogue's 8-token bucket +# The drafter's o_proj and down projection, on the copied split-K kernel (DRAFT_LOWM_OWN_SPLITK), also push their rows +# into the drafter boundary all-reduce's exchange (the canonical words its own push would store, at the call parity), +# which then runs prepushed 2 instead of pushing x itself; the copy releases its dependents only after its own grid +# wait and the pushes (the prepushed >= 2 contract). Bitwise. TP2: the 4-row first pass only (its 1-row linears run +# FlashInfer's direct kernel). False:. +DRAFT_AR_PUSH = True +# The drafter boundary all-reduce fed by its projection's push (DRAFT_AR_PUSH): its own taken / not-taken lines. +DRAFT_AR_PUSH_PATH = "m3_draft_ar_push" +# The same push from the drafter's direct ('d') low-M linears (the 1-row passes' o_proj and down at TP2), through this +# target's copy of FlashInfer's direct kernel (ops/m3_lowm_direct; its output FlashInfer's): the owning thread +# stores the CTA's column pairs into the exchange before the kernel's last-instruction release. Bitwise. Rubin TP2 +# world 2 (dpush-1 782685), per 1-row boundary: o_proj 9.22 -> 8.09 us, down 13.81 -> 12.72. False:. +DRAFT_AR_PUSH_DIRECT = True +# uGPU: a verify of these row counts computes the target logits shard as two half-vocabulary in-tree GEMMs, each on +# the green context of the die whose memory holds that half of a placed copy of the head (ops/gemm/m3_ugpu_head; the +# module's weight is untouched, so every other reader keeps ordinary memory); the logits are the in-tree head's bit for +# bit. Rubin TP2 head shard: 93.3 -> 79.5 us at +# 4 rows, 93.5 -> 80.6 / 80.7 at 12 / 16 (at 1 / 2 rows the half-shard GEMM is another kernel). In-model the decode +# step is ~45 us shorter at one request and ~90 us at two, and the gain holds at long context. +KILL_UGPU_HEAD = False +UGPU_HEAD_ROWS: Tuple[int, ...] = () if KILL_UGPU_HEAD else (4, 8, 12, 16) +# On a fused-verify step of up to this many verify rows the drafter's first fc is forked onto the aux stream right +# after the last captured target layer, overlapping the remaining target layers and the head (spec_worker). +FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover +# A sampled fused verify of up to this many verify rows on the in-tree head GEMM forks the drafter's first fc after the +# head instead of beside it (spec_worker): the fc's split-K grid otherwise holds the SMs first and the head GEMM starts +# late. 0:. Rubin TP2 c1 (locked nsys 768559): the head started 17.4 us after its input was ready. +FC_AFTER_HEAD_MAX_ROWS = SPEC_TAIL_MAX_TOKENS +# A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (spec_worker): per pass the +# draft input in one launch (ops/m3_spec_tail.draft_input), the drafter layer's in-tree attention, MLP and +# all-reduces, and one kernel for the draft head (its 12-bit coded copy, or the BF16 fragments when the coding does +# not decode bit for bit), the cross-rank argmax and the pass's bookkeeping (ops/m3_draft_step), at the +# configured draft length. The head kernel takes one row per request, at most 8. +# The fused draft loop streams the 12-bit coded copy of the draft head (when it decodes bit for bit); False keeps the +# BF16 fragments. +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +# 0: (the in-tree draft loop runs). Rubin TP2 per draft pass (head + argmax + exchange): 122.6 / 126.0 -> +# 102.8 / 102.8 us at b1 / b2 (BF16 fragments; the coded head is slower on sm_107: DRAFT_HEAD_CODED False). +DRAFT_STEP_MAX_REQUESTS = min(SPEC_TAIL_MAX_REQUESTS, 8) # capacity: fixed: the draft head's 8 rows +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) run FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm, the TP8 port) at the rows its TP2 table is certified for, up to this many rows; other calls +# keep the in-tree cuBLAS GEMM. 0:. Rubin TP2 (mclk locked, 4 rotating sets) at 1 row: qkv 13.18 -> 9.71, +# o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. +DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the certified low-M rows +# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (ops/m3_lowm_splitk): bit for bit the +# same output at every certified cell (spec/test_modeling_v2_m3_lowm_splitk.py); the seam for the drafter-boundary push +# epilogue and the down projection's SwiGLU prologue. False: (FlashInfer's kernel). +DRAFT_LOWM_OWN_SPLITK = True +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(m for m in rows if m <= DRAFT_LOWM_MAX_ROWS) + for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() +} +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (ops/triton_sparse_decode.MQ_MIN_TOKENS); the multi-query producer needs one KV head, so at the TP2 heads every +# identity step takes the per-row producer. +SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover +SPARSE_DECODE_KV_HEADS = 2 # capacity: fixed: the tp2 KV heads per rank +# The long-row top-16 selector takes engines of 129..LONG_SELECT_MAX_BLOCKS KV blocks of 128 tokens; it holds 2, 4 or +# 9 keys per thread by the row's valid-block count. At most 128 blocks, or past the maximum, the built-in selector +# runs; 0 turns the path off. Rubin TP2 (2 KV heads), 165,521 tokens, behind the scorer, ablation per layer: +# -0.9 us at 1 request and -1.2 us at 2 requests against the built-in selector. +LONG_SELECT_MIN_BLOCKS = 129 # capacity: fixed: past the built-in selector's 128 +LONG_SELECT_MAX_BLOCKS = 9216 # capacity: fixed: the select kernels' built blocks per row +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +# The long-row select releases the sparse decode after it right after its own grid wait (the decode stages q and the +# request length before its wait) up to this many query rows; past it the trigger stays at the select's exit. On Rubin +# TP2 the early-launched long-context sparse decode (SPARSE_DECODE_LC) crowds the select's CTAs off the SMs already at +# one request (4 rows): releasing at the exit is 336 us per decode step faster at 1 request and 350 at 2 (in-model, +# graph time, 165K-context mixed draws, locked memory clock), so the early release is off at every row count. +SELECT_EARLY_MAX_ROWS = 0 # capacity: fixed: measured early-release crossover +# Without the early release, the long-row select releases the sparse decode once every score of its row is loaded +# (1, before the exact rank of the candidates) instead of at its exit (0). The release still follows +# the select's own grid wait, which the sparse decode's pre-wait q / length / block-table reads rely on. The decode's +# 16-CTA cluster launch and prologue then overlap the rank: Rubin TP2 in-model, locked memory clock, AgentX draws, +# 2 instances vs the r4.3 stored A: replay period -162 us fixed, -94 at a 323K-token request. +SELECT_LONG_TRIGGER_PHASE = 1 +# The long-row select runs as ops/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per +# (query row, KV head) instead of one 1024-thread CTA, the same selection bit for bit (ties included), with 4, 11 or +# 24 keys per thread (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices. Rubin TP2 decode chain (2 KV heads, +# 4 rows, 57 layers behind the producer stream and the scorer, both orders x 3 rounds): -0.127 / -0.148 / -0.037 us +# per layer at 98K / 165K / 323K tokens against the 1024-thread kernel. False: (the 1024-thread kernel). +SELECT_LONG_DSL = True +LONG_SELECT_DSL_KEY_WIDTHS = (4, 11, 24) +# The MSA index scorer (ops/attention/m3_index_score) of a sparse decode step, up to INDEX_SCORE_MAX_REQUESTS requests +# (0 turns the path off; wider steps run the in-tree CuTe DSL scorer): its CTAs per request by request count, and up +# to INDEX_SCORE_EARLY_MAX_REQUESTS requests the early issue (Q and the first K tile right after the grid wait, one +# older block per CTA prefetched into L2 before it). The split is captured with the graph, so it is one value per +# request count over every context. Rubin TP2 (2 index heads), 165,521 tokens, behind a 58 MB PDL producer, ablation +# per layer: 1 request 3.27 us at 1536 CTAs (512: 3.76, in-tree 256-CTA scorer 6.07); 2 requests 4.97 us at 768 +# (in-tree 6.71). Within 0.2 us of each context's best split from 16K to 400K tokens except 2 requests at 400K (+0.6). +# The early issue reads the request's first page id from the MSA page table before the grid wait: exact because the +# page table is written only on the host side before the replay (attention/test_modeling_v2_m3_block_table_contract.py); +# the request length read there is a hint, every load and store is bounded by the one read after the wait. +# The batch-8 plan's 3- and 4-request steps take it too (Indexer k11 795543: Rubin TP2, 512 MB flush before each call, +# graph-timed, bitwise equal to the in-tree scorer; requests at 1.0 / 0.9 / 0.8 / 0.7 x the context): 3 requests at +# 424 CTAs each 6.08 / 9.04 / 14.20 us at 98K / 165K / 323K tokens (in-tree 6.65 / 9.44 / 14.49); 4 requests at 318 +# 7.41 / 10.82 / 16.80 (in-tree 7.60 / 10.97 / 16.96). 5..8 requests at about 1272 CTAs in total, 1272 // requests +# each (Indexer k11c 795819): 5 at 254 8.25 / 11.44 / 19.10 us (in-tree 8.66 / 12.10 / 19.65); 6 at 212; 7 at 181; +# 8 at 159 9.91 / 13.97 / 23.33 (10.33 / 14.47 / 23.74). Each within 0.03 us of its count's best split. +# 9..48 requests (W-IDX1, one split per bucket) win standalone (Indexer k14b 797357: 46 requests 52.29 -> 49.85 us per +# layer) but lose in-model (agg census c46 / c20: 52.13 -> 61.85 / 24.22 -> 29.06 us per layer, the in-tree scorer +# already at the single-stream bandwidth floor), so steps past 8 requests keep the in-tree scorer. +INDEX_SCORE_MAX_REQUESTS = 8 # capacity: fixed: measured in-model crossover +#: (largest request count of the bucket, CTAs per request) for 9..48 requests (W-IDX1; outside INDEX_SCORE). +INDEX_SCORE_WIDE_SPLITS = ((10, 127), (13, 97), (17, 74), (23, 55), (31, 41), (39, 32), (48, 26)) +INDEX_SCORE_SPLIT_K: Mapping[int, int] = MappingProxyType( + {1: 1536, 2: 768, 3: 424, 4: 318, 5: 254, 6: 212, 7: 181, 8: 159} + | {r: next(s for hi, s in INDEX_SCORE_WIDE_SPLITS if r <= hi) for r in range(9, 49)} +) +INDEX_SCORE_EARLY_MAX_REQUESTS = 48 # capacity: fixed: measured early-issue crossover + + +# Synth only (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS, fractional): the fused verify's forced accepted counts are +# drawn per request, keyed by the request's KV length at the verify (its prompt plus what it has accepted), instead of +# from the server-wide counter, so a request's forced counts do not depend on what the server ran before (A/B draws +# become history-independent). Production sampling never forces counts. False: (the shared counter, the in-tree +# sequence). +FORCE_AL_PER_REQUEST = True + +# The decode MXFP8 GEMVs (ops/mx_gemv, SM107 TP2 table), per projection (N, K): the rows each serves, on an +# activation its producer quantized (the MoE boundary for the QKV, the attention merge for the o_proj) and on a bf16 +# one (quantized in the GEMV's prologue, bit-identical to mxfp8_quantize). SM107 graph us per projection at 1..8 rows: +# 2.2..8.6 against 4.7..10.5 for the stock MXFP8 quantize + FlashInfer CuTe DSL GEMM (ops/mx_gemv/op.py). +# False puts every projection back on the stock decode-graph GEMM. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +MX_GEMV_TP2 = True +# MX-TAC: the decode-graph MXFP8 projections the GEMVs do not serve take one fixed per-bucket decision instead of the +# AutoTuner's per-start choice (whose L2-hot profiling left c2 bimodal at +-6%): FlashInfer's CuTe DSL quantize and +# CuTe DSL GEMM. Stock's candidates per bucket are quantize {TRT-LLM, CuTe DSL} x GEMM {CUTLASS, CuTe DSL}; on SM107 +# every TP2 projection at 1..16 rows runs fastest on this pair, cold and L2-hot (graph us, cold: 4.7..10.3 against +# 5.0..11.5 with the TRT-LLM quantize and 9.7..58 with CUTLASS), and all four candidates give the same bits (quantize +# bytes and scales, GEMM outputs). A constant is rank-agreed and start-independent by construction. +# False returns these projections to the AutoTuner. +MXFP8_DECODE_QUANT_TACTIC = 0 # MXFP8QuantizeRunner.CUTE_DSL +MXFP8_DECODE_GEMM_TACTIC = 0 # FlashInferMXFP8GemmRunner.CUTE_DSL +_ALL_ROWS = frozenset(range(1, TINY_MAX_TOKENS + 1)) +_GEMV_PROJECTIONS = ( + QKV_SPARSE, + QKV, + O_PROJ, + SHARED_GATE_UP, + SHARED_DOWN, + DENSE_GATE_UP, + DENSE_DOWN, +) +# B2: the same GEMVs at 9..16 rows (MMA width 16, the kernel's own wider configuration; contiguous k-tiles above 8 +# rows): standalone graph us at 16 rows 2.3..8.9 against 4.8..10.6 for the MX-TAC stock pair (b4p1 791712). Batch 8: +# at 17..32 rows (MMA width 32; rule 1b, a different K order than the generic GEMM). GEMM-AR standalone 794929 (iso +# latency vs the FlashInfer fallback on the same MXFP8 bytes), us per call at 24 / 32 rows: QKV-sparse -2.70 / -2.75, +# QKV -2.46 / -2.77, shared gate_up -3.12 / -3.12, shared down -0.94 / -1.01; dense gate_up +1.18 / +0.91, dense down +# +0.58 / +0.13 (the GEMV loses) and o_proj +0.11 / -0.18 (flat) keep 16 rows. Each projection is raised only as far +# as the kernel serves it (served_max_rows): the bf16 GEMV's in-kernel quantize declines dense gate_up / down past 8. +# MX_GEMV_B4 False: 9..32 rows decline to MX-TAC; B8_MX_GEMV_32 False: every GEMV to 16 rows (17..32-row MXFP8 +# inputs on FlashInfer's GEMM, exact). +MX_GEMV_B4 = True +B8_MX_GEMV_32 = True +GEMV_MAX_TOKENS = ( + (DECODE_MAX_TOKENS if B8_MX_GEMV_32 else B4_MAX_TOKENS) if MX_GEMV_B4 else TINY_MAX_TOKENS +) +_GEMV_B4_ROWS_ONLY = frozenset({DENSE_GATE_UP, DENSE_DOWN, O_PROJ}) + + +def _gemv_max(nk) -> int: + return min(GEMV_MAX_TOKENS, B4_MAX_TOKENS) if nk in _GEMV_B4_ROWS_ONLY else GEMV_MAX_TOKENS + + +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + { + nk: frozenset(range(1, served_max_rows(nk, False, TINY_MAX_TOKENS + 1, _gemv_max(nk)) + 1)) + for nk in _GEMV_PROJECTIONS + } + if MX_GEMV_TP2 + else {} +) +# W-G1 (agg wide decode, rule 1b): the MXFP8 GEMV also at 33..48 rows (MMA width 48 / 64) for the projections it wins +# on, after the decode-graph MXFP8 quantize the generic GEMM path runs anyway (the same activation bytes): standalone +# graph us per call at 48 rows, GEMV against the MX-TAC FlashInfer CuTe DSL GEMM on the same MXFP8 input, QKV-sparse +# 5.87 vs 10.82, QKV 5.53 vs 10.76, shared gate_up 4.52 vs 10.51, shared down 2.84 vs 4.74. o_proj and the dense MLP +# keep their rows (flat or losing at 24..48). A different K order than the GEMM (rule 1b). False: 33..48 rows back on +# the generic GEMM, exact. +WIDE_MX_GEMV_48 = True +WIDE_GEMV_MAX_TOKENS = 48 # capacity: fixed: the MXFP8 GEMV's widest build +_GEMV_W48_PROJECTIONS = (QKV_SPARSE, QKV, SHARED_GATE_UP, SHARED_DOWN) +_GEMV_W48_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + { + nk: frozenset( + range( + GEMV_MAX_TOKENS + 1, + served_max_rows(nk, False, GEMV_MAX_TOKENS + 1, WIDE_GEMV_MAX_TOKENS) + 1, + ) + ) + for nk in _GEMV_W48_PROJECTIONS + } + if MX_GEMV_TP2 and WIDE_MX_GEMV_48 and GEMV_MAX_TOKENS == DECODE_MAX_TOKENS + else {} +) +for _nk, _rows_w48 in _GEMV_W48_ROWS.items(): + MX_GEMV_ROWS[_nk] = MX_GEMV_ROWS.get(_nk, frozenset()) | _rows_w48 +MX_GEMV_BF16_MAX_TOKENS: Dict[Tuple[int, int], int] = ( + {nk: served_max_rows(nk, True, TINY_MAX_TOKENS + 1, _gemv_max(nk)) for nk in _GEMV_PROJECTIONS} + if MX_GEMV_TP2 + else {} +) +# W-G1: a bf16 activation (no MXFP8 hand-off) on rows the MXFP8 GEMV serves but the bf16 GEMV's in-kernel quantize does +# not (QKV / shared gate_up at 17..32 without the QKV hand-off, all four at 33..48) is quantized by the decode-graph +# MXFP8 quantize (the generic GEMM path's own) and runs the GEMV instead of that GEMM. +MX_GEMV_QUANT_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + { + nk: frozenset(r for r in MX_GEMV_ROWS.get(nk, ()) if r > MX_GEMV_BF16_MAX_TOKENS.get(nk, 0)) + for nk in _GEMV_W48_PROJECTIONS + } + if MX_GEMV_TP2 and WIDE_MX_GEMV_48 + else {} +) + +Key = Hashable + + +@dataclass(frozen=True) +class RowRange: + """The shapes one fast path takes, and the catalog entry whose cells certify each of them. + + ``keys`` are the covered shapes: row counts for most paths, (query rows per request, requests) for the attention + kernels, a variant name for thresholds. ``certifies`` maps a cell's spec of ``entry`` to the keys that cell + drives (empty for a cell of another variant). + """ + + path: str + keys: FrozenSet[Key] + entry: OpWrapper + certifies: Callable[[Mapping[str, Any]], Iterable[Key]] + + def covers(self, key: Key) -> bool: + return key in self.keys + + def certified(self) -> FrozenSet[Key]: + return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + + +def _rows(lo: int, hi: int) -> FrozenSet[int]: + return frozenset(range(lo, hi + 1)) + + +def _spec_rows(field: str, **match: Any) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + """A cell's ``spec[field]`` when the cell's other fields match ``match``.""" + + def certifies(spec: Mapping[str, Any]) -> Iterable[Key]: + if any(spec.get(name) != value for name, value in match.items()): + return () + return (spec[field],) + + return certifies + + +def _const(*keys: Key) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + return lambda spec: keys + + +def _sparse_decode_producer(spec: Mapping[str, Any]) -> Iterable[Key]: + if not spec["identity"]: + return () + rows = spec["decode_query_len"] * len(spec["seq_lens"]) + mq = rows >= SPARSE_DECODE_MQ_MIN_TOKENS and SPARSE_DECODE_KV_HEADS == 1 + return ("multi_query" if mq else "per_row",) + + +def index_score_schedule(requests: int) -> Tuple[int, bool]: + """The index scorer's (CTAs per request, early issue) at this many requests (a covered request count).""" + return INDEX_SCORE_SPLIT_K[requests], requests <= INDEX_SCORE_EARLY_MAX_REQUESTS + + +def _index_score_requests(spec: Mapping[str, Any]) -> Iterable[Key]: + """A cell certifies its request count when it drives the TP2 heads with that count's schedule.""" + requests = len(spec["seq_lens"]) + if spec.get("heads") != 2 or requests not in INDEX_SCORE_SPLIT_K: + return () + if (spec["split_k"], spec["early"]) != index_score_schedule(requests): + return () + return (requests,) + + +def _long_select_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The key widths a cell's rows reach: every width up to its blocks' when rows of mixed lengths are drawn.""" + blocks = spec["num_blocks"] + widest = 2 if blocks <= 2048 else 4 if blocks <= 4096 else 9 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_KEY_WIDTHS if w <= widest) + return (widest,) + + +def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The 384-thread select's key widths a cell's rows reach (every width up to its blocks' for mixed rows).""" + blocks = spec["num_blocks"] + widest = 4 if blocks <= 4 * 384 else 11 if blocks <= 11 * 384 else 24 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_DSL_KEY_WIDTHS if w <= widest) + return (widest,) + + +EMBED_FIRSTNORM = RowRange( + "embed_firstnorm", + _rows(1, EMBED_FIRSTNORM_MAX_TOKENS), + m3_embedding_firstnorm, + _spec_rows("rows"), +) +# One request decoding one token on an MXFP8 layer input: QKV GEMV + attention + o_proj in one launch. +# The identity-selection Triton producer + merge (every request within 16 blocks), certified at the TP2 heads on SM107. +SPARSE_DECODE_PRODUCER = RowRange( + _sparse_decode_kernel.PRODUCER_PATH, + _sparse_decode_kernel.PRODUCER_KEYS, + m3_sparse_decode, + _sparse_decode_producer, +) +# The long-context sparse decode at the TP2 heads (32 query / 2 KV heads): attention, split-K merge and the MXFP8 +# o_proj hand-off in one launch, for every decode step with a selected table. "mxfp8" is a step of at most +# MergeMxfp8.MAX_ROWS rows (the hand-off is written and published to m3_prequant.SLOT), "plain" a larger one. +# frozenset() routes every step back to the Triton producer + merge. +SPARSE_DECODE_LC = RowRange( + _sparse_decode_kernel.LC_PATH, + _sparse_decode_kernel.LC_KEYS, + m3_sparse_decode_lc, + lambda spec: ("mxfp8" if spec["mxfp8"] else "plain",), +) +# N1 at TP2: SPARSE_DECODE_LC records each step's selected block ids per sparse layer and, before its grid wait (while +# the select ranks), pulls into L2 the K/V pages + scales of the blocks the SAME layer selected one step back (65-68% of +# a layer's selected blocks recur one step later: MSA probe 778226). Same layer, not the next one as at TP4: the MoE +# between two sparse layers streams ~250 MB through the 126 MB L2, so an earlier issue would be evicted. Hints only; +# up to N1_MAX_ROWS query rows. False:. +SPARSE_DECODE_LC_PREFETCH_SELF_KV = True +N1_MAX_ROWS = 128 # capacity: fixed: the N1 prefetch's row bound +# Every token count of a sparse layer on an NVFP4 KV cache at the TP2 heads (prefill and decode): the producer takes the +# geometry, not a row range. frozenset() routes every sparse layer back to the native producer. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp2"}), m3_producer, _const("nvfp4_tp2")) +INDEX_SCORE = RowRange( + "m3_index_score", + _rows(1, INDEX_SCORE_MAX_REQUESTS), + m3_index_score, + _index_score_requests, +) +LONG_SELECT = RowRange( + "long_select", + frozenset(LONG_SELECT_KEY_WIDTHS) if LONG_SELECT_MAX_BLOCKS else frozenset(), + m3_select_long, + _long_select_widths, +) +LONG_SELECT_DSL = RowRange( + "long_select_dsl", + frozenset(LONG_SELECT_DSL_KEY_WIDTHS) + if LONG_SELECT_MAX_BLOCKS and SELECT_LONG_DSL + else frozenset(), + m3_select_long_dsl, + _long_select_dsl_widths, +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, PLAN_MC_AR_MAX_TOKENS) if MC_AR_TP2 else frozenset(), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=0), +) +MC_FINALIZE = RowRange( + "mc_finalize_allreduce_norm", + _rows(MC_FINALIZE_MIN_TOKENS, PLAN_MC_AR_MAX_TOKENS if B4_MC_FINALIZE else TINY_MAX_TOKENS), + m3_mc_finalize_allreduce_norm, + _spec_rows("m"), +) +ATTN_PUSH = RowRange( + "attn_push", + _rows(1, ATTN_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +DENSE_PUSH = RowRange( + "dense_push", + _rows(1, DENSE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +# The router of an unfused MoE of 9..48 rows (above the MoE engine's 8), a generation step's or a context step's: the +# TF32 tensor-core router (minimax_m3_router_tc, 16-CTA clusters, one reduction order; the arithmetic the built-in gate +# takes on SM107 at 9..16 generation rows), with the routed experts' NVFP4 activation from the same launch, as the agg +# canon routes these rows (2.9-3.5 us/call against ~9.8 us for the cuBLAS TF32 split-K router; closer to the fp32 +# product: max |dlogit| 1.7e-6 vs 6.7e-5). Devices that cannot co-schedule 16-CTA clusters keep the fp32 router. +ROUTER_TC_GEN_MIN_ROWS = 9 # capacity: fixed: past the MoE engine's 8 rows +# Up to 48 rows, as the agg canon runs it (v1.8.3, the router GEMV keys gone): the MoE of a 9..48-row step routes on +# it, with the routed experts' NVFP4 activation from the same kernel (the decode layer's MC finalize, the generic +# layer). GB300 tp4, behind a PDL producer, with the NVFP4 activation: 3.20 / 3.82 / 7.36 us at 8 / 12 / 48 tokens. +ROUTER_TC_GEN_MAX_ROWS = 48 # capacity: fixed: the router kernel's widest build +ROUTER_TC_GEN = RowRange( + "m3_router_tc", + _rows(ROUTER_TC_GEN_MIN_ROWS, ROUTER_TC_GEN_MAX_ROWS), + m3_router_tc, + _spec_rows("m"), # every cell checks the logits, with or without the NVFP4 activation +) +MOE_PUSH = RowRange( + "moe_push", + _rows(1, MOE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MOE_ENGINE = RowRange( + "m3_moe_engine", _rows(1, MOE_ENGINE_MAX_TOKENS), m3_moe_engine, _spec_rows("tokens") +) +# The GEMM paths are keyed (projection (N, K), rows): each projection runs its own split-K / ring / MMA +# configuration, so a row count is certified per projection. +MX_GEMV = RowRange( + "m3_mx_gemv", + frozenset((nk, m) for nk, rows in MX_GEMV_ROWS.items() for m in rows), + m3_mx_gemv, + lambda spec: ((spec["nk"], spec["m"]),), +) +MX_GEMV_BF16 = RowRange( + "m3_mx_gemv_bf16", + frozenset((nk, m) for nk, hi in MX_GEMV_BF16_MAX_TOKENS.items() for m in _rows(1, hi)), + m3_mx_gemv_bf16, + lambda spec: ((spec["nk"], spec["m"]),), +) +# (draft length, requests) of a greedy verify step. +SHARD_ACCEPT = RowRange( + "m3_shard_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SHARD_ACCEPT_MAX_REQUESTS + 1)), + m3_shard_accept, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) + +# Rows of a fused-draft-loop all-reduce boundary (drafter hidden 6144) on the multicast all-reduce. +DRAFT_MC_AR = RowRange( + "m3_draft_mc_allreduce_norm", + _rows(1, DRAFT_MC_AR_MAX_ROWS), + m3_mc_allreduce_norm, + lambda spec: (spec["m"],) if spec.get("drafter") else (), +) + +# (draft length, requests) of a sampled verify step. +SAMPLED_ACCEPT = RowRange( + "m3_sampled_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SAMPLED_VERIFY_MAX_REQUESTS + 1)), + m3_sampled_accept, + lambda spec: ((SHARD_ACCEPT_DRAFT_LEN, spec["nreq"]),), +) + + +# (query rows per request, requests) of a decode step whose dense-layer attention runs on PrimTS split-KV decode. +FULLCTX_DECODE = RowRange( + "m3_fullctx_decode", + frozenset(FULLCTX_DECODE_RECIPES) if FULLCTX_DECODE_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + +# (query rows per request, requests) of a generation-only drafter pass that runs on the full-context decode. +DRAFTER_FULLCTX = RowRange( + "m3_drafter_fullctx", + frozenset(FULLCTX_DECODE_RECIPES) if DRAFTER_FULLCTX_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + + +def _verify_rows(max_rows: int) -> FrozenSet[int]: + """The verify rows of 1..8 requests at the shard-accept draft length, up to ``max_rows``.""" + return frozenset( + n * (SHARD_ACCEPT_DRAFT_LEN + 1) + for n in range(1, 9) + if n * (SHARD_ACCEPT_DRAFT_LEN + 1) <= max_rows + ) + + +# Requests of a fused-verify step whose first drafter inputs are the fused ones; certified by the chunk-norm cells +# at those requests' verify rows (the norm is the only kernel of its own; the gather rows and the fc are in-tree). +FUSED_FIRST_INPUTS = RowRange( + "m3_fused_first_inputs", + _rows(1, FUSED_FIRST_INPUTS_MAX_REQUESTS), + m3_chunk_norm, + lambda spec: ( + (spec["m"] // (SHARD_ACCEPT_DRAFT_LEN + 1),) + if spec["m"] % (SHARD_ACCEPT_DRAFT_LEN + 1) == 0 + else () + ), +) +# Verify rows of a fused-verify step whose Eagle3 captures apply the fc_norm. +CAPTURE_NORM = RowRange( + "m3_capture_norm", _verify_rows(CAPTURE_NORM_MAX_ROWS), m3_capture_norm, _spec_rows("m") +) +# Verify rows of a first draft pass whose fc_norm runs as the chunk norm. +CHUNK_NORM = RowRange( + "m3_chunk_norm", _verify_rows(CHUNK_NORM_MAX_ROWS), m3_chunk_norm, _spec_rows("m") +) + +# ((N, K), rows) of a drafter BF16 linear on its low-M kernel. +DRAFT_LOWM = RowRange( + "m3_draft_lowm", + frozenset((nk, m) for nk, rows in DRAFT_LOWM_ROWS.items() for m in rows), + m3_draft_lowm, + lambda spec: ((spec["nk"], spec["m"]),), +) +# Verify rows of a step whose target logits come from the die-placed head (in-tree GEMM halves per die). +UGPU_HEAD = RowRange("m3_ugpu_head", frozenset(UGPU_HEAD_ROWS), m3_ugpu_head, _spec_rows("m")) + +# (draft length, requests) of a fused draft loop, and the rows of its passes' draft input: the first pass's +# requests x (draft length + 1), one per request on the later passes. +DRAFT_STEP = RowRange( + "m3_draft_step", + frozenset((DRAFT_STEP_DRAFT_LEN, n) for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1)), + m3_draft_step, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) +DRAFT_INPUT = RowRange( + "m3_draft_input", + frozenset( + rows + for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1) + for rows in (n, n * (DRAFT_STEP_DRAFT_LEN + 1)) + ), + m3_draft_input, + _spec_rows("m"), +) + +FAST_PATHS: Tuple[RowRange, ...] = ( + EMBED_FIRSTNORM, + FULLCTX_DECODE, + DRAFTER_FULLCTX, + SPARSE_DECODE_PRODUCER, + SPARSE_DECODE_LC, + PRODUCER, + INDEX_SCORE, + LONG_SELECT, + LONG_SELECT_DSL, + MC_ALLREDUCE, + MC_FINALIZE, + ATTN_PUSH, + DENSE_PUSH, + MOE_PUSH, + MOE_ENGINE, + ROUTER_TC_GEN, + MX_GEMV, + MX_GEMV_BF16, + SHARD_ACCEPT, + SAMPLED_ACCEPT, + DRAFT_MC_AR, + UGPU_HEAD, + DRAFT_STEP, + DRAFT_INPUT, + DRAFT_LOWM, + FUSED_FIRST_INPUTS, + CAPTURE_NORM, + CHUNK_NORM, +) + +# -------------------------------------------------------------------------------------------------------------------- +# Selection: pure functions of the step shape +# -------------------------------------------------------------------------------------------------------------------- + + +def plan_for(num_tokens: int) -> DecodePlan: + for plan in DECODE_PLANS: + if num_tokens <= plan.max_tokens: + return plan + raise AssertionError(f"no decode plan covers {num_tokens} tokens") + + +def is_decode_step(num_tokens: int) -> bool: + """Whether a pure generation step of ``num_tokens`` rows goes to the decode target.""" + return num_tokens <= DECODE_MAX_TOKENS + + +def moe_path(num_tokens: int) -> Optional[str]: + """The MoE kernel of a decode step: the engine, the unfinalized MoE + finalize all-reduce, or None (generic).""" + if MOE_ENGINE.covers(num_tokens): + return MOE_ENGINE.path + if MC_FINALIZE.covers(num_tokens): + return MC_FINALIZE.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# The dense layers' swigluoai activation as this target's C++ op (ops/gemm/m3_dense_swiglu, +# minimaxM3Sm107Tp2/swigluOai.cu, LongAttn c267035d2f: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, +# beta 1), PDL-launched with an early release so the down GEMV stages its weights during it; the down GEMV reads the +# activation only after its own grid wait. Indexer kdev17 783285 (TP2 chain): -0.69 us per dense layer vs the Triton +# kernel; LongAttn gate a 783511 (oracle 23/23). False runs the Triton kernel. +DENSE_SWIGLU_CPP_PATH = "m3_dense_swiglu" +# FlashInfer's CuTe DSL MXFP8 GEMM on the producer's MXFP8 activation (rows the GEMV does not serve). +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +# The MXFP8 quantize (tuned backend) and FlashInfer's CuTe DSL GEMM: every row the decode kernels do not serve. +GENERIC_GEMM_PATH = "mxfp8_gemm" + + +def linear_path(n: int, k: int, num_tokens: int, prequant: bool, low_m_gemm: bool) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. ``low_m_gemm`` is accepted for the call sites' + symmetry: no low-M GEMM configuration is certified at TP2, so it selects nothing. + """ + if prequant: + if num_tokens in MX_GEMV_ROWS.get((n, k), ()): + return MX_GEMV_PATH + return PREQUANT_GEMM_PATH + if num_tokens <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0): + return MX_GEMV_BF16_PATH + if num_tokens in MX_GEMV_QUANT_ROWS.get((n, k), ()): + # W-G1: the decode-graph MXFP8 quantize, then the MXFP8 GEMV. + return MX_GEMV_PATH + return GENERIC_GEMM_PATH + + +# -------------------------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A covered decode shape could not take its fast path under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() + + +def set_require(require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def missed(path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if _require[0]: + raise FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " + f"{reason}" + ) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def taken(path: str, rows: Optional[int], what: str = "") -> None: + """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line + every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" + at = "" if rows is None else f" at {rows} rows" + detail = f" ({what})" if what else "" + _log_once( + logger.info, + path, + f"taken: {rows}" + (f": {what}" if what else ""), + f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", + ) + + +def long_select_covers(num_blocks: int) -> bool: + """Whether ``LONG_SELECT`` takes rows of ``num_blocks`` KV blocks (the engine's per-row block width).""" + return LONG_SELECT_MIN_BLOCKS <= num_blocks <= LONG_SELECT_MAX_BLOCKS + + +def long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: + """The (PDL launch, early trigger, trigger phase) keys of the 384-thread select a decode step can launch, compiled + before warmup; none when ``LONG_SELECT_DSL`` covers nothing.""" + if not LONG_SELECT_DSL.keys: + return () + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + pdl = get_env_enable_pdl() + keys = [(pdl, False, SELECT_LONG_TRIGGER_PHASE)] + if SELECT_EARLY_MAX_ROWS: + keys.append((pdl, True, 0)) + return tuple(keys) + + +def long_select_fallback(num_blocks: int, rows: int) -> None: + """The in-tree selector is about to run instead of the long-row select: a miss when ``LONG_SELECT`` covers + ``num_blocks`` (the kernel declined operands its cells certify; raises under require), uncovered otherwise (rows of + at most 128 blocks, or wider than ``LONG_SELECT_MAX_BLOCKS``: the in-tree selector serves them in every mode).""" + if long_select_covers(num_blocks): + missed( + LONG_SELECT.path, + rows, + f"the long-row select declined {num_blocks}-block scores it covers", + ) + else: + span = f"{LONG_SELECT_MIN_BLOCKS}..{LONG_SELECT_MAX_BLOCKS}" + uncovered( + LONG_SELECT.path, f"{num_blocks} KV blocks, outside {span}: the built-in selector" + ) + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +# Kernel modules report from inside their calls through ops/m3_report: the same policy. +m3_report.install( + on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed +) + + +# -------------------------------------------------------------------------------------------------------------------- +# Target constants +# -------------------------------------------------------------------------------------------------------------------- + +# The generic boundary takes the low-latency multicast all-reduce + residual + RMSNorm up to this many tokens; with +# the shared + routed add folded in, up to its own maximum: the decode boundaries' certified rows. +MC_AR_MAX_TOKENS = GENERIC_MC_AR_MAX_TOKENS +MC_AR_FOLDED_MAX_TOKENS = GENERIC_MC_AR_FOLDED_MAX_TOKENS +# The shared expert folds into the routed engine from one token (dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# SMs left to the concurrent kernels while a persistent routed engine runs. +MOE_ENGINE_RESERVED_SMS = 24 # capacity: fixed: an SM split, not a size +# Id of the target's final hidden state in Eagle3's ``layers_to_capture``. +FINAL_HIDDEN_LAYER_ID = -1 # capacity: fixed: an Eagle3 layer id, not a size +# The trtllm-gen MoE runner's routing kind for MiniMax-M3 (sigmoid + correction bias, renormalized, scaled) and its +# gated activation (SwiGLU with the per-expert alpha / beta / clamp operands). +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) +INDEX_SCORE_PATH = INDEX_SCORE.path + +# -------------------------------------------------------------------------------------------------------------------- +# Streaming: also emit a response on decoding iterations 2..FIRST_CONTENT_EMIT_ITERS (RuntimeHooks), so a request's +# first content leaves one step after its first token. 1: the stock schedule. +FIRST_CONTENT_EMIT_ITERS = 2 + +# The runtime-value dump (TRTLLM_MODELING_V2_VALUE_DUMP=1): the metadata of every fast-path and entry decision a step +# takes, recorded on eager steps and captures (replay runs no Python), printed once at exit. +# -------------------------------------------------------------------------------------------------------------------- + +_VALUE_DUMP = os.environ.get("TRTLLM_MODELING_V2_VALUE_DUMP") == "1" +_VALUES: set = set() + + +def _record(*meta: Hashable) -> None: + if _VALUE_DUMP: + _VALUES.add(meta) + + +def _print_values() -> None: + if _VALUES: + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "0")) + for meta in sorted(_VALUES, key=repr): + print(f"[modeling_v2 value dump] rank {rank}: {meta}", flush=True) + + +if _VALUE_DUMP: + atexit.register(_print_values) + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing + sees only its KV cache algorithm; this check at construction refuses every other runtime. + + The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index + projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. + """ + sparse = model_config.sparse_attention_config + if ( + sparse is None + or getattr(sparse, "algorithm", None) != "minimax_m3" + or sparse.implementation != "msa" + ): + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs +# -------------------------------------------------------------------------------------------------------------------- + + +def _mxfp8_linear( + core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + _record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +_PINNED_RUNNERS: dict = {} + + +def _pinned_runner(cls, dtype): + runner = _PINNED_RUNNERS.get((cls, dtype)) + if runner is None: + runner = _PINNED_RUNNERS[(cls, dtype)] = cls(dtype) + return runner + + +def _pinned_mxfp8_gemm(act, act_scale, weight, weight_scale, out_dtype): + """The decode-graph MXFP8 GEMM at the target's fixed tactic (``MXFP8_DECODE_GEMM_TACTIC``), no AutoTuner: the + same choice on every rank and every start.""" + _record("pinned_mxfp8_gemm", tuple(weight.shape), act.shape[0], MXFP8_DECODE_GEMM_TACTIC) + runner = _pinned_runner(FlashInferMXFP8GemmRunner, out_dtype) + return runner.forward([act, act_scale, weight, weight_scale], tactic=MXFP8_DECODE_GEMM_TACTIC) + + +def _pinned_mxfp8_quantize(x): + """The decode-graph MXFP8 activation quantize at the target's fixed tactic (``MXFP8_DECODE_QUANT_TACTIC``).""" + _record("pinned_mxfp8_quantize", x.shape[0], MXFP8_DECODE_QUANT_TACTIC) + return _pinned_runner(MXFP8QuantizeRunner, x.dtype).forward( + [x], tactic=MXFP8_DECODE_QUANT_TACTIC + ) + + +_MX_TAC_ENGAGED: set = set() + + +def _note_mx_tac(kind: str, n: int, k: int, rows: int) -> None: + """MX-TAC's engagement line, once per (activation kind, projection, rows).""" + key = (kind, n, k, rows) + if key not in _MX_TAC_ENGAGED: + _MX_TAC_ENGAGED.add(key) + logger.info( + f"[M3 sm107 tp2] MX-TAC engaged: pinned MXFP8 decision (quantize tactic " + f"{MXFP8_DECODE_QUANT_TACTIC}, GEMM tactic {MXFP8_DECODE_GEMM_TACTIC}) " + f"for the [{n}, {k}] projection at {rows} rows, {kind} activation" + ) + + +def _gemv_variants(n: int, k: int, push: Optional[str], world: int): + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch routes to an [N, K] + projection's GEMV: an MXFP8 activation on the rows ``MX_GEMV_ROWS`` gives it, a bf16 one up to + ``MX_GEMV_BF16_MAX_TOKENS``; each pushing when the projection's boundary push covers the rows. ``push``: "attn" + (the o_proj into the attention boundary), "dense" (the dense MLP's down projection into the post-MLP boundary) + or "moe" (the shared expert's down projection into the MoE boundary).""" + push_rows = {"attn": ATTN_PUSH.keys, "dense": DENSE_PUSH.keys, "moe": MOE_PUSH.keys}.get( + push, frozenset() + ) + variants = [] + max_rows = max([DECODE_MAX_TOKENS, *(max(r) for r in MX_GEMV_ROWS.values() if r)]) + for rows in range(1, max_rows + 1): + for x_bf16, served in ( + (False, rows in MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + bucket = next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + variants.append((rows, x_bf16, (world, bucket, AR_PUSH_MODE))) + return variants + + +def _uncompiled_note() -> str: + key = mx_gemv.last_uncompiled + return "" if key is None else f" (its kernel {key} was not compiled before warmup)" + + +def _linear( + core, x: torch.Tensor, wk: str, sk: str, low_m_gemm: bool, push: Optional[tuple] = None +) -> Tuple[torch.Tensor, bool]: + """One MXFP8 projection ``x @ w[wk]^T``; returns (output, whether its GEMV pushed into ``push``). + + Outside decode CUDA graphs: the engine-tuned MXFP8 quantize + GEMM. Inside them the projection's (N, K) and the + step's rows pick one path (``linear_path``): + + 1. an activation its producer already quantized to MXFP8 (``m3_prequant.SLOT``, taken here): the tcgen05 decode + GEMV on the rows it serves, else the decode-graph MXFP8 GEMM on that activation; + 2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; + 3. W-G1: a bf16 activation on the rows the MXFP8 GEMV serves but the bf16 GEMV does not + (``MX_GEMV_QUANT_ROWS``): the decode-graph MXFP8 quantize, then the MXFP8 GEMV; + 4. otherwise the decode-graph MXFP8 quantize + GEMM. + + MX-TAC: on a decode graph whose MXFP8 projections the engine tunes (``tune_decode_graph_backends``), the GEMM + (and, on a bf16 activation, the quantize) the GEMVs do not serve take one fixed decision + (``MXFP8_DECODE_QUANT_TACTIC`` / ``MXFP8_DECODE_GEMM_TACTIC``: FlashInfer's CuTe DSL quantize and CuTe DSL GEMM) + instead of the AutoTuner's per-start choice (whose L2-hot profiling left c2 bimodal at +-6%). On SM107 every TP2 + projection at 1..16 rows runs fastest on this pair, cold and L2-hot (graph us, cold: 4.7..10.3 against 5.0..11.5 + with the TRT-LLM quantize and 9.7..58 with CUTLASS), and all four candidates give the same bits (quantize bytes + and scales, GEMM outputs). A constant is rank-agreed and start-independent by construction. + + A kernel declining rows its path covers is a miss; the chain then continues down the list. ``push`` (an + ``m3_mx_gemv`` ``ar`` tuple) is the boundary all-reduce exchange the GEMV writes its rows into. ``low_m_gemm`` + marks the projections the fused low-M MXFP8 GEMM may take (``LOW_M_GEMM_MAX_TOKENS``: no TP2 configuration is + built, so ``linear_path`` never selects it). + """ + w, s = core.w[wk], core.w[sk] + x = x.contiguous() + mode = _mx_mode(core.mxfp8) + if not mode.decode_graph: + return _mxfp8_linear(core, x, w, s, mode), False + use_pdl = get_env_enable_pdl() + rows = x.shape[0] + n, k = w.shape + prequant = m3_prequant.SLOT.take(x) + if prequant is None and x.dtype != torch.bfloat16: + path = GENERIC_GEMM_PATH + else: + path = linear_path(n, k, rows, prequant is not None, low_m_gemm) + _record("linear", (n, k), rows, path, push is not None) + if prequant is not None: + act, act_scale = prequant + if path == MX_GEMV_PATH: + output = m3_mx_gemv(act, act_scale, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] mxfp8 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if path == PREQUANT_GEMM_PATH: + uncovered(MX_GEMV_PATH, f"the [{n}, {k}] projection at rows the GEMV does not serve") + if mode.tuned_quantize_gemm: + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("mxfp8", n, k, rows) + return output, False + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_PATH: + # W-G1 (MX_GEMV_QUANT_ROWS): no MXFP8 hand-off, so the decode-graph quantize the generic GEMM path would run + # (the same bytes), then the GEMV in place of that GEMM. + act, act_scale = _pinned_mxfp8_quantize(x) + output = m3_mx_gemv(act, act_scale, w, s, True, None, use_pdl, None) + if output is not None: + taken(path, rows, f"[{n}, {k}] mxfp8 quantize + GEMV") + return output, False + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's quantized activation" + + _uncompiled_note(), + ) + if mode.tuned_quantize_gemm: + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("mxfp8", n, k, rows) + return output, False + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16(x, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] bf16 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, rows, f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note() + ) + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + if mode.tuned_quantize_gemm and x.dtype == torch.bfloat16: + act, act_scale = _pinned_mxfp8_quantize(x) + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("bf16", n, k, rows) + return output, False + return _mxfp8_linear(core, x, w, s, mode), False + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer boundaries: all-reduce + residual + RMSNorm +# -------------------------------------------------------------------------------------------------------------------- + + +def _can_mc(x, folded: bool = False) -> bool: + """Whether the generic boundary takes the multicast all-reduce + residual + RMSNorm: the rows + ``MC_AR_MAX_TOKENS`` / ``MC_AR_FOLDED_MAX_TOKENS`` (with the shared + routed add folded in) give it.""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (MC_AR_FOLDED_MAX_TOKENS if folded else MC_AR_MAX_TOKENS) + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + +def _mc_allreduce_norm(core, x, x2, residual, norm_w, *, quant=False, prepushed=0, capture=None): + """The multicast all-reduce + optional second addend + residual + RMSNorm (``m3_mc_allreduce_norm``), or None + without a workspace or for operands the kernel does not take (the same on every rank). With ``quant`` the normed + rows are also published, MXFP8-quantized, to ``m3_prequant.SLOT`` for the next MXFP8 projection; ``capture`` = + (out, fc_norm weight, eps) also writes the layer's Eagle3 capture norm of the updated residual.""" + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + if workspace is None: + return None + _record( + "m3_mc_allreduce_norm", x.shape[0], x2 is not None, quant, prepushed, capture is not None + ) + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm_w, + core.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + capture=capture, + ) + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm.""" + _record("m3_allreduce", x.shape[0]) + return m3_allreduce( + x, + core.allreduce, + AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ), + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The MoE: router, routed experts (trtllm-gen NVFP4), shared expert +# -------------------------------------------------------------------------------------------------------------------- + + +@dataclass +class _MoeScalars: + """The routed experts' kernel scalars of one layer, derived from the checkpoint's per-tensor scales once the + weights are real (``MiniMaxM3Tp2ModelingV2._derive_moe``).""" + + fc31_input_scale: torch.Tensor + fc2_input_scale: torch.Tensor + fc31_alpha: torch.Tensor + fc2_alpha: torch.Tensor + fc31_scale_c: torch.Tensor + act_alpha: torch.Tensor + act_beta: torch.Tensor + act_clamp: torch.Tensor + + +def _router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The fp32 router logits of MoE layer ``i`` when the MoE runs unfused outside ``ROUTER_TC_GEN``'s rows (the decode + steps' engines route in-kernel): ``F.linear(float(h), w)``, the built-in gate's generation-step arithmetic; the + experts then quantize their bf16 input themselves.""" + _record("m3_fp32_linear", h.shape[0] if h.dim() == 2 else 0) + return m3_fp32_linear(h, core.w[f"l{i}_router"]) + + +def _decode_router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The router logits of a pure generation step on the generic layer: ``ROUTER_TC_GEN``'s rows take the TF32 + tensor-core router when the device co-schedules 16-CTA clusters (decided at prebuild); every other step the fp32 + router (``_router_logits``).""" + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + ROUTER_TC_GEN.covers(rows) + and core.router_tc_cluster16 + and h.dtype == torch.bfloat16 + and h.stride(1) == 1 + and (h.stride(0) * 2) % 16 == 0 + ): + taken(ROUTER_TC_GEN.path, rows) + _record("m3_router_tc", rows) + return m3_router_tc(h, core.w[f"l{i}_router"], None, get_env_enable_pdl())[0] + return _router_logits(core, i, h) + + +def _router_tc_fp4(core, i: int, h: torch.Tensor) -> Optional[tuple]: + """(router logits, NVFP4 activation, its scales) of an unfused MoE on ``ROUTER_TC_GEN``'s rows, a generation or a + context step: the TF32 tensor-core router with the routed experts' NVFP4 activation quantization fused in (the agg + canon's router on these rows); None elsewhere (the caller routes on its fp32 router).""" + rows = h.shape[0] if h.dim() == 2 else 0 + if not ( + ROUTER_TC_GEN.covers(rows) + and core.router_tc_cluster16 + and h.dtype == torch.bfloat16 + and h.stride(1) == 1 + and (h.stride(0) * 2) % 16 == 0 + ): + return None + taken(ROUTER_TC_GEN.path, rows, "with the NVFP4 activation") + _record("m3_router_tc", rows, True) + return tuple( + m3_router_tc( + h, core.w[f"l{i}_router"], core.moe[i].fc31_input_scale.reshape(1), get_env_enable_pdl() + ) + ) + + +def _decode_routed_unfinalized(core, i: int, h: torch.Tensor): + """The routed experts of a pure generation step, unfinalized (``[gemm2 out, expert scales, expanded->permuted + index]``): routed on the fused tensor-core router where it covers the rows, else on the router logits with the + experts quantizing their bf16 input.""" + fused = _router_tc_fp4(core, i, h) + if fused is not None: + logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, logits, False) + return _routed_experts(core, i, h, None, _decode_router_logits(core, i, h), False) + + +def _triton_finalize(rows: int, routed, shared: Optional[torch.Tensor]) -> torch.Tensor: + """The finalized routed output (+ ``shared``) from the runner's unfinalized parts via ops/m3_moe_finalize + (``MOE_TRITON_FINALIZE``). Raises if the parts are not the trtllm-gen layout (the runner did not finalize them: + no fallback).""" + out = None + if isinstance(routed, (list, tuple)) and len(routed) == 3: + out = m3_moe_finalize(routed[0], routed[1], routed[2], shared) + if out is None: + raise RuntimeError( + f"the routed experts' unfinalized output at {rows} rows is not the trtllm-gen layout" + ) + taken( + "m3_moe_finalize", rows, "Triton finalize" + (" + shared add" if shared is not None else "") + ) + return out + + +def _triton_finalize_rows(core, h: torch.Tensor) -> bool: + """Whether a finalized MoE forward of ``h`` takes ``MOE_TRITON_FINALIZE``: at least its rows, in one runner call.""" + rows = h.shape[0] if h.dim() == 2 else 0 + return ( + MOE_TRITON_FINALIZE + and rows >= MOE_TRITON_FINALIZE_MIN_TOKENS + and rows <= core.moe_max_num_tokens + ) + + +def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): + """The trtllm-gen NVFP4 MoE runner over this rank's tensor-parallel slice of every expert, routing in-kernel. + + ``x_sf`` None means ``x`` is the bf16 activation, quantized here with the experts' global input scale. Returns + the finalized rows, or the unfinalized ``[gemm2 out, expert scales, expanded->permuted index]``. + """ + moe = core.moe[i] + if x_sf is None: + x, x_sf = fp4_quantize(x, moe.fc31_input_scale, 16, False, False) + rows = x.shape[0] + x_sf = x_sf.view(rows, -1).flatten() + _record("fp4_block_scale_moe_runner", rows, do_finalize) + outs = fp4_block_scale_moe_runner( + router_logits, + core.w[f"l{i}_router_bias"], + x, + x_sf.view(torch.float8_e4m3fn), + core.w[f"l{i}_e_w31"], + core.w[f"l{i}_e_s31"].view(torch.float8_e4m3fn), + None, + moe.act_alpha, + moe.act_beta, + moe.act_clamp, + core.w[f"l{i}_e_w2"], + core.w[f"l{i}_e_s2"].view(torch.float8_e4m3fn), + None, + moe.fc31_scale_c, + moe.fc31_alpha, + moe.fc2_alpha, + core.num_experts, + core.top_k, + None, + None, + core.moe_inter, + 0, + core.num_experts, + core.routed_scaling_factor, + ROUTING_MINIMAX2, + do_finalize, + act_type=ACT_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=core.max_num_tokens, + use_dp=False, + ) + return outs[0] if do_finalize else outs + + +def _split_chunk(rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +# -------------------------------------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------------------------------------- + + +def _nvfp4_layer(manager, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _producer_covered(core) -> bool: + """Whether ``PRODUCER`` covers this target's sparse layers: an NVFP4 KV cache at the TP2 heads (32 query, two KV + and two index heads). A covered layer that cannot take ``m3_producer`` is a miss.""" + return core.main_kv_is_nvfp4 and (core.heads, core.kv_heads, core.index_heads) == (32, 2, 2) + + +def _producer_layout(kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 pools are the layout ``m3_producer`` is built for: [pages, 2, 2, 128, 64] with + its scale pool [pages, 2, 2, 128, 8].""" + return ( + kv_data is not None + and kv_scale is not None + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 2, 128, 64) + and kv_data.stride()[2:] == (128 * 64, 64, 1) + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 2, 128, 8) + and kv_scale.stride()[2:] == (128 * 8, 8, 1) + and kv_scale.shape[0] == kv_data.shape[0] + ) + + +def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or data.shape[0] != scales.shape[0] + or tuple(data.shape[1:]) != (2, num_heads, 128, 64) + or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) + or data.dtype not in (torch.int8, torch.uint8) + or scales.dtype != torch.uint8 + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or cache.stride(1) < num_heads * cache.stride(2) + or cache.stride(0) < 2 * cache.stride(1) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + + +#: The step phases ("decode" / "prefill") whose producer engagement has been logged, once per process. +_PRODUCER_ENGAGED: set = set() + + +def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: + if covered: + missed(PRODUCER.path, rows, reason) + return None + + +def _sparse_qkv_producer( + core, i: int, packed, position_ids, kv_cache_manager, msa_idx_k_cache, msa_out_cache_loc +): + """The horizontal producer of a sparse layer: Q / index-Q (FP8) out, K / V / index-K written into the paged + caches (NVFP4 or FP8 main cache), for every step type (token-major). An NVFP4 cache at the TP2 heads takes + ``m3_producer`` (``PRODUCER``); a covered layer that cannot is a miss. Returns (q, index_q), or None for caches + neither kernel takes. + + ``kv_cache_manager``, ``msa_idx_k_cache`` (the per-layer index-K cache getter) and ``msa_out_cache_loc`` are the + calling target's ``step_args`` entries.""" + covered = _producer_covered(core) + rows = int(packed.shape[0]) + if packed.dtype != torch.bfloat16 or position_ids is None: + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + manager = kv_cache_manager + if manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = manager.get_buffers(i, kv_layout="HND") + index_k_cache = msa_idx_k_cache(i) + out_cache_loc = msa_out_cache_loc + layer_uses_nvfp4 = _nvfp4_layer(manager, i) + scale_buffers = ( + manager.get_block_scale_buffers(i, kv_layout="HND") if layer_uses_nvfp4 else None + ) + kv_quant_scale = core.inv_kv_scales[i] if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supported_main = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_fused_write_layout(buffers, scale_buffers, core.kv_heads) + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale.is_cuda + and kv_quant_scale.dtype == torch.float32 + and kv_quant_scale.numel() >= 3 + and kv_quant_scale.is_contiguous() + ) + else: + supported_main = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, core.kv_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= core.kv_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + cos_sin = core.rotary_cos_sin + supported_rope = ( + cos_sin.is_cuda + and cos_sin.dtype == torch.float32 + and cos_sin.is_contiguous() + and cos_sin.dim() == 3 + and tuple(cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main + or not supported_index + or not supported_rope + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < rows + ): + return _producer_miss( + covered and layer_uses_nvfp4, + rows, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the layout the " + "producer is built for", + ) + w = core.w + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + layer_uses_nvfp4 + and (core.heads, core.kv_heads, core.index_heads) == (32, 2, 2) + and _producer_layout(buffers, scale_buffers) + ): + _record("m3_producer", rows) + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + core.heads, + get_env_enable_pdl(), + PRODUCER_EARLY_TRIGGER, + PRODUCER_PREWAIT_STEP_INPUTS, + ) + phase = "decode" if rows <= DECODE_MAX_TOKENS else "prefill" + if phase not in _PRODUCER_ENGAGED: + _PRODUCER_ENGAGED.add(phase) + logger.info( + f"[M3 sm107 tp2] fast path engaged: m3_producer ({phase}, T={rows}, " + f"early_trigger={PRODUCER_EARLY_TRIGGER}, prewait_step_inputs={PRODUCER_PREWAIT_STEP_INPUTS})" + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool {tuple(scale_buffers.shape)} is not the " + "producer's layout", + ) + _record("m3_qkv_insert", rows, layer_uses_nvfp4) + q, index_q = m3_qkv_insert( + packed.contiguous(), + buffers.view(torch.uint8) if layer_uses_nvfp4 else buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.heads, + core.kv_heads, + core.index_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + ) + return q.flatten(1), index_q.flatten(1) + + +def _qk_norm_rope(core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool) -> torch.Tensor: + """The per-head Gemma RMSNorm of Q and K and the partial NeoX RoPE of a dense layer's packed [Q | K | V], in one + kernel: in place, or an FP8 copy (Q / K normed and roped, V cast) when ``out_fp8``.""" + qkv = qkv.contiguous() + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + args = ( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + 1.0, # factor: no YARN + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + _record("qk_norm_rope", qkv.shape[0], out_fp8) + if out_fp8: + return m3_qk_norm_rope_fp8(*args) + return fused_qk_norm_rope(*args) + + +def _dense_fp8_op_applies(qkv: torch.Tensor, position_ids) -> bool: + """Whether a dense layer's FP8 q/k/v come from this target's op (``minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8``, + else the BF16 fused_qk_norm_rope output cast to FP8).""" + return DENSE_QKV_FP8_VIA_BF16 and position_ids is not None and qkv.dtype == torch.bfloat16 + + +def _dense_qk_norm_rope_fp8( + core, i: int, qkv: torch.Tensor, position_ids, kv_dest: Optional[tuple] = None +) -> torch.Tensor: + """A dense layer's FP8 q/k/v: QK norm + RoPE with Q/K rounded to BF16 and then to E4M3 and V cast, the bits of the + BF16 fused_qk_norm_rope output cast to FP8 (K is stored in the FP8 cache), in one launch of this target's op + (``DENSE_QKV_FP8_VIA_BF16``); outside its supported inputs, the BF16 kernel and the cast. + + With ``kv_dest`` = (the layer's HND pool view, the step's int32 cache slots, live tokens, a [tokens, q heads, 128] + FP8 buffer) the op also stores the live tokens' K and V rows into the pool (the bytes ``write_msa_phase_kv`` + stores, one launch instead of about ten) and every row's Q into the buffer; ``kv_dest`` is passed only where the + op is known to run (``_dense_fp8_op_applies``).""" + if not _dense_fp8_op_applies(qkv, position_ids): + return fp8_e4m3_saturate(_qk_norm_rope(core, i, qkv, position_ids, False)) + _record("m3_dense_qk_norm_rope_fp8", qkv.shape[0], kv_dest is not None) + args = ( + qkv.contiguous(), + core.heads, + core.kv_heads, + core.kv_heads, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, # NeoX RoPE + position_ids.reshape(-1).contiguous().to(torch.int32), + True, # Gemma norms + ) + if kv_dest is not None: + return torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv(*args, *kv_dest) + return torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + + +def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: + """The dense MLP's SwiGLU-OAI activation: this target's C++ op (``m3_dense_swiglu``, minimaxM3Sm107Tp2/ + swigluOai.cu: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, beta 1, PDL-launched with an early + release so the down GEMV stages its weights during it) at the rows its cells certify, in any graph or eager step; + otherwise the Triton kernel (``m3_silu_and_mul``). Indexer kdev17 (TP2 chain): -0.69 us per dense layer vs the + Triton kernel. Every decline is logged once.""" + path = DENSE_SWIGLU_CPP_PATH + rows = gu.shape[0] + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + params = (core.swiglu_alpha, core.swiglu_limit, core.swiglu_beta) + act = None + if rows > max_rows: + uncovered(path, f"{rows} rows, past its certified {max_rows}") + elif params != (DENSE_ALPHA, DENSE_LIMIT, DENSE_BETA): + uncovered(path, f"activation parameters {params}") + else: + act = m3_dense_swiglu( + gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + ) + if act is None: + uncovered(path, "the gate_up layout") + else: + _record("m3_dense_swiglu", rows) + taken(path, rows, "dense swigluoai") + if act is None: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return act + + +def _fp8_index_qk_norm_rope( + core, i: int, idx_qk: torch.Tensor, position_ids, msa_idx_k_cache, msa_out_cache_loc +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + _record("m3_fp8_indexer_qk_norm_rope", rows) + return m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer helpers both targets' bodies call: the Eagle3 capture hook and the norm after a layer +# -------------------------------------------------------------------------------------------------------------------- + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + +def _next_norm(core, i: int) -> torch.Tensor: + """The norm the boundary after layer ``i`` applies: the next layer's input norm, the final norm after the last.""" + return core.w[f"l{i + 1}_in_norm"] if i + 1 < core.num_layers else core.w["final_norm"] + + +# -------------------------------------------------------------------------------------------------------------------- +# Kernel prebuilds the core runs before the executor's warmup +# -------------------------------------------------------------------------------------------------------------------- + + +def _in_tree_scorer(): + """The in-tree CuTe DSL index scorer's runner (importing it registers ``cute_dsl_minimax_m3_index_decode_score``); + imported on first use, as the CuTe DSL ops are.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLMiniMaxM3IndexDecodeScoreRunner, + ) + + return CuteDSLMiniMaxM3IndexDecodeScoreRunner + + +def _dummy_scorer_operands(dtype, heads, qlen, dim, page, device): + """One request of two blocks: real-shaped operands that load a compiled scorer module.""" + rows, blocks = 1, 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros(rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view(rows, blocks) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + return idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + + +#: Mixed (context + generation) steps run their generation rows on the built-in sparse decode +#: (``MsaDecodeFmha._run_sparse``). Its Triton producer and merge specialize on the split-K factor (16 / 8 / 4 / 2 / 1 +#: by the row count on SM107), on the row count's divisibility and on the 16-byte alignment of the request rows (a +#: mixed step's generation requests follow its context requests, so its block-table and length rows start mid-buffer), +#: so ``_prebuild_builtin_sparse_decode`` launches them per decode row count up to this many rows, both alignments. +BUILTIN_SPARSE_DECODE_PREBUILD_MAX_ROWS = ( # capacity: fixed: the prebuilt built-in decode rows + 64 # capacity: fixed: the mixed steps' generation rows prebuilt +) + + +def _prebuild_builtin_sparse_decode(query_lens, num_heads, num_kv_heads, device) -> list: + """Compile (and load) the built-in sparse decode's Triton producer and merge for every generation row count a mixed + step can have, so no step compiles one after the server is ready. Each launch mirrors the operands + ``MsaDecodeFmha._run_sparse`` passes for an NVFP4 layer (FP8 q, the HND pool's K / V planes, the scale pool's + planes, the head-major top-k view, BF16 output) on dummy values, on the current stream. Returns the row counts.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode, + ) + + page, head_dim, topk = 128, 128, MSA_REQUIRED_TOPK + pages = topk + pool = torch.zeros( + pages, 2, num_kv_heads, page, head_dim // 2, dtype=torch.uint8, device=device + ) + scales = torch.zeros( + pages, 2, num_kv_heads, page, head_dim // 16, dtype=torch.uint8, device=device + ) + global_scales = torch.ones(2, 4, dtype=torch.float32, device=device) + done = [] + for qlen in sorted(set(query_lens)): + for rows in range(qlen, BUILTIN_SPARSE_DECODE_PREBUILD_MAX_ROWS + 1, qlen): + requests = rows // qlen + q = torch.zeros(rows, num_heads, head_dim, dtype=torch.bfloat16, device=device).to( + torch.float8_e4m3fn + ) + indexes = ( + torch.arange(topk, dtype=torch.int32, device=device) + .repeat(rows, num_kv_heads, 1) + .contiguous() + ) + # One leading context request: its rows are skipped, so the generation rows start mid-buffer. + block_rows = torch.arange(topk, dtype=torch.int32, device=device).repeat( + requests + 1, 1 + ) + lengths = torch.full((requests + 1,), topk * page, dtype=torch.int32, device=device) + output = torch.empty(rows, num_heads, head_dim, dtype=torch.bfloat16, device=device) + for first in (0, 1): + minimax_m3_sparse_attn_decode( + q, + pool[:, 0], + pool[:, 1], + indexes.permute(1, 0, 2), + block_rows[first : first + requests], + lengths[first : first + requests], + sm_scale=head_dim**-0.5, + output=output, + decode_query_len=qlen, + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=global_scales[0, :1], + v_global_scale=global_scales[1, :1], + ) + done.append(rows) + return done + + +# -------------------------------------------------------------------------------------------------------------------- +# The core +# -------------------------------------------------------------------------------------------------------------------- + +# Engine states (compiled kernel + workspace) per (device, token count[, shared fold]), shared by every layer: a step +# runs the layers one after another on one stream. +_ENGINE_STATES: Dict[tuple, object] = {} + + +class MiniMaxM3Tp2ModelingV2(DecoderModel): + """The MiniMax-M3 text model of this target: its weights, its runtime resources, and the per-step dispatch.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_runtime(model_config) + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + self.mapping = mapping + # Whether the 9..16-row generation router takes the TF32 tensor-core kernel (set by prebuild_kernels). + self.router_tc_cluster16 = False + self.vocab_size = cfg.vocab_size + d = _weights.MODEL_WEIGHTS.dims(self) + self.num_layers = d.num_layers + self.hidden = d.hidden + self.heads = d.heads + self.kv_heads = d.kv_heads + self.head_dim = d.head_dim + self.index_heads = d.index_heads + self.index_dim = d.index_dim + self.q_size = d.heads * d.head_dim + self.kv_size = d.kv_heads * d.head_dim + self.moe_inter = d.moe_inter + self.num_experts = d.num_experts + self.top_k = int(cfg.num_experts_per_tok) + self.routed_scaling_factor = float(getattr(cfg, "routed_scaling_factor", 1.0)) + self.rms_eps = float(cfg.rms_norm_eps) + self.swiglu_alpha = float(getattr(cfg, "swiglu_alpha", 1.702)) + self.swiglu_beta = 1.0 # the (up + 1) offset of SwiGLU-OAI + self.swiglu_limit = float(getattr(cfg, "swiglu_limit", 7.0)) + self.max_num_tokens = model_config.max_num_tokens + self.moe_max_num_tokens = model_config.moe_max_num_tokens + self.sparse_layers = frozenset(_sparse_layer_ids(cfg)) + self.moe_layers = frozenset(_moe_layer_ids(cfg)) + mode = model_config.quant_config.quant_mode + self.main_kv_is_fp8 = bool(mode is not None and mode.has_fp8_kv_cache()) + self.main_kv_is_nvfp4 = bool(mode is not None and mode.has_fp4_kv_cache()) + # The sparse layers' QKV and index projections: one five-way packed projection, or the QKV projection and a + # separate BF16 index projection (fuse_qkv_index_projection=False) run on two streams. + self.fused_projection = bool( + getattr(model_config.sparse_attention_config, "fuse_qkv_index_projection", False) + ) + + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # The engine reads these two off the module tree: the MXFP8 GEMM state it tunes, and the opt-in to the + # automatic FlashInfer MXFP8 path on decode graphs. + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + + # The built-in fused all-reduce + residual + RMSNorm of the layer boundaries the multicast kernel does not take. + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + + # Streams: the shared / routed expert overlap and, when the MoE chunks its input, the chunking overlap. + self.aux_stream_moe = torch.cuda.Stream() + chunked = self.moe_max_num_tokens < model_config.max_num_tokens * mapping.dp_size + self.aux_stream_moe_chunk = torch.cuda.Stream() if chunked else None + self.moe_events = {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers} + self.moe_chunk_events = { + i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers + } + # The unfused sparse layers' index branch overlaps the QKV branch on its own stream. + self.aux_stream_attn = None if self.fused_projection else torch.cuda.Stream() + self.attn_events = ( + {} + if self.fused_projection + else {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.sparse_layers} + ) + + # Partial NeoX RoPE (rotary_dim of head_dim): the cos/sin table every attention kernel reads. + rope = RopeParams.from_config(cfg) + self.rope_theta = rope.theta + self.rotary_dim = int(rope.dim) + self.rotary_cos_sin = rope.create_rope_const_params(interleave=False)[1].reshape( + rope.max_positions, 2, -1 + ) + + # One MSA attention backend per layer (dense and sparse; the step metadata is the backend's). + quant_config = model_config.get_quant_config() + self.attn: List[MiniMaxM3MsaSparseAttention] = [] + for i in range(self.num_layers): + backend = create_attention( + model_config.attn_backend, + i, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=None, + quant_config=quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + q_scaling=1.0, + sparse_params=model_config.sparse_attention_config.to_sparse_params( + pretrained_config=cfg, layer_idx=i + ), + ) + backend.update_quant_config(quant_config) + if not isinstance(backend, MiniMaxM3MsaSparseAttention): + raise ValueError( + f"layer {i}: expected the MiniMax-M3 MSA backend, got {type(backend).__name__}" + ) + if i not in self.sparse_layers: + # The KV split trtllm-gen dense decode runs at, per request count. + backend.dense_decode_max_seq_len_by_batch = DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + self.attn.append(backend) + + # Built once the weights are real (``post_load_weights``): the KV (Q, K, V) scales, the MoE scalars, the + # decode engines and workspaces, the two targets. + self.kv_scales: Dict[int, torch.Tensor] = {} + self.inv_kv_scales: Dict[int, torch.Tensor] = {} + self.moe: Dict[int, _MoeScalars] = {} + self.mx_alpha: Optional[torch.Tensor] = None + self.moe_engines: Dict[int, Dict[Tuple[int, bool, bool], Optional[object]]] = {} + self.embed_owners: Dict[tuple, DecodeEmbedding] = {} + # Each sparse layer's selected block ids per (KV head, query row, slot) at its last long-context decode, which + # that decode prefetches into L2 before its grid wait (``SPARSE_DECODE_LC_PREFETCH_SELF_KV``, up to + # ``N1_MAX_ROWS`` query rows). + self.n1_ids: Dict[int, torch.Tensor] = ( + { + i: torch.full( + (self.kv_heads, N1_MAX_ROWS, MSA_REQUIRED_TOPK), + -1, + dtype=torch.int32, + device=torch.device("cuda", torch.cuda.current_device()), + ) + for i in sorted(self.sparse_layers) + } + if SPARSE_DECODE_LC_PREFETCH_SELF_KV + else {} + ) + self._built_for: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def sm_scale(self, i: int) -> float: + return (self.head_dim**-0.5) / float(self.attn[i].q_scaling) + + # ---------------------------------------------------------------------------------------------- after the load + + def post_load_weights(self) -> None: + """Derive the kernels' state from the loaded weights, then build both targets and every decode kernel. + + Idempotent per set of weight storages: the loader's post-load walk reaches this module directly and through + the shell, and a second call on the same storages does nothing. + """ + identity = ( + self.w["embed"].data_ptr(), + self.w["l0_in_norm"].data_ptr(), + self.w["final_norm"].data_ptr(), + ) + if self._built_for == identity: + return + device = self.w["embed"].device + derived = nn.Module() + derived.register_buffer( + "mx_alpha", torch.ones(1, dtype=torch.float32, device=device), persistent=False + ) + for i in range(self.num_layers): + # No KV calibration in the checkpoint: unit (Q, K, V) scales and their reciprocals. + derived.register_buffer( + f"l{i}_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + derived.register_buffer( + f"l{i}_inv_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + for i in sorted(self.moe_layers): + for name, value in dataclasses.asdict(self._derive_moe(i)).items(): + derived.register_buffer(f"l{i}_{name}", value, persistent=False) + self.derived = derived + self.mx_alpha = derived.mx_alpha + self.kv_scales = {i: getattr(derived, f"l{i}_kv_scales") for i in range(self.num_layers)} + self.inv_kv_scales = { + i: getattr(derived, f"l{i}_inv_kv_scales") for i in range(self.num_layers) + } + self.moe = { + i: _MoeScalars( + **{ + f.name: getattr(derived, f"l{i}_{f.name}") + for f in dataclasses.fields(_MoeScalars) + } + ) + for i in sorted(self.moe_layers) + } + self._build_embed_owners() + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self._built_for = identity + + def cache_derived_state(self) -> None: + """The staged loaders' entry to the same build (``post_load_weights``).""" + self.post_load_weights() + + def _derive_moe(self, i: int) -> _MoeScalars: + """The routed experts' kernel scalars from the checkpoint's per-tensor scales (trtllm-gen NVFP4 layout). + + Global input scales: the reciprocal of the largest per-expert input scale (gate and up must agree). Per expert: + ``alpha = 1 / (input_scale * (1 / weight_scale_2))`` for the gate (fc31_alpha), the up half (its GEMM2-input + rescale ``fc31_scale_c = fc2_input_scale * up_alpha``) and the down projection; SwiGLU-OAI's activation + constants per expert, with beta and the clamp in the pre-dequant GEMM output domain (divided by fc31_alpha). + """ + w = self.w + isc1, isc3, isc2 = w[f"l{i}_e_isc1"].data, w[f"l{i}_e_isc3"].data, w[f"l{i}_e_isc2"].data + if not torch.allclose(isc1, isc3): + raise ValueError(f"layer {i}: the routed experts' w1 and w3 input scales differ") + fc31_input_scale = isc1.max().reciprocal() + fc2_input_scale = isc2.max().reciprocal() + ws2_1, ws2_3, ws2_2 = ( + w[f"l{i}_e_ws2_1"].data, + w[f"l{i}_e_ws2_3"].data, + w[f"l{i}_e_ws2_2"].data, + ) + fc31_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_1)) + up_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_3)) + fc2_alpha = 1.0 / (fc2_input_scale * (1.0 / ws2_2)) + fc31_scale_c = fc2_input_scale * up_alpha + n = self.num_experts + device = fc31_alpha.device + act_alpha = torch.full((n,), self.swiglu_alpha, dtype=torch.float32, device=device) + act_beta = torch.full((n,), self.swiglu_beta, dtype=torch.float32, device=device) + act_clamp = torch.full((n,), self.swiglu_limit, dtype=torch.float32, device=device) + act_beta.div_(fc31_alpha) + act_clamp.div_(fc31_alpha) + return _MoeScalars( + fc31_input_scale=fc31_input_scale, + fc2_input_scale=fc2_input_scale, + fc31_alpha=fc31_alpha, + fc2_alpha=fc2_alpha, + fc31_scale_c=fc31_scale_c, + act_alpha=act_alpha, + act_beta=act_beta, + act_clamp=act_clamp, + ) + + def _build_embed_owners(self) -> None: + """The fused embedding + first RMSNorm's caller-owned buffers, one per (rows, id dtype) of a decode step.""" + table, weight = self.w["embed"], self.w["l0_in_norm"] + self.embed_owners = { + (rows, dtype): DecodeEmbedding(table, weight, self.rms_eps, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + + def projections(self) -> List[Tuple[str, str, Optional[str]]]: + """Every MXFP8 projection in layer order: (weight key, scale key, the boundary its GEMV pushes into).""" + out = [] + for i in range(self.num_layers): + out.append((f"l{i}_qkv_w", f"l{i}_qkv_s", None)) + out.append((f"l{i}_o_w", f"l{i}_o_s", "attn")) + if i in self.moe_layers: + out.append((f"l{i}_sh_gu_w", f"l{i}_sh_gu_s", None)) + out.append((f"l{i}_sh_dn_w", f"l{i}_sh_dn_s", "moe")) + else: + out.append((f"l{i}_mlp_gu_w", f"l{i}_mlp_gu_s", None)) + out.append((f"l{i}_mlp_dn_w", f"l{i}_mlp_dn_s", "dense")) + return out + + def decode_query_lens(self) -> List[int]: + """The query rows per request a decode step can have: 1, and 1 + each draft length speculation runs.""" + lens = {1} + spec = getattr(self.model_config, "spec_config", None) + if spec is not None: + for drafts in ( + getattr(spec, "max_draft_len", None), + getattr(spec, "max_total_draft_tokens", None), + *(getattr(spec, "draft_len_schedule", None) or {}).values(), + ): + if drafts: + lens.add(1 + int(drafts)) + return sorted(lens) + + def prebuild_kernels(self) -> Dict[str, object]: + """Compile (and load) every kernel a decode forward dispatches to that would otherwise compile on its first + call: each projection's decode GEMVs (every row count, activation and boundary push), the MoE engines of + ``MOE_ENGINE``'s token counts, the index scorers at the decode query lengths (this target's and the in-tree + one wider steps take), the 384-thread long-row select, the dense layers' full-context decode plans, and the + built-in sparse decode a mixed step's generation rows run. Runs on the current (main) stream before the + executor's warmup, in the same order on every rank, then synchronizes; after it, GEMV compiles are frozen and + the forward only looks kernels up (a missing one is a miss). Returns the seconds each part took and the keys. + """ + times: Dict[str, object] = {} + # The 9..16-row generation router: decided once, from the device, never inside a CUDA-graph capture. + # The built-in gate's conditions, mirrored exactly: 128 experts, a hidden size in 1024s, 16-CTA clusters. + cfg = self.model_config.pretrained_config + self.router_tc_cluster16 = ( + int(cfg.num_local_experts) == 128 + and int(cfg.hidden_size) % 1024 == 0 + and bool(torch.ops.trtllm.minimax_m3_router_tc_cluster16(int(cfg.hidden_size))) + ) + logger.info( + f"[M3 prebuild] rank {mpi_rank()}: TF32 tensor-core router for {ROUTER_TC_GEN_MIN_ROWS}.." + f"{ROUTER_TC_GEN_MAX_ROWS}-row generation steps: " + + ( + "engaged (16-CTA clusters)" + if self.router_tc_cluster16 + else "not used (no 16-CTA clusters)" + ) + ) + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + from .ops import m3_fullctx_decode as _fullctx_decode_kernel + from .ops import m3_select_long_dsl as _select_long_dsl_kernel + from .ops import m3_tokmap_guard + + device = self.w["embed"].device + # The one-request token-map guard of the stock generation preprocess (``PREPROCESS_TOKMAP_FIX``): every path + # that can run it (the drafter's full-context decode or its fallback, any trtllm-gen fallback), so it is + # installed whatever else is enabled. + m3_tokmap_guard.install( + torch.device("cuda", torch.cuda.current_device()), + DECODE_MAX_TOKENS, + lambda rows: taken(PREPROCESS_TOKMAP_PATH, rows, "one-request token map"), + lambda rows, reason: missed(PREPROCESS_TOKMAP_PATH, rows, reason), + ) + mx_gemv.freeze(False) + use_pdl = get_env_enable_pdl() + world = int(self.mapping.tp_size) + start = time.perf_counter() + for wk, sk, push in self.projections(): + weight, scale = self.w[wk], self.w[sk] + mx_gemv.ensure_compiled(weight, scale, use_pdl) + n, k = weight.shape + for rows, x_bf16, ar in _gemv_variants(n, k, push, world): + mx_gemv.precompile(weight, scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + torch.cuda.synchronize() + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for i in sorted(self.moe_layers): + self._prebuild_engines(i) + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = self._prebuild_index_scorer(self.decode_query_lens(), device) + torch.cuda.synchronize() + times["scorer_s"] = time.perf_counter() - start + start = time.perf_counter() + times["select_keys"] = [] + if long_select_dsl_keys(): + # The 384-thread long-row select (``SELECT_LONG_DSL``): one kernel per launch key. + times["select_keys"] = _select_long_dsl_kernel.prebuild(long_select_dsl_keys(), device) + torch.cuda.synchronize() + times["select_s"] = time.perf_counter() - start + start = time.perf_counter() + times["fullctx_plans"] = 0 + if FULLCTX_DECODE.keys: + # The dense layers' full-context decode: one compiled PrimTS plan per covered (rows, requests). + times["fullctx_plans"] = _fullctx_decode_kernel.prebuild( + device, + {key: FULLCTX_DECODE_RECIPES[key] for key in FULLCTX_DECODE.keys}, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), + use_pdl, + ) + torch.cuda.synchronize() + times["fullctx_s"] = time.perf_counter() - start + start = time.perf_counter() + times["builtin_decode_rows"] = _prebuild_builtin_sparse_decode( + self.decode_query_lens(), self.heads, self.kv_heads, device + ) + torch.cuda.synchronize() + times["builtin_decode_s"] = time.perf_counter() - start + mx_gemv.freeze(True) + logger.info( + "[M3 prebuild] decode kernels built before warmup: GEMVs %.1f s, MoE engines %.1f s, index scorer " + "%.1f s (%d variants), long-row select %.1f s (%d variants), full-context decode %.1f s (%d plans), " + "built-in sparse decode %.1f s (%d row counts)" + % ( + times["gemv_s"], + times["moe_s"], + times["scorer_s"], + len(times["scorer_keys"]), + times["select_s"], + len(times["select_keys"]), + times["fullctx_s"], + times["fullctx_plans"], + times["builtin_decode_s"], + len(times["builtin_decode_rows"]), + ) + ) + return times + + # ------------------------------------------------------------------------------------------ the MoE engines + + def _expert_operands(self, i: int): + """The routed experts' operands the engines read (trtllm-gen NVFP4 SwiGLU-OAI layout), and the input scale.""" + w, m = self.w, self.moe[i] + tensors = [ + w[f"l{i}_e_w31"].data, + w[f"l{i}_e_s31"].data, + w[f"l{i}_e_w2"].data, + w[f"l{i}_e_s2"].data, + m.fc31_alpha, + m.fc31_scale_c, + m.act_alpha, + m.act_beta, + m.act_clamp, + m.fc2_alpha, + ] + return tensors, m.fc31_input_scale + + def _shared_weights(self, i: int): + w = self.w + return ( + w[f"l{i}_sh_gu_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_gu_s"].data, + w[f"l{i}_sh_dn_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_dn_s"].data, + ) + + def engine(self, i: int, num_tokens: int, shared: bool) -> Optional[object]: + """Layer ``i``'s prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in (and the shared + expert with ``shared``). None after a miss: the build failed or was never made (a forward never builds one).""" + return self._lookup( + i, self.moe_engines[i], (num_tokens, True, shared), MOE_ENGINE.path, num_tokens + ) + + def _lookup(self, i: int, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {i}'s engine" + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, what) + else: + missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + missed(path, num_tokens, "its build failed (logged)") + return engine + + def _prebuild_engines(self, i: int) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of layer ``i`` can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in. Whether an engine is kept is agreed + across the TP group (one host collective per engine): a build that fails on one rank drops that engine on + every rank, so the ranks never mix an engine (and its boundary push) with the generic MoE.""" + engines = self.moe_engines.setdefault(i, {}) + for num_tokens in sorted(MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in engines: + engines[key] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build( + i, MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + f"[M3 dispatch] {MOE_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + + def _try_build(self, i: int, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(i, num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {i}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, i: int, num_tokens: int, shared: bool) -> object: + """The 1..8-token engine (router, routed experts and the shared expert in one launch, at the TP2 width 1536). + Rubin TP2 ablation chain (57 layers behind a 58.4 MB PDL producer, row = (P + MoE) - P, both orders): 31.3 / + 50.7 us/layer at 4 / 8 tokens against the built-in MoE's 42.4 / 55.4 (FlashInfer MXFP8 shared expert); bit for + bit the built-in block's shared.add_(routed). When the engine pushes its sum into the boundary all-reduce, + every CTA releases the PDL dependents right after its grid dependency wait instead of after the push + (``MOE_EARLY_AR_RELEASE``), so the all-reduce (which polls the pushed tagged words, not the grid) becomes + resident on the SMs the engine frees: scheduling only, the same bits (Rubin TP2 greedy early column, same-job + A,B + B,A: c1 -15 us).""" + tensors, scale = self._expert_operands(i) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + kwargs["shared_weights"] = self._shared_weights(i) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = M3MoeState( + device, + num_tokens=num_tokens, + num_ctas=sms - MOE_ENGINE_RESERVED_SMS, + shared=shared, + early_release=MOE_EARLY_AR_RELEASE, + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + # ------------------------------------------------------------------------------------------ the index scorers + + def _index_geometries(self) -> list: + """Every sparse layer's (index dtype, index heads per rank, index dim), deduplicated.""" + geoms = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + geom = (dtype, int(config.num_index_heads), int(config.sparse_index_dim)) + if geom not in geoms: + geoms.append(geom) + return geoms + + def _scorer_variants(self, query_lens) -> list: + """The scorer keys a decode step can take: every sparse layer's index geometry at each decode query length, + with each covered request count's schedule (``index_score_schedule``), where the kernel serves it.""" + runner = _index_score_kernel.IndexScoreRunner + keys = [] + for dtype, heads, dim in self._index_geometries(): + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=heads, + head_dim=dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for requests in sorted(INDEX_SCORE.keys): + split_k, early = index_score_schedule(requests) + key = (dtype, heads, int(qlen), dim, int(split_k), bool(early)) + if key not in keys: + keys.append(key) + return keys + + def _in_tree_scorer_variants(self, query_lens) -> list: + """The in-tree CuTe DSL scorer's keys (``CuteDSLMiniMaxM3IndexDecodeScoreRunner``: index dtype, index heads, + decode query length, index dim, its 256-CTA split) a decode step outside ``INDEX_SCORE`` can take: every query + length at which a step of more than ``INDEX_SCORE_MAX_REQUESTS`` requests fits in ``DECODE_MAX_TOKENS`` (e.g. + 3..8 requests decoding one row each), where that kernel serves the geometry.""" + in_tree = _in_tree_scorer() + widest = DECODE_MAX_TOKENS // (INDEX_SCORE_MAX_REQUESTS + 1) + keys = [] + for dtype, heads, dim in self._index_geometries(): + for qlen in sorted(set(query_lens)): + if qlen > widest or not in_tree.is_supported( + q_dtype=dtype, + num_heads=heads, + head_dim=dim, + page_size=in_tree.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + key = (dtype, heads, qlen, dim, in_tree.SPLIT_K) + if key not in keys: + keys.append(key) + return keys + + def _prebuild_index_scorer(self, query_lens, device) -> list: + """Compile every key of ``_scorer_variants`` and ``_in_tree_scorer_variants`` and launch each once on + real-shaped dummy operands (loading its module), on the current stream, so no decode step compiles a scorer. + The kernels' batch, token, page and block counts are symbolic: their compiles need no step metadata. Returns + the keys.""" + runner = _index_score_kernel.IndexScoreRunner + keys = self._scorer_variants(query_lens) + for key in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + q, k, bt, score, seq = _dummy_scorer_operands( + dtype, heads, qlen, dim, runner.SUPPORTED_PAGE_SIZE, device + ) + runner.kernel_cache[key](q, k, bt, score, seq) + in_tree = _in_tree_scorer() + in_tree_keys = self._in_tree_scorer_variants(query_lens) + for key in in_tree_keys: + dtype, heads, qlen, dim, split_k = key + in_tree._compile(dtype, heads, qlen, dim, split_k) + q, k, bt, score, seq = _dummy_scorer_operands( + dtype, heads, qlen, dim, in_tree.SUPPORTED_PAGE_SIZE, device + ) + in_tree.kernel_cache[key](q, k, bt, score, seq) + return keys + in_tree_keys + + # --------------------------------------------------------------------------------------------- the dispatch + + def _select_target(self, attn_metadata: AttentionMetadata): + """Which target runs this step: one host-side integer, fixed per CUDA-graph capture. A step with context rows + (a mixed step included) is prefill's; every pure generation step is decode's, whatever its width.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + capture: Callable[..., None] = _plain_capture, + capture_target: Optional[Callable[..., Optional[tuple]]] = None, + **kwargs, + ) -> torch.Tensor: + """``capture``: the Eagle3 hidden-state capture of a layer (the worker's, with its capture norm), and + ``capture_target``: where a decode boundary all-reduce writes a capture layer's capture norm itself + (``CAPTURE_NORM_IN_BOUNDARY``; the worker's, None without one).""" + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if kwargs.get("lora_params") is not None: + raise ValueError("this MiniMax-M3 target serves no LoRA adapters") + return self._select_target(attn_metadata).forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The targets +# -------------------------------------------------------------------------------------------------------------------- + + +#: The cores the breakable prefill graphs' eager attention boundary looks up by key: an ``eager_on_graph`` callable +#: replays with weak references to its arguments, so they may be tensors or scalars only. +_BCG_CORES: Dict[str, "MiniMaxM3Tp2ModelingV2"] = {} + + +def _msa_attention_live(core, i: int, q, k, v, idx_q, md, output) -> None: + """The MSA backend over a step's live rows (``md.num_tokens``; a padded step's pad rows of ``output`` are left as + allocated): the KV write unless a producer wrote the caches, the sparse layers' top-k block selection, the FMHA + over the context rows and any generation rows beside them.""" + live = int(md.num_tokens) + sparse = i in core.sparse_layers + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k is not None) + m3_msa_attention( + core.attn[i], + q[:live], + k[:live] if k is not None else None, + v[:live] if v is not None else None, + idx_q[:live] if idx_q is not None else None, + md, + output[:live], + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + + +def _prefill_attention_boundary( + packed, position_ids, q, k, v, idx_q, output, core_key: str, i: int +) -> None: + """The attention boundary of a prefill layer under a breakable prefill CUDA graph, run eagerly at every replay (as + the built-in model's MiniMax-M3 attention boundary is): everything that depends on the step's live row count. A + sparse layer passes its ``packed`` five-way projection: the horizontal producer runs here over the live rows + (K / V / index-K into the paged caches, Q / index-Q out), then the MSA backend; a dense layer passes q / k / v.""" + core = _BCG_CORES[core_key] + md = get_model_extra_attrs()["attention_metadata"]() + live = int(md.num_tokens) + if packed is not None: + horizontal = _sparse_qkv_producer( + core, + i, + packed[:live], + position_ids[..., :live] if position_ids is not None else None, + md.kv_cache_manager, + md.msa_idx_k_cache, + md.msa_out_cache_loc, + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + k = v = None + _msa_attention_live(core, i, q, k, v, idx_q, md, output) + + +_prefill_attention_eager = eager_on_graph(_prefill_attention_boundary) + +#: The drafters the breakable prefill graphs' eager draft attention boundary looks up by key (``_BCG_CORES``'s rule). +_BCG_DRAFTS: Dict[str, "M3Eagle3Draft"] = {} + + +def _draft_attention_boundary(qkv, output, draft_key: str) -> None: + """The drafter layer's attention under a breakable prefill CUDA graph, run eagerly at every replay (as the in-tree + drafter's attention boundary is): the backend's RoPE + KV append + causal attention over the step's live rows, + with the step's metadata read at the replay; ``output``'s pad rows are left as allocated.""" + draft = _BCG_DRAFTS[draft_key] + md = get_model_extra_attrs()["attention_metadata"]() + out = m3_draft_attention(draft.attn, qkv, md) + output[: out.shape[0]].copy_(out) + + +_draft_attention_eager = eager_on_graph(_draft_attention_boundary) + + +class PrefillTarget: + """A step with context rows: prefill chunks, and the mixed steps that carry generation rows beside them. Every + layer runs in its generic form; the attention core is the MSA backend's for every row. + + Under a breakable (piecewise) prefill CUDA graph the projections, norms, MoE and boundaries are captured per padded + row bucket and the attention boundary runs eagerly at every replay (``_prefill_attention_eager``: the live-row + horizontal producer and the MSA backend), as the built-in model's MiniMax-M3 attention does.""" + + def __init__(self, core: MiniMaxM3Tp2ModelingV2) -> None: + self.core = core + self._bcg_key = f"m3_tp2_core_{id(core)}" + _BCG_CORES[self._bcg_key] = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + Counts and per-step host scalars are read here; device tensors are the metadata's persistent buffers, held + by reference; ``msa_idx_k_cache`` is the metadata's per-layer index-K cache getter. ``metadata`` is the + prepared object itself: the MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_tokens=md.num_tokens, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + """``capture_target`` is the decode boundaries' (unused: no prefill boundary writes a capture norm).""" + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + _record("prefill", int(step["num_tokens"]), int(step["num_contexts"])) + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + hidden, residual = self._layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + return hidden + + def _swiglu_mlp( + self, h: torch.Tensor, prefix: str, low_m_gemm: bool, dense: bool = False + ) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI (the dense MLP's on this target's + C++ op, ``_dense_swiglu``), down, no reduction.""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + out, _ = _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm) + return out + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: + the MoE scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: on ``ROUTER_TC_GEN``'s rows the TF32 tensor-core router with the + experts' NVFP4 activation fused in (as the agg canon routes a 9..48-row step with context rows; one runner + call), else the fp32 router logits and the experts quantizing their bf16 input (chunked).""" + core = self.core + fused = _router_tc_fp4(core, i, h) + if fused is not None: + logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, logits, True) + return self._routed_chunks(i, h, _router_logits(core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool): + """The MoE block of layer ``i`` without its reduction. Returns ``(out, routed)``: with ``return_parts`` the + shared and routed parts (routed None when the engine folded the shared expert in), else their sum in + ``out``.""" + core = self.core + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, False) + taken(MOE_ENGINE.path, rows) + folded = m3_moe_engine(h.contiguous(), engine, None, None) + return folded, None + if not return_parts and _triton_finalize_rows(core, h): + routed_parts, shared_out = maybe_execute_in_parallel( + lambda: _routed_experts(core, i, h, None, _router_logits(core, i, h), False), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + return _triton_finalize(rows, routed_parts, shared_out), None + routed_out, shared_out = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the + # same collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + # Row ranges at this dispatch: 9..48 rows take the decode target's generic MC finalize boundary (as the agg + # canon's generic layer does on every step, context rows included: the MoE unfinalized, its finalize + + # shared expert in the multicast all-reduce + residual + next norm); the CTX stack starts at 4096 rows. + folded = core._decode._generic_mc_finalize(i, hidden, residual, _next_norm(core, i)) + if folded is not None: + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, folded[1], None) + return folded + out, routed = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step): + """Attention of layer ``i`` up to its o_proj (no reduction). Returns (o_proj output, pushed): a prefill step + arms no boundary push, so pushed is always False.""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v for an FP8 cache, on this target's op; else in + # place in bf16), the MSA backend (it writes the layer's cache), o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = ( + _dense_qk_norm_rope_fp8(core, i, qkv, position_ids) + if core.main_kv_is_fp8 + else _qk_norm_rope(core, i, qkv, position_ids, False) + ) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + if not core.fused_projection: + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target runs a breakable prefill CUDA graph only with the five-way projection " + "(fuse_qkv_index_projection)" + ) + return self._unfused_sparse_attention(i, position_ids, hidden, step) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + if is_in_breakable_cuda_graph(): + output = packed.new_empty( + (packed.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16 + ) + _prefill_attention_eager( + packed, position_ids, None, None, None, None, output, self._bcg_key, i + ) + return _linear(core, output, f"l{i}_o_w", f"l{i}_o_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection + and its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and + the FP8 indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the + o_proj. A layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): the MSA + backend (``m3_msa_attention``: the KV write unless a producer wrote the caches, the sparse layers' top-k block + selection, the FMHA over the context rows and any generation rows beside them).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + if is_in_breakable_cuda_graph(): + # The live row count is a host value of the step: read at every replay inside the eager boundary. + _prefill_attention_eager(None, None, q, k, v, idx_q, output, self._bcg_key, i) + else: + _msa_attention_live(core, i, q, k, v, idx_q, step["metadata"], output) + return output + + +class DecodeTarget: + """A pure generation step (no context rows), of any width. + + Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids (up to 8 requests verifying 1 + 3 rows), per layer: + attention (o_proj all-reduce deferred, its rows pushed into the boundary all-reduce up to 16 rows) -> attention + boundary (all-reduce + residual + post-attention RMSNorm) -> MoE (the 1..8-row engine pushing its sum into the + boundary all-reduce; at 9..32 rows the unfinalized MoE with the finalize in the boundary all-reduce) or dense MLP + (its down projection pushing into the post-MLP boundary) -> MoE / dense boundary (all-reduce + shared/routed add + + residual + the next layer's RMSNorm, emitting the next QKV's MXFP8 activation). An Eagle3 capture layer runs the + same decode layer and captures its feed-forward boundary's reduced residual sum (the boundary writing the capture + norm itself when the worker gives it a target). The first layer runs the generic layer when the fused embedding + does not take the step. A wider step (or one carrying embeddings) runs every layer in its generic form, with this + target's attention (the dense layers' full-context decode, the decode selection and sparse decode where the + producer wrote the caches).""" + + def __init__(self, core: MiniMaxM3Tp2ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # Set by a feed-forward boundary that wrote the layer's Eagle3 capture (see _layer's ``capture``). + self._capture_written = False + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + States the two context counts instead of reading them back: this target runs only when ``num_contexts`` is + zero, so reading them would ask the metadata to confirm what the routing already decided. + + Graph classes: counts, the generation span, the identity-selection flag, the worst-case score extent and the + cache length bound are host values fixed per capture (prepare() stages them before the forward); the block + table, lengths, cache slots, identity table and valid-block counts are the metadata's persistent buffers, held + by reference and refreshed in place every step; ``msa_idx_k_cache`` / ``msa_proxy_max_score_view`` are the + metadata's getters of the per-layer index-K cache and of the score scratch view. ``metadata`` is the prepared + object itself: the MSA backend's step (``m3_msa_attention``) and the full-context decode's K/V write + (``write_msa_phase_kv``) take it as their argument. + """ + return dict( + num_contexts=0, + num_ctx_tokens=0, + num_tokens=md.num_tokens, + num_seqs=md.num_seqs, + max_seq_len=md.max_seq_len, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + msa_decode_span=md.msa_decode_span, + msa_block_table=md.msa_block_table, + msa_seq_lens_cuda=md.msa_seq_lens_cuda, + msa_short_ok=md._msa_short_ok, + msa_short_table=md.msa_short_table, + msa_proxy_max_score_view=md.msa_proxy_max_score_view, + msa_worst_case_max_k_tiles=md.msa_worst_case_max_k_tiles, + msa_n_valid_blocks=md.msa_n_valid_blocks, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + # The decode kernels' rows (a per-capture constant): up to DECODE_MAX_TOKENS tokens of 1-D input ids. + kernels = ( + inputs_embeds is None and input_ids.ndim == 1 and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + _record("decode", num_tokens, plan.name) + hidden, residual = self._embed(input_ids, num_tokens) + else: + # Breakable (piecewise) CUDA graphs capture steps with context rows only (prefill_cuda_graph_backend): + # every pure generation step, the 33..192-row ones on this generic branch included, runs eagerly or on + # the ordinary decode CUDA graphs, so this guard never fires in a supported deployment. + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + _record("decode", int(step["num_tokens"]), "generic") + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + captured = spec_metadata is not None and spec_metadata.is_layer_capture(i) + if kernels and residual is not None and captured: + # The capture layer on the decode layer (its boundaries push); the capture reads the MoE / dense + # boundary's reduced residual sum, as the generic layer's fused boundary capture does. Given a target + # (``CAPTURE_NORM_IN_BOUNDARY``), the boundary all-reduce writes the capture norm itself. + target = ( + capture_target(spec_metadata, i, num_tokens) + if capture_target is not None + else None + ) + self._capture_written = False + hidden, residual = self._layer( + i, plan, position_ids, hidden, residual, step, capture=target + ) + taken("capture_layer_fast", num_tokens, f"layer {i}") + if target is not None and self._capture_written: + taken("capture_norm_in_boundary", num_tokens, f"layer {i}") + capture(spec_metadata, i, residual, None, written=True) + else: + capture(spec_metadata, i, residual, None) + elif not kernels or residual is None or captured: + hidden, residual = self._generic_layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + else: + hidden, residual = self._layer(i, plan, position_ids, hidden, residual, step) + return hidden + + def _embed(self, input_ids: torch.Tensor, num_tokens: int): + """(hidden, residual) of the fused embedding + first RMSNorm; (embedding, None) when it does not run, so the + first layer runs its own input norm. SM107: 1.55..1.65 us against 3.4..6.0 us for the stock embedding + + RMSNorm at 1..8 rows (the gather bitwise, the norm one ulp off on ~5e-6 of the elements: summation order).""" + core = self.core + if not EMBED_FIRSTNORM.covers(num_tokens): + uncovered(EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens") + return m3_embedding(input_ids, core.w["embed"]), None + owner = core.embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + missed( + EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 after the load)" + if owner is None + else "the input ids are not contiguous", + ) + return m3_embedding(input_ids, core.w["embed"]), None + _record("m3_embedding_firstnorm", num_tokens, input_ids.dtype) + taken(EMBED_FIRSTNORM.path, num_tokens, str(input_ids.dtype)) + return m3_embedding_firstnorm( + input_ids, owner.table, owner.weight, owner.eps, owner.raw, owner.output, owner.use_pdl + ) + + def _bucket_flags(self, num_tokens: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(self.core.mapping, self.core.hidden) + workspace = None if workspace is None else m3_mc_allreduce._bucket(workspace, num_tokens) + flags = None if workspace is None else workspace.flags_for(8) + return None if flags is None else (workspace, flags) + + @staticmethod + def _no_bucket(path: str, num_tokens: int) -> None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, num_tokens: int): + """The push tuple into the attention boundary's all-reduce (the o_proj's; also the dense down projection's into + the post-MLP boundary, the same exchange), or None (after a miss). + + None comes only from the token count or from workspace state created collectively in the same eager call on + every rank, so all ranks agree. (A lone non-pushing rank would still be safe: its all-reduce pushes the same + bf16 rows into the same set A.)""" + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(ATTN_PUSH.path, num_tokens) + return None + workspace, flags = found + return ( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + AR_PUSH_MODE, + ) + + def _moe_push_targets(self, num_tokens: int): + """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None + (after a miss). + + This decision must match across ranks: prepushed 3 reads set B of every rank, which a non-pushing rank never + fills. None comes only from the token count or from state created in the same eager call on every rank + (workspace, flags, the descriptor cache), and the engine it is armed for is agreed at build time + (``_prebuild_engines``), so it does.""" + hidden = self.core.hidden + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(MOE_PUSH.path, num_tokens) + return None + workspace, flags = found + key = (workspace.mc.data_ptr(), flags.data_ptr()) + desc = self._moe_push_descs.get(key) + if desc is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor") + return None + desc = self._moe_push_descs[key] = ar_push_desc( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + flags.device, + ) + set_b = workspace.mc.data_ptr() + 2 * 4 * workspace.max_tokens * workspace.world_size * ( + hidden // 2 + ) + return desc, ( + set_b, + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + 2, + ) + + def _boundary(self, use_mc: bool, x, residual, norm_w, x2=None, quant=False, capture=None): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce. + ``capture``: the Eagle3 capture the multicast kernel also writes (``_capture_written`` records it).""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant, capture=capture) + if fused is not None: + if capture is not None: + self._capture_written = True + taken( + MC_ALLREDUCE.path, + x.shape[0], + ("x + x2" if x2 is not None else "x") + (", MXFP8 copy" if quant else ""), + ) + return fused + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + missed( + MC_ALLREDUCE.path, + x.shape[0], + m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + ) + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _prepushed( + self, i: int, x, residual, norm_w, quant, prepushed, source, x2=None, capture=None + ): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed, capture=capture + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + if capture is not None: + self._capture_written = True + taken(MC_ALLREDUCE.path, x.shape[0], f"prepushed {prepushed}, {source}") + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step, capture=None): + """Layer ``i`` on the decode kernels. ``capture``: (out, fc_norm weight, eps) for the feed-forward boundary to + write the layer's Eagle3 capture norm (``CAPTURE_NORM_IN_BOUNDARY``); a boundary that writes it sets + ``_capture_written``.""" + core = self.core + num_tokens = hidden.shape[0] + # Padded rows: the boundary all-reduce polls and re-arms rows t < M = residual.shape[0], CUDA-graph padding + # rows included. Every producer pushes all rows of the same tensor: the o_proj / shared down / dense down GEMV + # t < m = x.shape[0] (its row count comes from the host shape, not a live count), the MoE engine every row of + # the num_tokens it was built for (= num_tokens here). The all-reduce checks x.sizes() == residual.sizes(), so + # no padded row is left with a stale flag. + # The GEMV dispatch (and with it every push into a boundary all-reduce) is a decode-graph one. + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + uncovered( + "decode_graph", + "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + ) + push = plan.attn_mc_ar and decode_graph and ATTN_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if push else None + hidden, attn_pushed = self._attention(i, position_ids, hidden, step, armed) + post_norm = core.w[f"l{i}_post_norm"] + if armed is not None and attn_pushed: + hidden, residual = self._prepushed( + i, hidden, residual, post_norm, False, AR_PUSH_MODE, "o_proj push" + ) + else: + if armed is not None: + missed( + ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden, residual = self._boundary(plan.attn_mc_ar, hidden, residual, post_norm) + + norm_w = _next_norm(core, i) + if i not in core.moe_layers: + # A dense layer: its MLP's down projection pushes into the post-MLP boundary's all-reduce, as the o_proj + # does into the attention one (the same exchange; the gate_up projection never pushes). SM107 TP2: the + # two non-capture dense layers' post-MLP hops end 5.5-5.8 us after the down GEMV unpushed vs 2.4-2.5 us + # pushed; bitwise equal to the unpushed boundary (the same bf16 words). + dense_push = plan.attn_mc_ar and decode_graph and DENSE_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if dense_push else None + out, pushed = self._swiglu_mlp(hidden, f"l{i}_mlp", True, armed, dense=True) + if armed is not None and pushed: + return self._prepushed( + i, + out, + residual, + norm_w, + False, + AR_PUSH_MODE, + "dense down push", + capture=capture, + ) + if armed is not None: + missed( + DENSE_PUSH.path, + num_tokens, + "the dense MLP's down projection did not push its rows into the armed boundary all-reduce", + ) + return self._boundary(plan.attn_mc_ar, out, residual, norm_w, capture=capture) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + moe_push = None + if moe_path(num_tokens) == MOE_ENGINE.path: + # The engine runs inside the MoE; looked up (or found missing) here first so a push is armed only for an + # engine that will take it. + engine = core.engine(i, num_tokens, shared=True) + if ( + engine is not None + and plan.moe_mc_ar + and decode_graph + and MOE_PUSH.covers(num_tokens) + ): + moe_push = self._moe_push_targets(num_tokens) + elif moe_path(num_tokens) == MC_FINALIZE.path and plan.moe_mc_ar: + if self._bucket_flags(num_tokens) is None: + self._no_bucket(MC_FINALIZE.path, num_tokens) + else: + return self._mc_finalize(i, hidden, residual, norm_w, quant, capture) + shared, routed, engine_pushed, shared_pushed = self._moe( + i, hidden, return_parts=True, push=moe_push + ) + if moe_push is not None: + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + if engine_pushed: + return self._prepushed( + i, shared, residual, norm_w, quant, 2, "MoE engine push", capture=capture + ) + else: + if engine_pushed and shared_pushed: + return self._prepushed( + i, + shared, + residual, + norm_w, + quant, + 3, + "MoE engine + shared push", + x2=routed, + capture=capture, + ) + if engine_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + missed( + MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary( + plan.moe_mc_ar, shared, residual, norm_w, x2=routed, quant=quant, capture=capture + ) + + def _mc_finalize(self, i: int, h, residual, norm_w, quant: bool, capture=None, what: str = ""): + """The MoE of layer ``i`` unfinalized (its shared expert beside it) and the finalize + all-reduce + residual + + next RMSNorm in one launch (``m3_mc_finalize_allreduce_norm``): returns (normed, updated residual). The + caller checked the rows and the multicast bucket (decided from shapes alone, so the same on every TP rank).""" + core = self.core + rows = h.shape[0] + routed, (shared, _) = maybe_execute_in_parallel( + lambda: _decode_routed_unfinalized(core, i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + fc2, scales, perm = routed + _record("m3_mc_finalize_allreduce_norm", rows, quant, capture is not None) + fused = m3_mc_allreduce.m3_mc_finalize_allreduce_norm( + fc2, + perm.view(rows, -1).to(torch.int32), + scales.reshape(rows, -1), + shared, + residual, + norm_w, + core.rms_eps, + m3_mc_allreduce.get_workspace(core.mapping, core.hidden), + use_pdl=get_env_enable_pdl(), + quant=quant, + capture=capture, + ) + if fused is None: + raise RuntimeError( + f"layer {i}: the MC finalize all-reduce rejected the unfinalized MoE outputs" + ) + if capture is not None: + self._capture_written = True + taken(MC_FINALIZE.path, rows, what or ("MXFP8 copy" if quant else "")) + return fused + + def _generic_mc_finalize(self, i: int, h, residual, norm_w): + """Past ``DECODE_MAX_TOKENS`` rows, up to ``GENERIC_MC_FINALIZE_MAX_TOKENS``: the generic layer's MoE boundary + as ``_mc_finalize``; None (from shapes alone, the same on every rank) keeps the generic boundary.""" + rows = h.shape[0] if h.dim() == 2 else 0 + cap = min(GENERIC_MC_FINALIZE_MAX_TOKENS, m3_mc_allreduce.FINALIZE_MAX_TOKENS) + if ( + not (DECODE_MAX_TOKENS < rows <= cap) + or h.dtype != torch.bfloat16 + or torch.compiler.is_compiling() + ): + return None + if self._bucket_flags(rows) is None: + return None + return self._mc_finalize(i, h, residual, norm_w, False, what="generic") + + def _swiglu_mlp( + self, + h: torch.Tensor, + prefix: str, + low_m_gemm: bool, + push: Optional[tuple] = None, + dense: bool = False, + ): + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI (the dense MLP's on this target's + C++ op, ``_dense_swiglu``), down, no reduction. Returns (output, whether the down projection's GEMV pushed into + ``push``).""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm, push) + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: the MoE + scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: on ``ROUTER_TC_GEN``'s rows the TF32 tensor-core router with the + experts' NVFP4 activation fused in (one runner call: the rows are within ``moe_max_num_tokens``), else the + fp32 router and the experts quantizing their bf16 input (chunked).""" + core = self.core + fused = _router_tc_fp4(core, i, h) + if fused is not None: + logits, x_fp4, x_sf = fused + return _routed_experts(core, i, x_fp4, x_sf, logits, True) + return self._routed_chunks(i, h, _decode_router_logits(core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool, push=None): + """The MoE block of layer ``i`` without its reduction. + + ``push`` = (the engine's ``arp`` descriptor, the shared down projection's ``ar`` tuple) the decode layer armed + for the MoE boundary's exchange sets, or None. Returns ``(out, routed, engine_pushed, shared_pushed)``: with + ``return_parts`` the shared and routed parts (routed None when the engine folded the shared expert in), else + their sum in ``out``. + """ + core = self.core + engine_arp, shared_ar = push if push is not None else (None, None) + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, engine_arp is not None) + taken(MOE_ENGINE.path, rows) + folded = m3_moe_engine(h.contiguous(), engine, None, engine_arp) + return folded, None, engine_arp is not None, False + if not return_parts and push is None and _triton_finalize_rows(core, h): + routed_parts, (shared_out, _) = maybe_execute_in_parallel( + lambda: _decode_routed_unfinalized(core, i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + return _triton_finalize(rows, routed_parts, shared_out), None, False, False + routed_out, (shared_out, shared_pushed) = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False, shared_ar), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the same + # collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out, False, shared_pushed + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None, False, shared_pushed + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + if x.shape[0] > DECODE_MAX_TOKENS: + taken( + MC_ALLREDUCE.path, + x.shape[0], + "generic " + ("x + x2" if x2 is not None else "x"), + ) + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _generic_layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + hidden, residual = self._generic_boundary(hidden, None, residual, core.w[f"l{i}_post_norm"]) + if i in core.moe_layers: + folded = self._generic_mc_finalize(i, hidden, residual, _next_norm(core, i)) + if folded is not None: + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, folded[1], None) + return folded + out, routed, _, _ = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = None): + """Attention of layer ``i`` up to its o_proj (no reduction). ``push`` is the attention boundary's exchange the + o_proj may write into. Returns (o_proj output, pushed).""" + core = self.core + if i not in core.sparse_layers: + return self._dense_attention(i, position_ids, hidden, step, push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _dense_attention(self, i: int, position_ids, hidden, step, push): + """Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v: an FP8 cache, or the FP8 pool the dense layers sit + in within a hybrid NVFP4 cache on a decode step), the attention core, o_proj. + + On a covered pure generation step (``_fullctx_plan``) the attention is this target's PrimTS split-KV decode + over the layer's full causal context, and the norm/RoPE op also stores the step's new K/V rows into the + layer's paged pool and the query compact, as the decode reads it (``DENSE_KV_INSERT_FUSED``; the bytes + ``write_msa_phase_kv`` stores). Otherwise the MSA backend.""" + core = self.core + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + plan = self._fullctx_plan(i, qkv.device, step) + kv_dest = None + if ( + plan is not None + and fp8 + and DENSE_KV_INSERT_FUSED + and _dense_fp8_op_applies(qkv, position_ids) + ): + kv_dest = ( + step["kv_cache_manager"].get_buffers(i, kv_layout="HND"), + step["msa_out_cache_loc"], + plan[1], + torch.empty( + (qkv.shape[0], core.heads, core.head_dim), + dtype=torch.float8_e4m3fn, + device=qkv.device, + ), + ) + if fp8: + fused = _dense_qk_norm_rope_fp8(core, i, qkv, position_ids, kv_dest) + if kv_dest is not None: + taken(DENSE_KV_INSERT_PATH, kv_dest[2]) + else: + fused = _qk_norm_rope(core, i, qkv, position_ids, False) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + if kv_dest is not None: + q = kv_dest[3].view(q.shape[0], -1) + o = None + if plan is not None: + o = self._fullctx_decode(i, q, k, v, step, plan, kv_written=kv_dest is not None) + if o is None: + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _fullctx_plan(self, i: int, device: torch.device, step): + """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs + on ``ops/m3_fullctx_decode`` (``FULLCTX_DECODE``): a covered (query rows per request, requests) shape over + the layer's FP8 pool. None otherwise (the MSA backend runs). + + Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 (29.3), (1, 1) + 15.7 (28.0), (4, 2) 24.5 (29.3), (1, 2) 20.2 (28.1); the stock split is capped at 16 cluster CTAs per (request, + KV head), i.e. 32 of 212 CTAs at one request. With the BF16 split-KV merge (``FULLCTX_DECODE_OWN_REDUCER``): + (1, 1) 14.5 vs 18.5, (4, 1) 17.3 vs 21.3, (1, 2) 18.5 vs 23.1, (4, 2) 23.3 vs 27.5.""" + path = FULLCTX_DECODE.path + span = step["msa_decode_span"] + if span is None or span.row_first != 0: + uncovered(path, "no decode span starting at row 0") + return None + nreq = int(step["num_seqs"]) + qlen = int(span.query_len) + if not FULLCTX_DECODE.covers((qlen, nreq)): + uncovered(path, f"{qlen} query rows x {nreq} requests is not a covered shape") + return None + if int(step["max_seq_len"]) > FULLCTX_DECODE_MAX_KV_LEN: + uncovered(path, "the cache length bound is beyond the plans' K/V capacity") + return None + num_tokens = nreq * qlen + pool, _ = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + if pool.dtype != torch.float8_e4m3fn: + missed(path, num_tokens, "the dense layer's KV pool is not FP8") + return None + from .ops import m3_fullctx_decode as _fullctx_decode_kernel + + own_reducer = FULLCTX_DECODE_OWN_REDUCER + if own_reducer and not _fullctx_decode_kernel.has_own_reducer(): + missed(path, num_tokens, "the BF16 split-KV merge op is not built") + return None + # Compiled before warmup (``prebuild_kernels``); the forward never compiles one. + plan = _fullctx_decode_kernel.plan_for(device, qlen, nreq, own_reducer) + if plan is None: + missed(path, num_tokens, "the PrimTS decode plan is not prebuilt") + return None + return plan, num_tokens + + def _fullctx_decode(self, i: int, q, k, v, step, plan_tokens, *, kv_written: bool): + """Dense layer ``i``'s decode attention on the PrimTS plan of ``_fullctx_plan``: the step's new K/V go into the + paged cache as the MSA backend writes them (``write_msa_phase_kv``; already there when ``kv_written``), then + every query row attends its request's full causal context.""" + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + write_msa_phase_kv, + ) + + from .ops import m3_fullctx_decode as _fullctx_decode_kernel + + core = self.core + plan, num_tokens = plan_tokens + nreq = int(step["num_seqs"]) + if not kv_written: + write_msa_phase_kv( + core.attn[i], + k[:num_tokens], + v[:num_tokens], + step["metadata"], + AttentionInputType.mixed, + token_offset=0, + ) + pool, subpages = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + k_pages, v_pages = _fullctx_decode_kernel.kv_page_views(pool, subpages) + _record("m3_fullctx_decode", num_tokens, nreq, kv_written) + plan.run( + fp8_e4m3_saturate(q[:num_tokens].reshape(num_tokens, core.heads, core.head_dim)), + k_pages, + v_pages, + step["msa_block_table"][:nreq], + step["msa_seq_lens_cuda"][:nreq], + core.sm_scale(i), + output[:num_tokens], + ) + taken(FULLCTX_DECODE.path, num_tokens) + return output + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step, push): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection and + its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and the FP8 + indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the o_proj. A + layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): a sparse + layer whose caches the producer wrote runs this target's selection and sparse decode when the generation span + starts at row 0; everything else the MSA backend (``m3_msa_attention``: the KV write unless a producer wrote the + caches, the sparse layers' top-k block selection, the FMHA).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + span = step["msa_decode_span"] + sparse = i in core.sparse_layers + if sparse and k_live is None and span is not None and span.row_first == 0: + self._decode_sparse_core(i, q_live, idx_q_live, step, out_live) + return output + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + def _index_scores(self, i: int, idx_q, step, num_tokens: int, qlen: int): + """The per-index-head block maxima of a non-identity decode step: this target's CuTe DSL scorer at the request + counts ``INDEX_SCORE`` covers (each count's CTAs per request and early issue, ``index_score_schedule``), the + in-tree CuTe DSL scorer beyond them. Rubin TP2 (2 index heads), 165,521 tokens, behind a 58 MB PDL producer, + per layer: 1 request 3.27 us at 1536 CTAs (in-tree 256-CTA scorer 6.07); 2 requests 4.97 us at 768 (in-tree + 6.71).""" + config = self.core.attn[i].m3_config + rows = int(step["num_seqs"]) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = step["msa_idx_k_cache"](i) + max_score = step["msa_proxy_max_score_view"]( + config.num_index_heads, step["msa_worst_case_max_k_tiles"], num_tokens + ) + if INDEX_SCORE.covers(rows): + split_k, early = index_score_schedule(rows) + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(split_k), + bool(early), + ) + if key not in _index_score_kernel.IndexScoreRunner.kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + missed( + INDEX_SCORE_PATH, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score", key) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + split_k, + early, + ) + taken(INDEX_SCORE_PATH, rows) + else: + uncovered( + INDEX_SCORE_PATH, + f"{rows} requests, past {INDEX_SCORE_MAX_REQUESTS}: the in-tree CuTe DSL scorer", + ) + _in_tree_scorer() # registers the op + _record("cute_dsl_minimax_m3_index_decode_score", rows, qlen) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + ) + return max_score + + def _select(self, i: int, max_score, step, num_tokens: int): + """The top-16 table of a non-identity decode step from the block maxima reduced to the KV head: the long-row + select (the 384-thread kernel, ``SELECT_LONG_DSL``, when compiled) on rows of ``LONG_SELECT_MIN_BLOCKS.. + LONG_SELECT_MAX_BLOCKS`` blocks, else the built-in selector (``m3_select_blocks``).""" + config = self.core.attn[i].m3_config + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + group = config.num_index_heads // config.num_kv_heads + max_score_kv = ( + max_score.view(config.num_kv_heads, group, max_score.shape[1], max_score.shape[2]).amax( + dim=1 + ) + if group > 1 + else max_score + ) + from .ops import m3_select_long_dsl as _select_long_dsl_kernel + + num_blocks = int(max_score_kv.shape[1]) + table = None + dsl = False + if long_select_covers(num_blocks): + early_trigger = num_tokens <= SELECT_EARLY_MAX_ROWS + if LONG_SELECT_DSL.keys: + # Compiled before warmup (prebuild_kernels); a missing kernel falls back to the 1024-thread one. + key = _select_long_dsl_kernel.kernel_key( + get_env_enable_pdl(), early_trigger, SELECT_LONG_TRIGGER_PHASE + ) + dsl = key in _select_long_dsl_kernel.kernel_cache + if not dsl: + missed( + LONG_SELECT_DSL.path, + num_tokens, + f"the 384-thread select {key} was not compiled before warmup", + ) + _record("long_select", num_blocks, num_tokens, dsl) + table = long_select_blocks( + max_score_kv, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + use_pdl=get_env_enable_pdl(), + early_trigger=early_trigger, + trigger_phase=SELECT_LONG_TRIGGER_PHASE, + dsl=dsl, + ) + if table is not None: + taken((LONG_SELECT_DSL if dsl else LONG_SELECT).path, num_tokens) + else: + long_select_fallback(num_blocks, num_tokens) + _record("m3_select_blocks", num_blocks, num_tokens) + table = m3_select_blocks( + max_score_kv, + MSA_REQUIRED_TOPK, + n_valid_blocks, + config.init_blocks, + config.local_blocks, + True, + ) + return table + + def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step of a sparse layer whose caches + the producer already wrote. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise the + index scorer writes the per-index-head block maxima, reduced to the KV head, and the long-row selector takes the + top 16 (the built-in selector beyond its shapes). On an NVFP4 layer the decode (``m3_sparse_decode``: the + long-context decode or the Triton producer + merge) prefetches into L2, before its grid wait, the K/V pages of + the blocks the same layer selected one step back (``SPARSE_DECODE_LC_PREFETCH_SELF_KV``: 65-68% of a layer's + selected blocks recur one step later; same layer, not the next one as at TP4: the MoE between two sparse layers + streams ~250 MB through the 126 MB L2, so an earlier issue would be evicted; hints only). + """ + from .ops import m3_attn_lc as _attn_lc_kernel + + core = self.core + backend = core.attn[i] + num_tokens = int(q.shape[0]) + rows = int(step["num_seqs"]) + qlen = int(step["msa_decode_span"].query_len) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + num_kv_heads = int(backend.indexer.config.num_kv_heads) + # The index scorers are compiled and loaded by prebuild_kernels, so no step compiles one and the identity can + # be taken from the first decode step. + identity = step["msa_short_ok"] and step["msa_short_table"] is not None + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + else: + max_score = self._index_scores(i, idx_q, step, num_tokens, qlen) + topk = self._select(i, max_score, step, num_tokens).permute(1, 0, 2) + manager = step["kv_cache_manager"] + k_paged, v_paged = msa_paged_kv(manager, i) + nvfp4 = {} + if manager.is_nvfp4_layer(i): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(backend, core.kv_scales[i]) + scales = manager.get_block_scale_buffers(i, "HND") + nvfp4 = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + # The long-context decode (if it runs) takes the layer's N1 id buffer; any other route leaves the slot unused, + # and it is cleared after the call either way. + n1_ids = core.n1_ids.get(i) + if n1_ids is not None and nvfp4 and num_tokens <= N1_MAX_ROWS: + _attn_lc_kernel.N1_IDS.ids = n1_ids + _record("m3_sparse_decode", num_tokens, qlen, identity) + try: + m3_sparse_decode( + q.view(num_tokens, core.heads, core.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=core.sm_scale(i), + output=output.view(num_tokens, core.heads, core.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4, + ) + finally: + _attn_lc_kernel.N1_IDS.ids = None + + +_ENGAGED: set = set() + + +def _log_engaged(path: str, mapping) -> None: + """One info line per fast path and process (every GEN rank logs its own): the run's engagement evidence.""" + if path not in _ENGAGED: + _ENGAGED.add(path) + rank = getattr(mapping, "tp_rank", "?") if mapping is not None else "?" + logger.info(f"[M3 sm_107 tp2] fast path engaged: {path} (tp rank {rank})") + + +class M3SpecStepFold: + """The target's speculative step-input staging (the in-tree ``SpecStepFold`` protocol). + + Between two decode graphs of a one-model Eagle3 step, the spec sampler stores the step's outputs into its per-slot + store and the engine gathers them back into its flat input buffers for the next step. This object does both in + one launch each (``m3_spec_tail.slot_scatter`` / ``step_gather``): integer copies only, the same values the + engine's and the sampler's own copies write. + + The forward of a decode-only step also does both at its tail (``store_in_forward``, one ``m3_fold_tail`` launch + inside the step graph, this target's C++ op): it stores the step's outputs for the rows ``begin_step`` armed and + stages the next step's inputs for those rows in that order. The arming is a device-side row count and slot list + staged before the replay, so a captured graph whose step was not armed writes nothing. The sampler then skips its + store for exactly the armed rows (``sampler_rows_stored``), and the next step skips its gathers when its rows are + those rows (``begin_step``). The sampler copies the store to the host on its D2H side stream and records + ``store_copies_done`` after the copies; the forward's tail waits on it before it rewrites the store. + + * ``store`` -- the spec sampler's per-slot store (``bind_sampler_store``). + * ``step_inputs`` -- the engine's (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) device buffers + (``bind_engine_buffers``). + """ + + # Read by the shell's forward (the tail runs only when the fold stores in the forward). + in_forward = True + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[tuple] = None + # Device arming read by the forward's tail: ctrl[0] = armed rows (0 = write nothing), slots[:rows]. + self._ctrl: Optional[torch.Tensor] = None + self._slots: Optional[torch.Tensor] = None + # The rows (sequence slots) the next forward stores and stages for, as begin_step armed them. + self._armed: Optional[tuple] = None + # The rows the previous forward staged the next step's inputs for (confirmed by sampler_rows_stored). + self._staged_for: Optional[tuple] = None + # Output buffers the tail ran on: the captured graphs' static outputs (by address) and the last eager step's. + self._graph_outputs: set = set() + self._eager_outputs: Optional[torch.Tensor] = None + self._logged = False + # The (external) event the sampler records on its D2H side stream after copying the store to the host, waited + # on by the forward's tail before it rewrites the store (see SpecSamplerBase._copy_outputs_to_host). + self.store_copies_done: Optional[torch.cuda.Event] = None + + def bind_engine_buffers( + self, input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ): + self.step_inputs = (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) + # Recorded once now so a captured wait has a completed record to wait on before the first copy. + self.store_copies_done = torch.cuda.Event(external=True) + self.store_copies_done.record() + logger.info( + "[M3 host] spec store host copies on the D2H side stream, awaited by the forward's tail" + ) + logger.info("[M3 host] spec fold tail: one C++ launch (no Triton) at the forward's tail") + self._ctrl = torch.zeros(1, dtype=torch.int32, device=input_ids.device) + self._slots = torch.zeros(max(max_batch_size, 1), dtype=torch.long, device=input_ids.device) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + if self._ctrl is None: + return False + rows = tuple(seq_slots) if seq_slots is not None else None + staged = rows is not None and rows == self._staged_for + self._staged_for = None + armed = rows if rows is not None and 0 < len(rows) <= self._slots.numel() else None + # Read only by this step's forward, so both uploads may wait for the step-input pack's flush. + if armed is not None: + slots = torch.tensor(armed, dtype=torch.long, pin_memory=prefer_pinned()) + copy_h2d_if_changed(self._slots[: len(armed)], slots, pack=True) + ctrl = torch.tensor( + [len(armed) if armed is not None else 0], dtype=torch.int32, pin_memory=prefer_pinned() + ) + copy_h2d_if_changed(self._ctrl, ctrl, pack=True) + self._armed = armed + return staged + + # The launches the store wait may precede. store_copies_done is an external event: under graph capture its wait + # is a node of its own, so the kernel launched right after it loses its programmatic (PDL) edge and with it the + # overlap of its pre-wait part. Only a launch that waits for its predecessor's completion anyway may follow it: + # the fold tail (a plain launch, no griddepcontrol). Waiting before the last draft_step instead cost +13 / +17 + # us/step. + STORE_WAIT_BEFORE = frozenset({"fold_tail"}) + + def _wait_store_copies(self, before: str) -> None: + """Order the store's rewrite after the sampler's host copies of it; ``before`` names the launch that follows. + Under graph capture a launch outside STORE_WAIT_BEFORE fails the capture.""" + if self.store_copies_done is None: + return + if before not in self.STORE_WAIT_BEFORE and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"M3SpecStepFold: the store wait would precede {before!r}, a launch outside " + f"{sorted(self.STORE_WAIT_BEFORE)}; an external-event wait node drops the next kernel's PDL edge" + ) + torch.cuda.current_stream().wait_event(self.store_copies_done) + + def store_in_forward(self, outputs: dict) -> None: + """At the forward's tail: store the armed rows' outputs and stage the next step's inputs (one launch).""" + if self._ctrl is None or self.store is None or self.step_inputs is None: + return + store = self.store + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + new_tokens = outputs.get("new_tokens") + next_new_tokens = outputs.get("next_new_tokens") + new_tokens_lens = outputs.get("new_tokens_lens") + next_draft_tokens = outputs.get("next_draft_tokens") + if any( + t is None or not isinstance(t, torch.Tensor) + for t in (new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens) + ): + return + rows = new_tokens.shape[0] + # The previous step's host copies of the store run on the sampler's side stream. + self._wait_store_copies(before="fold_tail") + _record("m3_fold_tail", rows) + if not m3_fold_tail.fold_tail( + self._ctrl, + self._slots, + min(rows, self._slots.numel()), + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + next_draft_tokens, + store.next_draft_tokens, + new_tokens_lens, + store.new_tokens_lens, + input_ids, + draft_tokens, + position_id_offsets, + kv_lens_offsets, + next_new_tokens.shape[1], + cpp=True, + ): + return + if torch.cuda.is_current_stream_capturing(): + self._graph_outputs.add(new_tokens.data_ptr()) + else: + self._eager_outputs = new_tokens + + def gather( + self, + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + if self.step_inputs is None: + return False + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + return m3_spec_tail.step_gather( + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len, + tokens_per_gen_step, + draft_token_buffer_width, + input_ids, + num_tokens, + draft_tokens, + num_draft_tokens, + position_id_offsets, + kv_lens_offsets, + first_previous_row, + ) + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list) -> bool: + armed, self._armed = self._armed, None + new_tokens = outputs.get("new_tokens") + ran = isinstance(new_tokens, torch.Tensor) and ( + new_tokens is self._eager_outputs or new_tokens.data_ptr() in self._graph_outputs + ) + self._eager_outputs = None + stored = ( + ran + and armed is not None + and num_skip == 0 + and tuple(seq_slots) == armed + and new_tokens.shape[0] == len(armed) + ) + self._staged_for = armed if stored else None + if stored and not self._logged: + self._logged = True + logger.info("[M3 host] spec step stored and next step staged in the forward tail") + return stored + + def scatter( + self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) -> bool: + store = self.store + if store is None: + return False + return m3_spec_tail.slot_scatter( + seq_slots, + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + new_tokens_lens, + store.new_tokens_lens, + next_draft_tokens, + store.next_draft_tokens, + ) + + +class _DraftForward: + """``draft_model.model`` for the in-tree worker: the drafter layer's forward, called as + ``model(input_ids=, position_ids=, hidden_states=, attn_metadata=, spec_metadata=, all_rank_num_tokens=)`` and + returning (hidden, hidden to carry to the next pass). ``d2t`` None: the draft vocabulary is the target's.""" + + d2t = None + + def __init__(self, draft: "M3Eagle3Draft") -> None: + self._draft = draft + + def __call__(self, **kwargs): + return self._draft.forward_layer(**kwargs) + + +# The drafter's projections by weight-table name, for the low-M linears' engagement and prebuild. +_DRAFT_LINEARS = ("fc", "qkv", "o", "gu", "dn") + + +class M3Eagle3Draft(nn.Module): + """The one-layer BF16 Eagle3 drafter (``speculative_model``): its weights (``weights.DRAFT_WEIGHTS``), its + vocab-parallel BF16 head and logits processor, the drafter layer's TRTLLM attention backend (layer index = the + target's layer count, FP8 KV) and its TP all-reduce. + + What the worker reads: ``model`` (the layer forward), ``apply_eagle3_fc`` (the per-layer fc_norms + fc over the + captured target hidden states), ``lm_head`` and ``logits_processor``; the fused draft loop reads the layer's parts + (``draft_input_weights``, ``attention``, ``mlp``, ``ar_norm``). The drafter embeds with the target's table. + + The BF16 projections (fc, qkv, o_proj, gate_up, down; not the head) run FlashInfer's fixed-tactic low-M kernels + (``ops/gemm/m3_draft_lowm``, the TP8 port) at the rows ``DRAFT_LOWM`` covers, compiled before the warmup + (``prebuild_linears``); every other call is the in-tree cuBLAS GEMM. Rubin TP2 (mclk locked, 4 rotating sets) at + 1 row: qkv 13.18 -> 9.71, o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. The + split-K ('s') tactics run this target's copy of FlashInfer's kernel (``ops/m3_lowm_splitk``, bit for bit the + same output); with ``DRAFT_AR_PUSH`` the o_proj and down projection of the fused draft loop also push their rows + into the drafter boundary all-reduce's exchange (``arm_push`` / ``take_pushed``), as do the direct ('d') 1-row + linears through this target's copy of the direct kernel (``ops/m3_lowm_direct``; o_proj 9.22 -> 8.09 us, + down 13.81 -> 12.72 us per 1-row boundary). + """ + + def __init__( + self, + core: MiniMaxM3Tp2ModelingV2, + draft_config: ModelConfig, + start_layer_idx: int, + capture_layers: int, + ) -> None: + super().__init__() + cfg = draft_config.pretrained_config + _ensure_draft_vocab_size(cfg) + eagle = getattr(cfg, "eagle_config", None) or {} + mapping = draft_config.mapping + self.draft_config = draft_config + # The engine's draft weight mapper binds to these two. + self.model_config = draft_config + self.config = cfg + self.capture_layers = int(capture_layers) + self.layer_idx = int(start_layer_idx) + self.dtype = cfg.torch_dtype + self.hidden = int(cfg.hidden_size) + self.hidden_size_in = int(getattr(cfg, "target_hidden_size", cfg.hidden_size)) + self.rms_eps = float(cfg.rms_norm_eps) + self.return_hidden_post_norm = bool( + eagle.get("return_hidden_post_norm", False) or getattr(cfg, "norm_output", False) + ) + if ( + int(cfg.num_hidden_layers) != 1 + or eagle.get("eh_proj_before_attn", False) + or eagle.get("norm_before_fc", False) + or not getattr(cfg, "fc_norm", False) + or getattr(cfg, "kv_lora_rank", None) + or self.capture_layers <= 1 + or self.hidden_size_in != self.hidden + or cfg.draft_vocab_size != cfg.vocab_size + or getattr(cfg, "attention_bias", False) + or getattr(cfg, "mlp_bias", False) + or getattr(cfg, "bias", False) + or getattr(cfg, "tie_word_embeddings", False) + ): + raise ValueError( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle3 drafter with per-layer " + "fc_norms over " + "several captured layers, the target's hidden size and vocabulary and no biases" + ) + tp = mapping.tp_size + self.heads = cfg.num_attention_heads // tp + self.kv_heads = (cfg.num_key_value_heads + tp - 1) // tp + self.head_dim = int( + getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + ) + self.q_size = self.heads * self.head_dim + self.kv_size = self.kv_heads * self.head_dim + self.w = _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + self.lm_head = LMHead( + cfg.draft_vocab_size, + cfg.hidden_size, + dtype=cfg.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=draft_config.lm_head_gather_output, + reduce_output=False, + use_custom_cublas_mm=False, + quant_config=None, + ) + self.logits_processor = LogitsProcessor() + self.allreduce = AllReduce( + mapping=mapping, strategy=draft_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + quant_config = draft_config.get_quant_config() + self.attn = create_attention( + draft_config.attn_backend, + self.layer_idx, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, rope=RopeParams.from_config(cfg) + ), + quant_config=quant_config, + skip_create_weights_in_init=draft_config.skip_create_weights_in_init, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): + raise ValueError( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + self._bcg_key = f"m3_tp2_draft_{id(self)}" + _BCG_DRAFTS[self._bcg_key] = self + # The drafter embeds with the target's table, read off the core on every call (the engine materializes the + # core's weights after construction). + self._core = [core] + self.model = _DraftForward(self) + # The low-M linears: (N, K, M, PDL) of every call prebuild_linears compiled, and (N, K, M) of every call that + # logged its engagement line. + self._lowm_built: set = set() + self._lowm_engaged: set = set() + # DRAFT_AR_PUSH: the weight whose next low-M call also pushes its rows into a boundary all-reduce (data + # pointer, the bucket's multicast words, its flags, rank, world, max tokens), set by arm_push; and the data + # pointer of the output the last armed call pushed, 0 if none (read once by take_pushed). Host state set and + # read in one forward's issue order. + self._push_arm: Optional[tuple] = None + self._pushed = 0 + + @property + def embed_table(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def load_weights_from_target_model(self, target_model: nn.Module) -> None: + """Nothing to alias: the drafter reads the target's embedding table through the core and owns its head.""" + + # ------------------------------------------------------------------------------------------- the projections + + def linear(self, x: torch.Tensor, key: str) -> torch.Tensor: + """``x @ w[key]^T``: the low-M kernel where ``DRAFT_LOWM`` covers the call, else the in-tree cuBLAS GEMM.""" + w = self.w[key] + out = self._lowm_linear(w, x) + return out if out is not None else m3_bf16_linear(x, w) + + def _lowm_linear(self, weight: torch.Tensor, x: torch.Tensor) -> Optional[torch.Tensor]: + """x [M, K] @ weight^T through the low-M kernel when ``DRAFT_LOWM`` covers it; None for the in-tree GEMM. + A projection none of whose row counts is covered never reaches the low-M checks.""" + n, k = weight.shape + if x.dim() != 2 or not any(nk == (n, k) for nk, _ in DRAFT_LOWM.keys): + return None + m = int(x.shape[0]) + if not DRAFT_LOWM.covers(((n, k), m)): + uncovered(DRAFT_LOWM.path, f"the drafter's {n}x{k} linear past its tactic table's rows") + return None + if not x.is_contiguous(): + # The first pass's fc may read a column window of the captured hidden states; the in-tree GEMM takes + # strides. + uncovered(DRAFT_LOWM.path, "a strided activation (the in-tree GEMM reads it in place)") + return None + use_pdl = get_env_enable_pdl() + if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + missed(DRAFT_LOWM.path, m, "the activation or weight is not BF16") + return None + if (n, k, m, use_pdl) not in self._lowm_built: + missed( + DRAFT_LOWM.path, m, f"the {n}x{k} kernel at {m} rows was not compiled before warmup" + ) + return None + kind, tactic = tactic_for(n, k, m) + if (n, k, m) not in self._lowm_engaged: + self._lowm_engaged.add((n, k, m)) + logger.info( + f"[M3 sm_107 tp2] fast path engaged: m3_draft_lowm {n}x{k} at {m} rows (device {x.device.index})" + ) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + ar = None + if ( + self._push_arm is not None + and self._push_arm[0] == weight.data_ptr() + and ( + (kind == "s" and DRAFT_LOWM_OWN_SPLITK) + or (kind == "d" and DRAFT_AR_PUSH_DIRECT and tactic[1] % 2 == 0) + ) + ): + ar = self._push_arm[1:] + self._pushed = out.data_ptr() + _record("m3_draft_lowm", n, k, m, kind, ar is not None) + return m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, DRAFT_LOWM_OWN_SPLITK, ar) + + def prebuild_linears(self) -> int: + """Compile (and run once) every covered row count of every drafter projection, outside capture, on the + current stream, before the executor's warmup. Returns the number of calls compiled.""" + from .ops.m3_lowm_direct import compile_direct_push + from .ops.m3_lowm_splitk import compile_push + + use_pdl = get_env_enable_pdl() + built = 0 + shapes = {tuple(self.w[key].shape): self.w[key] for key in _DRAFT_LINEARS} + for (n, k), weight in sorted(shapes.items()): + if weight.dtype != torch.bfloat16: + continue + for _, m in sorted(key for key in DRAFT_LOWM.keys if key[0] == (n, k)): + if (n, k, m, use_pdl) in self._lowm_built: + continue + kind, tactic = tactic_for(n, k, m) + x = torch.zeros(m, k, dtype=torch.bfloat16, device=weight.device) + out = torch.empty(m, n, dtype=torch.bfloat16, device=weight.device) + m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, DRAFT_LOWM_OWN_SPLITK) + if DRAFT_AR_PUSH and DRAFT_LOWM_OWN_SPLITK and kind == "s": + # The push variant a boundary all-reduce may arm (o_proj / down): compiled only, never run here + # (its pushes need the all-reduce's exchange and a matching call). + compile_push(x, weight.t(), out, use_pdl, tactic) + if DRAFT_AR_PUSH and DRAFT_AR_PUSH_DIRECT and kind == "d" and tactic[1] % 2 == 0: + # The direct kernel's push variant (ops/m3_lowm_direct), compiled only. + compile_direct_push(x, weight.t(), out, use_pdl, tactic) + self._lowm_built.add((n, k, m, use_pdl)) + built += 1 + return built + + def arm_push(self, key: str, workspace, rows: int) -> None: + """Arm ``w[key]``'s next covered low-M call to push into ``workspace``'s bucket for ``rows``.""" + bucket = m3_mc_allreduce._bucket(workspace, rows) + flags = None if bucket is None else bucket.flags_for(8) + self._pushed = 0 + self._push_arm = None + if bucket is not None and flags is not None and bucket.max_tokens == 8: + self._push_arm = ( + self.w[key].data_ptr(), + bucket.mc.data_ptr(), + flags.data_ptr(), + bucket.rank, + bucket.world_size, + bucket.max_tokens, + ) + + def take_pushed(self, x: torch.Tensor) -> bool: + """Whether the armed call pushed its rows; disarms. ``x`` is what the boundary all-reduce is about to reduce: + it must be the pushing call's own output (the prepushed all-reduce reads only the pushed words, never ``x``), + so an operation between that call and the all-reduce raises instead of silently reducing the wrong rows.""" + pushed, self._pushed, self._push_arm = self._pushed, 0, None + if pushed and x.data_ptr() != pushed: + raise RuntimeError( + "the drafter's pushed rows are not the boundary all-reduce's input (an operation ran in between)" + ) + return bool(pushed) + + # ------------------------------------------------------------------------------------------------- the layer + + def fc(self, hs: torch.Tensor) -> torch.Tensor: + return self.linear(hs, "fc") + + def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The captured target hidden states [T, layers x H] to the drafter's input width: each layer's fc_norm, + then the fc (the in-tree ``Eagle3ForCausalLM.apply_eagle3_fc``).""" + hidden_states = hidden_states.to(self.dtype) + if hidden_states.shape[-1] != self.hidden: + chunks = hidden_states.chunk(self.capture_layers, dim=-1) + hidden_states = torch.cat( + [ + flashinfer_rmsnorm(c, self.w[f"fc_norm{i}"], self.rms_eps) + for i, c in enumerate(chunks) + ], + dim=-1, + ) + hidden_states = self.fc(hidden_states) + return hidden_states + + def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: + """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is + False, where the caller fuses the reduction). The backend's generation phase runs this target's full-context + decode with the fused RoPE + FP8 K/V append when the shell installed it (``ops/m3_drafter_attn``).""" + qkv = self.linear(x, "qkv") + _record("m3_draft_attention", int(md.num_tokens)) + if is_in_breakable_cuda_graph(): + # A context step's draft pass is captured with the target's: the attention runs in an eager break, on the + # step's metadata (the drafter's: the runtime's draft swap installs it on the target's object). + if get_model_extra_attrs()["attention_metadata"]() is not md: + raise RuntimeError( + "MiniMax-M3 drafter under a breakable prefill CUDA graph: its metadata is not the step's" + ) + o = qkv.new_empty((qkv.shape[0], self.q_size), dtype=torch.bfloat16) + _draft_attention_eager(qkv, o, self._bcg_key) + else: + o = m3_draft_attention(self.attn, qkv, md) + out = self.linear(o, "o") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool, cpp_silu: bool = False) -> torch.Tensor: + """gate_up, the plain SwiGLU, down (all-reduced unless ``reduce`` is False). ``cpp_silu`` (the fused draft + loop): the SwiGLU on this target's C++ op (``ops/spec/m3_draft_silu``, minimaxM3Sm107Tp2/siluAndMul.cu), bit + for bit the in-tree Triton silu_and_mul in one launch either way (no Triton on the GEN decode path).""" + gu = self.linear(h, "gu") + act = None + if cpp_silu: + act = m3_draft_silu(gu) + if act is None: + uncovered( + DRAFT_SILU_PATH, "the drafter SwiGLU input layout (the Triton op runs it)" + ) + else: + taken(DRAFT_SILU_PATH, int(gu.shape[0]), "drafter MLP SwiGLU, one C++ launch") + if act is None: + act = m3_silu_and_mul(gu, None, None, None) + out = self.linear(act, "dn") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def ar_norm(self, x: torch.Tensor, residual: torch.Tensor, norm_key: str): + """One draft-layer boundary in the in-tree structure: the projection's plain all-reduce (BF16 out), then the + fused residual add + RMSNorm (FP32 add, norm of the FP32 sum). Returns (normed, residual sum).""" + return flashinfer_fused_add_rmsnorm( + m3_allreduce(x, self.allreduce, None), residual, self.w[norm_key], self.rms_eps + ) + + def forward_layer( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + all_rank_num_tokens: Optional[List[int]] = None, + ): + """The drafter layer (the in-tree ``Eagle3DraftModel.forward`` for one layer): the embedded tokens and the + fc'd hidden states normed and concatenated, attention, the post-attention norm with the residual, the MLP and + the final norm. Returns (normed output, the hidden state the next pass carries).""" + previous = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + embeds = m3_embedding(input_ids, self.embed_table).to(self.dtype) + residual = hidden_states + h = flashinfer_rmsnorm(hidden_states, self.w["hidden_norm"], self.rms_eps) + e = flashinfer_rmsnorm(embeds, self.w["in_norm"], self.rms_eps) + h = torch.cat([e, h], dim=-1) + h = self.attention(h, attn_metadata, reduce=True) + h, residual = flashinfer_fused_add_rmsnorm( + h, residual, self.w["post_norm"], self.rms_eps + ) + h = self.mlp(h, reduce=True) + spec_metadata.maybe_capture_hidden_states(self.layer_idx, h, residual) + h, to_save = flashinfer_fused_add_rmsnorm( + h, residual, self.w["final_norm"], self.rms_eps + ) + return (h, h) if self.return_hidden_post_norm else (h, to_save) + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous + + +def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, capture_norm) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "tp_size", 1) > 1 + ) + + +class _M3ShardDraftLogits(nn.Module): + """W2b: the drafter's logits processor on an in-tree draft step: this rank's BF16 vocab shard of the drafter head + (no all-gather, no fp32 cast), which the worker's greedy top-1 reads (``greedy_sample_draft_with_tp_gather``). + Rows the caller has not selected (``return_context_logits`` False) take the original processor.""" + + def __init__(self, original: nn.Module): + super().__init__() + self._m3_original = [original] # not a submodule: the target shares it + + def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits: bool = False): + if not return_context_logits: + return self._m3_original[0]( + hidden_states, lm_head, attn_metadata, return_context_logits + ) + return lm_head.apply_linear(hidden_states, None) + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify and fused draft loop, at the in-tree numerics. + + * Verify: on decode-only batches whose verify rows the fused kernels cover (``tail_eligible``), the target's BF16 + lm_head computes this rank's vocabulary shard of the logits (the in-tree head's own local GEMM, as two + half-vocabulary GEMMs on the die that holds each half at ``UGPU_HEAD``'s rows) and one kernel replaces the logits + all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. + * Draft: the drafter's first pass may take the per-layer fc_norms in one launch (or folded into the hidden-state + capture) and run its fc beside the target head GEMM (forked right after the last captured target layer at most + ``FC_EARLY_MAX_ROWS`` verify rows; after a sampled verify's head at most ``FC_AFTER_HEAD_MAX_ROWS``). On a + decode-only step ``DRAFT_STEP`` covers, the draft loop is fused: each pass's draft input in one launch, the + drafter layer's attention and MLP with their all-reduces fused with the residual add and the next RMSNorm (the + multicast all-reduce at ``DRAFT_MC_AR``'s rows, fed by the projection's push), and one kernel for the BF16 head, + the cross-rank argmax and the pass's bookkeeping. + + Every other batch takes the in-tree paths. Methods whose docstring says "Mirrors" copy the named in-tree method + with the fused branches inserted. + """ + + def __init__(self, spec_config, mapping, use_separate_draft_kv_cache: bool = False) -> None: + super().__init__( + spec_config, mapping=mapping, use_separate_draft_kv_cache=use_separate_draft_kv_cache + ) + # Per-step state, reset at the top of every _forward_impl: the first draft pass's gather rows when this step + # takes the fused first drafter inputs, and the fc output a fork computed beside the target head. + self._m3_first_gather_ids: Optional[torch.Tensor] = None + self._m3_fc_hidden: Optional[torch.Tensor] = None + self._m3_target_lm_head: Optional[nn.Module] = None + # (fc_norm0 weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple] = None + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + # Armed per step by arm_capture_norm: the drafter whose first fc is forked at the last captured target layer, + # and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + # The 12-bit coded copy of the target head shard the fused verify streams, built before the warmup. + # The die-placed target head (uGPU), built before the warmup, or None. + self._m3_ugpu_head = None + self._m3_force_vec_warm = False + # The streamed copy of the drafter's BF16 head shard the fused draft loop reads, built before the warmup, the + # head weight version it was built from, and the replaced copies a captured graph may still read. + self._m3_draft_head = None + self._m3_draft_head_version: Optional[int] = None + self._m3_draft_head_graveyard: list = [] + # Per step: the next step's input tokens [nreq, D + 1], column 0 written by the fused verify (or the seed) and + # columns 1.. by the fused draft passes; set as _m3_next_new_tokens once the fused draft loop filled it. + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + # Per step: the first draft inputs are the fused ones (accepted tokens as they lie, fc output of every row). + self._m3_fused_first_inputs = False + + def capture_target(self, spec_metadata, layer_id: int, num_tokens: int) -> Optional[tuple]: + """Where this step's capture norm writes layer ``layer_id``'s capture, for a boundary all-reduce that writes it + itself (CAPTURE_NORM_IN_BOUNDARY): (the capture buffer's [num_tokens, H] column block, the layer's fc_norm + weight, eps), or None (no capture norm armed, not a captured layer, or more rows than the buffer). The + boundary writes every row it reduces; rows past num_capture_tokens are unread.""" + capture_norm = self._m3_capture_norm + if capture_norm is None or spec_metadata.hidden_states is None: + return None + if num_tokens > spec_metadata.hidden_states.shape[0]: + return None + hidden_size = spec_metadata.hidden_size + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + out = spec_metadata.hidden_states[ + :num_tokens, i * hidden_size : (i + 1) * hidden_size + ] + return out, capture_norm[0][i], capture_norm[1] + return None + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None, written: bool = False + ) -> None: + """The Eagle3 capture with this step's capture norm (``written``: the layer's boundary all-reduce already wrote + it, see ``capture_target``); on a step armed for the early fc fork (``arm_capture_norm``), the last captured + layer then forks the drafter's first fc onto the aux stream, where it overlaps the remaining target layers + and the head (joined by ``_m3_target_accept``).""" + if not written: + _eagle3_capture(spec_metadata, layer_id, hidden_states, residual, self._m3_capture_norm) + drafter = self._m3_fc_early_drafter + layers = spec_metadata.layers_to_capture + if drafter is not None and layers and layer_id == layers[-1]: + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = self._m3_fc_fork( + spec_metadata, spec_metadata.num_capture_tokens, drafter + ) + + def prebuild(self, lm_head: nn.Module, draft: Optional[M3Eagle3Draft]) -> None: + """From the loaded weights, before the executor's warmup, on every rank: the die-placed copy of the target head + (uGPU), the fused draft loop's head copy and the drafter's low-M + linears.""" + start = time.perf_counter() + if self.prebuild_ugpu_head(lm_head): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] uGPU target head copy placed per die before warmup: %.1f s" + % (time.perf_counter() - start) + ) + start = time.perf_counter() + if self.prebuild_draft_head(draft): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + start = time.perf_counter() + built = draft.prebuild_linears() if draft is not None else 0 + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] drafter low-M GEMMs built before warmup: %d calls, %.1f s" + % (built, time.perf_counter() - start) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Mirrors Eagle3OneModelWorker._forward_impl, with the fused verify branch.""" + runtime_draft_len = spec_metadata.runtime_draft_len + # skip the draft forward if the runtime draft length is 0 + if runtime_draft_len == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._m3_first_gather_ids = None + self._m3_next_new = None + self._m3_next_new_tokens = None + self._m3_fused_first_inputs = False + if logits is None: + # tail_eligible() skipped the gathered target logits: verify from this rank's logits shard. + accepted_tokens, num_accepted_tokens, self._m3_first_gather_ids = ( + self._m3_target_accept( + hidden_states, + spec_metadata, + batch_size, + runtime_draft_len, + draft_model, + kv_lens=getattr(attn_metadata, "kv_lens_cuda", None), + ) + ) + else: + self._execute_guided_decoder_if_present(logits) + + # Sample and accept tokens. ``input_ids`` is required by the relaxed-acceptance path (scans for + # thinking-phase tokens); ignored otherwise. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, attn_metadata, spec_metadata + ) + if self.draft_only_eligible(attn_metadata, spec_metadata, draft_model): + self._m3_seed_first_pass(num_accepted_tokens, spec_metadata, batch_size) + + # Mamba hybrid models need state updates after token acceptance (a no-op on this target's cache manager). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None) + if self.sa_enhancer is not None and sa_manager is not None: + self.sa_enhancer.extend_and_prepare( + sa_manager=sa_manager, + request_ids=spec_metadata.request_ids, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_gens=num_gens, + num_contexts=num_contexts, + max_draft_len=runtime_draft_len, + ) + + if self._uses_external_shared_target_kv: + next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + ) + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here so the post-loop restore (below) + # can put attn_metadata back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + # prepare next new tokens to support overlap scheduler + if self._m3_next_new_tokens is not None: + # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). + next_new_tokens = self._m3_next_new_tokens + else: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + # Commit model-owned target state only after draft preparation has succeeded; SpecWorkerBase aborts the + # still-pending snapshots on any exception above this point. + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices[:batch_size], + num_contexts, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """The fused first draft inputs on a step that verified with the fused kernels (or was seeded for them), + else the in-tree ones.""" + fused = self._m3_first_drafter_inputs( + position_ids, + accepted_tokens, + input_ids.shape[0], + attn_metadata.num_contexts, + spec_metadata, + draft_model, + ) + self._m3_fused_first_inputs = fused is not None + if fused is not None: + fused.update(attn_metadata=attn_metadata, spec_metadata=spec_metadata) + return fused + return super().prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + def _m3_structural_ok(self, attn_metadata, spec_metadata) -> bool: + """The batch conditions shared by the fused verify and the fused first drafter inputs: a decode-only plain + Eagle3 step without attention DP, guided decoding, relaxed acceptance or a suffix-automaton enhancer, of at + most WIDE_VERIFY_MAX_TOKENS verify rows. Host state only, so the TP ranks agree.""" + if self.is_mtp_eagle or self.use_dynamic_tree: + return False + if self._uses_external_shared_target_kv or self.sa_enhancer is not None: + return False + if ( + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp + ): + return False + if getattr(self.spec_config, "use_relaxed_acceptance_for_thinking", False): + return False + if attn_metadata.num_contexts != 0 or attn_metadata.num_seqs == 0: + return False + draft_len = spec_metadata.runtime_draft_len + if draft_len == 0 or attn_metadata.num_seqs * (draft_len + 1) > WIDE_VERIFY_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Whether this step verifies with the fused kernels instead of the gathered target logits. + + Decode-only Eagle3 batches (``_m3_structural_ok``) over an unpadded vocab-sharded bias-free target head, all + greedy at a request count ``SHARD_ACCEPT`` covers, or sampled under the temperature + top-p contract (see + _m3_sampled_verify_ok) at a request count ``SAMPLED_ACCEPT`` covers, with nothing else reading the target + logits. Every input is host state shared by all TP ranks, so the ranks agree. + """ + if not self._m3_structural_ok(attn_metadata, spec_metadata): + return False + if getattr(spec_metadata, "enable_penalty", False): + return False + if not _sharded_head(lm_head) or getattr(lm_head, "bias", None) is not None: + return False + draft_len = spec_metadata.runtime_draft_len + if spec_metadata.is_all_greedy_sample: + key = (draft_len, attn_metadata.num_seqs) + if not SHARD_ACCEPT.covers(key): + uncovered( + SHARD_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + else: + if not self._m3_sampled_verify_ok(spec_metadata): + return False + key = (draft_len, attn_metadata.num_seqs) + if not SAMPLED_ACCEPT.covers(key): + uncovered( + SAMPLED_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + self._m3_target_lm_head = lm_head + _log_engaged( + SHARD_ACCEPT.path if spec_metadata.is_all_greedy_sample else SAMPLED_ACCEPT.path, + self.mapping, + ) + return True + + def draft_only_eligible(self, attn_metadata, spec_metadata, draft_model) -> bool: + """Whether a step verified by the in-tree path can still take the fused first drafter inputs: up to + FIRST_DRAFTER_INPUTS_MAX_REQUESTS requests, the structural conditions of tail_eligible, without the + target-sampling ones.""" + if attn_metadata.num_seqs > FIRST_DRAFTER_INPUTS_MAX_REQUESTS: + uncovered( + "m3_first_drafter_inputs", f"in-tree verify at {attn_metadata.num_seqs} requests" + ) + return False + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None + + def _m3_seed_first_pass(self, num_accepted_tokens, spec_metadata, batch_size: int) -> None: + """The first draft pass's gather rows (the last accepted row of each request) from the in-tree acceptance, as + _m3_target_accept leaves them.""" + k1 = spec_metadata.runtime_draft_len + 1 + n = num_accepted_tokens[:batch_size].to(torch.int64) + rows = spec_metadata.batch_indices_cuda[:batch_size].to(torch.int64) + self._m3_first_gather_ids = rows * k1 + n - 1 + + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], + ): + """The in-tree first draft inputs of a decode-only step, with the per-layer fc_norm RMSNorms of + ``apply_eagle3_fc`` in one kernel (unless the capture already applied them) before the fc; the input ids are + the accepted tokens as they lie. None to use the in-tree path.""" + if ( + self._m3_first_gather_ids is None + or self.is_mtp_eagle + or num_contexts != 0 + or draft_model is None + ): + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + ): + return None + if not self._m3_first_inputs_covered(num_tokens, spec_metadata): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not draft_model.w["fc_norm0"]: + if torch.cuda.is_current_stream_capturing(): + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), + ) + self._m3_fc_norm_w = cache + hidden = self._m3_fc_hidden + self._m3_fc_hidden = None + if hidden is None or hidden.shape[0] != num_tokens: + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + _record("m3_chunk_norm", num_tokens) + hs = m3_spec_tail.chunk_norm(hs, cache[1], draft_model.rms_eps) + _log_engaged(CHUNK_NORM.path, self.mapping) + hidden = draft_model.fc(hs) + input_ids = accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].reshape(-1) + _log_engaged(FUSED_FIRST_INPUTS.path, self.mapping) + return {"input_ids": input_ids, "position_ids": position_ids, "hidden_states": hidden} + + def _m3_first_inputs_covered(self, num_tokens: int, spec_metadata) -> bool: + """Whether the fused first drafter inputs of ``num_tokens`` verify rows are covered: their request count + (``FUSED_FIRST_INPUTS``), and the chunk norm's rows (``CHUNK_NORM``) unless the captures already applied the + fc_norm.""" + k1 = spec_metadata.runtime_draft_len + 1 + if num_tokens % k1 or not FUSED_FIRST_INPUTS.covers(num_tokens // k1): + uncovered( + FUSED_FIRST_INPUTS.path, + f"{num_tokens} verify rows at draft length {k1 - 1} (the in-tree first drafter inputs run)", + ) + return False + if self._m3_capture_norm is None and not CHUNK_NORM.covers(num_tokens): + uncovered( + CHUNK_NORM.path, f"{num_tokens} verify rows (the in-tree first drafter inputs run)" + ) + return False + return True + + def _m3_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + after: Optional[torch.cuda.Event] = None, + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (after the event + ``after`` when given, else after everything queued on the current stream); the result is picked up by + _m3_first_drafter_inputs. Returns the stream to join, or None when the fc path does not apply (the first + drafter inputs then compute it).""" + self._m3_fc_hidden = None + if draft_model is None: + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + or self._m3_fc_norm_w is None + or not self._m3_first_inputs_covered(num_tokens, spec_metadata) + ): + return None + aux = self._m3_aux_stream + if aux is None: + aux = self._m3_aux_stream = torch.cuda.Stream() + if after is not None: + aux.wait_event(after) + else: + aux.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(aux): + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + _record("m3_chunk_norm", num_tokens) + hs = m3_spec_tail.chunk_norm( + hs, self._m3_fc_norm_w[1], draft_model.rms_eps, use_pdl=False + ) + _log_engaged(CHUNK_NORM.path, self.mapping) + self._m3_fc_hidden = draft_model.fc(hs) + return aux + + def _m3_fc_join(self, fork: torch.cuda.Stream) -> None: + """Join the fc fork onto the current stream and keep its output's memory alive for that stream. + + _m3_fc_hidden is allocated in fork (aux) stream order and read on the current stream. Without the record, + dropping it returns its block to the aux stream's pool, where later aux work can take it while the current + stream's reads are still queued (eager warm-up / mixed steps). Nothing is recorded under graph capture, as in + maybe_execute_in_parallel: a captured graph relies on the next fork's aux.wait_stream(main) to order the + reuse.""" + main = torch.cuda.current_stream() + main.wait_stream(fork) + hidden = self._m3_fc_hidden + if ( + hidden is not None + and hidden.is_cuda + and not (torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()) + ): + hidden.record_stream(main) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model: Optional[M3Eagle3Draft] + ) -> None: + """Before every target forward: whether this step's Eagle3 captures apply the drafter's fc_norm themselves. + + On a fused-verify step of ``CAPTURE_NORM``'s rows whose first drafter inputs are the fused ones and whose + fc_norm weights are already stacked (by an earlier fused first pass), the capture writes + rmsnorm(h + r) * fc_norm[i] (one kernel per captured layer instead of add + copy, or the layer's boundary + all-reduce itself, ``capture_target``) and the fused first drafter inputs read it without their own + chunk_norm. A fused-verify step of at most ``FC_EARLY_MAX_ROWS`` verify rows is also armed to fork the + drafter's first fc at the last captured layer (``capture_hidden_states``). + """ + self._m3_capture_norm = None + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = None + if draft_model is None: + return + if not self.tail_eligible(attn_metadata, spec_metadata, lm_head): + return + rows = attn_metadata.num_seqs * (spec_metadata.runtime_draft_len + 1) + # The fork reads every captured column at the last captured layer, so a final-state capture, which is written + # after the model returns, keeps the fc at its stock place. + if ( + rows <= FC_EARLY_MAX_ROWS + and FINAL_HIDDEN_LAYER_ID not in spec_metadata.layers_to_capture + ): + self._m3_fc_early_drafter = draft_model + if not CAPTURE_NORM.covers(rows): + uncovered(CAPTURE_NORM.path, f"{rows} verify rows (the in-tree capture runs)") + return + # The normed captures are read only by the fused first drafter inputs; the in-tree ones norm again. + if not FUSED_FIRST_INPUTS.covers(attn_metadata.num_seqs): + return + fc_norm_w = self._m3_fc_norm_w + if ( + fc_norm_w is not None + and fc_norm_w[0] is draft_model.w["fc_norm0"] + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.rms_eps) + _log_engaged(CAPTURE_NORM.path, self.mapping) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata) -> bool: + """Whether sampled verify rows can use the fused sampled op: only under advanced_sampling_mode=no_topk, the + deploy-time guarantee that no row uses top-k or min-p (that op applies neither).""" + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + return mode is not None and mode.skips_top_k and not mode.skips_top_p + + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + kv_lens: Optional[torch.Tensor] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path. ``kv_lens``: the + requests' KV lengths at this verify (the per-request forced-acceptance key, FORCE_AL_PER_REQUEST).""" + force_total = 0 + force_vec = None + if self.force_num_accepted_tokens != 0.0: + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # The draw pool is built by an H2D copy and the draw kernels JIT-compile per block width, so both + # happen here, before the first captured step (a scratch counter keeps the real draw sequence). + self._ensure_force_accept_rng_state(hidden_states.device) + if not self._m3_force_vec_warm: + scratch = torch.zeros(1, dtype=torch.int64, device=hidden_states.device) + dummy = torch.empty(64, dtype=torch.int32, device=hidden_states.device) + for block in (16, 32, 64): + m3_force_vec( + self._force_accept_rng_pool, + scratch, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + m3_force_vec_req( + self._force_accept_rng_pool, + dummy, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + self._m3_force_vec_warm = True + if not is_warmup: + force_total = min(int(self.force_num_accepted_tokens) + 1, draft_len + 1) + force_vec = self._m3_force_vec( + batch_size, force_total, draft_len + 1, hidden_states.device, kv_lens + ) + draft_tokens = spec_metadata.draft_tokens.view(batch_size, draft_len) + if draft_tokens.dtype != torch.int32: + draft_tokens = draft_tokens.to(torch.int32) + # Column 0 receives the last accepted token, the next step's first input (the fused draft loop adds the rest). + self._m3_next_new = torch.empty( + (batch_size, draft_len + 1), dtype=torch.int32, device=hidden_states.device + ) + rows = hidden_states[: batch_size * (draft_len + 1)] + lm_head = self._m3_target_lm_head + batch_idx = spec_metadata.batch_indices_cuda[:batch_size] + # The Eagle3 fc on the aux stream: the fork made at the last captured layer, else one beside the head. A + # sampled verify's head (the in-tree GEMM up to FC_AFTER_HEAD_MAX_ROWS verify rows) holds + # every SM while it streams: its fc is forked after it instead (overlapping the sampled acceptance), where its + # many split-K CTAs do not delay the head's start. Rubin TP2 c1: the head started 17.4 us after its input was + # ready with the fc beside it. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + head_first = ( + fork is None + and rows.shape[0] <= FC_AFTER_HEAD_MAX_ROWS + and not spec_metadata.is_all_greedy_sample + ) + if fork is None and not head_first: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the head's own local GEMM (or its per-die split), whose + # all-gather the fused kernels replace (they exchange per-row records across ranks, not logits). + ugpu = self._m3_ugpu_head + if ugpu is not None and UGPU_HEAD.covers(rows.shape[0]): + if m3_ugpu_head_op.current(ugpu, lm_head.weight): + taken(UGPU_HEAD.path, rows.shape[0]) + _record("m3_ugpu_head", rows.shape[0]) + logits = m3_ugpu_head_op.m3_ugpu_head(rows, ugpu) + else: + missed( + UGPU_HEAD.path, + rows.shape[0], + "the placed head copy is not of the head's current weight (prebuild)", + ) + logits = lm_head.apply_linear(rows, None) + else: + logits = lm_head.apply_linear(rows, None) + head_done = None + if head_first: + head_done = torch.cuda.Event() + head_done.record() + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept_kernel.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = SHARD_ACCEPT.path + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept_kernel.sampled_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + spec_metadata.temperatures[:n_rows], + spec_metadata.top_ks[:n_rows], + spec_metadata.top_ps[:n_rows], + None, + seeds, + offsets, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = "spec_tail" + if head_done is not None: + # Enqueued after the sampled acceptance (which starts first) but waits on the head only. + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model, after=head_done) + if fork is not None: + self._m3_fc_join(fork) + if result is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, None, "the fused verify's exchange workspace") + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + return result + + def prebuild_ugpu_head(self, lm_head: nn.Module) -> bool: + """Build the die-placed copy of the target head shard for the uGPU verify head (``UGPU_HEAD``) before the + executor's warmup, on every rank (the module's weight is not changed): a verify of those rows computes the + logits shard as two half-vocabulary in-tree GEMMs, each on the green context of the die whose memory holds + that half, bit for bit the in-tree head's logits. Rubin TP2 head shard: 93.3 -> 79.5 us at 4 rows, 93.5 -> + 80.6 / 80.7 at 12 / 16; in-model the decode step is ~45 us shorter at one request and ~90 us at two. False + (nothing changed) when the head or device does not qualify (see m3_ugpu_head.prepare_head).""" + self._m3_ugpu_head = None + if not UGPU_HEAD.keys or not _sharded_head(lm_head): + return False + self._m3_ugpu_head = m3_ugpu_head_op.prepare_head(lm_head) + return self._m3_ugpu_head is not None + + def _m3_force_vec( + self, + batch_size: int, + base_total: int, + max_total: int, + device, + kv_lens: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + With ``kv_lens`` (FORCE_AL_PER_REQUEST): each request draws from the fixed-seed pool at its own key, its KV + length at this verify (its prompt plus what it has accepted so far), so its draws do not depend on the + server's history or on the other requests; every TP rank holds the same lengths, so the ranks agree. Without + lengths: exactly as _apply_force_accepted_tokens draws for a decode-only batch (one shared counter bump per + step, the (counter, slot) hash), the generic path's sequence. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: + return None + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + block = max(16, 1 << (batch_size - 1).bit_length()) + if FORCE_AL_PER_REQUEST and kv_lens is not None: + keys = kv_lens[:batch_size] + if keys.dtype != torch.int32: + keys = keys.to(torch.int32) + return m3_force_vec_req( + self._force_accept_rng_pool, + keys, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + # One launch: bump the shared counter and draw, with _apply_force_accepted_tokens' exact arithmetic. + return m3_force_vec( + self._force_accept_rng_pool, + self._force_accept_rng_counter, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[M3Eagle3Draft]) -> bool: + """Build the copy of the drafter's BF16 head shard that the fused draft loop streams (rebuilt only when the + head weight tensor or its version changed; a replaced copy is kept alive), on every rank. False when the + drafter is not one the loop serves.""" + old = self._m3_draft_head + if draft_model is None or not self._m3_drafter_supported(draft_model): + self._m3_draft_head = None + return False + weight = draft_model.lm_head.weight + # Idempotent: the same weight tensor at the same version keeps the built copy, so a repeated + # cache_derived_state never re-places it under a captured CUDA graph. + if ( + old is not None + and old.weight is weight + and self._m3_draft_head_version == weight._version + ): + return True + if old is not None: + # A rebuilt copy (new weight values): keep the old one alive, since a graph captured earlier may still + # read its (placed) memory until the graphs are recaptured. + self._m3_draft_head_graveyard.append(old) + self._m3_draft_head = m3_draft_step_kernel.prepare_head( + draft_model.lm_head, self.mapping.tp_rank + ) + self._m3_draft_head_version = weight._version + return True + + def _m3_drafter_supported(self, draft_model: M3Eagle3Draft) -> bool: + """The drafter the fused draft loop runs (``M3Eagle3Draft`` validates its shape at construction) over this + worker's tensor-parallel group with a bias-free vocab-sharded BF16 head without padding.""" + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if self._d2t is not None: + return False + lm_head = draft_model.lm_head + return _sharded_head(lm_head) and m3_draft_step_kernel.head_supported(lm_head) + + def _m3_draft_ar_norm( + self, + draft_model: M3Eagle3Draft, + x: torch.Tensor, + residual: torch.Tensor, + norm_key: str, + pushed: bool = False, + ): + """One drafter boundary of the fused draft loop: the multicast all-reduce + residual + RMSNorm with the + drafter's stock rounding (``DRAFT_MC_AR``: opts bit 16, normalizing the unrounded FP32 sum as FlashInfer + fused_add_rmsnorm does; Rubin TP2 10.0..11.0 -> 4.6..5.3 us per boundary), else the in-tree structure + (``M3Eagle3Draft.ar_norm``). ``pushed``: the boundary's projection already pushed its rows into the exchange + (prepushed 2). The choice depends on shapes and a collectively created workspace only, so the TP ranks + agree. Returns (normed, residual sum).""" + rows = x.shape[0] + path = DRAFT_MC_AR.path + if DRAFT_MC_AR.covers(rows): + workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) + out = None + if workspace is not None: + _record("m3_draft_mc_ar", rows, pushed) + out = m3_mc_allreduce.m3_mc_allreduce_norm( + x.contiguous(), + residual.contiguous(), + draft_model.w[norm_key], + draft_model.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + drafter_norm=True, + prepushed=2 if pushed else 0, + ) + if out is not None: + _log_engaged(path, self.mapping) + if pushed: + taken(DRAFT_AR_PUSH_PATH, rows, "the projection pushed its rows (prepushed 2)") + elif DRAFT_AR_PUSH: + uncovered( + DRAFT_AR_PUSH_PATH, + "the boundary's projection did not push (the AR pushes x)", + ) + return out + if pushed: + raise RuntimeError( + "the drafter's rows were pushed into a boundary all-reduce that did not run" + ) + missed(path, rows, "the multicast all-reduce workspace or operand layout") + return draft_model.ar_norm(x, residual, norm_key) + + def _m3_arm_draft_push(self, draft_model: M3Eagle3Draft, key: str, rows: int) -> None: + """DRAFT_AR_PUSH: let projection ``key``'s next low-M call push its rows into the boundary all-reduce that + ``_m3_draft_ar_norm`` then runs (the multicast path only; shapes and a collectively created workspace decide + it, so the TP ranks agree).""" + if not (DRAFT_AR_PUSH and DRAFT_MC_AR.covers(rows)): + return + workspace = m3_mc_allreduce.get_workspace(self.mapping, draft_model.w[key].shape[0]) + if workspace is not None: + draft_model.arm_push(key, workspace, rows) + + def _m3_fused_loop_ok( + self, attn_metadata, spec_metadata, batch_size: int, draft_model: Optional[M3Eagle3Draft] + ) -> bool: + """Whether this step drafts with the fused loop: a decode-only step the fused verify or the seed prepared + (``_m3_first_gather_ids``) with the fused first draft inputs, a supported drafter, a (draft length, requests) + ``DRAFT_STEP`` covers, the prebuilt head and the exchange workspace. Host state only, so the TP ranks + agree.""" + if ( + self._m3_first_gather_ids is None + or not self._m3_fused_first_inputs + or draft_model is None + ): + return False + path = DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + uncovered( + path, "the drafter or its attention metadata is not one the fused loop serves" + ) + return False + draft_len = spec_metadata.runtime_draft_len + if not DRAFT_STEP.covers((draft_len, batch_size)): + uncovered(path, f"draft length {draft_len} at {batch_size} requests") + return False + head = self._m3_draft_head + if head is None or head.weight is not draft_model.lm_head.weight: + missed( + path, + batch_size, + "the drafter head's streamed copy is not built for its current weight (prebuild)", + ) + return False + if m3_draft_step_kernel.workspace(self.mapping, head.wf.device) is None: + capture_missed(path, None, "the draft-step exchange workspace") + return False + _log_engaged(path, self.mapping) + return True + + def _forward_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ): + """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop (W2b: drafting from + the head's local shard where ``_m3_wide_draft_logits_ok``).""" + if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): + return self._m3_fused_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + # W2b: on a step the fused loop declines, the drafter head returns this rank's BF16 vocab shard (not the + # gathered fp32 logits), so the greedy top-1 runs the fused shard argmax (greedy_sample_draft_with_tp_gather). + swap = self._m3_wide_draft_logits_ok(spec_metadata, draft_model) + original = draft_model.logits_processor if swap else None + if swap: + draft_model.logits_processor = _M3ShardDraftLogits(original) + try: + return super()._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + finally: + if swap: + draft_model.logits_processor = original + + def _m3_wide_draft_logits_ok(self, spec_metadata, draft_model: Optional[nn.Module]) -> bool: + """Whether the in-tree draft loop may draft from the head's local shard (W2b): ``WIDE_SPEC_TAIL``, an Eagle3 + drafter with an unpadded vocab-sharded head under plain TP, argmax drafting (no rejection sampling) and no + guided decoder (which reads full logits). Host state only, so the TP ranks agree.""" + if not WIDE_SPEC_TAIL or draft_model is None or self.is_mtp_eagle: + return False + if self.guided_decoder is not None or spec_metadata.wants_advanced_draft_sampling: + return False + if self.mapping is None or self.mapping.enable_attention_dp: + return False + lm_head = getattr(draft_model, "lm_head", None) + return ( + lm_head is not None + and getattr(draft_model, "logits_processor", None) is not None + and _sharded_head(lm_head) + and getattr(lm_head, "bias", None) is None + ) + + def greedy_sample_draft_with_tp_gather( + self, logits: torch.Tensor, spec_metadata=None, mapping_lm_head_tp=None + ) -> torch.Tensor: + """W2b (``WIDE_SPEC_TAIL``): the in-tree draft loop's greedy top-1 over this rank's vocab shard of the drafter + logits, from the fused shard argmax + cross-rank key exchange (ops/m3_shard_accept at draft length 0, at + most 48 rows per launch) instead of the per-rank max, the all-gather of (index, value) pairs and the pick. The + same tokens (keys order BF16 values exactly; ties take the lowest global index, as the pick over ranks and the + first-index max within a shard). Plain TP only (no LM-head TP in ADP); else, or when the logits are not a + unit-stride BF16 shard, or before the exchange workspace exists during a capture, the in-tree sampler.""" + tokens = None + if ( + WIDE_SPEC_TAIL + and mapping_lm_head_tp is None + and logits.dim() == 2 + and logits.dtype == torch.bfloat16 + and logits.stride(-1) == 1 + and self._draft_logits_are_sharded(logits, spec_metadata) + ): + m = logits.shape[0] + empty = torch.empty((m, 0), dtype=torch.int32, device=logits.device) + out = m3_shard_accept_kernel.shard_accept( + logits, + self.mapping.tp_rank * logits.shape[-1], + self.mapping, + empty, + m3_shard_accept_kernel._local_rows(m, logits.device), + 0, + ) + if out is not None: + tokens = out[0][:, 0] + _log_engaged("m3_draft_shard_top1", self.mapping) + if tokens is None: + return super().greedy_sample_draft_with_tp_gather( + logits, spec_metadata, mapping_lm_head_tp + ) + return tokens + + def _m3_fused_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model: M3Eagle3Draft, + draft_kv_cache_manager, + batch_size: int, + num_accepted_tokens, + ): + """Mirrors Eagle3OneModelWorker._forward_linear_draft_loop for a decode-only Eagle3 step, with each pass's + glue fused (see ``_m3_fused_loop_ok``): + + * the draft input -- the embedding lookup, the layer's input norm and hidden norm and their concat -- in one + launch (``m3_spec_tail.draft_input``; the second pass gathers the last accepted rows there); + * the drafter layer's attention and MLP (its SwiGLU on this target's C++ op), each followed by the all-reduce + with the residual add and the next RMSNorm (the post-attention norm, then the drafter's final norm): the + multicast boundary fed by the o_proj / down projection's push where covered (``_m3_draft_ar_norm``); + * one kernel for the BF16 head over the head rows (pass 0: the last accepted row of each request), the + cross-rank argmax and the bookkeeping the in-tree loop does between passes: the draft token, the next pass's + input id and position, the draft KV lengths (and pass 0's sequence lengths) (``m3_draft_step``). + + The host metadata transitions are the in-tree loop's. The drafts land in columns 1.. of the step's next-input + tokens (column 0 is the last accepted token), which ``_forward_impl`` returns as they lie. + """ + draft_len = spec_metadata.runtime_draft_len + table = draft_model.embed_table + dev = table.device + head = self._m3_draft_head + gather_ids = self._m3_first_gather_ids + ids = inputs["input_ids"] + hidden = inputs["hidden_states"] + pos = inputs["position_ids"] + if ids.dtype != torch.int32: + ids = ids.to(torch.int32) + if pos.dtype != torch.int32: + pos = pos.to(torch.int32) + num_acc = num_accepted_tokens[:batch_size] + if num_acc.dtype != torch.int32: + num_acc = num_acc.to(torch.int32) + next_new = self._m3_next_new + if next_new is None: + # Seeded from the in-tree acceptance: the last accepted token is the first-pass input at the gather row. + next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=dev) + next_new[:, 0] = ids[gather_ids] + draft_tokens = next_new[:, 1:] + w = draft_model.w + eps = draft_model.rms_eps + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) as md: + for i in range(draft_len): + x, residual = m3_spec_tail.draft_input( + ids, + table, + hidden, + gather_ids if i == 1 else None, + w["in_norm"], + eps, + w["hidden_norm"], + eps, + ) + rows = x.shape[0] + self._m3_arm_draft_push(draft_model, "o", rows) + h = draft_model.attention(x, md, reduce=False) + h, residual = self._m3_draft_ar_norm( + draft_model, h, residual, "post_norm", pushed=draft_model.take_pushed(h) + ) + self._m3_arm_draft_push(draft_model, "dn", rows) + h = draft_model.mlp(h, reduce=False, cpp_silu=True) + h, to_save = self._m3_draft_ar_norm( + draft_model, h, residual, "final_norm", pushed=draft_model.take_pushed(h) + ) + if draft_model.return_hidden_post_norm: + to_save = h + ids_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + pos_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + if not m3_draft_step_kernel.draft_step( + head, + h, + batch_size, + gather_ids if i == 0 else None, + self.mapping, + draft_tokens, + i, + ids_next, + pos, + pos_next, + md.kv_lens_cuda, + num_acc if i == 0 else None, + md._seq_lens_cuda if i == 0 else None, + draft_len, + ): + raise RuntimeError( + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" + ) + ids, pos, hidden = ids_next, pos_next, to_save + if i == 0: + # The draft step set the device sequence lengths to 1 and rolled the KV lengths back. + md._seq_lens[:batch_size].fill_(1) + md.on_update() + has_kv_cache = md.kv_cache_manager is not None + if has_kv_cache and hasattr(md, "host_request_types"): + md.host_request_types[: md.num_contexts].fill_(1) + if has_kv_cache: + md.num_contexts = 0 + self._prepare_flash_mla_generation_layout(md, 0, batch_size) + md.update_for_spec_dec() + md.use_spec_decoding = False + else: + md.update_for_spec_dec() + self._m3_next_new_tokens = next_new + return draft_tokens + + +def _build_speculation(shell, model_config: ModelConfig): + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise ValueError( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + _set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return draft_config, draft, worker, use_separate + + +# -------------------------------------------------------------------------------------------------------------------- +# Runtime hooks and the shell +# -------------------------------------------------------------------------------------------------------------------- + + +def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: + """The step-loop hooks this target opts into. + + * ``pack_step_inputs``: a decode step's small input copies go out as one upload and one scatter at the graph + replay instead of a chain of copies the next graph waits on; the scatter is parallel (a program per copy), issued + by the decode graph's own prologue (this target's one-launch C++ op, ``m3_step_prologue``) from a window copied + on a side stream while the previous graph runs, so no operation is left between the step graphs and the replay + does not start with a host copy (its copy-engine -> compute turnaround was 10.3 us at c1). Sampled speculative + batches stage their Philox seed/offset uploads, and the accepted-draft counts, into the same pack; the sampler's + slot ids stay in a device buffer re-uploaded only when they change. + * ``skip_unchanged_h2d``: steady decode re-uploads slot indices, query lengths and page tables that rarely change. + * ``kv_lens_update_once``: a speculative overlap step re-derives the MSA KV-length metadata once, after the + forward's KV-length correction (the first pass is rewritten in full by the second). + * ``spec_offsets_op``: the overlap scheduler's speculative step correction (position_ids += offsets, kv_lens += + offsets) as this target's one plain launch (``m3_spec_offsets``) instead of two ATen adds at the head of the + decode graph: the same int32 adds. + * ``loop_on_execution_stream``: the overlap loop's own stream work and the forward on one stream, so step N+1's + inputs wait on graph N in program order rather than through two cross-stream event hops. + * ``merge_tuned_tactics_across_ranks``: tensor-parallel peers capture the same GEMM tactics (rank-local profiling + noise otherwise splits them, and the slowest rank sets every step). + * ``skip_greedy_rng_upload``: all-greedy speculative batches never read the Philox seeds. + * ``spec_step_fold``: the spec sampler's slot stores and the engine's next-step input gathers; the decode forward + stores its outputs and stages the next step's inputs at its own tail (``M3SpecStepFold``). + * ``sampler_d2h_side_stream``: the spec sampler copies its per-slot store to the host on a side stream, overlapping + the next step's input staging. + * ``iteration_timing_one_event``: the iteration log times the device step with one event per boundary on a side + stream (the reported step times are the same). + * MSA: identity selections for short sequences (and short/long graph families), the one-launch KV-length update, + unchanged block-table copies skipped, pool views built once, block-id tables copied from the page-index views + (the stock builder makes a Python list of the whole padded row: on the Vera CPU 282 us per call at one request + and 561 us at two, against 8 and 13 us). + """ + return RuntimeHooks( + first_content_emit_iters=FIRST_CONTENT_EMIT_ITERS, + mxfp8_pinned_tactics=True, + late_compile_detect=True, + pack_step_inputs=True, + pack_parallel_scatter=True, + pack_in_graph=True, + pack_early_copy=True, + pack_prologue_op=m3_step_prologue, + skip_unchanged_h2d=True, + kv_lens_update_once=True, + spec_offsets_op=m3_spec_offsets, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + pack_spec_rng_uploads=True, + spec_small_h2d=True, + spec_step_fold=spec_fold, + sampler_d2h_side_stream=True, + iteration_timing_one_event=True, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=_msa_kv_update_kernel.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + block_ids_from_views=True, + ), + ) + + +class ModelingV2MiniMaxM3Nvfp4Sm107Tp2( + DecoderModelForCausalLM[MiniMaxM3Tp2ModelingV2, PretrainedConfig] +): + """MiniMax-M3 NVFP4 / sm_107 / tp2, text only, with one-model Eagle3 or no speculation. + + The shell owns the vocab-parallel BF16 ``lm_head`` and the logits processor (``DecoderModelForCausalLM``); with + speculation, the drafter (``M3Eagle3Draft``, its weights from the speculative checkpoint through + ``load_draft_weights``) and the worker (``M3Eagle3Worker``) that verifies and drafts each step. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + assert torch.cuda.get_device_capability() == _SM, ( + f"this target is built for sm_{_SM[0]}{_SM[1]}, the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + # Under modeling_v2='require' a covered decode shape that cannot take its fast path raises. + set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + model_config = _text_model_config(model_config) + cfg = model_config.pretrained_config + spec_config = getattr(model_config, "spec_config", None) + if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: + # The one-model Eagle3 layer has no NVFP4 trtllm-gen decode kernel: keep its modules and the shared + # draft cache on FP8 while the target stays NVFP4. + model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 + super().__init__( + MiniMaxM3Tp2ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self.spec_config = spec_config + self.draft_model: Optional[M3Eagle3Draft] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[M3Eagle3Worker] = None + self.use_separate_draft_kv_cache = False + self.layer_idx = FINAL_HIDDEN_LAYER_ID + self._spec_fold: Optional[M3SpecStepFold] = None + if spec_config is not None: + ( + self.draft_config, + self.draft_model, + self.spec_worker, + self.use_separate_draft_kv_cache, + ) = _build_speculation(self, model_config) + self._spec_fold = M3SpecStepFold() + self.runtime_hooks = _runtime_hooks(self._spec_fold) + if mpi_rank() == 0: + logger.info( + f"[M3 sm_107 tp2] routing: modeling_v2 target __tp2, decode rows 1..{WIDE_MAX_TOKENS} (decode layer " + f"1..{DECODE_MAX_TOKENS}; fast-path cells up to 128 rows; 129..{WIDE_MAX_TOKENS} take the stock-exact " + "paths)" + ) + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights: Dict, weight_mapper=None) -> None: + """The drafter's checkpoint (``speculative_model``), loaded by the engine after the target's.""" + _weights.DRAFT_WEIGHTS.load(self.draft_model, weights) + + def set_guided_decoder(self, guided_decoder) -> bool: + """Mirrors ``SpecDecOneEngineForCausalLM.set_guided_decoder``: guided decoding with speculation is installed + on the worker.""" + if hasattr(self.spec_worker, "set_guided_decoder"): + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + def cache_derived_state(self) -> None: + """Build the targets and every decode kernel once the weights are real, before the executor's warmup; with + speculation also the drafter's full-context decode, the fused verify's coded target head, the die-placed + target head copy, the fused draft loop's head copy and the drafter linears' low-M kernels.""" + super().cache_derived_state() + self.model.post_load_weights() + if self.draft_model is not None and DRAFTER_FULLCTX.keys: + path = DRAFTER_FULLCTX.path + recipes = {k: FULLCTX_DECODE_RECIPES[k] for k in DRAFTER_FULLCTX.keys} + installed = m3_drafter_attn.install( + self.draft_model, + recipes, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and m3_fullctx_decode_kernel.has_own_reducer(), + get_env_enable_pdl(), + fused_rope_mode=(DRAFTER_ROPE_MODE if DRAFTER_ROPE_FUSED else None), + on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), + on_taken=lambda rows: taken(path, rows), + on_uncovered=lambda reason: uncovered(path, reason), + on_missed=lambda rows, reason: missed(path, rows, reason), + ) + logger.info(f"[M3 prebuild] {path} installed on {installed} drafter attention layer(s)") + if self.spec_worker is not None: + self.spec_worker.prebuild(self.lm_head, self.draft_model) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + """One engine step: the target forward, then the logits (or, with speculation, the worker's verify and + draft, and on a decode step the spec fold's store of its outputs). Text only.""" + multimodal_params = kwargs.pop("multimodal_params", None) or [] + if any(p is not None and getattr(p, "multimodal_data", None) for p in multimodal_params): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + worker = self.spec_worker + capture = _plain_capture + capture_target = None + if worker is not None: + capture = worker.capture_hidden_states + capture_target = worker.capture_target + if spec_metadata is not None: + worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + capture=capture, + capture_target=capture_target, + **kwargs, + ) + if getattr(attn_metadata, "csa2_remote_tail_mode", None) == "source": + # A context worker stopped at the encoder/decoder boundary: a disposable token for the context-only + # completion, and no draft worker. + return self.logits_processor.forward(hidden_states, self.lm_head, attn_metadata, False) + if spec_metadata is not None and spec_metadata.is_layer_capture(FINAL_HIDDEN_LAYER_ID): + capture(spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + if worker is None: + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], self.lm_head, attn_metadata, True + ) + spec_input_ids = input_ids if input_ids is not None else kwargs.get("orig_input_ids") + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if spec_input_ids is not None: + spec_input_ids = spec_input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens) + outputs = worker( + input_ids=spec_input_ids, + position_ids=spec_position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + spec_fold = self._spec_fold + if ( + spec_fold is not None + and spec_fold.in_forward + and attn_metadata.num_contexts == 0 + and isinstance(outputs, dict) + ): + # The decode forward stores its outputs into the spec sampler's store and stages the next step's inputs + # at its tail (one launch inside the step graph). + spec_fold.store_in_forward(outputs) + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/__init__.py new file mode 100644 index 000000000000..b8c7426e3933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Activation ops of this target: the gated MLPs' fused SwiGLU.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..daea97df9d12 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SiluAndMul(OpWrapper): + """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with + ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain + SwiGLU ``silu(gate) * up``. + + The MXFP8 dense MLPs (layers 0-2, intermediate 6144 per rank) and the MoE's shared expert (intermediate 1536 per + rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 9216 per + rank) runs it plain. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=1536, oai=True)), + Cell(why="shared expert, an 8-row decode verify", spec=dict(m=8, inter=1536, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=6144, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=6144, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=9216, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=9216, oai=False), + ), + ) + + note = """ + The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two + differ only by the fast sigmoid's last bits (the test's band). + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + gate, up = x.float().chunk(2, dim=-1) + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + a = 1.0 if alpha is None else alpha + b = 0.0 if beta is None else beta + return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" + + +m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/__init__.py new file mode 100644 index 000000000000..a7f7fee7a234 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention ops of this target: MSA block selection, the sparse decode and the KV-length update.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..789ae79db88e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per +request count (modeling's ``DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 16 +KV_HEADS = 1 +HEAD_DIM = 128 +#: The split this target hands trtllm-gen at one request (modeling's DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model +#: maximum the in-tree step passes otherwise. +CAPPED = 16384 +MODEL_MAX = 1 << 20 + + +def dense_attention_reference( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first + ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's + tokens in order). Returns fp32 [total_q, heads, head_dim].""" + batch, kv_heads, max_len, _ = keys.shape + heads = q.shape[1] + group = heads // kv_heads + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + pos = torch.arange(max_len, device=q.device) + with true_fp32_matmul(): + for b in range(batch): + full = int(seq_lens[b]) + rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] + limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) + valid = pos[None, :] < limit[:, None] # [qlen, max_len] + for h in range(kv_heads): + qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] + logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + probs = torch.softmax(logits, dim=-1) + out[ + b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group + ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) + return out + + +def gather_paged_kv( + kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool + (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + page = pool.shape[2] + pages = (int(seq_lens.max()) + page - 1) // page + slots = block_table[:, :pages].long() + k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] + v = pool[slots * subpages + 1] + batch, _, kv_heads, _, d = k.shape + + def order(t): + return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) + + return order(k), order(v) + + +class M3DenseDecode(OpWrapper): + """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len + ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. + + Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into + ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: + every CTA still covers its share of each request's actual length, which may exceed it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="two requests keep the model maximum (no cap at two requests)", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), + ), + ) + + note = """ + The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped + one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV + band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + minimax_m3_trtllm_gen_dense_decode( + q, + kv_cache_manager, + layer_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + max_seq_len=max_seq_len, + max_num_requests=max_num_requests, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). + + ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + if keys.dtype == torch.float8_e4m3fn: + q = fp8_e4m3_saturate(q) + return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> None: + # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. + assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" + # Each verify row attends a causal prefix of at least one key. + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + + +m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..c2b23aa18750 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention on the TRTLLM backend: RoPE, KV-cache append and causal attention.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class M3DraftAttention(OpWrapper): + """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the + backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's + NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` + ``num_tokens``); returns [num_tokens, q heads x 128] bf16. + + ``backend`` (``TrtllmAttention``) and ``metadata`` are the runtime's objects for the drafter layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="first draft pass of a 2-request verify (2 x 4 rows)", spec=dict(reqs=2, qlen=4)), + Cell(why="a later draft pass, one row per request", spec=dict(reqs=2, qlen=1)), + Cell(why="a context chunk's draft rows", spec=dict(reqs=1, context=512)), + ) + + note = """ + `reference` is exact (float64) causal softmax attention: Q rotated by the backend's NeoX RoPE at each token's + position, over the drafter layer's K / V rows read back from its FP8 cache (dequantized by the backend's + ``kv_scale_quant_orig``). What remains between the two is the FMHA's own arithmetic (Q's rounding to the + kernel's input precision, its P rounding and accumulation order), which the test bounds. + """ + + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors grouped KV heads (query head h reads KV head h // (heads / kv heads); two per rank at tp2) and + un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + # [kv heads, tokens, dim]: each KV head's rows in token order across the pages. + keys = ( + pool[pages, 0] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + values = ( + pool[pages, 1] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + keys = keys.repeat_interleave(group, dim=0) # [heads, tokens, dim] + values = values.repeat_interleave(group, dim=0) + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] + rows = q[first : first + n] + x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,hkd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.einsum( + "thk,hkd->thd", torch.softmax(logits, dim=-1), values + ) + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) + + +m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..2fecad0e9afa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q +out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3Fp8IndexerQkNormRope(OpWrapper): + """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and + writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, + then RoPE on its leading ``rotary_dim`` channels (NeoX). + + The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; + the fused layers' producer writes the index cache itself. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, two index heads per rank", spec=dict(m=1, heads=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, heads=2)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=2)), + ) + + note = """ + `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 + before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of + values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. + """ + + def __call__( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + index_k_cache, + out_cache_loc, + num_heads_q, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + position_ids, + ) + + def reference( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, + then the E4M3 cast.""" + out = fused_qk_norm_rope( + qk.clone(), + num_heads_q, + 1, + 0, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) + return ( + q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), + k.to(torch.float8_e4m3fn), + ) + + +m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py new file mode 100644 index 000000000000..bc3cd23b289f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' full-context decode attention at TP2 on Rubin (``ops/m3_fullctx_decode``): PrimTS split-KV +decode at the recipe modeling's ``FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) + +from .. import m3_fullctx_decode as kernel +from .m3_dense_decode import dense_attention_reference, gather_paged_kv + +__all__ = ["m3_fullctx_decode"] + +HEADS = kernel.HEADS +KV_HEADS = kernel.KV_HEADS +HEAD_DIM = kernel.HEAD_DIM +#: The plan capacity (the per-request K/V length bound PrimTS sizes its policy for): dispatch's +#: FULLCTX_DECODE_MAX_KV_LEN, which the forward's prebuilt plans use. +MAX_KV_LEN = 1 << 21 + + +class M3FullCtxDecode(OpWrapper): + """Decode attention of every query row over its request's full causal context in the paged FP8 pool, written + into ``output`` (BF16 [T, heads, head_dim]) in place. + + ``recipe`` is ``modeling's FULLCTX_DECODE_RECIPES[(qlen, requests)]``: the PrimTS Q tile and the split-KV fanout. + ``own_reducer`` merges the splits on the BF16 op (modeling's ``FULLCTX_DECODE_OWN_REDUCER``) rather than PrimTS's + FP16 reducer and a cast. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over an AgentX-length (165k) context", + spec=dict(kv_lens=(165525,), qlen=4), + ), + Cell( + why="one request, one draft-pass row, over a 165k context", + spec=dict(kv_lens=(165522,), qlen=1), + ), + Cell( + why="two requests verifying 1 + 3 rows at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165525, 120011), qlen=4), + ), + Cell( + why="two requests, one row each, at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165522, 120008), qlen=1), + ), + Cell( + why="one request verifying 1 + 3 rows over a short (4.1k) context: most splits are empty", + spec=dict(kv_lens=(4104,), qlen=4), + ), + ) + + note = """ + Same operands and arithmetic as the stock trtllm-gen call (E4M3 q, E4M3 K/V, P quantized to E4M3 per split, + fp32 accumulation): the split differs, so the result differs from stock as a stock split change does (the P + quantization is per split). Deterministic, and the same bits eager and in a CUDA graph. Held to the fp32 + reference within the FP8-KV band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + shape = (decode_query_len, int(seq_lens.shape[0])) + kernel.prebuild(q.device, {shape: recipe}, max_kv_len, own_reducer, get_env_enable_pdl()) + plan = kernel.plan_for(q.device, *shape, own_reducer) + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + k_pages, v_pages = kernel.kv_page_views(pool, subpages) + plan.run(fp8_e4m3_saturate(q), k_pages, v_pages, block_table, seq_lens, sm_scale, output) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V, q rounded to E4M3 first as the kernel runs it.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + return dense_attention_reference( + fp8_e4m3_saturate(q), keys, values, seq_lens, decode_query_len, sm_scale + ) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> None: + pool, _ = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + assert pool.dtype == torch.float8_e4m3fn, "an FP8 E4M3 pool" + assert tuple(pool.shape[1:]) == (KV_HEADS, kernel.PAGE, HEAD_DIM), ( + "the TP2 dense-layer pool geometry" + ) + assert q.shape[1:] == (HEADS, HEAD_DIM), "32 query heads of 128 per rank" + assert int(q.shape[0]) == int(seq_lens.shape[0]) * decode_query_len, ( + "q covers exactly the batch" + ) + if torch.cuda.is_current_stream_capturing(): + return # the length checks read the device: eager calls only + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + assert bool((seq_lens <= max_kv_len).all()), "a request longer than the plan capacity" + + +m3_fullctx_decode = M3FullCtxDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..be1c89b97261 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py @@ -0,0 +1,239 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score``.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_index_score"] + +#: The wide target's 9..48-request buckets (modeling's INDEX_SCORE_WIDE_SPLITS, repeated here because the ops do not +#: import the dispatch; test_modeling_v2_m3_dispatch checks that every covered request count is certified, so they +#: cannot drift). +_WIDE_SPLITS = ((10, 127), (13, 97), (17, 74), (23, 55), (31, 41), (39, 32), (48, 26)) + + +def _wide_split(requests: int) -> int: + return next(s for hi, s in _WIDE_SPLITS if requests <= hi) + + +def _wide_lens(requests: int) -> tuple: + """Short, mid and edge contexts (one 165521-token request at 12 / 20 / 32 / 46 requests).""" + base = (129, 4000, 20000, 65536, 300, 1000, 8193, 16384) + lens = [base[i % len(base)] for i in range(requests)] + if requests in (12, 20, 32, 46): + lens[0] = 165521 + return tuple(lens) + + +PAGE = 128 +HEAD_DIM = 128 + + +def compiled_scorer(key: tuple): + """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss + (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" + from ..m3_index_score import IndexScoreRunner + + if key not in IndexScoreRunner.kernel_cache: + IndexScoreRunner._compile(*key) + return IndexScoreRunner.kernel_cache[key] + + +class M3IndexScore(OpWrapper): + """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the + positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of + ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is + written). FP8 or BF16 q and index-K, fp32 products and sums. + + ``early`` (the target's kernel at every request count it covers) issues the Q and first K tiles right after the + grid wait from a pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise + every older block of the CTA is prefetched. Neither, nor ``split_k``, changes a score. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over 165521 tokens (AgentX ISL) at TP2: 1536 CTAs, early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=2, split_k=1536, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most of the 1536 CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=2, split_k=1536, early=True), + ), + Cell( + why="one request of 400000 tokens: 3125 blocks, so the 1536 CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=2, split_k=1536, early=True), + ), + Cell( + why="one request decoding one row over 65536 tokens", + spec=dict(seq_lens=(65536,), qlen=1, heads=2, split_k=1536, early=True), + ), + Cell( + why="two requests verifying 4 rows each over 165521 and 20000 tokens: 768 CTAs each, early issue", + spec=dict(seq_lens=(165521, 20000), qlen=4, heads=2, split_k=768, early=True), + ), + Cell( + why="two requests, one of 300 tokens and one of 262144 (2048 blocks: the 768 CTAs loop)", + spec=dict(seq_lens=(300, 262144), qlen=4, heads=2, split_k=768, early=True), + ), + Cell( + why="three requests verifying 4 rows each over 165521, 148968 and 132416 tokens: 424 CTAs each", + spec=dict(seq_lens=(165521, 148968, 132416), qlen=4, heads=2, split_k=424, early=True), + ), + Cell( + why="three requests, 129 / 323584 / 20000 tokens (2528 blocks: the 424 CTAs loop; the short one drains)", + spec=dict(seq_lens=(129, 323584, 20000), qlen=4, heads=2, split_k=424, early=True), + ), + Cell( + why="four requests verifying 4 rows each over 165521, 148968, 132416 and 115864 tokens: 318 CTAs each", + spec=dict( + seq_lens=(165521, 148968, 132416, 115864), qlen=4, heads=2, split_k=318, early=True + ), + ), + Cell( + why="four requests, 300 / 400000 / 129 / 65536 tokens (3125 blocks: the 318 CTAs loop)", + spec=dict(seq_lens=(300, 400000, 129, 65536), qlen=4, heads=2, split_k=318, early=True), + ), + Cell( + why="five requests over 165521 down to 99312 tokens (0.6 x): 254 CTAs each", + spec=dict( + seq_lens=(165521, 148968, 132416, 115864, 99312), + qlen=4, + heads=2, + split_k=254, + early=True, + ), + ), + Cell( + why="six requests, 129 / 400000 / 300 / 20000 / 65536 / 1000: 212 CTAs (3125 blocks loop; short drain)", + spec=dict( + seq_lens=(129, 400000, 300, 20000, 65536, 1000), + qlen=4, + heads=2, + split_k=212, + early=True, + ), + ), + Cell( + why="seven requests over 323584 down to 129414 tokens (0.4 x): 181 CTAs each, all loop", + spec=dict( + seq_lens=(323584, 291225, 258867, 226508, 194150, 161792, 129433), + qlen=4, + heads=2, + split_k=181, + early=True, + ), + ), + Cell( + why="eight requests, 165521 down to 49656 tokens (0.3 x) and one of 129: 159 CTAs each", + spec=dict( + seq_lens=(165521, 148968, 132416, 115864, 99312, 82760, 49656, 129), + qlen=4, + heads=2, + split_k=159, + early=True, + ), + ), + ) + tuple( + Cell( + why=f"{r} requests at {_wide_split(r)} CTAs each (W-IDX1): " + + ("one long request" if r in (12, 20, 32, 46) else "short and mid contexts"), + spec=dict(seq_lens=_wide_lens(r), qlen=4, heads=2, split_k=_wide_split(r), early=True), + ) + for r in range(9, 49) + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..6ec67e73edd4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MiniMax-M3 MSA attention backend's step: KV write, sparse top-k block selection, FMHA.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + +from .m3_sparse_decode_lc import dequant_nvfp4_pages + + +class M3MsaAttention(OpWrapper): + """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` + [T, heads * 128] (bf16, the step's live rows). + + * ``k`` / ``v`` given (the dense layers): written into the layer's paged main cache first (NVFP4 with + ``kv_scale_orig_quant``, else FP8 / BF16 as the pool is). None: a producer already wrote the caches. + * ``sparse``: the indexer scores the layer's index-K cache (already written) with ``idx_q`` and publishes the + top-k blocks; the FMHA attends those blocks. Otherwise the FMHA attends the whole page table. + + The FMHA dispatches on the step's metadata (context rows: the MSA prefill FMHA; generation rows: the MSA decode + FMHA, the dense layers' trtllm-gen decode at their capped KV split). ``backend`` and ``metadata`` are the runtime's + objects for this layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dense layer, a prefill chunk, FP8 pool", + spec=dict(layer="dense", phase="context", tokens=1024), + ), + Cell( + why="dense layer, a decode step of 2 requests", + spec=dict(layer="dense", phase="generation", reqs=2), + ), + Cell( + why="sparse layer, a prefill chunk beyond 16 blocks, NVFP4 pool", + spec=dict(layer="sparse", phase="context", tokens=4096), + ), + Cell( + why="sparse layer, a mixed step (context + generation rows)", + spec=dict(layer="sparse", phase="mixed", tokens=512, reqs=2), + ), + ) + + note = """ + `reference` is exact (float64) softmax attention over the cache rows the step attends, read back from the paged + cache the call wrote (or a producer wrote before it) and dequantized as the FMHA dequantizes them: FP8 rows as + E4M3, NVFP4 rows as E2M1 x E4M3 block scale x the per-tensor ``kv_scale_quant_orig``. A sparse layer's top-k blocks + are re-selected from ``idx_q`` and the index-K cache with the backend's rule (per-block max over the causal + keys, the query's own block forced in, ties to the lower block). On an FP8 pool the FMHA takes Q as a plain E4M3 + cast, which the reference applies too. What remains between the two is the FMHA's own arithmetic (its P + rounding and accumulation order), which the test bounds. + """ + + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors grouped KV heads (query + head h reads KV head h // (heads / kv heads); two per rank at tp2), each KV head selecting its own top-k + blocks from its own index heads against the rank's single index-K head. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + assert config.num_index_heads % kv_heads == 0 + idx_group = config.num_index_heads // kv_heads + + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + qpos = cached[row] + torch.arange(n, device=q.device) + causal = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + for g in range(kv_heads): + if nvfp4: + keys = dequant_nvfp4_pages( + pool[:, 0].view(torch.uint8), + scale_pool[:, 0], + pages, + quad_scales=False, + kv_head=g, + ) + values = dequant_nvfp4_pages( + pool[:, 1].view(torch.uint8), + scale_pool[:, 1], + pages, + quad_scales=True, + kv_head=g, + ) + keys, values = keys * k_global, values * v_global + else: + keys = pool[pages, 0, g].to(torch.float64) + values = pool[pages, 1, g].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + allowed = causal + if sparse: + # One index-K head per rank, shared by every index head (the cache is [pages, 1, page, dim]). + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = _selected_blocks( + idx64[first : first + n, g * idx_group : (g + 1) * idx_group], + idx_keys, + causal, + qpos, + page, + config, + ) + allowed = causal & chosen.repeat_interleave(page, dim=1)[:, :length] + h0, h1 = g * group, (g + 1) * group + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = ( + torch.einsum("thd,kd->thk", q64[first + lo : first + hi, h0:h1], keys) + * sm_scale + ) + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi, h0:h1] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + (ids >= n_valid - config.local_blocks) & (ids < n_valid) + ) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen + + +m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py new file mode 100644 index 000000000000..ff9a05c0233b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the TP2 heads, with its constant loads +staged before the PDL grid-dependency wait and an early dependents trigger +(``kernels/minimaxM3Sm107Tp2/producerPw2.cu``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +__all__ = ["m3_producer"] + +HEAD_DIM = 128 +ROTARY_DIM = 64 +PAGE = 128 +#: Heads per rank at TP2 (MiniMax-M3: 64 query, 4 KV and 4 index heads in all; one index-K head per token). +HEADS_Q = 32 +HEADS_KV = 2 +HEADS_INDEX = 2 +#: Packed heads per token: q, k, v, index-q and the index-k head. +HEADS_PACKED = HEADS_Q + 2 * HEADS_KV + HEADS_INDEX + 1 +#: NVFP4 block: 16 values share one E4M3 scale. +FP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _norm_rope( + x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(FP4_BLOCK) + + +def scale_index(within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return (within_page // 4) * 32 + cols * 4 + within_page % 4 + + +class M3Producer(OpWrapper): + """The NVFP4 QKV + indexer producer of a sparse layer: ``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert`` at the + TP2 heads (32 query, two KV and two index heads), head dim 128 and rotary dim 64, with the same arguments and + outputs bit for bit. + + ``packed`` [T, 39 * 128] BF16 holds each token's 32 q heads, 2 k heads, 2 v heads, 2 index-q heads and the + index-k head. Every head but v gets the Gemma RMSNorm with its weight and NeoX RoPE on its first 64 dims + (``rotary_cos_sin[position]``). Returns (q FP8 [T, 32, 128], index-q FP8 [T, 2, 128]); writes the token's slot + ``out_cache_loc[t]`` (skipped when negative or past the pool): each KV head's k and v as NVFP4 into its block of + the HND pools (16-value blocks, E4M3 scales against ``kv_quant_scale[1]`` / ``[2]``, the V scales in the 4-token + swizzle) and index-k as FP8. ``early`` triggers the dependent kernel on entry (its outputs are read only after the + dependent's own grid wait); ``prewait_step_inputs`` reads the position ids, the slots and the RoPE rows before the + kernel's own grid wait (none of them may be written by a PDL predecessor). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token (c1 without speculation)", spec=dict(tokens=1, scale=1.0)), + Cell(why="a c1 verify: 4 tokens", spec=dict(tokens=4, scale=3.0)), + Cell(why="a c4 verify: 16 tokens, small activations", spec=dict(tokens=16, scale=1e-3)), + Cell( + why="a c8 verify: 32 tokens, large activations, a skipped and an out-of-pool slot", + spec=dict(tokens=32, scale=40.0, bad_slots=True), + ), + Cell(why="the widest decode step: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="a prefill chunk of 1000 tokens, early trigger off", + spec=dict(tokens=1000, scale=1.0, early=False), + ), + ) + + note = """ + The kernel is the native producer's arithmetic, compiled with the library's --use_fast_math, so it is held to that + kernel bit for bit on every output and cache byte. The torch reference states the math in fp32: the FP8 outputs are + held to it within one E4M3 step, and the NVFP4 K/V, decoded, within one E2M1 step of their block scale. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + q, index_q = torch.ops.trtllm.minimax_m3_sm107_tp2_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + prewait_step_inputs, + ) + return q, index_q + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, 32, 128], "index_q" [T, 2, 128], per token and KV head "k" + and "v" [T, 2, 128], and "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl``, + ``early`` and ``prewait_step_inputs`` do not change them.""" + t = packed.shape[0] + x = packed.float().view(t, HEADS_PACKED, HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = rope[:, 0], rope[:, 1] + + def heads(first, count, weight): + return torch.stack( + [_norm_rope(x[:, h], weight, cos, sin, eps) for h in range(first, first + count)], + dim=1, + ) + + k0 = HEADS_Q + v0 = k0 + HEADS_KV + iq0 = v0 + HEADS_KV + return { + "q": heads(0, HEADS_Q, q_weight), + "k": heads(k0, HEADS_KV, k_weight), + "v": x[:, v0 : v0 + HEADS_KV], + "index_q": heads(iq0, HEADS_INDEX, index_q_weight), + "index_k": _norm_rope(x[:, iq0 + HEADS_INDEX], index_k_weight, cos, sin, eps), + } + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> None: + assert heads_q == HEADS_Q, f"the producer is built for {HEADS_Q} query heads, not {heads_q}" + t = packed.shape[0] + # The RoPE row is read at the position with no bound. + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + # Two tokens writing one slot race. + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" + + +m3_producer = M3Producer() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..e253c70a53aa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3QkNormRopeFp8(OpWrapper): + """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded + once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place + entry. + + The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps + of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp2 heads (32 q, 2 kv)", spec=dict(m=1, q=32, kv=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, q=32, kv=2)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=32, kv=2)), + ) + + note = """ + `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from + fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. + """ + + def __call__( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + + def reference( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" + out = fused_qk_norm_rope( + qkv.clone(), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return out.to(torch.float8_e4m3fn) + + +m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..21d361cac90c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into +the paged caches.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3QkvInsert(OpWrapper): + """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the + partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and + index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. + + ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool + [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 + [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp2 + heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), + Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), + ) + + note = """ + The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` + returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + common = ( + num_heads, + num_kv_heads, + num_index_heads, + head_dim, + rotary_dim, + eps, + q_norm, + k_norm, + index_q_norm, + index_k_norm, + rotary_cos_sin, + position_ids, + ) + if kv_scale is not None: + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common + ) + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, index_k_cache, out_cache_loc, *common + ) + + def reference( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of + Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as + packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" + m = packed.shape[0] + sizes = [ + num_heads * head_dim, + num_kv_heads * head_dim, + num_kv_heads * head_dim, + num_index_heads * head_dim, + head_dim, + ] + q, k, v, iq, ik = packed.float().split(sizes, dim=-1) + cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) + half = rotary_dim // 2 + + def norm_rope(x, w, heads): + x = x.reshape(m, heads, head_dim) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) + x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] + c, s = cos.unsqueeze(1), sin.unsqueeze(1) + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) + + return { + "q": norm_rope(q, q_norm, num_heads), + "index_q": norm_rope(iq, index_q_norm, num_index_heads), + "k": norm_rope(k, k_norm, num_kv_heads), + "v": v.reshape(m, num_kv_heads, head_dim), + "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], + } + + +m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..76841cca732d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SelectBlocks(OpWrapper): + """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), + with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the + row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` + backs the result head-major so ``permute(1, 0, 2)`` is contiguous. + + The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="eight decode rows of 17..128 blocks", + spec=dict(total_q=8, n_blocks=128, mixed=True), + ), + Cell( + why="one decode row at the 128-block edge", + spec=dict(total_q=1, n_blocks=128, mixed=False), + ), + Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), + ) + + note = """ + Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. + """ + + def __call__( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + return torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output + ) + + def reference( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + """Plain torch selection; ``head_major_output`` changes the backing only.""" + heads, n_blocks, total_q = max_score_kv.shape + out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) + blocks = torch.arange(n_blocks, device=max_score_kv.device) + for t in range(total_q): + nv = int(n_valid_blocks[t]) + for h in range(heads): + s = max_score_kv[h, :, t].clone().double() + s[blocks >= nv] = float("-inf") + forced = (blocks < min(init_blocks, nv)) | ( + (blocks >= nv - local_blocks) & (blocks < nv) + ) + s[forced] = float("inf") + order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] + ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) + out[t, h, : ids.numel()] = ids + return out + + +m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..66f7f33e8452 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). +MAX_BLOCKS = 9216 +#: Blocks the selection keeps per (query, KV head). +TOPK = 16 +#: Scores the forcing substitutes: initial blocks rank above local ones, and both above any finite score. +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 + +# M3's checkpoint at TP2 (two KV heads per rank): no forced initial blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=2, init_blocks=0, local_blocks=1) + + +class M3SelectLong(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query's top-16 blocks. + + ``scores`` is the per-KV-head block maxima [num_kv_heads, num_blocks, + total_q] the CuTe DSL index scorer wrote (amax over the KV head's index + heads). Per (query, KV head) row the op keeps the ``n_valid_blocks[query]`` + leading blocks (clamped to ``[0, num_blocks]``), forces the last + ``local_blocks`` of them (counted from the *unclamped* valid count) to + ``LOCAL_SCORE`` and the first ``init_blocks`` to ``INIT_SCORE`` -- local wins + where both apply -- and selects the 16 largest, ties to the smaller block. + A pick whose forced score is -inf is emitted as -1. The ids are written + ascending with a -1 tail. + + Same contract as ``torch.ops.trtllm.minimax_m3_select_blocks`` (the in-tree + ``select_blocks_from_maxscore``) and bit-identical to it, except on a row + where more than 2048 valid blocks tie at the 16th-largest score: there the + in-tree op fills the tied slots in atomic order (its staging holds 2048 + candidates), while this op breaks the tie to the smallest blocks as the + contract says. Up to 2048 tied blocks both take the smallest. The test checks + both against ``reference``. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the TP2 engine's 5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + note = """ + Certified at two KV heads per rank (TP2, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`). The kernel takes any head count and + forcing; other forcing is not driven here. + + `scores` may be strided (the op reads it by stride); the target hands over + the scorer's contiguous view, which is what the cells drive. Rows of at + most 128 blocks are within the kernel's range, but the target keeps those on + the in-tree op, so they are not certified here. + + Ordering is by the descending-score key `orderedKey`: a positive NaN ranks + above +inf, +0 above -0, and a negative NaN below -inf -- which, unlike + -inf, is emitted as a block id rather than -1. The cells drive positive + NaN only. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sm107_tp2_select_long( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and + ``trigger_phase`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() # [total_q, num_blocks] + forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where( + in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max + ) + k = min(TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) + ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> None: + """The inputs this op takes and answers wrongly. + + Negative forcing counts are accepted and push the forced ranges past the + row instead of failing; the in-tree op this one stands in for refuses + them, so a caller relying on this one would silently diverge from it. + An ``out`` overlapping ``scores`` is overwritten while it is still read. + """ + assert init_blocks >= 0 and local_blocks >= 0, ( + f"init_blocks ({init_blocks}) and local_blocks ({local_blocks}) must be non-negative; the " + "op accepts negative counts and selects as if the forced ranges were elsewhere" + ) + storage = scores.untyped_storage() + assert not (storage.data_ptr() <= out.data_ptr() < storage.data_ptr() + storage.nbytes()), ( + "out must not alias scores: the op writes it while scores are still read" + ) + + +m3_select_long = M3SelectLong() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py new file mode 100644 index 000000000000..205e304d5f8e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long_dsl.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_select_long import _M3, MAX_BLOCKS, M3SelectLong + + +class M3SelectLongDsl(OpWrapper): + """``m3_select_long``'s contract, cells and reference, run by ``ops/m3_select_long_dsl``: the same selection + bit for bit as the 1024-thread ``minimax_m3_sm107_tp2_select_long`` (also on rows where more than 2048 blocks tie + at the 16th-largest score), from one 384-thread CTA per row. Returns ``out``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the TP2 engine's 5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + """``m3_select_long``'s reference: the selection is the same bit for bit.""" + return M3SelectLong.reference( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> None: + """``m3_select_long``'s domain.""" + M3SelectLong.is_valid( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + + note = """ + Certified at two KV heads per rank (TP2, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`), and bit for bit against the + 1024-thread `m3_select_long` on every cell. The kernel takes any head count + and forcing; other forcing is not driven here. + + The kernel is compiled per (PDL launch, early trigger, trigger phase) with + symbolic shapes and strides; the target compiles its keys before warmup + (`ops/m3_select_long_dsl.prebuild`), and a call outside them compiles + on first use. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + from ..m3_select_long_dsl import select_long_dsl + + if scores.shape[1] > MAX_BLOCKS: + raise RuntimeError( + f"m3_select_long_dsl supports at most {MAX_BLOCKS} blocks, got {scores.shape[1]}" + ) + select_long_dsl( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + +m3_select_long_dsl = M3SelectLongDsl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..98f04538e784 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This target's block-sparse decode attention (Triton, with the one-launch long-context route).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..triton_sparse_decode import minimax_m3_sparse_attn_decode +from .m3_sparse_decode_lc import PAGE, sparse_attention_reference + + +class M3SparseDecode(OpWrapper): + """``minimax_m3_sparse_attn_decode`` as ``M3Attention._decode_sparse_core`` calls it. + + Attention of each query token over the blocks ``topk_idx`` selected for it + (``[num_kv_heads, T, topk]``, ascending ids with a -1 tail), each block + truncated at the token's causal extent, written into ``output`` in place. + The target passes FP8 q, the packed NVFP4 sparse-layer cache (two KV heads) with its block and per-tensor + scales, and BF16 ``output`` [T, 32, 128]. Returns ``output``. + + Which kernel runs is decided inside, per step: a selected (non-identity) table on the TP2 geometry goes to the + one-launch ``m3_sparse_decode_lc`` kernel while modeling's ``SPARSE_DECODE_LC`` covers it; an identity table + (every request within 16 blocks, staged by the metadata), or a selected table with that range killed, runs the + Triton per-row producer, then the LSE merge. (The multi-query producer needs one KV head: never at TP2.) Every + route also publishes the MXFP8 o_proj hand-off to ``m3_prequant.SLOT``, bit for bit + ``torch.ops.trtllm.mxfp8_quantize(output, True)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="identity selection, decode without speculation at batch 8: the Triton producer + merge", + spec=dict( + seq_lens=(1, 128, 129, 700, 2048, 1500, 64, 0), decode_query_len=1, identity=True + ), + ), + Cell( + why="identity selection, Eagle3 verify at one request (T=4): the Triton producer + merge", + spec=dict(seq_lens=(1900,), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 2 (T=8), a padding request", + spec=dict(seq_lens=(2048, 0), decode_query_len=4, identity=True), + ), + Cell( + why="selected blocks, Eagle3 verify at one request (T=4): the one-launch long-context kernel", + spec=dict(seq_lens=(40000,), decode_query_len=4, identity=False), + ), + Cell( + why="selected blocks, Eagle3 verify at batch 2 (T=8): the one-launch long-context kernel", + spec=dict(seq_lens=(9000, 40000), decode_query_len=4, identity=False), + ), + ) + + note = """ + Certified only as the target calls it: two KV heads of 16 query heads each, head_dim 128, FP8 q, packed NVFP4 + cache, BF16 output, `num_topk_chunks` left to the dispatcher. The function also takes BF16 / FP8 caches, other + head geometries and an explicit split-K factor (the in-tree tests drive those); none is certified here. + + The MXFP8 hand-off of the long-context route is certified on its own entry (`m3_sparse_decode_lc`); the Triton + merge's (`_store_mxfp8_head`) is held to `torch.ops.trtllm.mxfp8_quantize` by this entry's test. + """ + + def __call__( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + kv_scale=kv_scale, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + num_topk_chunks=num_topk_chunks, + identity_blocks=identity_blocks, + ) + return output + + def reference( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + """``output`` as the call is supposed to write it: float64 block-sparse attention over the + dequantized NVFP4 blocks (``sparse_attention_reference``), rounded once to ``output``'s dtype. + + Only the packed NVFP4 cache is mirrored -- the target's only configuration -- one KV head at a time: KV head + ``h`` serves query heads ``[h * group, (h + 1) * group)`` with its own top-k row. ``kv_scale``, + ``num_topk_chunks`` and ``identity_blocks`` do not change what is computed. + """ + assert k_block_scale is not None, "mirrors the NVFP4 cache only" + kv_heads = k_paged.shape[1] + group = q.shape[1] // kv_heads + heads = [] + for h in range(kv_heads): + heads.append( + sparse_attention_reference( + q[:, h * group : (h + 1) * group], + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[h], + block_table, + seq_lens, + decode_query_len, + sm_scale, + kv_head=h, + ) + ) + return torch.cat(heads, dim=1).to(output.dtype) + + def is_valid( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) + + +m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py new file mode 100644 index 000000000000..2733eb1e7524 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py @@ -0,0 +1,474 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch long-context sparse decode over the selected NVFP4 blocks (16-CTA clusters), SM107 TP2 geometry. + +Also the plain-torch block-sparse attention and MXFP8 hand-off mirrors the +Triton sparse decode entry (``m3_sparse_decode``) reuses: both compute the same +thing over the same cache. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The TP2 geometry the kernel is compiled for: 16 query heads on each of 2 KV heads per rank. +HEADS_PER_KV = 16 +KV_HEADS = 2 +NUM_HEADS = HEADS_PER_KV * KV_HEADS +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +#: NVFP4 elements per E4M3 block scale. +SF_VEC = 16 +#: MXFP8 elements per UE8M0 scale, and the o_proj width the hand-off is laid out for. +MX_BLOCK = 32 +WIDTH = NUM_HEADS * HEAD_DIM + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, + block_scale: torch.Tensor, + pages: torch.Tensor, + quad_scales: bool, + kv_head: int = 0, +) -> torch.Tensor: + """Dequantize KV head ``kv_head`` of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, kv_head].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, kv_head].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, + kv_head: int = 0, +) -> torch.Tensor: + """Block-sparse decode attention over KV head ``kv_head`` of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages( + k, k_block_scale, pages, quad_scales=False, kv_head=kv_head + ).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True, kv_head=kv_head).reshape( + -1, HEAD_DIM + ) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, + sf [num_rows_alloc * WIDTH // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) + num_sf_k_tiles = WIDTH // 128 + sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +class M3SparseDecodeLc(OpWrapper): + """Sparse decode attention, split-K merge and MXFP8 hand-off of one TP2 decode step in one launch. + + One cluster of 16 CTAs per (query token, KV head): CTA ``r`` attends the + token's ``r``-th selected block of that head (``topk[h, tok, r]``), and the + cluster merges the 16 partials over distributed shared memory into the + head's 16 BF16 rows of ``out[tok]``. With ``q8`` / ``sf`` it also writes + the row's MXFP8 copy and swizzled UE8M0 scales for the o_proj, bitwise + ``torch.ops.trtllm.mxfp8_quantize(out, True)``. All operands are the byte + views ``m3_attn_lc.m3_lc_attn`` hands over (E4M3 ``q``, packed E2M1 pages, + E4M3 block scales). Returns ``out``. + + Per block: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16 + (exact), QK and PV on FP16 tensor cores with FP32 accumulation, the softmax + probabilities rounded to FP16 for PV, log2-domain softmax with + ``ex2.approx``. ``linear_softmax`` picks the denominator recurrence (the two + agree to rounding); SM107 keeps the log-domain recurrence of the stock + Triton producer. ``prewait`` stages the request's block-table row in shared + memory before the grid dependency wait; it does not change a value. + ``n1_ids`` (int32 [2, rows, 16], dense rows, initialized to -1) holds the + layer's selected block ids one decode step back: each CTA prefetches its + slot's block into L2 before the grid dependency wait and overwrites the slot + with this step's id (-1 past the token's block count); rows at or past + ``rows`` skip both. A hint: it does not change a value. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="Eagle3 verify at one request, 165k context: T=4, log-domain recurrence, MXFP8 hand-off", + spec=dict( + seq_lens=(165521,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="Eagle3 verify at two requests, long and medium context: T=8, no hand-off", + spec=dict( + seq_lens=(165521, 98765), + decode_query_len=4, + linear_softmax=False, + mxfp8=False, + prewait=False, + ), + ), + Cell( + why="Eagle3 verify at two requests: T=8 with the MXFP8 hand-off", + spec=dict( + seq_lens=(165521, 12345), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="two short requests with -1 tails and a zero-length padding request: T=8", + spec=dict( + seq_lens=(130, 0), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="decode without speculation at two requests: T=2", + spec=dict( + seq_lens=(165521, 5000), + decode_query_len=1, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="the wide target's largest hand-off step: 32 requests, T=128 (one full scale tile), mixed lengths " + "with tile edges and a zero-length padding request, MXFP8 hand-off", + spec=dict( + seq_lens=tuple((165521, 20000, 1281, 129)[i % 4] for i in range(31)) + (0,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="the linear denominator and pre-wait staging variants, driven at T=8", + spec=dict( + seq_lens=(165521, 20000), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + prewait=True, + ), + ), + ) + + note = """ + Certified: the eight instantiated variants (`linear_softmax` x `prewait` x + hand-off on / off). The hand-off cells are certified bit-exactly against + `mxfp8_handoff_reference` applied to the op's own BF16 rows, which equals + `torch.ops.trtllm.mxfp8_quantize(out, True)`; the BF16 rows themselves within + the band the test derives. `prewait` and the hand-off do not change the BF16 + rows. + + `topk` is the head-major [2, T, 16] table; it is read for the first + `min(16, ceil(kv_len / 128))` entries of each token; a -1 among them skips + that CTA's block, entries past them are never read. `use_pdl` only changes + the launch attribute. The hand-off assumes `num_sf_k_tiles = 32` (a + 4096-wide o_proj input) and buffers with at least as many rows as the step; + `MERGE_MXFP8` sizes them at 128 rows. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + n1_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sm107_tp2_sparse_decode_lc( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + decode_query_len, + sm_scale, + linear_softmax, + use_pdl, + prewait, + n1_ids, + ) + return out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + n1_ids: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 4096] the op is supposed to write: ``sparse_attention_reference`` per KV head in + float64, rounded once. ``linear_softmax``, ``use_pdl``, ``prewait`` and ``n1_ids`` do not change what is + computed; the hand-off buffers are checked separately with ``mxfp8_handoff_reference``. + """ + q3 = q.view(-1, NUM_HEADS, HEAD_DIM) + attn = torch.cat( + [ + sparse_attention_reference( + q3[:, h * HEADS_PER_KV : (h + 1) * HEADS_PER_KV], + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk[h], + block_table, + seq_lens, + decode_query_len, + sm_scale, + kv_head=h, + ) + for h in range(KV_HEADS) + ], + dim=1, + ) + return attn.reshape(out.shape).to(out.dtype) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + n1_ids: Optional[torch.Tensor] = None, + ) -> None: + """The inputs this op takes and answers wrongly (``m3_attn_lc.lc_applies`` is the target's form of the + layout half). + + The kernel addresses a page's token rows at a fixed 64-byte (data) and + 8-byte (scale) pitch and loads them 16 / 8 bytes at a time, so only the + page and KV-head strides are read from the views. It reads the selected + block ids, the block table and the hand-off buffers unchecked. Reads + ``topk`` and ``seq_lens`` back to the host: a test-only cost. + """ + for name, t in (("k", k), ("v", v)): + assert ( + t.dim() == 4 and t.shape[1] == KV_HEADS and t.shape[2:] == (PAGE, HEAD_DIM // 2) + ), f"{name} must be the [pages, 2, 128, 64] byte view of the rank's KV heads" + assert t.stride(3) == 1 and t.stride(2) == HEAD_DIM // 2, ( + f"{name} rows must be dense 64-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0 and t.stride(1) % 16 == 0, ( + f"{name} pages and heads must be 16-byte aligned" + ) + assert v.stride() == k.stride(), "k and v share their strides" + for name, t in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + assert ( + t.dim() == 4 + and t.shape[1] == KV_HEADS + and t.stride(3) == 1 + and t.stride(2) == HEAD_DIM // SF_VEC + ), f"{name} rows must be dense 8-byte rows: the kernel ignores the row stride" + assert t.data_ptr() % 8 == 0 and t.stride(0) % 8 == 0 and t.stride(1) % 8 == 0, ( + f"{name} pages and heads must be 8-byte aligned" + ) + assert q.dim() == 3 and q.shape[1:] == (NUM_HEADS, HEAD_DIM), "q must be [T, 32, 128]" + assert q.data_ptr() % 16 == 0 and q.stride(0) % 16 == 0 and q.stride(1) % 16 == 0, ( + "q rows are loaded 16 bytes at a time and must be 16-byte aligned" + ) + assert k_global_scale.numel() == 1 and v_global_scale.numel() == 1, ( + "global scales are single values" + ) + total_q = q.shape[0] + assert topk.dim() == 3 and topk.shape[0] == KV_HEADS and topk.shape[1] == total_q, ( + "topk is the head-major [2, T, 16] table" + ) + assert ( + total_q % decode_query_len == 0 and seq_lens.numel() >= total_q // decode_query_len + ), ( + "every query token must belong to a request: seq_lens is read at tok // decode_query_len unchecked" + ) + seq = seq_lens.tolist() + num_cols = block_table.shape[1] + for h in range(KV_HEADS): + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk[h, tok, : min(TOPK, nblocks)] + assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( + f"head {h} token {tok}: selected block ids must be -1 or below the request's block count; " + "the block table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + f"num_sf_k_tiles must be {WIDTH // 128}: the hand-off rows are {WIDTH} wide and a smaller " + "value overlaps them silently" + ) + assert ( + q8.numel() >= total_q * WIDTH + and sf is not None + and sf.numel() >= ((total_q + 127) // 128 * 128 * (WIDTH // MX_BLOCK)) + ), "the hand-off buffers must hold every row of the step: they are written unchecked" + + +m3_sparse_decode_lc = M3SparseDecodeLc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..43166b472eca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The MSA pool's page: one sparse block is one 128-token KV page. +TOKENS_PER_BLOCK = 128 + + +class MsaKvUpdate(OpWrapper): + """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. + + ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its + own buffers at the head of every captured speculative decode step. One + single-CTA launch writes, for ``batch`` requests and ``total_q`` query + tokens (all int32): + + * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, + ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); + * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` + attends up to, through ``block_table``; + * ``n_valid_blocks[t]``, when given, the block count top-k selection is + bounded by. + + It returns None; the cells read the four buffers. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation at one request: qo = 1, every buffer written", + spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="decode without speculation at the captured batch 8", + spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", + spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), + ), + Cell( + why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", + spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), + ), + Cell( + why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", + spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), + ), + Cell( + why="positions crossing a page boundary: the slot moves to the previous page of the request", + spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), + ), + ) + + note = """ + The kernel covers any batch up to the block size's loop (it strides over + 256 threads), and any page size; only `tokens_per_block = 128`, the MSA + pool's page, is certified. + + The op does not bounds-check the request index a token reads + (`q_batch_row`) or the block-table column it derives: both come from the + metadata's own staging, which keeps them in range by construction, and an + out-of-range value reads whatever memory sits there. `is_valid` states both. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.minimax_m3_msa_kv_update( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + + def reference( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], + n_valid_blocks[:total_q])`` as the op is supposed to write them. + + The in-tree metadata's own torch re-derivation + (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook + call), statement for statement, returning new tensors rather than + writing the buffers. The buffer arguments only say which outputs exist: + a ``None`` buffer gives a ``None`` result. + """ + page = int(tokens_per_block) + kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) + cu_kv = None + if cu_kv_lens is not None: + cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) + cu_kv[1:] = torch.cumsum(kv_true, 0) + qbr = q_batch_row[:total_q].to(torch.long) + qo_dev = seq_lens[:batch] + pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(block_table.shape[1]) - 1 + ) + slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + n_valid = None + if n_valid_blocks is not None: + n_valid = torch.div( + (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" + ).to(torch.int32) + return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid + + def is_valid( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + """The inputs this op takes and answers wrongly. + + Every operand but the block table is read as a dense run from its first + element, so a strided view reads the wrong elements. The request index of + each token and the table row it selects are not bounds-checked. Reads + ``q_batch_row`` back to the host: a test-only cost. + """ + for name, t in ( + ("kv_lens", kv_lens), + ("kv_lens_staged", kv_lens_staged), + ("seq_lens", seq_lens), + ("q_batch_row", q_batch_row), + ("q_intra", q_intra), + ("msa_seq_lens", msa_seq_lens), + ("out_cache_loc", out_cache_loc), + ("cu_kv_lens", cu_kv_lens), + ("n_valid_blocks", n_valid_blocks), + ): + assert t is None or t.is_contiguous(), ( + f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" + ) + rows = q_batch_row[:total_q] + assert total_q == 0 or ( + int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) + ), ( + "q_batch_row must index a request below batch (and a block-table row): the op reads the " + "request's lengths and table row unchecked" + ) + + +msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/__init__.py new file mode 100644 index 000000000000..5da86fe24bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective-communication entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..7f480855c170 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams + + +class M3AllReduce(OpWrapper): + """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the + plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: + ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; + returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the + communicator's per call. + + The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps + beyond the multicast kernel's rows. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="tp2, a prefill chunk", spec=dict(world=2, m=2048, hidden=6144)), + Cell(why="tp2, an 8-row generation step", spec=dict(world=2, m=8, hidden=6144)), + Cell( + why="tp2, plain sum of eight drafter rows", + spec=dict(world=2, m=8, hidden=6144, plain=True), + ), + ) + + note = """ + The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then + adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum + to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. + """ + + def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): + return comm(x, all_reduce_params=params) + + def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): + """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" + total = torch.stack([x.float() for x in xs]).sum(0) + if params is None: + return total.to(xs[0].dtype) + updated = (params.residual.float() + total).to(xs[0].dtype) + h = updated.float() + normed = ( + h + * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) + * params.norm_weight.float() + ) + return normed.to(xs[0].dtype), updated + + +m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..b624ae64e59e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,319 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode layer boundary, 2 ranks on SM107.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer + fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; + only the stored residual is rounded to BF16. + + The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are + summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic + and the normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + pre = total.to(torch.bfloat16).float() + residual.float() + updated = pre.to(torch.bfloat16) + u = pre.double() if drafter else updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +class M3McAllreduceNorm(OpWrapper): + """One TP group's boundary: ``updated = bf16(residual + bf16(sum_r bf16(x_r + x2_r)))``, ``normed = + RMSNorm(updated)``: bitwise the stock all-reduce + fused residual RMSNorm (+ mxfp8_quantize) of SM107 TP2. + + Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls + its own view ``ws_uc`` (re-arming the words it read) and rotates the per-CTA call parities in ``flags``; the + workspace is the ``McAllReduceWorkspace`` bucket of ``maxt`` tokens. All ranks must make the same calls in the + same stream order. ``prepushed`` 2 means a producer already pushed ``x`` (and launched this grid after its + pushes); 3 that two producers pushed ``x`` into exchange set A and ``x2`` into set B. Returns + ``[normed, updated]`` or, with ``quant``, also the MXFP8 copy of ``normed`` (E4M3 rows, swizzled UE8M0). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The tp2 decode boundaries (hidden 6144, 2 ranks, eps 1e-6): every decode step fits the 8-token bucket, epilogue + # option 7 at 1, 2, 7, 8 tokens and 15 at 3..6 (ops/m3_mc_allreduce.AR_OPTS_BY_M; every option is + # bit-identical). + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token at the attention boundary, x only, option 7", + spec=dict(m=1, maxt=8, opts=7, x2=False, quant=False, prepushed=0), + ), + Cell( + why="MoE boundary at 4 rows (one request verifying 1 + 3): shared + routed summed in-kernel, MXFP8 copy", + spec=dict(m=4, maxt=8, opts=15, x2=True, quant=True, prepushed=0), + ), + Cell( + why="MoE boundary at 8 rows (two requests), MXFP8 copy for the next QKV", + spec=dict(m=8, maxt=8, opts=7, x2=True, quant=True, prepushed=0), + ), + Cell( + why="attention boundary whose 8 rows the o_proj GEMV already pushed (prepushed 2)", + spec=dict(m=8, maxt=8, opts=7, x2=False, quant=False, prepushed=2), + ), + # Every other decode row count (modeling's MC_ALLREDUCE), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} rows: a boundary row count between the ones above", + spec=dict( + m=m, + maxt=8, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + x2=m % 2 == 0, + quant=m % 3 == 0, + prepushed=0, + ), + ) + for m in range(1, 9) + if m not in (1, 4, 8) + ), + # B3, the batch-4 decode steps (modeling's MC_ALLREDUCE past 8 rows): every row count 9..16 on the 16-token + # bucket (option 15), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} rows on the 16-token bucket (batch-4 decode step)", + spec=dict(m=m, maxt=16, opts=15, x2=m % 2 == 0, quant=m % 3 == 0, prepushed=0), + ) + for m in range(9, 17) + ), + # The batch-8 decode steps: every row count 17..32 on the 48-token bucket (option 15), the second addend and + # the MXFP8 copy alternating. + *( + Cell( + why=f"{m} rows on the 48-token bucket (batch-8 decode step)", + spec=dict(m=m, maxt=48, opts=15, x2=m % 2 == 0, quant=m % 3 == 0, prepushed=0), + ) + for m in range(17, 33) + ), + # The Eagle3 drafter boundaries (modeling's DRAFT_MC_AR, SpecTail): option 16 (the drafter's stock rounding) on + # 23 / 31 at every draft-pass row count 1..8 (pass 0: 4 x requests, later passes: requests), x only. + *( + Cell( + why=f"Eagle3 drafter boundary at {m} rows (opts {23 if m in (1, 2, 7, 8) else 31})", + spec=dict( + m=m, + maxt=8, + opts=23 if m in (1, 2, 7, 8) else 31, + x2=False, + quant=False, + prepushed=0, + drafter=True, + ), + ) + for m in range(1, 9) + ), + # Every other row count the o_proj pushes (modeling's ATTN_PUSH). + *( + Cell( + why=f"{m} rows prepushed by the o_proj GEMV", + spec=dict( + m=m, + maxt=8, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + x2=False, + quant=False, + prepushed=2, + ), + ) + for m in range(1, 8) + ), + # B6, the batch-4 decode steps (modeling's ATTN_PUSH past 8 rows): every row count 9..16 the o_proj pushes into + # the 16-token bucket (option 15). + *( + Cell( + why=f"{m} rows prepushed by the o_proj GEMV on the 16-token bucket (batch-4 decode step)", + spec=dict(m=m, maxt=16, opts=15, x2=False, quant=False, prepushed=2), + ) + for m in range(9, 17) + ), + ) + + note = """ + The epilogue option bits (1 early residual / weight load, 2 st.async cluster exchange, 4 early parity rotation, + 8 max carveout) are bit-identical by construction; only the four (bucket, option) pairs of the sm_103 kernel are + built. SM107 TP2 against the stock AUTO all-reduce + fused residual RMSNorm + mxfp8_quantize, rows 1..24, x alone + and x + x2: normed, updated and the MXFP8 copy (data and the scale bytes of valid rows) bitwise equal. `prepushed` + 1 (pushed, grid wait kept) is accepted by the op but no caller passes it. `pdl` changes only launch overlap and is + driven off in the cells. + + `normed` is gated to one BF16 rounding step: the kernel's RMS statistic is a warp / cluster tree of FP32 sums + and its `rsqrtf` is approximate, so a row whose scaled value lands next to a BF16 midpoint can round either + way. `updated` and the MXFP8 copy of the op's own `normed` are bit-exact. + """ + + def __call__( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_sm107_tp2_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + None, + None, + maxt, + quant, + prepushed, + opts, + cap_out, + cap_w, + cap_eps, + ) + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant, drafter) + + def is_valid( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> None: + h = residual.shape[-1] + # The op checks norm_w's element count, not its layout: a strided view is read as packed memory. + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + # The workspace is addressed as [set][parity][maxt][world][H / 2] words and never bounds-checked: a smaller + # buffer (another bucket's, or one built for fewer ranks) is polled and re-armed out of range. + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + # Flags are per (token, cluster CTA) of the bucket; the op checks the count, not that they belong to the + # workspace. Flags of another bucket hold another call history, and the parity they select is not the one the + # peers push into: the poll then waits on the other half forever or reads last call's rows. + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + + +m3_mc_allreduce_norm = M3McAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..9861ed2f0513 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize + multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode MoE boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_mc_allreduce_norm import boundary + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class M3McFinalizeAllreduceNorm(OpWrapper): + """The MoE boundary with the routed experts' finalize folded in. + + Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, + ``scales`` fp32 or bf16), then the group's all-reduce, residual add and RMSNorm exactly as + ``m3_mc_allreduce_norm``: same workspace, flags and call-order contract. Returns ``[normed, updated]`` or, with + ``quant``, also the MXFP8 copy of ``normed``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Not taken at tp2 (modeling's MC_FINALIZE is empty: the decode step is at most 8 rows); the cells are the tp4 + # target's, on the SM107 kernel. The tp4 target takes this path from 9 tokens (top-4 of 128 experts, fp32 routing + # weights, the shared expert always present): the 16- and 48-token buckets, epilogue option 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="the smallest finalize call (9 tokens, 16-token bucket) with the next QKV's MXFP8 copy", + spec=dict(m=9, maxt=16, opts=15, topk=4, quant=True, dropped=0), + ), + Cell( + why="a full 16-token bucket with dropped expert slots", + spec=dict(m=16, maxt=16, opts=15, topk=4, quant=False, dropped=5), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=24, maxt=48, opts=15, topk=4, quant=False, dropped=0), + ), + Cell( + why="the 48-token bucket full, MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, topk=4, quant=True, dropped=3), + ), + # Every other row count of the unfinalized MoE (MC_FINALIZE): the bucket and options the target picks. + *( + Cell( + why=f"{m} tokens: a finalize row count between the ones above", + spec=dict( + m=m, + maxt=16 if m <= 16 else 48, + opts=15, + topk=4, + quant=m % 2 == 0, + dropped=m % 4, + ), + ) + for m in range(9, 49) + if m not in (9, 16, 24, 48) + ), + ) + + note = """ + BF16 routing weights and the 8-token bucket are accepted by the op (option 7 there) but the target only enters + this path at 9 tokens and up, with the built-in MoE's fp32 weights; neither is a cell. The rest is as in + `m3_mc_allreduce_norm`'s note. + """ + + def __call__( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_sm107_tp2_mc_allreduce_norm( + fc2, + shared, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + 0, + opts, + cap_out, + cap_w, + cap_eps, + ) + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> None: + h = residual.shape[-1] + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + # The kernel gathers fc2[perm] without a bound: an index past the FC2 rows reads whatever follows them. + lo, hi = int(perm.min()), int(perm.max()) + assert lo >= -1 and hi < fc2.shape[0], ( + f"perm spans [{lo}, {hi}]; entries must be -1 (dropped) or a row of fc2 (< {fc2.shape[0]})" + ) + + +m3_mc_finalize_allreduce_norm = M3McFinalizeAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/__init__.py new file mode 100644 index 000000000000..e9f4d8e20b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedding ops of this target: the generic token-embedding gather.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..71fce6b253cb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="an 8-row decode verify, int64 ids", + spec=dict(rows=8, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ) + + note = """ + An exact gather: `reference` must match bit for bit. + """ + + def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) + + def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] + return rows * valid.unsqueeze(-1).to(table.dtype) + + +m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/__init__.py new file mode 100644 index 000000000000..4c6fc92cd49c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 decode embedding gather + first-layer RMSNorm as one CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py new file mode 100644 index 000000000000..78563b4db40d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. + +128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original +K3 wrapper this target does not promise FlashInfer's bit-identical reduction. +""" + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +NORM_THREADS = 128 +NORM_VEC = 8 +NORM_WARPS = 4 # capacity: fixed: the CTA's warps + + +@cute.kernel +def m3_embed_norm_kernel( + ids: cute.Tensor, # int32 or int64 [N] + table: cute.Tensor, # bf16 [V, H] + weight: cute.Tensor, # bf16 [H] + raw: cute.Tensor, # bf16 [N, H], out: the rows + out: cute.Tensor, # bf16 [N, H], out: the normed rows + vocab: cutlass.Int64, + eps: cutlass.Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + hidden: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 + ) + + tok = cutlass.Int64(ids[bidx]) + valid = (tok >= cutlass.Int64(0)) & (tok < vocab) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + + gX = cute.local_tile(table, tiler_mn, (row, 0)) + gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) + gY = cute.local_tile(out, tiler_mn, (bidx, 0)) + w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 + ) + thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + tXgR = thr_copy_O.partition_D(gR) + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. + tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) + if valid: + cute.copy(copy_atom_load, tXgX, tXrX) + cute.copy(copy_atom_load, tWgW, tWrW) + cute.copy(copy_atom_store, tXrX, tXgR) + + # flashinfer's RMSNormKernel arithmetic, in its order (see the header). + x = tXrX.load().to(cutlass.Float32) + x_sq = x * x + sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) + for i in cutlass.range_constexpr(5): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + if lane == 0: + reduction_buffer[0, warp] = sum_sq + cute.arch.barrier() + total = cutlass.Float32(0.0) + if lane < NORM_WARPS: + total = reduction_buffer[0, lane] + for i in cutlass.range_constexpr(5): + total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) + mean_sq = total / cutlass.Float32(hidden) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + w = tXrW.load().to(cutlass.Float32) + y = x * rstd * (w + cutlass.Float32(0.0)) + tXrO.store(y.to(out.element_type)) + cute.copy(copy_atom_store, tXrO, tXgO) + + +@cute.jit +def m3_embed_norm( + ids: cute.Tensor, + table: cute.Tensor, + weight: cute.Tensor, + raw: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int64, + eps: cutlass.Float32, + n_tokens: cutlass.Constexpr[int], + hidden: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. + blocks = hidden // (NORM_VEC * NORM_THREADS) + tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), + stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip + m3_embed_norm_kernel( + ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden + ).launch( + grid=[n_tokens, 1, 1], + block=[NORM_THREADS, 1, 1], + smem=NORM_WARPS * 4, + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py new file mode 100644 index 000000000000..e0caa5e350dc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host side of the MiniMax-M3 decode embedding + first-layer RMSNorm (m3_embedding_firstnorm_kernel.py): H 6144, +1..32 rows (one CTA per row: the decode layer's 1..8 rows and the batch-4 / batch-8 plans' 9..32), into caller-owned +buffers. + +Compiled on the first call per (rows, id dtype, table shape, PDL), which must happen outside CUDA-graph capture +(:class:`DecodeEmbedding` compiles at construction). +""" + +import threading +from typing import Dict, Tuple + +import cuda.bindings.driver as cuda_driver +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +MAX_ROWS = 32 # capacity: fixed: the kernel's built rows (decode steps) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def arg(tensor): + return from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _check(ids, table, weight, raw, out) -> None: + rows = ids.shape[0] if ids.dim() == 1 else 0 + if not ( + 1 <= rows <= MAX_ROWS + and ids.dtype in (torch.int32, torch.int64) + and ids.is_contiguous() + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN,) + and all( + t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() + for t in (raw, out) + ) + ): + raise ValueError( + f"m3_embed_norm: expects ids int32/int64 [1..{MAX_ROWS}], a contiguous bf16 table [V, 6144], " + "bf16 weight [6144] and contiguous bf16 raw / out [rows, 6144]" + ) + + +def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) + fn = _compiled[key] = cute.compile( + _m3_embed_norm_jit, + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + _check(ids, table, weight, raw, out) + fn = _compile(ids, table, weight, eps, raw, out, use_pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return out, raw + + +class DecodeEmbedding: + """Caller-owned buffers of one (rows, id dtype) decode embedding: ``raw`` (the gathered rows, the first residual) + and ``output`` (their RMSNorm), over ``table`` and the first layer's plain (non-Gemma) input norm ``weight`` / + ``eps``. Built (and compiled) before CUDA-graph capture.""" + + def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): + assert not torch.cuda.is_current_stream_capturing() + assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN + assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) + assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( + table, + weight, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(eps) + self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) + + def __call__(self, ids): + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return m3_embed_norm( + ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/__init__.py new file mode 100644 index 000000000000..9cc374bdf1cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMM entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..3ed3cc97a229 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3Bf16Linear(OpWrapper): + """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. + + The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), + o_proj (this rank's input columns), gate_up and down projections run it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=4608, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=4096)), + Cell(why="drafter gate_up, one row", spec=dict(m=1, n=18432, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=9216)), + ) + + note = """ + cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x, w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return (x.float() @ w.float().t()).to(x.dtype) + + +m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py new file mode 100644 index 000000000000..f1b619a13564 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense MLP's swigluoai activation as one C++ launch with PDL and an early release (DENSE_SWIGLU_CPP).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# The dense layers' swigluoai parameters, which the kernel hard-codes (modeling_minimaxm3._build_swiglu_oai_dense_mlp). +LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 + + +class M3DenseSwiglu(OpWrapper): + """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, + rounded once to bf16: trtllm::minimax_m3_sm107_tp2_swiglu_oai (minimaxM3Sm107Tp2/swigluOai.cu), bit for bit the + Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the + grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights + during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even + half width; ``__call__`` returns None otherwise or for other activation parameters (the caller runs the Triton + op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP2 dense MLP shard: gate | up 12288 wide (d = 6144), at the decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP2 dense MLP at 1 row", spec=dict(rows=1, d=6144)), + Cell(why="the TP2 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=6144)), + Cell(why="the TP2 dense MLP at 8 rows", spec=dict(rows=8, d=6144)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the + alpha-scaled ex2.approx sigmoid, div.full), checked over every bf16 gate and every bf16 up pattern. + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float] = LIMIT, + alpha: Optional[float] = ALPHA, + beta: Optional[float] = BETA, + use_pdl: bool = True, + ) -> Optional[torch.Tensor]: + if ( + (limit, alpha, beta) != (LIMIT, ALPHA, BETA) + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 4 != 0 + or x.data_ptr() % 4 != 0 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp2_swiglu_oai(x, out, bool(use_pdl)) + return out + + def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_dense_swiglu = M3DenseSwiglu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py new file mode 100644 index 000000000000..94f5e112e592 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py @@ -0,0 +1,169 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 x BF16 -> BF16 low-M GEMM of the Eagle3 drafter's linears at fixed FlashInfer tactics (SM10x).""" + +from __future__ import annotations + +from typing import Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: (N, K) of the TP2 drafter shards -> ((max rows, kind, tactic), ...), rows ascending. Kind "d" is FlashInfer's direct +#: SIMT kernel with tactic (block, outputs per block, rows per block); kind "s" is its split-K tensor-core kernel with +#: tactic (mma_m, mma_n, split_k, ab_stages). A row count is served by the first entry whose max rows holds it; the +#: rows modeling's DRAFT_LOWM covers are only the ones measured (1, and 4 where a tactic beats cuBLAS at no larger error +#: vs fp64). Rubin SM107, mclk locked, 4 rotating weight sets, CUDA-graph timing, us cuBLAS -> low-M at 1 / 4 rows: +#: qkv 13.18 -> 9.71 / 13.31 -> 11.72, o 7.17 -> 5.33 / 7.12 -> 5.78, gate_up 22.27 -> 17.70 / -, down 12.52 -> 9.76 / +#: 12.59 -> 10.13, fc 21.50 -> 17.77 / -. +TACTICS: Mapping[Tuple[int, int], tuple] = { + # qkv + (4608, 12288): ((1, "d", (256, 2, 1)), (4, "d", (64, 4, 4))), + # o_proj + (6144, 4096): ( + (1, "d", (128, 2, 1)), + (4, "s", (64, 8, 2, 10)), + (8, "s", (64, 8, 2, 10)), + (16, "s", (64, 16, 2, 10)), + ), + # gate_up + (18432, 6144): ((1, "d", (256, 4, 1)),), + # down + (6144, 9216): ( + (1, "d", (128, 2, 1)), + (4, "s", (64, 8, 2, 10)), + (16, "s", (64, 16, 2, 10)), + (32, "s", (64, 32, 2, 6)), + ), + # the Eagle3 fc over the three captured hidden states + (6144, 18432): ((1, "d", (384, 2, 1)), (16, "s", (64, 16, 2, 10))), +} +#: Rows each (N, K) is certified at (a subset of what its table would serve). +ROWS: Mapping[Tuple[int, int], Tuple[int, ...]] = { + (4608, 12288): (1, 4), + (6144, 4096): (1, 4, 8, 16), + (18432, 6144): (1,), + (6144, 9216): (1, 4, 16, 32), + (6144, 18432): (1, 16), +} + + +def tactic_for(n: int, k: int, m: int) -> Optional[Tuple[str, tuple]]: + """(kind, tactic) of an [m, k] x [n, k]^T call, or None when the table does not serve it.""" + for rows, kind, tactic in TACTICS.get((n, k), ()): + if m <= rows: + return kind, tactic + return None + + +class M3DraftLowM(OpWrapper): + """``out`` [M, N] BF16 = ``x`` [M, K] BF16 @ ``weight`` [N, K]^T BF16, FP32 accumulation. + + ``kind`` / ``tactic`` pick FlashInfer's kernel: "d" runs ``run_direct_dense`` (SIMT, one dot product per output in + K order), "s" runs ``run_splitk_dense`` (tcgen05 MMA, K split over ``split_k`` CTAs and summed in FP32). Both + round once to BF16. The kernel is compiled on its first call per (M, N, K, tactic, PDL); the target compiles every + covered call before warmup (``M3Eagle3Draft.prebuild_linears``). + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter qkv, 1 rows: d tactic", spec=dict(nk=(4608, 12288), m=1)), + Cell(why="drafter qkv, 4 rows: d tactic", spec=dict(nk=(4608, 12288), m=4)), + Cell(why="drafter o_proj, 1 rows: d tactic", spec=dict(nk=(6144, 4096), m=1)), + Cell(why="drafter o_proj, 4 rows: s tactic", spec=dict(nk=(6144, 4096), m=4)), + Cell(why="drafter gate_up, 1 rows: d tactic", spec=dict(nk=(18432, 6144), m=1)), + Cell(why="drafter down, 1 rows: d tactic", spec=dict(nk=(6144, 9216), m=1)), + Cell(why="drafter down, 4 rows: s tactic", spec=dict(nk=(6144, 9216), m=4)), + Cell(why="drafter fc, 1 rows: d tactic", spec=dict(nk=(6144, 18432), m=1)), + ) + + note = """ + The table is the TP2 drafter's: its shapes are the per-rank shards of the Eagle3 layer at 2-way tensor + parallelism (64 query and 4 KV heads, intermediate 18432) and the replicated fc, at the b1 decode rows (1 row + in draft passes 1-2, 4 in pass 0). Certified at those shapes and rows only; FlashInfer's kernels take other + shapes, which are not measured here. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + if kind == "d" and ar is not None: + # The boundary push (modeling's DRAFT_AR_PUSH_DIRECT): this target's copy of the direct kernel, whose + # output is FlashInfer's, also stores the rows into the boundary all-reduce's exchange. + from ..m3_lowm_direct import DirectTactic, run_direct_dense_push + + run_direct_dense_push(x, weight.t(), out, use_pdl, DirectTactic(*tactic), ar) + elif kind == "d": + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + DirectTactic, + run_direct_dense, + ) + + run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) + elif own_splitk: + # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # it also pushes the rows into a drafter boundary all-reduce (modeling's DRAFT_AR_PUSH). + from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push + + if ar is not None: + run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) + else: + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + else: + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + return out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """bf16(x @ weight^T) summed in fp64 and rounded once; ``out``, ``kind``, ``tactic`` and ``use_pdl`` only + choose where and how the op writes.""" + return (x.double() @ weight.double().t()).to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> None: + """The kernels index ``x`` and ``out`` densely: a strided view is read as if it were packed.""" + assert x.dim() == 2 and x.is_contiguous(), "x must be a contiguous [M, K] matrix" + assert weight.dim() == 2 and weight.is_contiguous(), ( + "weight must be a contiguous [N, K] matrix" + ) + assert out.is_contiguous() and tuple(out.shape) == (x.shape[0], weight.shape[0]), ( + "out must be a contiguous [M, N] matrix" + ) + assert ar is None or ( + len(ar) == 5 and ((own_splitk and kind == "s") or (kind == "d" and tactic[1] % 2 == 0)) + ), "ar (the boundary all-reduce push) runs only on this target's split-K or direct copy" + + +m3_draft_lowm = M3DraftLowM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..9702cc0778a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class M3Fp32Linear(OpWrapper): + """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the + fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies + (the engine leaves torch's default).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), + Cell( + why="a generation step above the tensor-core router's 48 rows", + spec=dict(m=64, e=128, k=6144), + ), + ) + + note = """ + Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against + the true fp32 reference. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x.to(torch.float32), w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return x.float() @ w.t() + + +m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..96352e76d700 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 x MXFP8 GEMV (tcgen05, block-scaled, CuTe DSL) on an activation its producer already quantized.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_dequant +from ..mx_gemv import op as _gemv + +__all__ = ["PROJECTIONS", "gemv_is_valid", "m3_mx_gemv", "served_cells", "served_max_rows"] + +# (N, K) of this target's MXFP8 projections at tp2, the GEMV's tuned shapes. +QKV_SPARSE = (4992, 6144) # the fused QKV + index projection (index values disabled) +QKV = (4608, 6144) +O_PROJ = (6144, 4096) +SHARED_GATE_UP = (3072, 6144) +SHARED_DOWN = (6144, 1536) +DENSE_GATE_UP = (12288, 6144) +DENSE_DOWN = (6144, 6144) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "shared gate_up": SHARED_GATE_UP, + "shared down": SHARED_DOWN, + "dense gate_up": DENSE_GATE_UP, + "dense down": DENSE_DOWN, +} + + +def served_cells(named: tuple[Cell, ...], x_bf16: bool, max_m: int) -> tuple[Cell, ...]: + """A cell for every (projection, rows) the GEMV serves up to ``max_m`` rows that ``named`` does not list. + + Certification is per projection: a projection's split-K, ring and MMA width differ, so each (N, K, M) the + decode dispatch can send here is its own cell. + """ + listed = {(cell.spec["nk"], cell.spec["m"]) for cell in named} + return tuple( + Cell(why=f"{name}, {m} rows", spec=dict(nk=nk, m=m)) + for name, nk in PROJECTIONS.items() + for m in range(1, max_m + 1) + if (nk, m) not in listed and _gemv.serves(nk[0], nk[1], m, x_bf16) + ) + + +def served_max_rows(nk: tuple[int, int], x_bf16: bool, lo: int, hi: int) -> int: + """The largest ``m`` in ``lo..hi`` such that the GEMV serves an (N, K) projection at every row count ``lo..m``; + ``lo - 1`` when it declines ``lo``. The decode dispatch raises a projection's rows past the batch-1/2 step only + as far as this (the bf16 GEMV's in-kernel quantize declines the K 6144, split 2 projections at MMA width 16).""" + m = lo - 1 + while m < hi and _gemv.serves(nk[0], nk[1], m + 1, x_bf16): + m += 1 + return m + + +def gemv_reference( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """bf16(x @ dequant(weight)^T) for fp32 activation values x [M, K], summed in fp64 and rounded once.""" + return (x.double() @ mxfp8_dequant(weight, weight_scale).double().t()).to(torch.bfloat16) + + +def gemv_is_valid(weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + n, k = weight.shape + # The kernel reads both from data_ptr() as dense buffers at the layout it was compiled for. + assert weight.dtype == torch.float8_e4m3fn and weight.is_contiguous(), ( + "weight must be contiguous e4m3 [N, K]" + ) + assert weight_scale.numel() == n * k // 32, ( + "weight_scale must be the 128x4-swizzled UE8M0 buffer [N * K / 32]" + ) + + +class M3MxGemv(OpWrapper): + """bf16 [M, N] = xq [M, K] (e4m3, 128x4-swizzled UE8M0 scales ``xsf``) x weight [N, K]^T (e4m3, swizzled + ``weight_scale``), M <= 48; None when the shape is not served (the caller falls back). + + ``trigger_early``: release the PDL dependents before the epilogue. ``cfg``: a (split, ring[, contig, nacc, a_tmem, + st_async]) tuning override; default the table's. ``use_pdl``: default the environment's. ``ar`` = (multicast + workspace address, flags address, rank, world, bucket tokens, mode): also push the rows into the next boundary + all-reduce's exchange (then ``op.last_pushed`` is True). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The projections' (N, K) at the token counts each build differs in: the MMA width (8 columns up to the TP2 decode + # step's 8 rows) is one, so one cell per projection pins the shape, and the rest are every served row count. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 4)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 8 rows: two requests verifying 1 + 3", + spec=dict(nk=QKV_SPARSE, m=8), + ), + Cell( + why="dense-attention QKV, 4 rows: one request verifying 1 + 3", spec=dict(nk=QKV, m=4) + ), + Cell(why="o_proj, one token (split 2)", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 8 rows", spec=dict(nk=O_PROJ, m=8)), + Cell(why="shared gate_up, 4 rows (split 4)", spec=dict(nk=SHARED_GATE_UP, m=4)), + Cell(why="shared down, one token (split 2, K 1536)", spec=dict(nk=SHARED_DOWN, m=1)), + Cell(why="dense gate_up, 8 rows (N 12288)", spec=dict(nk=DENSE_GATE_UP, m=8)), + Cell(why="dense down, 4 rows (K 6144)", spec=dict(nk=DENSE_DOWN, m=4)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling's MX_GEMV_ROWS). + CELLS = CELLS + served_cells(CELLS, x_bf16=False, max_m=48) + + note = """ + The kernel sums in fp32 in split-K order where `reference` sums in fp64: the bf16 outputs agree within the band the + test derives (SM107: at most as many misrounded elements as the stock FlashInfer CuTe DSL GEMM on the same + operands, rand / wide-scale / cancellation / tie / zero inputs). Not covered: the all-reduce push (`ar`), which + needs a multicast exchange; `cfg` overrides beyond the table (the target passes none). Shapes outside the table + and M > 48 return None; the target sends at most 8 rows. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv(xq, xsf, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The dequantized product; the launch options are ignored.""" + return gemv_reference(mxfp8_dequant(xq, xsf), weight, weight_scale) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + m, k = xq.shape + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous(), ( + "xq must be contiguous e4m3 [M, K]" + ) + # A linear-layout [M, K / 32] scale buffer is read as the swizzled one: wrong scales, no error. + assert xsf.numel() >= (m + 127) // 128 * 128 * (k // 32), ( + "xsf must be the 128x4-swizzled scale buffer (rows padded to 128), not the linear one" + ) + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv = M3MxGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..0398fe64b3d5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mx_gemv_bf16.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 GEMV (CuTe DSL) on a bf16 activation, quantized to MXFP8 in the kernel's prologue.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_quant_dequant +from ..mx_gemv import op as _gemv +from .m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + gemv_is_valid, + gemv_reference, + served_cells, +) + +__all__ = ["m3_mx_gemv_bf16"] + + +class M3MxGemvBf16(OpWrapper): + """bf16 [M, N] = mxfp8_quantize(x) x weight [N, K]^T for a contiguous, 16-byte aligned bf16 ``x`` [M, K]; the + quantization is mxfp8_quantize's (UE8M0 scale per 32 = the smallest power of two >= amax / 448). None when the + shape is not served. The other arguments as ``m3_mx_gemv``'s. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The prologue quantizes one 32-element block per thread and round: the projections at the token counts the + # target sends here (no producer-quantized activation), every TP2 decode row count. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token", spec=dict(nk=QKV_SPARSE, m=1)), + Cell(why="dense-attention QKV, 8 rows", spec=dict(nk=QKV, m=8)), + Cell(why="o_proj, 4 rows (the attention output in bf16)", spec=dict(nk=O_PROJ, m=4)), + Cell(why="shared gate_up, 8 rows", spec=dict(nk=SHARED_GATE_UP, m=8)), + Cell(why="shared down, 4 rows (K 1536)", spec=dict(nk=SHARED_DOWN, m=4)), + Cell( + why="dense gate_up, 8 rows: 24 K tiles per split-2 CTA", + spec=dict(nk=DENSE_GATE_UP, m=8), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling's MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=32) + + note = """ + As `m3_mx_gemv`. The in-kernel quantization is mxfp8_quantize's bit for bit except where the kernel's + rcp.approx(448) and an exact 1/448 put a block's amax on different sides of a power of two (not observed on the + cells). Not covered: `ar`, `cfg` overrides. Returns None for rows beyond 8 rounds of the prologue per CTA. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv_bf16(x, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The MXFP8 round trip of x times the dequantized weight; the launch options are ignored.""" + return gemv_reference(mxfp8_quant_dequant(x), weight, weight_scale) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv_bf16 = M3MxGemvBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..3bfcbc625016 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned + +from .m3_mxfp8_linear import mxfp8_reference + + +class M3Mxfp8Gemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight + [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key + and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The decode dispatch hands it the boundary / merge MXFP8 activation on rows the GEMV does not serve: the + # unsplit shared down projection at 17..32 rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=1536)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=1536)), + ) + + note = """ + Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + """The dequantized product; ``tune`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..9eadb1e0d20a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned + +_BLOCK = 32 + + +def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" + pad_r = (rows + 127) // 128 * 128 + pad_c = (cols + 3) // 4 * 4 + s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) + return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] + + +def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, k = x.shape + s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) + factor = torch.pow(2.0, (s - 127).to(torch.float32)) + return ( + x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) + ).view(rows, k) + + +def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: + """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" + with true_fp32_matmul(): + return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=4992, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=4096, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=6144, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=3072, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=4096, path="tuned")), + ) + + note = """ + Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` + quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and + multiplies dequantized operands in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + if tuned_quantize_gemm: + return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + if flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, output_dtype + ) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are + ignored.""" + m, k = x.shape + blocks = x.float().view(m, k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(-1, keepdim=True) + exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) + scale = torch.pow(2.0, exp) + q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale + with true_fp32_matmul(): + out = q.view(m, k) @ _dequant(weight, weight_scale).t() + return (out * alpha.float()).to(output_dtype) + + +m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py new file mode 100644 index 000000000000..f8deb858c843 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, +and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). + +The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values +``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 +marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a +per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: +per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 +code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved +it decodes to the BF16 weight bit for bit. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_target_head"] + +#: Verify rows per call (requests x (draft length + 1)). +MAX_ROWS = 8 +#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). +CTA_WARPS = 16 +#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. +TARGET_SHARD = 50016 +HIDDEN = 6144 +#: Codes per window: 0..14 decode, 15 marks an exception. +WINDOW = 15 +#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. +MAX_BASE = 254 - WINDOW +#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). +ROWS_PER_CHUNK = 4096 + + +def _pad_rows(weight: torch.Tensor) -> torch.Tensor: + n = weight.shape[0] + tiles = (n + 15) // 16 + if tiles * 16 == n: + return weight + return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) + + +def _fragment_view(w: torch.Tensor) -> torch.Tensor: + """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" + n, k = w.shape + # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) + return ( + w.view(n // 16, 2, 8, k // 16, 2, 4, 2) + .permute(0, 3, 2, 5, 4, 1, 6) + .reshape(n // 16, k // 16, 32, 8) + ) + + +def _exponents(weight: torch.Tensor) -> torch.Tensor: + return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF + + +def _window_base(hist: torch.Tensor) -> int: + """The base of the 15-exponent window holding the most weights (the lowest such base).""" + window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) + return int(window.argmax()) + + +def encode(weight: torch.Tensor, base: Optional[int] = None): + """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's + own window. + + Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile + << 16 | k), exc_w fp32 [E, 2] ({w, 0})). + """ + w = _pad_rows(weight) + n, k = w.shape + tiles, ks = n // 16, k // 16 + assert ks % 4 == 0, "K must be a multiple of 64" + bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF + expo = (bits >> 7) & 0xFF + if base is None: + base = _window_base(torch.bincount(expo.flatten(), minlength=256)) + code = expo - base + exc = (code < 0) | (code >= WINDOW) + code = torch.where(exc, torch.full_like(code, WINDOW), code) + sm = ((bits >> 8) & 0x80) | (bits & 0x7F) + sm = torch.where(exc, torch.zeros_like(sm), sm) + # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. + rows, cols = exc.nonzero(as_tuple=True) + order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) + rows, cols = rows[order], cols[order] + exc_w = torch.stack( + [ + w[rows, cols].float(), + torch.zeros((rows.numel(),), device=w.device), + ], + dim=1, + ) + exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) + exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) + exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( + torch.int32 + ) + # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. + smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) + cf = ( + _fragment_view(code.to(torch.uint8).view(n, k)) + .view(tiles, ks // 4, 4, 32, 8) + .to(torch.int32) + ) + # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. + cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) + for p in range(4): + cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) + cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 + # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. + seg01 = ( + smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) + ) # [.., 2, 32, 16] + seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) + packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] + return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() + + +def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): + """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" + n, _ = weight.shape + assert rows_per_chunk % 16 == 0 + hist = torch.zeros(256, dtype=torch.int64, device=weight.device) + for r in range(0, n, rows_per_chunk): + hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) + base = _window_base(hist) + packed, offs, poss, ws = [], [], [], [] + n_exc = 0 + for r in range(0, n, rows_per_chunk): + p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) + packed.append(p) + offs.append(off[:-1] + n_exc) + poss.append(pos) + ws.append(ew) + n_exc += pos.numel() + exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) + return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) + + +def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: + """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" + tiles, ks = (n + 15) // 16, k // 16 + p = packed.view(tiles, ks // 4, 3, 32, 16) + sm = ( + p[:, :, :2] + .permute(0, 1, 3, 2, 4) + .reshape(tiles, ks // 4, 32, 4, 8) + .permute(0, 1, 3, 2, 4) + .to(torch.int32) + ) + cw = ( + p[:, :, 2] + .contiguous() + .view(torch.int32) + .view(tiles, ks // 4, 32, 4) + .permute(0, 1, 3, 2) + .long() + & 0xFFFFFFFF + ) + codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( + torch.int32 + ) + bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) + bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) + frag = bits.reshape(tiles, ks, 32, 8) + # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). + w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) + return w.to(torch.int16).view(torch.bfloat16)[:n] + + +def decode( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + n: int, + k: int, + row0: int = 0, + row1: Optional[int] = None, +) -> torch.Tensor: + """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). + ``row0`` must be a multiple of 16.""" + row1 = n if row1 is None else row1 + assert row0 % 16 == 0 + t0, t1 = row0 // 16, (row1 + 15) // 16 + unit_bytes = (k // 64) * 3 * 32 * 16 + dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() + counts = (exc_off[1:] - exc_off[:-1]).long() + tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) + rows = tile_of * 16 + (exc_pos >> 16).long() + cols = (exc_pos & 0xFFFF).long() + sel = (rows >= row0) & (rows < row1) + dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) + return dec + + +def round_trip_ok( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + weight: torch.Tensor, + rows_per_chunk: int = ROWS_PER_CHUNK, +) -> bool: + """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent + window below the BF16 infinity exponent.""" + n, k = weight.shape + if not 0 <= base <= MAX_BASE or k % 64: + return False + if ( + packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 + or exc_off.numel() != (n + 15) // 16 + 1 + ): + return False + for r0 in range(0, n, rows_per_chunk): + r1 = min(n, r0 + rows_per_chunk) + dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) + if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): + return False + return True + + +class M3TargetHead(OpWrapper): + """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit + the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in + another summation order). + + ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` + (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are + the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and + # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with + # a shard whose rows are not a multiple of 16. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{m} verify row(s) of the TP4 target head shard", + spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), + ) + for m in range(1, MAX_ROWS + 1) + ), + Cell( + why="4 rows with outliers and exact zeros in many tiles (the exception lists)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), + ), + Cell( + why="2 rows over a shard of 4000 rows (250 tiles), K 1024", + spec=dict(m=2, n=4000, k=1024, exceptions=True), + ), + ) + + note = """ + The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception + adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs + from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); + integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp + order: the result is deterministic. + """ + + def __call__( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_target_head( + w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl + ) + return out + + def reference( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. + + ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" + weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + return logits.to(torch.bfloat16) + + def is_valid( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> None: + # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count + # stores a partial sum. + assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + assert out.data_ptr() != x.data_ptr(), "out must not alias x" + + +m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py new file mode 100644 index 000000000000..c734b2a3172e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_ugpu_head.py @@ -0,0 +1,168 @@ +# 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. +"""The target lm_head shard GEMM of a verify step on a two-die GPU, each die reading only its own memory. + +A copy of the head shard [N, K] (BF16) is placed once as one virtual range whose rows [0, S) live in locality domain 0 +and rows [S, N) in domain 1 (:meth:`LocalityMap.place_split`; S on a 2 MiB granule). A covered verify computes +``F.linear`` of each row half of the copy on a green-context stream holding that domain's SMs, and concatenates: each +logit is the in-tree head GEMM's own logit (the same cuBLAS GEMM per output column) at the row counts the cells +certify. The module's own weight is never re-placed: a kernel reading die-placed memory from all SMs runs 1.5-1.7x +slower (SpecTail 769262), so every other reader (uncovered row counts, prefill logits, warmup, fallbacks) keeps the +ordinary weight, and only this op reads the placed copy. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Optional + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger + +__all__ = ["UgpuHead", "m3_ugpu_head", "prepare_head"] + +TARGET_SHARD = 100032 # the TP2 target head shard rows (vocabulary 200064 / 2) +HIDDEN = 6144 +_ENGAGED: set = set() + + +@dataclass +class UgpuHead: + """A placed head: ``weight`` the die-placed copy (read only by :data:`m3_ugpu_head`), ``source`` the module's own + weight it copies and ``version`` that weight's version counter at the copy (a later in-place change makes the copy + stale: see :func:`current`), ``split`` the first domain-1 row, ``streams`` the two domains' green-context + streams.""" + + weight: torch.Tensor + source: torch.Tensor + version: int + split: int + streams: List[torch.cuda.ExternalStream] + + +def current(head: Optional["UgpuHead"], weight: torch.Tensor) -> bool: + """Whether ``head`` is a placed copy of ``weight`` as it is now (same tensor, not modified since the copy).""" + return head is not None and head.source is weight and weight._version == head.version + + +def _split_row(n: int, row_bytes: int, granule: int) -> Optional[int]: + """The granule-aligned row nearest the middle (the halves' bytes differ by < one granule), or None.""" + import math + + step = granule // math.gcd(granule, row_bytes) # rows per granule-aligned boundary + s = round(n / 2 / step) * step + return s if 0 < s < n else None + + +def prepare_head(lm_head: torch.nn.Module) -> Optional[UgpuHead]: + """A die-placed copy of ``lm_head.weight`` for :data:`m3_ugpu_head` (the module is not changed), or None when the + head does not take the plain ``F.linear`` path this op reproduces, or the device has no two locality domains. + Before capture.""" + from tensorrt_llm._torch.modules import low_m_gemm + + from ..m3_locality import LocalityMap + + w = getattr(lm_head, "weight", None) + if ( + w is None + or w.dtype != torch.bfloat16 + or w.dim() != 2 + or not w.is_contiguous() + or getattr(lm_head, "bias", None) is not None + or getattr(lm_head, "use_cute_dsl_bf16_gemm", False) + or getattr(lm_head, "use_custom_cublas_mm", False) + or low_m_gemm.LOW_M_GEMM_ACTIVE + or type(getattr(lm_head, "quant_method", None)).__name__ != "UnquantizedLinearMethod" + ): + return None + try: + lm = LocalityMap.get(w.device) + except (ImportError, RuntimeError) as err: + logger.warning(f"[M3 sm_107 tp2] uGPU head: no locality map ({err}); the head stays as is") + return None + s = _split_row(w.shape[0], w.shape[1] * w.element_size(), lm.granularity) + if s is None: + return None + with torch.no_grad(): + placed = lm.place_split(w.detach(), s) + return UgpuHead( + weight=placed, source=w, version=w._version, split=s, streams=lm.domain_streams() + ) + + +class M3UgpuHead(OpWrapper): + """``logits`` [M, N] (BF16) = x [M, K] (BF16) times the placed head W [N, K] transposed, as + ``cat(F.linear(x, W[:S]), F.linear(x, W[S:]))`` with the halves on the two domains' green-context streams + (forked from and joined back into the caller's stream).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The verify rows of one to four requests at draft length 3 (the TP2 GEN batch is 1: every E2E step is 4 rows). + CELLS: tuple[Cell, ...] = ( + Cell( + why="4 verify rows: one request, draft length 3 (the b1 point)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN), + ), + Cell( + why="8 verify rows: two requests, draft length 3", + spec=dict(m=8, n=TARGET_SHARD, k=HIDDEN), + ), + Cell(why="12 verify rows: three requests", spec=dict(m=12, n=TARGET_SHARD, k=HIDDEN)), + Cell( + why="16 verify rows: four requests (the c4 point)", + spec=dict(m=16, n=TARGET_SHARD, k=HIDDEN), + ), + ) + + note = """ + Each output column is computed by cuBLAS over the same weight row and the same input as the in-tree head GEMM; at + the certified row counts the half-N call picks the full call's kernel, so the logits are the in-tree head's bit + for bit (the cells check it on SM107). At 1 and 2 rows cuBLAS picks a different kernel for the half shard: not + covered. Placement only moves pages between the dies; every other reader of the weight is unchanged. + """ + + def __call__(self, x: torch.Tensor, head: UgpuHead) -> torch.Tensor: + key = (x.device, x.shape[0]) + if key not in _ENGAGED: + _ENGAGED.add(key) + logger.info( + f"[M3 sm_107 tp2] fast path engaged: m3_ugpu_head ({x.shape[0]} rows, split {head.split} of " + f"{head.weight.shape[0]}, on {x.device})" + ) + main = torch.cuda.current_stream(x.device) + fork = torch.cuda.Event() + fork.record(main) + parts = [] + for stream, w in zip(head.streams, (head.weight[: head.split], head.weight[head.split :])): + stream.wait_event(fork) + with torch.cuda.stream(stream): + parts.append(F.linear(x, w)) + for stream in head.streams: + done = torch.cuda.Event() + done.record(stream) + main.wait_event(done) + for p in parts: + p.record_stream(main) + return torch.cat(parts, dim=1) + + def reference(self, x: torch.Tensor, head: UgpuHead) -> torch.Tensor: + """The in-tree head GEMM over the whole (placed) weight.""" + return F.linear(x, head.weight) + + +m3_ugpu_head = M3UgpuHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/long_select.py new file mode 100644 index 000000000000..79aaf6674300 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/long_select.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. + +Same contract and bit-identical output as ``torch.ops.trtllm.minimax_m3_select_blocks`` on per-KV-head scores +[num_kv_heads, n_blocks, total_q]: init / local forcing, per-query valid-block count, ties to the smaller block index, +forced -inf picks emitted as -1, ascending block ids with -1 tail padding. + +One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups ``block % 16`` is a lower +bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared memory +and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every candidate +against the whole row (re-read). + +The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long`` (this target's copy of ``minimax_m3_select_long`` +with a choice of where it releases the sparse decode), called through this target's ``m3_select_long`` entry, or with +``dsl`` the same selection from one 384-thread CTA per row (``m3_select_long_dsl``, ops/m3_select_long_dsl). +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long +from ..ops.attention.m3_select_long_dsl import m3_select_long_dsl + + +def long_select_blocks( + max_score_kv: torch.Tensor, + *, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + dsl: bool = False, +) -> Optional[torch.Tensor]: + """Top-16 select over per-KV-head scores [num_kv_heads, n_blocks, total_q] with 129..9216 blocks. + + Returns the head-major backed [total_q, num_kv_heads, 16] int32 table, or None outside this kernel's shapes + (callers keep the reference op). ``early_trigger`` releases the dependent sparse decode right after the select's + grid wait (modeling's ``SELECT_EARLY_MAX_ROWS``); otherwise ``trigger_phase`` 1 releases it once every score of the + row is loaded and 0 at the select's exit (modeling's ``SELECT_LONG_TRIGGER_PHASE``). ``dsl`` runs the 384-thread + kernel (modeling's ``SELECT_LONG_DSL``; the same table). + """ + if ( + max_score_kv.dtype != torch.float32 + or not max_score_kv.is_cuda + or max_score_kv.dim() != 3 + or not 128 < max_score_kv.shape[1] <= MAX_BLOCKS + or n_valid_blocks.dtype != torch.int32 + or not n_valid_blocks.is_cuda + or n_valid_blocks.shape[0] < max_score_kv.shape[2] + ): + return None + num_kv_heads, _, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, 16), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + return (m3_select_long_dsl if dsl else m3_select_long)( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_attn_lc.py new file mode 100644 index 000000000000..ddfbb0e4b2a9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_attn_lc.py @@ -0,0 +1,164 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 long-context sparse decode attention in one launch. + +Replaces the generic chain's split-K producer (``_gqa_sparse_decode_kernel``, one block per chunk) and the LSE +merge (``_merge_topk_attn_out_kernel``, with its MXFP8 copy for the o_proj) once the top-k block table exists. + +Grid: one cluster of 16 CTAs per (query token, KV head); CTA ``r`` attends the token's ``r``-th selected block of +that head (the Triton producer's chunk ``r`` at 16 chunks), keeps the partial in shared memory, and after a cluster +barrier merges head ``r`` of the KV head's group from all 16 partials over DSMEM, writing the BF16 row and, +optionally, its MXFP8 copy and swizzled UE8M0 scales, bitwise ``torch.ops.trtllm.mxfp8_quantize(out, True)``. + +Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +log-domain recurrence. Accumulation order inside the MMAs and the reductions differs, so results match to rounding, +not bit for bit. + +Schedule: the request length and the query rows (FP8 widened to FP16 in shared memory) are read before the PDL grid +wait, which overlaps them with the select ahead when it releases this grid early (modeling's ``SELECT_EARLY_MAX_ROWS``); +the split-K partials travel by ``st.async`` into each receiving CTA's own mbarrier instead of a cluster-wide barrier. +Neither changes a value or the merge order. + +Geometry: 16 query heads per KV head, two KV heads per rank (TP2), head_dim 128, page 128, top-k 16. The kernel +is ``torch.ops.trtllm.minimax_m3_sm107_tp2_sparse_decode_lc``, called through this target's ``m3_sparse_decode_lc`` +entry. +""" + +import torch + +from tensorrt_llm.logger import logger + +from ..ops.attention.m3_sparse_decode_lc import m3_sparse_decode_lc + + +class N1Slot: + """The sparse layer's N1 id buffer (modeling's ``SPARSE_DECODE_LC_PREFETCH_SELF_KV``), set by the layer around its + decode call and taken by the long-context decode; empty means no prefetch.""" + + def __init__(self) -> None: + self.ids = None + + def take(self): + ids, self.ids = self.ids, None + return ids + + +# int32 [2, N1_MAX_ROWS, 16]: the layer's selected block ids per (KV head, query row, slot) at its previous decode step. +N1_IDS = N1Slot() + + +def lc_applies(q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output) -> bool: + """Whether the step's operands fit the kernel: TP2 geometry, packed NVFP4 pages, FP8 q, BF16 [T, 32, 128] out.""" + if q.dim() != 3 or tuple(q.shape[1:]) != (32, 128) or q.dtype != torch.float8_e4m3fn: + return False + if q.stride(2) != 1 or q.stride(1) % 16 or q.stride(0) % 16 or q.data_ptr() % 16: + return False + if ( + int(k_paged.shape[1]) != 2 + or topk_idx.dim() != 3 + or int(topk_idx.shape[0]) != 2 + or int(topk_idx.shape[1]) != int(q.shape[0]) + or int(topk_idx.shape[-1]) != 16 + or topk_idx.dtype != torch.int32 + or topk_idx.stride(2) != 1 + ): + return False + if v_paged.stride() != k_paged.stride(): + return False + for t in (k_paged, v_paged): + if ( + t.stride(3) != 1 + or t.stride(2) != 64 + or int(t.shape[3]) != 64 + or t.data_ptr() % 16 + or t.stride(0) % 16 + or t.stride(1) % 16 + ): + return False + for t in (k_block_scale, v_block_scale): + if ( + t is None + or int(t.shape[1]) != 2 + or t.stride(3) != 1 + or t.stride(2) != 8 + or t.data_ptr() % 8 + or t.stride(0) % 8 + or t.stride(1) % 8 + ): + return False + if v_block_scale.stride() != k_block_scale.stride(): + return False + return ( + output.dtype == torch.bfloat16 + and output.dim() == 3 + and output.is_contiguous() + and tuple(output.shape) == tuple(q.shape) + ) + + +def m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + linear_softmax=False, + mxfp8=None, + use_pdl=True, + prewait=False, + n1_ids=None, +): + """Drop-in for ``minimax_m3_sparse_attn_decode`` on the TP2 NVFP4 shape; callers check ``lc_applies`` first. + + ``topk_idx`` is the head-major [2, total_q, 16] table the Triton producer takes. ``mxfp8`` is ``(q8, sf)`` from + ``MERGE_MXFP8.buffers(4096, dev)`` to also emit the o_proj hand-off. ``n1_ids`` is the layer's N1 id buffer + (``N1_IDS``): the decode prefetches last step's blocks into L2 and records this step's. + """ + total_q = int(q.shape[0]) + q8, sf = ( + (mxfp8[0].view(torch.uint8), mxfp8[1].view(torch.uint8)) + if mxfp8 is not None + else (None, None) + ) + m3_sparse_decode_lc( + q.view(torch.uint8), + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + output.view(total_q, 4096), + q8, + sf, + 4096 // 128, + int(decode_query_len), + float(sm_scale), + bool(linear_softmax), + bool(use_pdl), + bool(prewait), + n1_ids, + ) + key = (total_q, q8 is not None, n1_ids is not None) + if key not in _engaged: + _engaged.add(key) + logger.info( + f"[M3 sm_107 tp2] sparse decode Lc engaged: T={total_q}, kv_heads=2, mxfp8_handoff={q8 is not None}, " + f"n1_self_kv={n1_ids is not None}" + ) + + +_engaged = set() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py new file mode 100644 index 000000000000..c61ab4acbf14 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 draft pass for the MiniMax-M3 drafter's vocab-sharded BF16 head (TP decode, <= 8 requests). + +One kernel replaces a draft pass's head and the draft loop's per-pass bookkeeping, at the in-tree precision: + +* The drafter's BF16 lm_head shard is kept a second time, built by ``prepare_head`` before the executor's warmup: in + the lossless 12-bit coding of the target head (``ops.gemm.m3_target_head``; 75% of the bytes, kept only when it + decodes to the BF16 shard bit for bit), else in ``mma.m16n8k16`` A-fragment order (``bf16_fragments``, the same + values). The kernel streams it with BF16 MMAs and FP32 accumulation over the BF16 head input rows; a coded + exception decodes to +0 and its exact product x * w is added at its tile's flush. +* Each logit is rounded to BF16 (the in-tree logits are the BF16 head GEMM's output) and keyed + ``ordered(value) << 32 | ~global_index`` with -0 as +0, so the larger value wins and ties take the lowest index, as + the in-tree ``torch.max`` / ``torch.argmax`` over the gathered logits (``greedy_sample_draft_with_tp_gather``). +* The last CTA exchanges each row's key across the TP ranks through a multicast mapping (Lamport slots, see + ``m3_shard_accept``), then stores the draft token, the next pass's input id and position, and advances the draft KV + lengths as the in-tree draft loop does between passes. + +The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu) and called through the +catalog entry ``ops.spec.m3_draft_step``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok +from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments +from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op +from . import m3_report +from .m3_shard_accept import AcceptWorkspace, create_workspace + +_ENGAGED: set = set() + + +# The fused draft loop streams the 12-bit coded copy of the draft head when it decodes bit for bit; off at TP2 (the +# coded head is slower on sm_107: the BF16 fragments stream). +HEAD_CODED = False +# The head copy is placed per locality domain and each pass runs as two launches, one per domain (placement only; +# 105.6 -> 86.3 / 86.6 us per pass at b1 / b2 on Rubin TP2). +UGPU = True +# The head streams an equal share of each launch's (tile, chunk) space per CTA instead of whole 16-row tiles (two +# zero-filled addends per boundary tile, so the sum does not depend on their order). +EQ_CHUNKS = True + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the 12-bit coded stream + with its exception lists (``coded``), or the fragment-ordered BF16 copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + coded: bool = False + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + exp_base: int = 0 + # BF16 fragments placed per locality domain (UGPU below): tiles [0, split_tile) live in domain 0 and + # the rest in domain 1; 0 = not placed (one launch over all SMs). + split_tile: int = 0 + # Equal-chunk split (EQ_CHUNKS below): zero-filled shared-tile partials [tiles, 32, 4] fp32 and + # arrival counts [tiles] int32 (the kernel leaves them zero); None = whole tiles per CTA. + bnd: Optional[torch.Tensor] = None + bnd_cnt: Optional[torch.Tensor] = None + + +def head_supported(lm_head) -> bool: + """A bias-free BF16 head whose local shard the kernel can own: whole 16-row tiles, K a multiple of 16, at most + MAX_TILES_PER_CTA tiles per CTA of a one-CTA-per-SM grid.""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + if getattr(lm_head, "bias", None) is not None: + return False + n, k = weight.shape + if n % 16 or k % 16: + return False + return (n // 16 + _num_sms() - 1) // _num_sms() <= MAX_TILES_PER_CTA + + +def prepare_head(lm_head, tp_rank: int) -> DraftHead: + """The streamed copy of ``lm_head``'s local shard (global ids from tp_rank * rows): its 12-bit coding when + ``HEAD_CODED``, K is a multiple of 64 and the coding decodes to the shard bit for bit, else (fail + closed) the fragment-ordered BF16 copy. + Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_draft_head``).""" + weight = lm_head.weight + offset = tp_rank * weight.shape[0] + with torch.no_grad(): + w = weight.detach() + if HEAD_CODED and w.shape[1] % 64 == 0: + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return DraftHead( + weight=weight, + wf=packed, + vocab_offset=offset, + coded=True, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + exp_base=base, + bnd=None, # the coded head (off on sm_107) keeps whole tiles + bnd_cnt=None, + ) + del packed, exc_off, exc_pos, exc_w + wf = bf16_fragments(w) + wf, split_tile = _place_per_domain(wf) + bnd, bnd_cnt = _eq_buffers(wf.shape[0], wf.device) + return DraftHead( + weight=weight, wf=wf, vocab_offset=offset, split_tile=split_tile, bnd=bnd, bnd_cnt=bnd_cnt + ) + + +def _eq_buffers(n_tiles: int, device): + """The equal-chunk split's zero-filled shared-tile partials and arrival counts, or (None, None) when + EQ_CHUNKS is off.""" + if not EQ_CHUNKS: + return None, None + return ( + torch.zeros(n_tiles, 32, 4, dtype=torch.float32, device=device), + torch.zeros(n_tiles, dtype=torch.int32, device=device), + ) + + +def _place_per_domain(wf: torch.Tensor): + """(wf, split tile): the BF16 fragments split-mapped over the two locality domains at the granule-aligned tile + nearest half (locality.LocalityMap.place_split; its domain streams built here, before capture), or (wf, 0) when + UGPU is off. Placement only: the streamed values are the same.""" + if not UGPU: + return wf, 0 + from .m3_locality import LocalityMap + + lm = LocalityMap.get(wf.device) + tile_bytes = wf[0].numel() * wf.element_size() + step = 1 + while (step * tile_bytes) % lm.granularity: + step += 1 + n_tiles = wf.shape[0] + split_tile = max(step, round(n_tiles / 2 / step) * step) + per = lm.sms_per_domain + if split_tile >= n_tiles or max(split_tile, n_tiles - split_tile) > per * MAX_TILES_PER_CTA: + return wf, 0 + lm.domain_streams() + return lm.place_split(wf, split_tile), split_tile + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace ([2][MAX_TOKENS][world] slots), created collectively on the first eager call; a + CUDA-graph capture that finds it missing is a miss.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("m3_draft_step", None, "the draft-step exchange workspace") + return None + _WS = create_workspace(mapping, device, MAX_TOKENS) + return _WS + + +def draft_step( + head: DraftHead, + x: torch.Tensor, + rows: int, + x_rows: Optional[torch.Tensor], + mapping, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, +) -> bool: + """One fused greedy draft pass over ``rows`` requests (see the module docstring); False when the exchange + workspace is missing during a capture. Collective over the TP group. + + Args: + head: the prepared head of this rank (``prepare_head``), coded or BF16 fragments. + x: [*, K] BF16 head input rows (the drafter's final-norm output); row ``x_rows[b]`` (int64) or ``b`` is + request b's. + draft_out: int32 [rows, >= step + 1], unit column stride: column ``step`` receives the tokens. + ids_next / pos_next: int32 [rows]: the tokens, and ``pos_src[x_rows[b] or b] + 1`` (pos_src int32). + kv_lens: int32 draft KV lengths: ``-= draft_len - num_acc`` with ``seq_lens = 1`` on pass 0 when ``num_acc`` + is given, ``+= 1`` on later passes. + """ + ws = workspace(mapping, x.device) + if ws is None: + return False + if head.split_tile: + _draft_step_split( + head, + x, + rows, + x_rows, + ws, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + ) + return True + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + _num_sms(), + True, + None, + head.exc_off, + head.exc_pos, + head.exc_w, + head.exp_base, + 0, + 0, + 0, + head.bnd, + head.bnd_cnt, + ) + return True + + +def _draft_step_split( + head, + x, + rows, + x_rows, + ws, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, +) -> None: + """One pass as two launches, one per locality domain on that domain's green-context stream, each streaming its own + domain's tiles of the split-mapped fragments; the keys' max and the last-CTA exchange and bookkeeping run once over + both (the last of the 2 x SMs-per-domain CTAs). Forked from and joined back to the current stream; no PDL on the + forks.""" + from .m3_locality import LocalityMap + + lm = LocalityMap.get(x.device) + per = lm.sms_per_domain + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + streams = lm.domain_streams() + n_tiles = head.wf.shape[0] + if x.device not in _ENGAGED: + _ENGAGED.add(x.device) + logger.info( + f"[M3 sm_107 tp2] fast path engaged: m3_draft_step_ugpu (split tile {head.split_tile} of {n_tiles}, " + f"{per} CTAs per domain, on {x.device})" + ) + for st, lo, hi in ((streams[0], 0, head.split_tile), (streams[1], head.split_tile, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + per, + False, + None, + None, + None, + None, + 0, + lo, + hi, + 2 * per, + head.bnd, + head.bnd_cnt, + ) + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py new file mode 100644 index 000000000000..a4c7b5ccb8fd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py @@ -0,0 +1,413 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's decode attention on Rubin: ``M3DrafterFullCtxFmha``. + +The drafter layer is a plain ``TrtllmAttention`` over its own FP8 paged cache (the full context, every draft pass). +On SM 10.7 its generation phase runs ``thop.attention``'s in-tree trtllm-gen runner, whose split is capped at 8 CTAs +per (request, KV head): 45-46 us per pass at 165k tokens on 212 SMs, for 85 MB. + +This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append +into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context +PrimTS split-KV decode of ``ops/m3_fullctx_decode`` at the same pinned recipes as the dense layers. + +With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead +(``minimax_m3_sm107tp2_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page +pool and block table taken from the KV cache pool pointers and mapping the preprocess reads. + +It serves generation-only steps of the covered (query rows per request, requests) shapes; the target installs it +first in the drafter layer's FMHA list (``install``), so every other step falls through to the stock libraries. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + FlashInferTrtllmGenFmha, +) +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode + +from . import m3_fullctx_decode as fullctx + +# The fused RoPE + FP8 K/V launch also takes a one-request multi-token pass (the drafter's first pass at one +# request), where the in-tree preprocess needs the token-map guard: bitwise the guarded preprocess + cast. +ROPE_ONE_REQUEST_MULTI = True +# With PDL, that launch takes the maximum shared-memory carveout so the drafter decode's large-smem CTAs co-reside with +# it (timing only). +ROPE_CARVEOUT = True + + +def fwd_has_no_mrope(fwd) -> bool: + return getattr(fwd, "mrope_position_deltas", None) is None + + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + +class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): + """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + + def __init__( + self, + attn: "TrtllmAttention", + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + pdl: bool = False, + on_rope_taken=None, + ) -> None: + super().__init__(attn) + self._recipes = recipes + self._max_kv_len = max_kv_len + self._own_reducer = own_reducer + self._fused_rope_mode = fused_rope_mode + self._pdl = pdl + self._on_rope_taken = on_rope_taken + self._one_request_multi = ROPE_ONE_REQUEST_MULTI + self._on_taken = on_taken + self._on_uncovered = on_uncovered + self._on_missed = on_missed + + def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: + nreq = int(metadata.num_generations) + if nreq <= 0 or int(q.shape[0]) % nreq: + return None + return int(q.shape[0]) // nreq, nreq + + def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: + if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: + return False + if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: + return False + attn = self.attn + if ( + attn.num_heads != fullctx.HEADS + or attn.num_kv_heads != fullctx.KV_HEADS + or attn.head_dim != fullctx.HEAD_DIM + or not QuantMode(attn.quant_mode).has_fp8_kv_cache() + ): + self._on_uncovered("not the TP2 FP8 drafter geometry") + return False + # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's + # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. + if ( + forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant + or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig + ): + self._on_uncovered("per-step KV cache scales") + return False + if int(metadata.max_seq_len) > self._max_kv_len: + self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") + return False + shape = self._shape(q, metadata) + if shape is None or shape not in self._recipes: + self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") + return False + # Compiled before warmup (``install``); the forward never compiles one. + if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: + self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") + return False + return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) + + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + attn = params.attn + width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim + rotary_dim = int(attn.rope_params.dim) + return ( + self._fused_rope_mode is not None + and params.qkv_input is not None + and attention_input.dim() == 2 + and attention_input.dtype == torch.bfloat16 + and attention_input.is_contiguous() + and int(attention_input.shape[1]) == width + and not params.is_cross + and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) + and attn.rotary_cos_sin is not None + and attn.rotary_cos_sin.dtype == torch.float32 + and fwd_has_no_mrope(params.fwd) + and 0 < rotary_dim <= attn.head_dim + and rotary_dim % 16 == 0 + and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) + # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo + # runs for several requests only) and reads the workspace's; with ROPE_ONE_REQUEST_MULTI + # the fused launch (which places rows by request = token / q_len, no map) takes it too, else it keeps + # the in-tree path (and its token-map guard). + and ( + self._one_request_multi + or not (self._multi_token(params) and int(params.batch_size) == 1) + ) + ) + + @staticmethod + def _multi_token(params: FmhaParams) -> bool: + return ( + params.spec_decoding_generation_lengths is not None + and int(params.attn.predicted_tokens_per_seq) > 1 + ) + + def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" + if not self._fused_rope_applies(params, attention_input): + return False + attn = params.attn + meta = params.meta + qlen, nreq = int(params.input_seq_length), int(params.batch_size) + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + if plan is None: + return False + # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary + # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the + # block table is the pool's [requests, 2, blocks] K / V sub-page indices. + pool_index, layer_in_pool = ( + int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] + ) + tpb = int(params.tokens_per_block) + sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim + base = ( + int(meta.host_kv_cache_pool_pointers[pool_index, 0]) + + layer_in_pool * int(params.kv_factor) * sub_page_bytes + ) + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) + ) + tables = meta.kv_cache_block_offsets[ + pool_index, params.seq_offset : params.seq_offset + nreq + ] + multi = self._multi_token(params) + if not torch.cuda.is_current_stream_capturing(): + # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess + # path checks; the op itself writes V at the table's V sub-page. + tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(tables.shape[-1], device=tables.device) + in_use = cols[None, :] < tpb_held[:, None] + offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + q = torch.empty( + int(params.num_tokens), + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=attention_input.device, + ) + torch.ops.trtllm.minimax_m3_sm107tp2_drafter_rope_fp8_kv( + attention_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + int(attn.rope_params.dim), + tpb, + self._fused_rope_mode, + self._pdl, + ROPE_CARVEOUT, + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + tables[:, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + if self._on_rope_taken is not None: + self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") + self._on_taken(params.num_tokens) + return True + + def run_generation(self, params: FmhaParams) -> None: + attention_input = params.qkv_input if params.qkv_input is not None else params.query_input + if attention_input is None: + raise RuntimeError("the drafter decode requires QKV or query input.") + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + if self._run_fused_rope(params, attention_input): + return + # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 + # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). + fp8_context_fmha = False + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + _fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + _window_left, + is_multi_token_gen, + ) = thop.trtllm_gen_generation_preprocess( + attention_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + fp8_context_fmha, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + ) + qlen, nreq = params.input_seq_length, params.batch_size + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the + # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row + # pass is never multi-token. + if plan is None or (is_multi_token_gen and qlen == 1): + raise RuntimeError( + f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" + ) + rows = params.num_tokens * attn.num_heads * attn.head_dim + q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) + q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) + # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 + # of the pool (the two views have the same page count, as the kernel requires). + if not torch.cuda.is_current_stream_capturing(): + # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. + lengths = params.spec_decoding_generation_lengths + if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): + raise RuntimeError( + f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" + ) + # Only the blocks a request holds are laid out; the table's tail entries are unset. + tpb = int(params.tokens_per_block) + held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(block_tables.shape[-1], device=block_tables.device) + in_use = cols[None, :] < held[:, None] + offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. + # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + block_tables[:nreq, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + self._on_taken(params.num_tokens) + + +def install( + draft_model: torch.nn.Module, + recipes, + max_kv_len: int, + own_reducer: bool, + pdl: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + on_rope_taken=None, +) -> int: + """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention + backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it + was installed on.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + count = 0 + device = next(draft_model.parameters()).device + # The backend is not a module: each attention module of the drafter holds it as ``attn``. + for module in draft_model.modules(): + backend = getattr(module, "attn", None) + if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): + continue + if not bool((backend.kv_cache_scaling_factor == 1.0).all()): + on_uncovered("a scaled FP8 KV cache") + continue + fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) + libs = backend._fmha_manager.fmha_libs + if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): + libs.insert( + 0, + M3DrafterFullCtxFmha( + backend, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode=fused_rope_mode, + pdl=pdl, + on_rope_taken=on_rope_taken, + ), + ) + count += 1 + return count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py new file mode 100644 index 000000000000..1736115f1386 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The speculative step's store and next-step staging at the tail of the forward (one Triton launch). + +For each of the first ``n`` rows (``n`` read from device memory, so a captured graph whose step was not armed writes +nothing), with ``s = slots[i]``: + +* the spec sampler's per-slot store, exactly as ``m3_spec_tail.slot_scatter`` writes it (zero-padded to each store + width): ``new[s][c]``, ``next_new[s][c]``, ``draft[s][c]`` and ``lens[s]``; +* the engine's next-step inputs, exactly as ``m3_spec_tail.step_gather`` writes them from that store for a step whose + rows are these rows in this order (no new tokens, first previous row 0): ``ids[i*T + t] = next_new[s][t]``, + ``draft_ids[i*D + d] = draft[s][d]``, ``pos_off[i*T + t] = lens[s]``, ``kv_off[i] = lens[s] - T``. + +Integer copies only: the values are the ones the sampler's and the engine's own copies write. + +With ``cpp`` the same writes run as this target's C++ op (trtllm::minimax_m3_sm107_tp2_fold_tail, +minimaxM3Sm107Tp2/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fold_tail_kernel( + ctrl, + slots, + src_new, + sn_s0, + sn_s1, + SN_W: tl.constexpr, + dst_new, + dn_s0, + dn_s1, + DN_W: tl.constexpr, + src_next, + sx_s0, + sx_s1, + SX_W: tl.constexpr, + dst_next, + dx_s0, + dx_s1, + DX_W: tl.constexpr, + src_draft, + sd_s0, + sd_s1, + SD_W: tl.constexpr, + dst_draft, + dd_s0, + dd_s1, + DD_W: tl.constexpr, + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + T: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(ctrl) + if i < n: + s = tl.load(slots + i).to(tl.int64) + c = tl.arange(0, BLOCK) + new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) + tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) + nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) + tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) + drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) + tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) + ln = tl.load(src_lens + i) + tl.store(dst_lens + s, ln) + # The next step's inputs read the store's padded rows: zero past each source width. + tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) + tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) + tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) + tl.store(kv_off + i, ln - T) + + +def fold_tail( + ctrl: torch.Tensor, + slots: torch.Tensor, + rows: int, + src_new: torch.Tensor, + dst_new: torch.Tensor, + src_next: torch.Tensor, + dst_next: torch.Tensor, + src_draft: torch.Tensor, + dst_draft: torch.Tensor, + src_lens: torch.Tensor, + dst_lens: torch.Tensor, + ids: torch.Tensor, + draft_ids: torch.Tensor, + pos_off: torch.Tensor, + kv_off: torch.Tensor, + tokens_per_step: int, + cpp: bool = False, +) -> bool: + """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs + are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width].""" + two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) + one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) + if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): + return False + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) + or slots.dtype != torch.long + ): + return False + t_step = tokens_per_step + d_step = t_step - 1 + if ( + rows <= 0 + or slots.numel() < rows + or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) + or ids.numel() < rows * t_step + or pos_off.numel() < rows * t_step + or draft_ids.numel() < rows * d_step + or kv_off.numel() < rows + ): + return False + widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) + widths += (src_draft.shape[1], dst_draft.shape[1], t_step) + if cpp and max(widths) <= 32: + torch.ops.trtllm.minimax_m3_sm107_tp2_fold_tail( + ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, + src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, + ) # fmt: skip + return True + block = triton.next_power_of_2(max(widths)) + _fold_tail_kernel[(rows,)]( + ctrl, + slots, + src_new, + src_new.stride(0), + src_new.stride(1), + src_new.shape[1], + dst_new, + dst_new.stride(0), + dst_new.stride(1), + dst_new.shape[1], + src_next, + src_next.stride(0), + src_next.stride(1), + src_next.shape[1], + dst_next, + dst_next.stride(0), + dst_next.stride(1), + dst_next.shape[1], + src_draft, + src_draft.stride(0), + src_draft.stride(1), + src_draft.shape[1], + dst_draft, + dst_draft.stride(0), + dst_draft.stride(1), + dst_draft.shape[1], + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + t_step, + d_step, + block, + num_warps=1, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_force_vec.py new file mode 100644 index 000000000000..cc0ce28d9639 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_force_vec.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's +``force_num_accepted_tokens``), in one launch. Reached through the ``ops/spec`` entries ``m3_force_vec`` and +``m3_force_vec_req``. + +Triton. ``force_vec_kernel``: one CTA bumps the shared step counter and draws each request's count from the fixed-seed +pool with the in-tree (counter, slot) hash, exactly as ``_apply_force_accepted_tokens`` does. ``force_vec_req_kernel``: +one draw per request from its own key (its KV length at this verify), no shared counter. +""" + +import triton +import triton.language as tl + + +# One compiled variant for every batch and accept count: specializing them (== 1, divisible by 16) would compile +# mid-serving the first time a step reaches a new one, past the worker's warm-up. +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_kernel( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK: tl.constexpr, +): + c = tl.load(counter) + 1 + slots = tl.arange(0, BLOCK) + tl.store(counter, c) + idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask + live = slots < batch_size + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) + + +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_req_kernel( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK: tl.constexpr, +): + # One draw per request from its own key (its KV length at this verify): no shared counter, so a request's forced + # counts depend only on its own prompt and progress, not on what else the server ran before or beside it. + slots = tl.arange(0, BLOCK) + live = slots < batch_size + k = tl.load(keys + slots, mask=live, other=0).to(tl.int64) + idx = (k * key_stride) & pool_mask + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py new file mode 100644 index 000000000000..b6705bdb46b7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 full-context FP8 decode attention on Rubin: ``fullctx_decode``. + +The dense layers (0-2) attend every token of their FP8 E4M3 (unscaled) paged KV cache. On SM 10.7 the stock call +(flashinfer trtllm-gen, ``minimax_m3_trtllm_gen_dense_decode``) caps its split at 16 cluster CTAs per (request, KV +head): at TP2 (2 KV heads) one request runs on 32 CTAs of 212 and takes 28-30 us for 85 MB at 165k tokens. + +This runs the same arithmetic -- FP8 Q (the same E4M3 rounding of q), FP8 K/V, P quantized to E4M3 per split -- on +the vendored PrimTS decode kernel (``attention/backends/prims_ts``, SM 10.7 through the sm_100f family) at a pinned +split-KV recipe, which fans one request out over 96-144 CTAs. The split partials are merged by +``torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce`` (``cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/ +fullCtxReduce.cu``): unnormalized fp32 accumulation over the active splits, one final division, one BF16 rounding, +the merge form of stock trtllm-gen. Only the split changes against stock (modeling's ``FULLCTX_DECODE_RECIPES``). + +A plan is compiled per (device, query rows per request, requests) by ``prebuild`` before the executor's warmup; the +forward only reads plans (``plan_for``). The plan is built from PrimTS's own config, compile and workspace helpers at +an explicit recipe; the PrimTS wrapper's planning policy is not used, so no PrimTS state is replaced. With +``own_reducer`` off the PrimTS reducer merges into FP16 and the result is cast to BF16. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Mapping, Optional, Tuple + +import torch + +HEADS = 32 +KV_HEADS = 2 # capacity: fixed: the tp2 KV heads per rank +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim +PAGE = 128 +# The per-request K/V length PrimTS's config selection sizes its policy for (Q/KV tiles, launch mode); the pinned +# recipes were measured with it. The compiled capacity (``max_kv_len``) may be larger. +POLICY_KV_LEN = ( + 1 << 20 +) # capacity: fixed: the PrimTS policy's K/V length bound (the plans' compile shape) +HEADS_PER_KV = HEADS // KV_HEADS # capacity: fixed: the tp2 query heads per KV head + +_plans: Dict[Tuple[int, int, int, bool], "FullCtxPlan"] = {} +_lock = threading.Lock() + + +class FullCtxPlan: + """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" + + def __init__( + self, + device: torch.device, + qlen: int, + nreq: int, + max_kv_len: int, + recipe: Mapping, + own_reducer: bool, + pdl: bool, + ): + from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec + from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( + get_max_active_clusters_for_cluster_size, + make_decode_config, + ) + + self._tsdec = tsdec + self.qlen = qlen + self.nreq = nreq + self.pdl = pdl + self.recipe = dict(recipe) + splits = int(self.recipe["splits"]) + fp8 = tsdec._cutlass_dtype("float8_e4m3fn") + # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer + # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents + # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an + # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the + # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own + # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge + # (minimax_m3_sm107_tp2_fullctx_reduce) waits as its first instruction. + cfg = make_decode_config( + headdim=HEAD_DIM, + args={"use_pdl": pdl, **self.recipe.get("config", {})}, + seq_len_q=qlen, + seq_len_kv=min(max_kv_len, POLICY_KV_LEN), + batch_size=nreq, + num_heads_q=HEADS, + num_heads_kv=KV_HEADS, + q_dtype=fp8, + k_dtype=fp8, + v_dtype=fp8, + o_dtype=tsdec._cutlass_dtype("float16"), + qkv_layout="pagedKv", + num_tokens_per_page=PAGE, + storage_tokens_per_page=PAGE, + split_kv_mode="gmem_reduction_with_separate_kernel", + splits_kv=splits, + max_splits_kv=splits, + min_loop_iters_per_split=1, + mask_type="causal" if qlen > 1 else "dense", + auto_tuner=True, + split_kv=True, + ) + _check_layout(cfg, qlen, splits) + with torch.cuda.device(device): + max_active = get_max_active_clusters_for_cluster_size(1) + spec = tsdec._decode_launch_spec_from_config( + cfg, + batch_size=nreq, + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + seq_len_q=qlen, + max_active_clusters=int(max_active), + ) + compile_spec = tsdec._make_decode_compile_spec( + spec, + device_index=_index(device), + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + page_size=PAGE, + max_kv_len=max_kv_len, + seq_len_q=qlen, + q_dtype_key="float8_e4m3fn", + k_dtype_key="float8_e4m3fn", + v_dtype_key="float8_e4m3fn", + output_dtype_key="float16", + use_packed_q=False, + kv_prefix_mode="dynamic", + kv_lengths_mode="dynamic", + ) + self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) + layout = tsdec._make_decode_workspace_layout( + spec.scratch_shapes, + torch.float16, + use_separate_reduction_kernel=True, + use_split_kv=True, + ) + self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) + self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) + self.workspace.split_kv_counter.zero_() + self.workspace.cu_seqlens_q.zero_() + self.workspace.attention_sinks.zero_() + rows = HEADS_PER_KV * qlen + if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): + raise RuntimeError( + f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " + f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" + ) + self.own_reducer = own_reducer + self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) + self.tile_kv = int(cfg.tile_size_kv) + self.insts_kv = int(cfg.num_insts_kv) + self.causal = qlen > 1 + shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) + self.staging = torch.empty(shape, dtype=torch.float16, device=device) + + def run( + self, + q: torch.Tensor, + k_pages: torch.Tensor, + v_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + ) -> None: + # The kernel takes a compact query; a column view of the fused QKV projection is copied once. + q = q.reshape(self.staging.shape).contiguous() + runtime = self._tsdec._prepare_decode_runtime_unchecked( + q, + (k_pages, v_pages), + None, + output_dtype=torch.float16, + bmm1_scale=sm_scale, + bmm2_scale=1.0, + out=self.staging, + ) + ws = self.workspace + self._main( + runtime.q, + runtime.k_cache, + runtime.v_cache, + runtime.k_sf_cache, + runtime.v_sf_cache, + runtime.out, + seq_lens, + ws.cu_seqlens_q, + block_table, + block_table, # Dense decode constexpr-elides the membership table. + ws.partial_o, + ws.partial_stats, + ws.split_kv_counter, + ws.attention_sinks, + runtime.num_physical_pages, + runtime.k_page_stride, + runtime.k_head_stride, + runtime.k_token_stride, + runtime.v_page_stride, + runtime.v_head_stride, + runtime.v_token_stride, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + if self.own_reducer: + torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce( + ws.partial_o, + ws.partial_stats, + seq_lens, + output, + self.qlen, + self.q_tokens_per_cta, + self.tile_kv, + self.insts_kv, + self.causal, + self.pdl, + ) + return + self._primts_reducer( + runtime.out, + seq_lens, + ws.cu_seqlens_q, + ws.partial_o, + ws.partial_stats, + ws.attention_sinks, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + output.view(self.staging.shape).copy_(self.staging) + + +def _check_layout(cfg, qlen: int, splits: int) -> None: + """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. + + Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is + ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding + window, attention sinks or packed Q. + """ + problems = [] + if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): + problems.append("not split-KV with the separate GMEM reducer") + if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: + problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") + if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: + problems.append("cluster-SMEM reduction or a persistent scheduler") + if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: + problems.append("packed Q, a sliding window or attention sinks") + if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): + problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") + if qlen > 1: + if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): + problems.append("multi-token Q without the grouped [token, head] tile layout") + elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): + problems.append( + f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" + ) + if problems: + raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + + +def has_own_reducer() -> bool: + """Require the BF16 split-KV merge op supplied by this target build.""" + torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce + return True + + +def plan_for( + device: torch.device, qlen: int, nreq: int, own_reducer: bool = True +) -> Optional[FullCtxPlan]: + """The prebuilt plan for this shape, or None.""" + return _plans.get((_index(device), qlen, nreq, own_reducer)) + + +def prebuild( + device: torch.device, + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + pdl: bool, +) -> int: + """Compile the plan of every (query rows per request, requests) in ``recipes``. Call before warmup.""" + built = 0 + for (qlen, nreq), recipe in recipes.items(): + key = (_index(device), qlen, nreq, own_reducer) + with _lock: + if key not in _plans: + _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) + built += 1 + return built + + +def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: + """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" + return pool[0::subpages_per_slot], pool[1::subpages_per_slot] + + +def _index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py new file mode 100644 index 000000000000..8088d79f014f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py @@ -0,0 +1,670 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with four schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. +* ``prestage_old`` (with ``early_issue``, ``INDEX_SCORE_PRESTAGE_OLD``): when the CTA's first block holds only + tokens older than this step (block < (seq_lens hint - decode query length) / 128), its K tile is TMA-loaded into its + stage before the wait; Q (and a first block that is not all old) is still loaded after it. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer``) through ``IndexScoreRunner`` +and called through this target's ``m3_index_score`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, griddepcontrol_wait + +# With the early issue, a CTA whose first block holds only tokens older than this step also TMA-loads that block's +# index-K tile before the grid wait (bit-identical scores; replay period -47 us per step on Rubin TP2). +INDEX_SCORE_PRESTAGE_OLD = True + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + prestage_old: bool = False, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K + # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's + # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. + self.prestage_old = prestage_old and early_issue + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # Otherwise the dependent block selector is released as the CTAs exit. + early_q = False + early_k = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + seqlen_hint = seq_lens[batch_id] + hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + if cutlass.const_expr(self.prestage_old): + # Stage the first K tile now when every token of the block predates this step. + if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: + early_k = True + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + if not early_k: + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # The length that bounds this grid's work is read after the wait (the + # pre-wait reads above are prefetch / early-issue hints only, clamped to + # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens + # patches every step, and num_blocks bounds an unmasked block_table load + # and an unmasked score store, so a stale length is an out-of-bounds write + # rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + prestage_old=early and INDEX_SCORE_PRESTAGE_OLD, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py new file mode 100644 index 000000000000..3e88c7fb1ce2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Locality domains ("uGPU") of a multi-die GPU (Rubin VR200: 2 domains of 100 SMs + 12 remainder SMs) for the +streaming decode kernels of this target: the SM -> domain map, per-domain weight placement, and the CTA -> SM +assignment a persistent kernel uses to read only its own die's memory. + +A kernel launched with one CTA per SM reads ``%smid`` once and takes its logical block index from a host-built +permutation (:meth:`LocalityMap.logical_map`), so every per-CTA structure (work plan row, epochs, roles) follows its +die. Weights that are split by output rows into per-domain halves (:meth:`LocalityMap.split_rows`) live in one virtual +range whose halves are physically in the two domains. Placement only: the arithmetic is unchanged, so results are +bitwise those of the unplaced weights. + +The SM -> domain map is attested by the driver: the SMs are split by locality domain +(``cuDevSmResourceSplit`` with ``CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID``), and a ``%smid`` probe runs on a green +context of each split. Allocations use driver VMM with ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN`` and are checked +with ``CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL``; any failure raises (fail closed).""" + +from __future__ import annotations + +import threading +from typing import Dict, List, Sequence, Tuple + +import torch + +_LOCK = threading.Lock() +_MAPS: Dict[int, "LocalityMap"] = {} +_PROBE_SRC = r""" +extern "C" __global__ void smid_probe(int* out) { + if (threadIdx.x == 0) { + unsigned s; + asm volatile("mov.u32 %0, %%smid;" : "=r"(s)); + out[blockIdx.x] = (int)s; + } +} +""" + + +def _drv(): + from cuda.bindings import driver + + return driver + + +def _ck(res, what: str): + """cuda-python returns (CUresult, *values): raise on an error, return the values.""" + err = res[0] if isinstance(res, tuple) else res + if int(err) != 0: + raise RuntimeError(f"locality: {what} failed: {err}") + if isinstance(res, tuple): + return res[1] if len(res) == 2 else res[1:] + return None + + +class _Mapping: + """One virtual range backed by per-domain physical allocations, exposed to torch through + ``__cuda_array_interface__``. Lives as long as a tensor that views it (the tensor holds this object).""" + + def __init__(self, device: int, pieces: Sequence[Tuple[int, int]]): + """``pieces``: (bytes, domain) in address order; each a multiple of the allocation granularity.""" + d = _drv() + self.device = device + self.size = sum(n for n, _ in pieces) + self.va = _ck(d.cuMemAddressReserve(self.size, 0, 0, 0), "cuMemAddressReserve") + self.handles = [] + off = 0 + try: + for nbytes, dom in pieces: + prop = _alloc_prop(device, dom) + h = _ck(d.cuMemCreate(nbytes, prop, 0), f"cuMemCreate(domain {dom})") + self.handles.append((h, nbytes)) + _ck(d.cuMemMap(int(self.va) + off, nbytes, 0, h, 0), "cuMemMap") + off += nbytes + acc = d.CUmemAccessDesc() + acc.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + acc.location.id = device + acc.flags = d.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _ck(d.cuMemSetAccess(self.va, self.size, [acc], 1), "cuMemSetAccess") + except Exception: + self.close() + raise + self.__cuda_array_interface__ = { + "shape": (self.size,), + "typestr": "|u1", + "data": (int(self.va), False), + "version": 3, + "strides": None, + "stream": None, + } + + def close(self): + d = _drv() + off = 0 + for h, n in self.handles: + d.cuMemUnmap(int(self.va) + off, n) + d.cuMemRelease(h) + off += n + self.handles = [] + if self.va is not None: + d.cuMemAddressFree(self.va, self.size) + self.va = None + + def __del__(self): + try: + self.close() + except Exception: + pass + + +_LIBCUDA = None + + +def _libcuda(): + global _LIBCUDA + if _LIBCUDA is None: + import ctypes + + _LIBCUDA = ctypes.CDLL("libcuda.so.1") + return _LIBCUDA + + +def _alloc_prop(device: int, domain: int): + d = _drv() + prop = d.CUmemAllocationProp() + prop.type = d.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.requestedHandleTypes = d.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE + prop.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN + prop.location.localized.deviceId = device + prop.location.localized.localityDomainId = domain + return prop + + +class LocalityMap: + """The locality domains of one device. Build with :meth:`get` before CUDA-graph capture.""" + + def __init__(self, device: int): + d = _drv() + _ck(d.cuInit(0), "cuInit") + dev = _ck(d.cuDeviceGet(device), "cuDeviceGet") + attr = d.CUdevice_attribute + self.device = device + self.n_domains = int( + _ck( + d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, dev), + "domain count", + ) + ) + self.sms_per_domain = int( + _ck( + d.cuDeviceGetAttribute( + attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_MULTIPROCESSOR_COUNT, dev + ), + "SMs per domain", + ) + ) + self.num_sms = int( + _ck(d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev), "SMs") + ) + if self.n_domains != 2: + raise RuntimeError( + f"locality: {self.n_domains} locality domains (this layer handles 2)" + ) + self._domain_sms = self._probe(d, dev) + dom = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for s in sms: + if dom[s] != -1: + raise RuntimeError(f"locality: SM {s} probed in two domains") + dom[s] = k + self.sm_domain: List[int] = dom + self.remainder: List[int] = [s for s in range(self.num_sms) if dom[s] < 0] + rt = _alloc_prop(device, 0) + self.granularity = int( + _ck( + d.cuMemGetAllocationGranularity( + rt, d.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM + ), + "granularity", + ) + ) + + @classmethod + def get(cls, device) -> "LocalityMap": + idx = torch.device(device).index if not isinstance(device, int) else device + idx = torch.cuda.current_device() if idx is None else idx + with _LOCK: + m = _MAPS.get(idx) + if m is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "LocalityMap probes the device: build it before CUDA-graph capture" + ) + m = _MAPS[idx] = cls(idx) + return m + + def sms(self, domain: int) -> List[int]: + return list(self._domain_sms[domain]) + + @property + def sm_table(self) -> torch.Tensor: + """int32 [num_sms] on the device: (domain << 16) | rank within the domain's sorted SMs; -1 = remainder SM. + For kernels that pick work by %smid themselves.""" + t = getattr(self, "_sm_table", None) + if t is None: + v = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for i, s in enumerate(sms): + v[s] = (k << 16) | i + t = self._sm_table = torch.tensor(v, dtype=torch.int32, device=f"cuda:{self.device}") + return t + + def _probe(self, d, dev) -> List[List[int]]: + """Split the SMs by locality domain (strict, no backfill) and record %smid of a probe kernel launched on a + green context of each part: the driver's own SM -> domain assignment.""" + from cuda.bindings import nvrtc + + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + # the probe module, compiled for this device + major = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev + ), + "cc", + ) + minor = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev + ), + "cc", + ) + prog = _ck( + nvrtc.nvrtcCreateProgram(_PROBE_SRC.encode(), b"smid_probe.cu", 0, [], []), + "nvrtcCreateProgram", + ) + opts = [f"--gpu-architecture=sm_{major}{minor}".encode()] + err = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)[0] + if int(err) != 0: + raise RuntimeError(f"locality: probe compile failed: {err}") + n = _ck(nvrtc.nvrtcGetCUBINSize(prog), "cubin size") + cubin = b" " * n + _ck(nvrtc.nvrtcGetCUBIN(prog, cubin), "cubin") + out: List[List[int]] = [] + nb = 4 * self.sms_per_domain # more blocks than SMs: every SM of the part shows up + torch.cuda.set_device(self.device) + buf = torch.full((nb,), -1, dtype=torch.int32, device=f"cuda:{self.device}") + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate(desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM), + "green ctx", + ) + st = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + ctx = _ck(d.cuCtxFromGreenCtx(g), "ctx from green ctx") + _ck(d.cuCtxPushCurrent(ctx), "push ctx") + try: + mod = _ck(d.cuModuleLoadData(cubin), "probe module") + fn = _ck(d.cuModuleGetFunction(mod, b"smid_probe"), "probe fn") + buf.fill_(-1) + torch.cuda.synchronize(self.device) + import ctypes + + arg = ctypes.c_void_p(buf.data_ptr()) + _ck( + d.cuLaunchKernel( + fn, nb, 1, 1, 32, 1, 1, 0, st, ((arg,), (ctypes.c_void_p,)), 0 + ), + "probe launch", + ) + _ck(d.cuStreamSynchronize(st), "probe sync") + sms = sorted(set(int(x) for x in buf.tolist() if x >= 0)) + d.cuModuleUnload(mod) + finally: + d.cuCtxPopCurrent() + d.cuStreamDestroy(st) + d.cuGreenCtxDestroy(g) + if len(sms) != self.sms_per_domain: + raise RuntimeError( + f"locality: domain {k} probed {len(sms)} SMs, expected {self.sms_per_domain}" + ) + out.append(sms) + return out + + def logical_map(self, per_domain: int, extra: int) -> Tuple[torch.Tensor, List[int]]: + """The SM -> logical block index permutation for a kernel of ``2 * per_domain + extra`` CTAs, one per SM: + logical blocks [d * per_domain, (d + 1) * per_domain) run on domain d's first ``per_domain`` SMs; logical + blocks 2 * per_domain .. run on the rest (the leftover domain SMs, alternating domains, then the remainder + SMs). Returns (int32 [num_sms] on the device, -1 = no block; the domain of each extra block, -1 remainder).""" + if 2 * per_domain + extra != self.num_sms: + raise ValueError( + f"locality: {2 * per_domain} + {extra} CTAs must cover the {self.num_sms} SMs" + ) + if per_domain > self.sms_per_domain: + raise ValueError(f"locality: {per_domain} CTAs per domain > {self.sms_per_domain} SMs") + lmap = [-1] * self.num_sms + for k in range(2): + for i, s in enumerate(self._domain_sms[k][:per_domain]): + lmap[s] = k * per_domain + i + left = [self._domain_sms[k][per_domain:] for k in range(2)] + rest = [x for pair in zip(*left) for x in pair] + self.remainder + ext_dom = [] + for i, s in enumerate(rest): + lmap[s] = 2 * per_domain + i + ext_dom.append(self.sm_domain[s]) + assert sorted(lmap) == list(range(self.num_sms)) + return torch.tensor(lmap, dtype=torch.int32, device=f"cuda:{self.device}"), ext_dom + + def domain_streams(self) -> List[torch.cuda.ExternalStream]: + """One non-blocking stream per domain on a green context holding exactly that domain's SMs (kept alive with + this map): work launched on stream d runs only on domain d's SMs. Fork / join with events from the caller's + stream; graph capture follows the forks. Built once, before capture.""" + st = getattr(self, "_streams", None) + if st is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.domain_streams: build before CUDA-graph capture") + d = _drv() + dev = _ck(d.cuDeviceGet(self.device), "cuDeviceGet") + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + self._green = [] + st = [] + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate( + desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM + ), + "green", + ) + s_ = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + self._green.append((g, s_)) + st.append(torch.cuda.ExternalStream(int(s_), device=f"cuda:{self.device}")) + self._streams = st + return st + + def domain_of(self, ptr: int) -> int: + """The locality domain ordinal of a device address (-1: not placed in one domain). Through libcuda + directly: cuda-python's attribute helper does not know this attribute.""" + import ctypes + + d = _drv() + lib = _libcuda() + out = ctypes.c_int(-9) + attr = int(d.CUpointer_attribute.CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL) + err = lib.cuPointerGetAttribute(ctypes.byref(out), ctypes.c_int(attr), ctypes.c_uint64(ptr)) + if err != 0: + raise RuntimeError( + f"locality: cuPointerGetAttribute(LOCALITY_DOMAIN_ORDINAL) failed: {err}" + ) + return int(out.value) + + def place_split(self, t: torch.Tensor, split_row: int) -> torch.Tensor: + """Split-mapped placement: a copy of ``t`` (contiguous, rows along dim 0) with the SAME shape and layout, in one + virtual range whose rows [0, split_row) are physically in domain 0 and rows [split_row, N) in domain 1. The + split must fall on an allocation granule (rows [0, split_row) a multiple of 2 MiB on VR200); the tail is padded + to a granule in domain 1. Stock consumers read it as a plain tensor; no extra memory once ``t`` is dropped.""" + if not t.is_contiguous(): + raise ValueError("locality.place_split: needs a contiguous tensor") + row_bytes = t[0].numel() * t.element_size() if t.dim() > 1 else t.element_size() + head = split_row * row_bytes + total = t.numel() * t.element_size() + if not 0 < head < total or head % self.granularity: + raise ValueError( + f"locality.place_split: rows [0, {split_row}) = {head} B must be a granule multiple " + f"({self.granularity} B) inside the {total} B tensor" + ) + tail = -(-(total - head) // self.granularity) * self.granularity + m = _Mapping(self.device, [(head, 0), (tail, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + out = flat[:total].view(t.dtype).view(t.shape) + out.copy_(t) + out._locality_mapping = m + for q, k in ( + (out.data_ptr(), 0), + (out.data_ptr() + head - 1, 0), + (out.data_ptr() + head, 1), + (out.data_ptr() + total - 1, 1), + ): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.place_split: byte {q:#x} is not in domain {k}") + return out + + def replace_param_(self, param: torch.nn.Parameter, split_row: int) -> torch.nn.Parameter: + """Re-place a parameter in place with :meth:`place_split` (swaps ``param.data``; the old block returns to the + caching allocator, so call ``torch.cuda.empty_cache()`` between large re-placements).""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.replace_param_: not during CUDA-graph capture") + param.data = self.place_split(param.data, split_row) + return param + + def split_rows(self, t: torch.Tensor, dim: int = 1) -> torch.Tensor: + """Place ``t`` (contiguous, 1-byte elements, rows along ``dim`` split into two equal halves) as + [2, *t.shape[:dim], t.shape[dim] // 2, *t.shape[dim + 1:]]: half k (rows [k h, (k + 1) h) of every leading + index) in domain k's memory, one virtual range. Each half's byte size must be a multiple of the allocation + granularity (2 MiB on VR200). The copy is exact; the caller drops ``t`` afterwards.""" + if t.element_size() != 1 or not t.is_contiguous(): + raise ValueError("locality.split_rows: needs a contiguous 1-byte tensor") + n = t.shape[dim] + if n % 2: + raise ValueError(f"locality.split_rows: {n} rows don't split in halves") + half_bytes = t.numel() // 2 + if half_bytes % self.granularity: + raise ValueError( + f"locality.split_rows: a half is {half_bytes} B, not a multiple of the {self.granularity} B granularity" + ) + m = _Mapping(self.device, [(half_bytes, 0), (half_bytes, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + shape = (2,) + tuple(t.shape[:dim]) + (n // 2,) + tuple(t.shape[dim + 1 :]) + out = flat.view(torch.uint8).view(shape) + src = t.view(torch.uint8) + for k in range(2): + out[k].copy_(src.narrow(dim, k * (n // 2), n // 2)) + out._locality_mapping = m # the tensor keeps its mapping alive + for k in range(2): + p = out[k].data_ptr() + for q in (p, p + half_bytes - 1): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.split_rows: byte {q:#x} is not in domain {k}") + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_direct.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_direct.py new file mode 100644 index 000000000000..30b905712b1f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_direct.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +"""Register-prefetch BF16 GEMM for low-M, long-K decode shapes. + +This target's copy of the installed FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_direct.py`` with one +addition, the drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH_DIRECT): after its one store of each output, +the owning thread also stores the CTA's column pairs as the boundary all-reduce's canonical exchange words (kernels/ +m3_mc_allreduce) at the call parity, before the kernel's (unchanged, last) release of its dependents. The plain +kernel and its output are FlashInfer's; only the push variant is called here (``run_direct_dense_push``). + +The kernel keeps a complete output dot product inside one CTA and reuses each +prefetched B value across several public-M rows. It is intentionally a +separate autotuner runner from the Blackwell tensor-core split-K kernel: the +two algorithms have different useful shape regions and tactic spaces. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import torch as _torch +from cutlass import Int32, const_expr +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +_VECTOR_WIDTH = 8 # capacity: fixed: a vector load width +_SUPPORTED_BLOCK_SIZES = (32, 64, 96, 128, 192, 256, 384) +_SUPPORTED_OUTPUTS_PER_BLOCK = (1, 2, 4) +_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound +_COMPILE_OPTIONS = "--ptxas-options -maxrregcount=64" + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.s32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + llvm.inline_asm( + None, + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dataclasses.dataclass(frozen=True) +class DirectTactic: + """One direct-kernel specialization.""" + + block_size: int + outputs_per_block: int + rows_per_block: int + + +def _default_rows_per_block(m: int) -> int: + if m <= 8: + return m + return next(rows for rows in (8, 4, 2, 1) if m % rows == 0) + + +def validate_tactic(tactic: DirectTactic, m: int, n: int, k: int) -> None: + """Reject a direct tactic that cannot serve ``(m, n, k)``.""" + if tactic.block_size not in _SUPPORTED_BLOCK_SIZES: + raise ValueError(f"unsupported block_size={tactic.block_size}") + if tactic.outputs_per_block not in _SUPPORTED_OUTPUTS_PER_BLOCK: + raise ValueError(f"unsupported outputs_per_block={tactic.outputs_per_block}") + if not 1 <= m <= _MAX_M: + raise ValueError(f"direct GEMM requires 1 <= M <= {_MAX_M}, got {m}") + if not 1 <= tactic.rows_per_block <= m or m % tactic.rows_per_block: + raise ValueError(f"rows_per_block={tactic.rows_per_block} must divide M={m}") + if n <= 0 or n % tactic.outputs_per_block: + raise ValueError(f"N={n} must be divisible by outputs_per_block={tactic.outputs_per_block}") + k_tile = tactic.block_size * _VECTOR_WIDTH + if k <= 0 or k % k_tile: + raise ValueError(f"K={k} must be divisible by {k_tile}") + + +def default_tactic(m: int, n: int, k: int) -> DirectTactic: + """Choose the measured register-prefetch fallback tactic.""" + block_size = next( + (block for block in (256, 192, 128, 96, 64, 32) if k % (block * _VECTOR_WIDTH) == 0), + None, + ) + if block_size is None: + raise ValueError("direct GEMM requires a supported 16-byte K tiling") + outputs_per_block = next(outputs for outputs in (2, 1) if n % outputs == 0) + tactic = DirectTactic( + block_size, + outputs_per_block, + _default_rows_per_block(m), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +def autotune_tactics(m: int, n: int, k: int) -> list[DirectTactic]: + """Enumerate the compact tactic space used by FlashInfer autotuning. + + Block sizes cover every configuration exercised in the H100/B200 sweep; + output grouping spans the measured 1/2/4-column choices. Row tiling stays + at the occupancy-oriented default to keep JIT cost bounded. + """ + try: + default = default_tactic(m, n, k) + except ValueError: + return [] + tactics = [default] + for block_size in _SUPPORTED_BLOCK_SIZES: + for outputs_per_block in _SUPPORTED_OUTPUTS_PER_BLOCK: + tactic = DirectTactic( + block_size, + outputs_per_block, + default.rows_per_block, + ) + try: + validate_tactic(tactic, m, n, k) + except ValueError: + continue + tactics.append(tactic) + return list(dict.fromkeys(tactics)) + + +def prefer_direct_bf16_gemm_sm100(m: int, n: int, k: int) -> bool: + """Return the conservative B200 no-autotune crossover heuristic. + + The three bands are a compact fit to a warm/cold sweep over M=1..16,24,32, + 18 N values, and 11 K values. This is deliberately not a blanket rule for + K=8192: direct wins only where public M and N leave the tensor-core path + with too little independent output work. + """ + return k == 8192 and ((m == 1 and n <= 4608) or (m <= 4 and n <= 512) or (m <= 8 and n <= 256)) + + +class DirectDenseGemmKernel: + """K-specialized direct GEMM with whole-mainloop vector prefetch.""" + + def __init__( + self, + *, + element_type, + num_rows: int, + k_extent: int, + tactic: DirectTactic, + use_pdl: bool, + ar_push: bool = False, + ) -> None: + validate_tactic(tactic, num_rows, tactic.outputs_per_block, k_extent) + self.ar_push = ar_push + self.element_type = element_type + self.num_rows = num_rows + self.rows_per_block = tactic.rows_per_block + self.k_extent = k_extent + self.block_size = tactic.block_size + self.outputs_per_block = tactic.outputs_per_block + self.vector_width = _VECTOR_WIDTH + self.use_pdl = use_pdl + self.num_warps = tactic.block_size // cute.arch.WARP_SIZE + self.num_k_tiles = k_extent // (tactic.block_size * _VECTOR_WIDTH) + + @cute.jit + def __call__( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64 = cutlass.Int64(0), + ar_flags: cutlass.Int64 = cutlass.Int64(0), + ar_rank: Int32 = Int32(0), + ar_world: Int32 = Int32(1), + ar_maxt: Int32 = Int32(8), + ) -> None: + n = cute.size(gB, mode=[0]) + copy_a = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.ALWAYS, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.CopyG2ROp(), + self.element_type, + num_bits_per_copy=self.vector_width * self.element_type.width, + load_cache_mode=cute.nvgpu.LoadCacheMode.STREAMING, + ) + self.kernel(gA, gB, gC, copy_a, copy_b, ar_ws, ar_flags, ar_rank, ar_world, ar_maxt).launch( + grid=[ + cute.ceil_div(n, self.outputs_per_block), + self.num_rows // self.rows_per_block, + 1, + ], + block=[self.block_size, 1, 1], + smem=self.rows_per_block * self.outputs_per_block * self.num_warps * 4, + stream=stream, + use_pdl=self.use_pdl, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + copy_a: cute.CopyAtom, + copy_b: cute.CopyAtom, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ) -> None: + tidx, _, _ = cute.arch.thread_idx() + block_idx, block_m, _ = cute.arch.block_idx() + warp_idx = cute.arch.warp_idx() + + num_rows: cutlass.Constexpr = self.rows_per_block + outputs_per_block: cutlass.Constexpr = self.outputs_per_block + vector_width: cutlass.Constexpr = self.vector_width + block_size: cutlass.Constexpr = self.block_size + num_warps: cutlass.Constexpr = self.num_warps + num_k_tiles: cutlass.Constexpr = self.num_k_tiles + + acc = cute.make_rmem_tensor( + cute.make_layout((num_rows, outputs_per_block), stride=(outputs_per_block, 1)), + cutlass.Float32, + ) + acc.fill(0.0) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + n_base = block_idx * outputs_per_block + m_base = block_m * num_rows + gA_vec = cute.logical_divide(gA, (None, vector_width)) + gB_vec = cute.logical_divide(gB, (None, vector_width)) + tA_all = cute.logical_divide(gA_vec, (None, (None, block_size))) + tB_all = cute.logical_divide(gB_vec, (None, (None, block_size))) + tA = tA_all[None, (None, (tidx, None))] + + b_regs = cute.make_rmem_tensor( + cute.make_layout( + (outputs_per_block, num_k_tiles, vector_width), + stride=(num_k_tiles * vector_width, vector_width, 1), + ), + self.element_type, + ) + for ni in cutlass.range_constexpr(outputs_per_block): + tB = tB_all[n_base + ni, (None, (tidx, None))] + for k_tile in cutlass.range_constexpr(num_k_tiles): + cute.copy(copy_b, tB[None, k_tile], b_regs[ni, k_tile, None]) + + a_regs = cute.make_rmem_tensor( + cute.make_layout((num_k_tiles, vector_width), stride=(vector_width, 1)), + self.element_type, + ) + for mi in cutlass.range_constexpr(num_rows): + for k_tile in cutlass.range_constexpr(num_k_tiles): + cute.copy( + copy_a, + tA[m_base + mi, None, k_tile], + a_regs[k_tile, None], + ) + for k_tile in cutlass.range_constexpr(num_k_tiles): + for vi in cutlass.range_constexpr(vector_width): + a_value = a_regs[k_tile, vi].to(cutlass.Float32) + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = acc[mi, ni] + a_value * b_regs[ni, k_tile, vi].to( + cutlass.Float32 + ) + + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni]) + + smem_layout = cute.make_layout( + (num_rows, outputs_per_block, num_warps), + stride=(outputs_per_block * num_warps, num_warps, 1), + ) + smem = cutlass.utils.SmemAllocator() + partials = smem.allocate_tensor(cutlass.Float32, smem_layout, byte_alignment=16) + with cute.arch.elect_one(): + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + partials[mi, ni, warp_idx] = acc[mi, ni] + + cute.arch.sync_threads() + if tidx == 0: + for mi in cutlass.range_constexpr(num_rows): + for ni in cutlass.range_constexpr(outputs_per_block): + total = cutlass.Float32(0.0) + for warp in cutlass.range_constexpr(num_warps): + total = total + partials[mi, ni, warp] + gC[m_base + mi, n_base + ni] = total.to(self.element_type) + if const_expr(self.ar_push): + acc[mi, ni] = total + if const_expr(self.ar_push): + # Push the CTA's column pairs of every row into the boundary all-reduce's exchange (prepushed 2), at the + # (token, column / (N / 8)) call parity; the release below follows the pushes (the prepushed contract). + n_cols = cute.size(gC, mode=[1]) + half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(n_cols // 2) + for mi in cutlass.range_constexpr(num_rows): + t = m_base + mi + for pi in cutlass.range_constexpr(outputs_per_block // 2): + col = n_base + 2 * pi + slot = cutlass.Int64(t * 8 + col // (n_cols // 8)) + par = _ld_volatile_s32(ar_flags + cutlass.Int64(4) * slot) & Int32(1) + row = cutlass.Int64(t * ar_world + ar_rank) + off = ( + cutlass.Int64(par) * half + + row * cutlass.Int64(n_cols // 2) + + cutlass.Int64(col // 2) + ) + word = _bf16x2_word(acc[mi, 2 * pi], acc[mi, 2 * pi + 1]) + _st_global_b32(ar_ws + off * cutlass.Int64(4), word) + + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +def _from_dlpack_static(tensor: _torch.Tensor): + # K is specialized and the row stride must retain its 16-byte divisibility + # for the verifier to accept vectorized G2R copies. + return from_dlpack(tensor, assumed_align=32) + + +def _make_compile_repr_tensors(dtype, m: int, n: int, k: int): + return tuple( + _from_dlpack_static(tensor) + for tensor in ( + _torch.empty((m, k), dtype=dtype, device="cuda"), + _torch.empty((n, k), dtype=dtype, device="cuda"), + _torch.empty((m, n), dtype=dtype, device="cuda"), + ) + ) + + +@functools.cache +def _get_compiled_direct_kernel( + dtype, + m: int, + n: int, + k: int, + tactic: DirectTactic, + use_pdl: bool, +): + if dtype != _torch.bfloat16: + raise ValueError(f"direct GEMM supports BF16; got {dtype}") + kernel = DirectDenseGemmKernel( + element_type=cutlass.BFloat16, + num_rows=m, + k_extent=k, + tactic=tactic, + use_pdl=use_pdl, + ) + tensors = _make_compile_repr_tensors(dtype, m, n, k) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + return cute_ext.compile(kernel, *tensors, stream, options=_COMPILE_OPTIONS) + + +def _validate_runtime_tensors(a, b, out, tactic: DirectTactic): + if any(not isinstance(tensor, _torch.Tensor) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("direct GEMM accepts only 2D tensors") + if a.device.type != "cuda" or b.device != a.device or out.device != a.device: + raise ValueError("a, b, and out must be on the same CUDA device") + if a.dtype != _torch.bfloat16 or b.dtype != a.dtype or out.dtype != a.dtype: + raise ValueError("a, b, and out must share BF16 dtype") + if not a.is_contiguous() or not b.T.is_contiguous() or not out.is_contiguous(): + raise ValueError("direct GEMM requires row-major A/out and column-major B") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + validate_tactic(tactic, m, n, k) + return m, n, k + + +def run_direct_dense(a, b, out, pdl: bool, tactic: DirectTactic): + """Run direct ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + m, n, k = _validate_runtime_tensors(a, b, out, tactic) + compiled = _get_compiled_direct_kernel(a.dtype, m, n, k, tactic, pdl) + tensors = tuple(_from_dlpack_static(tensor) for tensor in (a, b.T, out)) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + compiled(*tensors, stream) + return out + + +__all__ = [ + "DirectTactic", + "autotune_tactics", + "default_tactic", + "prefer_direct_bf16_gemm_sm100", + "run_direct_dense", + "run_direct_dense_push", + "compile_direct_push", +] + + +@functools.cache +def _get_compiled_direct_push(dtype, m: int, n: int, k: int, tactic: DirectTactic, use_pdl: bool): + assert tactic.outputs_per_block % 2 == 0, "the push writes column pairs" + kernel = DirectDenseGemmKernel( + element_type=cutlass.BFloat16, + num_rows=m, + k_extent=k, + tactic=tactic, + use_pdl=use_pdl, + ar_push=True, + ) + tensors = _make_compile_repr_tensors(dtype, m, n, k) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + return cute_ext.compile( + kernel, + *tensors, + stream, + cutlass.Int64(0), + cutlass.Int64(0), + Int32(0), + Int32(1), + Int32(8), + options=_COMPILE_OPTIONS, + ) + + +def compile_direct_push(a, b, out, pdl: bool, tactic) -> None: + """Compile (only) the push variant for these operands' shapes.""" + m, n, k = _validate_runtime_tensors(a, b, out, DirectTactic(*tactic)) + _get_compiled_direct_push(a.dtype, m, n, k, DirectTactic(*tactic), pdl) + + +def run_direct_dense_push(a, b, out, pdl: bool, tactic: DirectTactic, ar: tuple): + """``run_direct_dense`` that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange.""" + m, n, k = _validate_runtime_tensors(a, b, out, tactic) + compiled = _get_compiled_direct_push(a.dtype, m, n, k, tactic, pdl) + tensors = tuple(_from_dlpack_static(tensor) for tensor in (a, b.T, out)) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + ws, flags, rank, world, maxt = ar + compiled( + *tensors, + stream, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py new file mode 100644 index 000000000000..027541ea6990 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py @@ -0,0 +1,1249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +This target's copy of the installed FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the +Eagle3 drafter's low-M linears, modeling's ``DRAFT_LOWM_OWN_SPLITK``): the same kernel and the same split order at every +tactic, bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's +final cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, +the SwiGLU prologue that follow. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 # capacity: fixed: a kernel tile constant + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 # capacity: fixed: an alignment + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 # capacity: fixed: pipeline stages +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 # capacity: fixed: pipeline stages + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= _MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES + base = SplitKTactic(mma_m, mma_n, split_k, min_stages) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(min_stages, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 # capacity: fixed: a split index, not a size + + +# The drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH): the exchange words of +# ops/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.s32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + llvm.inline_asm( + None, + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + ar_push: bool = False, + ) -> None: + self.ar_push = ar_push + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64 = cutlass.Int64(0), + ar_flags: cutlass.Int64 = cutlass.Int64(0), + ar_rank: Int32 = Int32(0), + ar_world: Int32 = Int32(1), + ar_maxt: Int32 = Int32(8), + ar_tokens: Int32 = Int32(0), + ar_n: Int32 = Int32(0), + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel( + a, b, c, bias, ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_tokens, ar_n + ).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + ar_ws: cutlass.Int64, # ar_push: the boundary all-reduce's multicast exchange words (int32) + ar_flags: cutlass.Int64, # ar_push: its per-(token, cluster CTA) call flags (parity = flag & 1) + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, # the bucket's tokens + ar_tokens: Int32, # rows pushed (public M) + ar_n: Int32, # public N (the hidden size): the exchange row width + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + # ar_push: the owner's BF16 tile (kernel-M = output column, kernel-N = token), M-major. + sD = cute_ext.allocate( + mC.element_type, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=16, + ) + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sD, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl and not self.ar_push): + # ar_push: the dependent is a prepushed all-reduce, which skips its grid wait: it is released only after + # this grid's own grid wait and its pushes, in the owner's _store_tile (the prepushed >= 2 contract). + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + 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) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sD: cute.Tensor, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + self._store_tile(thr_t2r, rAcc, rD, gD_epi, c_dtype) + if cutlass.const_expr(self.ar_push): + self._push_tile( + thr_t2r, + rD, + sD, + epi_tid, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + @cute.experimental.jit + def _store_tile(self, thr_t2r, rAcc, rD, gD_epi, c_dtype: cutlass.Constexpr): + """The owner rank's reduced FP32 tile: one cast to ``c_dtype`` and the store to global memory.""" + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + @cute.experimental.jit + def _push_tile( + self, + thr_t2r, + rD, + sD, + epi_tid: cutlass.Int32, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output + would store them (ops/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" + cute_ext.partition_and_copy(thr_t2r, rD, sD) + cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + pairs = self.cta_m // 2 + col0 = bidx * self.cta_m + half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) + for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): + word = epi_tid + w0 + t = word // pairs + p = word % pairs + if t < ar_tokens: + n = col0 + 2 * p + par = _ld_volatile_s32( + ar_flags + cutlass.Int64(4) * cutlass.Int64(t * 8 + n // (ar_n // 8)) + ) & Int32(1) + off = ( + cutlass.Int64(par) * half + + cutlass.Int64(t * ar_world + ar_rank) * cutlass.Int64(ar_n // 2) + + cutlass.Int64(n // 2) + ) + _st_global_b32( + ar_ws + off * cutlass.Int64(4), + _bf16x2_word( + sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + ), + ) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +import torch as _torch # noqa: E402 (as in FlashInfer) + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ) + + +@cute.experimental.jit +def _bmm_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ar_tokens: Int32, + ar_n: Int32, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D 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, b, out, and bias must share BF16 or FP16 dtype") + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a, + b, + bias, + out, + pdl: bool, + tactic: SplitKTactic, +): + """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +@functools.cache +def _get_compiled_push( + dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] +): + kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, ar_push=True) + compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + z64, z32 = cutlass.Int64(0), Int32(0) + return cute_ext.compile( + _bmm_push, + kernel, + *compile_tensors[:3], + stream, + z64, + z64, + z32, + Int32(1), + Int32(8), + z32, + Int32(6144), + ) + + +def compile_push(a, b, out, pdl: bool, tactic) -> None: + """Compile (only) the push variant of ``run_splitk_dense_push`` for these operands' shapes and layouts.""" + cute_tensors = _to_cute_swap(a, b, out, None) + _get_compiled_push(a.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) + + +def run_splitk_dense_push(a, b, out, pdl: bool, tactic: SplitKTactic, ar: tuple): + """``run_splitk_dense`` without bias that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange; + ``ar`` = (multicast words address, flags address, rank, world, bucket tokens) of the 8-token bucket.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_push(a.dtype, tactic, pdl, cute_tensors[4]) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + ws, flags, rank, world, maxt = ar + compiled( + *cute_tensors[:3], + stream, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + Int32(out.shape[0]), + Int32(out.shape[1]), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py new file mode 100644 index 000000000000..76b73fc9ce91 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py @@ -0,0 +1,392 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-latency TP all-reduce + residual + RMSNorm for decode (M <= 8 tokens) over a multicast mapping. + + partial = bf16(x + x2) (optional second input, e.g. shared + routed experts) + updated = bf16(residual + sum_r partial_r) (ranks summed in FP32, rank order) + normed = bf16(updated * rsqrt(mean(updated^2) + eps) * w) + +One cluster of CLUSTER CTAs per token; CTA c owns columns [c * H / CLUSTER, ...), one 16-byte chunk (8 BF16) +per thread. Each CTA pushes its slice of this rank's partial through the multicast mapping (one store lands in +every rank's buffer), releases dependents, polls its own buffer until all ranks' slices arrived, sums them, adds +the residual and reduces the RMS statistic across the cluster through DSMEM. + +Exchange protocol as in the K3 sandwich: int32 words, 0x80000000 = empty (pushed -0.0 halves become +0.0), each +CTA keeps a call counter whose parity selects the buffer half, readers re-arm the words they read. All ranks of +the TP group must issue the same calls in the same stream order. + +The workspace holds two exchange sets. With prepushed 3 two producers pushed this rank's x (set A) and x2 (set B), +e.g. the routed experts' finalize and the shared expert's down projection, and the kernel forms bf16(x + x2) per +rank from the two slots, exactly as the unpushed path does. + +Producers that push (the o_proj / shared-expert GEMV ``ar`` tuple, the routed-expert engine's ``arp`` descriptor) are +pointed at a bucket's exchange by its ``mc`` address, ``flags_for(8)`` counters, ``rank``, ``world_size`` and +``max_tokens``; ``get_workspace`` and ``_bucket`` select it. + +The kernel is ``torch.ops.trtllm.minimax_m3_mc_allreduce_norm``, reached through the ``ops/comm`` entries +``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``; this module owns the workspaces, the bucket and +epilogue-option choice and the MXFP8 publication. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op +from ..ops.comm.m3_mc_finalize_allreduce_norm import ( + m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, +) + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +BUCKETS = (8, 16, 48) # token capacity of each workspace (grid = bucket * cluster CTAs) +# The MoE finalize (cluster kernel) stops at FINALIZE_MAX_TOKENS (= MAX_TOKENS here). +FINALIZE_MAX_TOKENS = 48 # capacity: fixed: the finalize kernel's built rows +EMPTY_WORD = -2147483648 + +# Attention / MoE boundary epilogue options (each bit-identical): 1 residual + norm weight loaded at entry (before +# the poll); 2 cluster RMS exchange by st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated +# right after the poll (no trailing __syncthreads); 8 max-smem carveout so the next kernel's CTAs co-reside early. +# The 8-token bucket, by token count 1..8 (GB300 TP4 tuning). +AR_OPTS_BY_M = (7, 7, 15, 15, 15, 15, 7, 7) +# The 16- and 48-token buckets. +AR_OPTS_BY_BUCKET = {16: 15, 48: 15} +# The MoE-finalize all-reduce, per bucket. +AR_FIN_OPTS_BY_BUCKET = {8: 7, 16: 15, 48: 15} + + +# Option 2 is also the safe ordering of the cluster RMS exchange: its stores into peer CTAs' shared memory follow a +# cluster barrier arrive / wait pair. Without it they run before any cluster barrier, when a peer need not have +# started, so every call keeps it. +_OPT_CLUSTER_BARRIER = 2 + + +def _ar_opts(m: int, bucket: int, finalize: bool = False) -> int: + """Epilogue options for an m-row call on the given bucket; every option is bit-identical. A bucket missing from + the tables gets 0, and every result keeps _OPT_CLUSTER_BARRIER.""" + if finalize: + opts = AR_FIN_OPTS_BY_BUCKET.get(bucket, 0) + elif bucket == 8: + opts = AR_OPTS_BY_M[m - 1] + else: + opts = AR_OPTS_BY_BUCKET.get(bucket, 0) + return opts | _OPT_CLUSTER_BARRIER + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags + + +class McBuckets: + """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" + + def __init__(self, workspaces): + self.workspaces = workspaces + self.hidden = workspaces[0].hidden + + def for_tokens(self, m: int): + for ws in self.workspaces: + if m <= ws.max_tokens: + return ws + return None + + +def _bucket(workspace, m): + return workspace.for_tokens(m) if isinstance(workspace, McBuckets) else workspace + + +def decline_reason(x, residual, workspace, cluster=8, x2=None) -> str: + """Why ``m3_mc_allreduce_norm`` returned None for these operands (for the dispatch's miss message).""" + if workspace is None: + return "no multicast workspace: it is created collectively on the first eager decode step" + bucket = _bucket(workspace, x.shape[0]) + if bucket is None: + return f"no workspace bucket holds {x.shape[0]} tokens (buckets {BUCKETS})" + if bucket.flags.get(cluster) is None: + return f"no {cluster}-CTA call counters for the {bucket.max_tokens}-token bucket" + return ( + f"operands outside the kernel: x {tuple(x.shape)} {x.dtype} contiguous={x.is_contiguous()}, " + f"residual contiguous={residual.is_contiguous()}, hidden {bucket.hidden}" + + ("" if x2 is None else f", x2 {tuple(x2.shape)} contiguous={x2.is_contiguous()}") + ) + + +def m3_mc_allreduce_norm( + x, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + x2=None, + quant=False, + prepushed=0, + drafter_norm=False, + capture=None, +): + """Return (normed, updated_residual) or None when unsupported. ``drafter_norm``: the Eagle3 drafter boundary + (opts bit 16: normalize the unrounded FP32 sum, as the drafter's stock FlashInfer fused_add_rmsnorm does); 8-token + bucket only, x only, no MXFP8 copy. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture`` = (out [>= M, H] view, fc_norm weight + [H], eps): the boundary also writes the Eagle3 capture of the updated residual, bit for bit + m3_spec_tail.capture_norm (CAPTURE_NORM_IN_BOUNDARY; not with ``drafter_norm``). + + A None makes the caller take another collective, so it must be the same on every TP rank. It is: it depends only + on the shapes, dtypes and layouts of the operands (the same code builds them on every rank) and on whether the + workspace's flags exist, which ``get_workspace`` creates collectively, before any capture, for cluster 8.""" + workspace = _bucket(workspace, x.shape[0]) + if workspace is None: + return None + if ( + x.dim() != 2 + or x.shape[0] > MAX_TOKENS + or x.dtype != torch.bfloat16 + or x.shape[1] != workspace.hidden + or not x.is_contiguous() + or not residual.is_contiguous() + or (x2 is not None and (not x2.is_contiguous() or x2.shape != x.shape)) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + if drafter_norm and ( + workspace.max_tokens != 8 or x2 is not None or quant or prepushed not in (0, 2) + ): + return None + outs = _allreduce_norm_op( + x, + x2, + residual, + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and x.shape[1] % 128 == 0), + int(prepushed), + _ar_opts(x.shape[0], workspace.max_tokens) | (16 if drafter_norm else 0), + *_capture_args(capture, drafter_norm), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def m3_mc_finalize_allreduce_norm( + fc2, + perm, + scales, + shared, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + quant=False, + capture=None, +): + """MoE finalize (sum_k scale * fc2[perm] + shared) + all-reduce + residual + RMSNorm in one kernel. + Returns (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture``: as ``m3_mc_allreduce_norm``.""" + m = residual.shape[0] + workspace = _bucket(workspace, m) + if workspace is None: + return None + if ( + residual.dim() != 2 + or m > MAX_TOKENS + or residual.shape[1] != workspace.hidden + or fc2.dtype != torch.bfloat16 + or not fc2.is_contiguous() + or fc2.dim() != 2 + or perm.dtype != torch.int32 + or perm.dim() != 2 + or perm.shape[0] != m + or not perm.is_contiguous() + or scales.dtype not in (torch.float32, torch.bfloat16) + or not scales.is_contiguous() + or (shared is not None and (shared.shape != residual.shape or not shared.is_contiguous())) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _finalize_allreduce_norm_op( + fc2, + perm, + scales, + shared, + residual.contiguous(), + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and residual.shape[1] % 128 == 0), + _ar_opts(m, workspace.max_tokens, finalize=True), + *_capture_args(capture, False), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def _capture_args(capture, drafter_norm: bool) -> tuple: + """(cap_out, cap_w, cap_eps) for the op, from ``capture`` = (out, fc_norm weight, eps) or None.""" + if capture is None: + return (None, None, 0.0) + assert not drafter_norm, "the capture epilogue is not built for the drafter boundary" + out, weight, eps = capture + return (out, weight, float(eps)) + + +_WORKSPACE: dict = {} +_AGREE_COMMS: dict = {} + + +def _tp_comm(mapping): + """A communicator over mapping's TP group, created once (collective over that group only).""" + key = tuple(mapping.tp_group) + comm = _AGREE_COMMS.get(key) + if comm is None: + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + comm = mapping.tp_group_pg + else: + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + _AGREE_COMMS[key] = comm + return comm + + +def tp_agree(mapping, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of mapping's TP group (one blocking host collective). + + For a fast path whose availability is decided from rank-local state (a kernel build that can fail on one rank): + every rank must take the same path, so that they issue the same collectives with the same exchange protocol and + the same kernels. Call it once per decision, at build time, at the same point on every rank, outside CUDA-graph + capture; never per step. With one rank, or attention DP (whose ranks run different batches), it returns ``ok``. + """ + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return ok + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("tp_agree must run outside CUDA-graph capture") + from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded + + return _mnnvl_workspace_all_succeeded(_tp_comm(mapping), bool(ok)) + + +def tp_agree_built(mapping, built, what: str): + """``built`` when it is not None on every TP rank, else None on every rank (see ``tp_agree``).""" + if tp_agree(mapping, built is not None) or built is None: + return built + from tensorrt_llm.logger import logger + + logger.warning(f"{what}: built here but not on another TP rank; every rank falls back") + return None + + +def get_workspace(mapping, hidden: int): + """Process-wide workspace (created collectively on first eager use); None while capturing before creation. + + The first eager call for a hidden size (the warmup forward, before any CUDA-graph capture) creates one + McAllReduceWorkspace per bucket in BUCKETS on every TP rank (MPI collective over the TP group) plus each + bucket's 8-CTA-cluster call counters; later calls, captured or not, return the same McBuckets. The workspaces + live until process exit. + """ + ws = _WORKSPACE.get(hidden) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + return None + wss = [McAllReduceWorkspace.create(mapping, hidden, b) for b in BUCKETS] + for w in wss: + w.flags_for(8) + ws = McBuckets(wss) + _WORKSPACE[hidden] = ws + return ws diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/__init__.py new file mode 100644 index 000000000000..2ed197d5c475 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/__init__.py @@ -0,0 +1,16 @@ +# 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. +"""MiniMax-M3 MoE layer of the sm_107 TP2 target for low-batch decode as a CuTe DSL kernel: ``m3_moe`` for 1..8 tokens +(op.py) at the TP2 routed and shared intermediate of 1536 per rank.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..084f54951209 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3936 @@ +# 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. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP2 rank (routed and shared intermediate 1536): +``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 # capacity: fixed: the checkpoint's hidden size +# routed intermediate per rank: 1536 at TP2 (this target: FC2 K = 6 full 256-K tiles), 768 at TP4, 384 at TP8 (FC2 K +# = 1.5 256-K tiles: the last tile's A box is half out of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 +# MMAs / scale atoms are issued; no padding, no extra bytes) +I_TP = _cfg("i_tp", 768) +TWO_I = 2 * I_TP +assert I_TP % 128 == 0 and 128 <= I_TP <= 1536, I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 # capacity: fixed: the CTA size +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# TAIL_TRIGGER (with LATE_TRIGGER, no all-reduce push): the CTAs that own more FC2 blocks than the rest (the first +# H / 32 mod NCTA, which end last) release the dependents right after the grid dependency, so the dependents launch +# when the other CTAs tear down, onto the SMs they leave, instead of after the last of the heavy CTAs. +TAIL_TRIGGER = bool(_cfg("tail_trigger", 1)) +NHEAVY = (H // 32) % NCTA if NCTA < H // 32 else 0 +# SH_TAKE: the last NHEAVY shared-expert CTAs, after their shared work, each run the routed body as an FC2-only CTA for +# one FC2 output block (the second block of routed CTA c < NHEAVY, which then owns one). The plan carries their rows +# (PLAN_COLS columns). Each block keeps its tiles, scales and finalize, so the output is unchanged. +SH_TAKE = bool(_cfg("sh_take", 0)) and SHARED and 0 < NHEAVY <= 24 +SH_HELPER0 = 24 - NHEAVY # first helper shared index (NSH = 24) +PLAN_COLS = NCTA + 24 if SH_TAKE else NCTA +# EARLY_AR (scheduling only): when the finalize pushes into the multicast all-reduce, EVERY CTA releases the +# dependents right after the grid dependency, so the all-reduce (which reads only the pushed tagged words, without a +# grid wait) and, through its own early release, the next layer's QKV GEMV become resident on each SM this kernel +# frees, instead of after the last CTA's teardown. The arithmetic is unchanged. +EARLY_AR = bool(_cfg("early_ar", 1)) +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +# the shared intermediate per rank: 1536 (TP2), 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate +# positions: its gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP2: 1 x (64 + +# 64), full K in one accumulator, the order of the stock tcgen05 MXFP8 GEMM; TP4: 2 x (32 + 32), the order of +# m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. +SH_I = _cfg("sh_i", 768) +SH_ROWS = SH_I // 24 +# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) +SH_SPLIT = 128 // (2 * SH_ROWS) +SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N +SH_KB = H // 32 # capacity: fixed: a kernel tile constant (gate_up K blocks) +SH_DKB = SH_I // 32 # down K blocks +SH_DKT = SH_I // 128 # down k-tiles +SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the +# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction +# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) +) # (in-chain: a win at M <= 2, not 4..8) +if SHARED: + assert SH_I in (384, 768, 1536), SH_I + assert SH_PRETILE or SH_SPLIT == 2, ( + "the un-pretiled shared stream is the 32-row (TP4) layout only" + ) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # capacity: fixed: pipeline stages (6 stages per unit) +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) +# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) +KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 + + +def _nkb2(t2: int) -> int: + """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" + return NUM_KB if t2 < KT2 - 1 else KB2_LAST + + +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy +# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few +# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. +HB_REPL = _cfg("hb_repl", 1) +assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) +HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +# Hang guard: every cross-CTA spin (the kernel needs all its CTAs resident at once) traps after this long, so a lost +# co-residency or a missing producer surfaces as a device exception instead of a silent hang. The timer is read only +# inside a spin iteration; the arithmetic is unchanged. +HANG_NS = _cfg("hang_ms", 2000) * 1000000 +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # capacity: fixed: a warp role index +S_WARP = 6 # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +# (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path). +E0_WARP = MMA_WARP # capacity: fixed: a warp role index +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and NRW0 <= 3 and bool(_cfg("e0_off", 1)) +# B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing +# warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the +# two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. +if B1_AFTER_TOPK and XSTAGE and EARLY0 and not E0_OFFLOAD: + raise ValueError( + "m3_moe: b1_after_topk=1 with xstage, early0 and no early-stage offload deadlocks (warp 4 waits on x_bar[1] " + "before TOPK2_BAR, the B1 build waits on TOPK2_BAR before x_bar[1]); set b1_after_topk=0 or early0=0" + ) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 # capacity: fixed: the SM's tensor-memory columns +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# TP2 at 8 tokens: the 98 KB of hidden states fit no 5-stage ring; the router GEMV then reads them from global (the +# kernel's non-XSTAGE path, same math) +if XSTAGE and M_MAX * H * 2 > STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA): + XSTAGE = False +assert not XSTAGE or M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 +_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 +_SH_FIX = ( + SH_GT * SH_NB * 128 + + SH_DKT * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + _SH_XCH_BYTES + + 2 * SH_ROWS * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * SH_NB * 128 +_SB_XC = _SB_BH + SH_DKT * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + _SH_XCH_BYTES +_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = SH_NB +SH_SFA = SH_ACC_DN + 16 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +# pre-tiled stages beyond the ring that warp 0 prefetches into L2, per lane (TP2: 72 stages) +SH_PF = max(2, (SH_STAGES - SH_NST + 31) // 32) if SHARED else 2 +# TP2 (SH_GT 48 gate_up k-tiles, SH_DKT 12) needs more than 256 columns: allocate the next power of two +SH_TMEM = 256 if SH_SFB_DN + 4 * SH_DKT <= 256 else 512 +# SH_TAKE helpers hand the shared allocation to the routed body, whose teardown frees TMEM_COLS +assert not SH_TAKE or SH_TMEM == TMEM_COLS, ( + "SH_TAKE needs the shared TMEM allocation to equal the routed one" +) +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 32 * SH_PF, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 4 * SH_DKT <= SH_TMEM, ( + SH_NST, + _SB_END, + BIG, + ) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _trap_after(elapsed, *, loc=None, ip=None): + """trap when elapsed (ns, int64) exceeds HANG_NS.""" + _llvm.inline_asm( + None, + [cutlass.Int64(elapsed).ir_value(loc=loc, ip=ip)], + f"{{ .reg .pred p; setp.gt.s64 p, $0, {HANG_NS}; @p trap; }}", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _hang_tick(t0): + """One spin iteration of a cross-CTA wait: the wait's start (latched on the first call; t0 = 0 before it), after + trapping once the wait has lasted longer than HANG_NS.""" + now = _gtime() + t1 = cutlass.select_(t0 == cutlass.Int64(0), now, t0) + _trap_after(now - t1) + return t1 + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + # (a partial last k-tile has no atoms past I_TP) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(SH_PF): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(SH_PF) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SH_NB, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if cutlass.const_expr(SH_SPLIT == 1): + # TP2: one K slice; quadrant qd = rows 32 qd .. (gate positions 0..63, then up 0..63), no partials: + # the owner sum is (0 + acc), then bf16, at _SB_GU word (u SH_ROWS + p) 8 + t = (32 qd + l) 8 + t + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + elif cutlass.const_expr(SH_SPLIT == 2): + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + else: + # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are + # accumulator columns 8 qd + t + av = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_GU + qd * cutlass.Int32(8), + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(1): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd == cutlass.Int32(0): + # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + for s_ in cutlass.range_constexpr(1, SH_SPLIT): + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + lane * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(SH_ROWS * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + t0h = cutlass.Int64(0) + while pend != cutlass.Int32(0): + t0h = _hang_tick(t0h) + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(SH_DKT): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + # SH_TAKE helpers keep the TMEM and the epoch for the routed body that follows (its teardown releases both) + not_helper = cutlass.Boolean(True) + if cutlass.const_expr(SH_TAKE): + not_helper = sid < cutlass.Int32(SH_HELPER0) + if (warp == 1) & not_helper: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if (tidx == 0) & not_helper: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if cutlass.const_expr(SH_TAKE): + is_routed = (bx < cutlass.Int32(NCTA)) | (bx >= cutlass.Int32(NCTA + SH_HELPER0)) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if (warp == MMA_WARP) & ( + bx < cutlass.Int32(NCTA) + ): # SH_TAKE helpers reuse the shared allocation + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(PLAN_COLS) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(LATE_TRIGGER and TAIL_TRIGGER and NHEAVY > 0): + late = bx < cutlass.Int32(NHEAVY) + if cutlass.const_expr(SH_TAKE): + late = bx >= cutlass.Int32(NCTA + SH_HELPER0) + if ar_off & late: + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(EARLY_AR): + if ar_on: + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + t0h = cutlass.Int64(0) + while done == cutlass.Int32(0): + t0h = _hang_tick(t0h) + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + t0h = cutlass.Int64(0) + while (cv - cnt_target) < cutlass.Int32(0): + t0h = _hang_tick(t0h) + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + for rr in cutlass.range_constexpr(HB_REPL): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4) + + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + hb_off = cutlass.Int32(0) + if cutlass.const_expr(HB_REPL > 1): + # this CTA's copy of the records + hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + t0h = cutlass.Int64(0) + while done == cutlass.Int32(0): + t0h = _hang_tick(t0h) + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview( + (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) + ), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + t0h = cutlass.Int64(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + t0h = _hang_tick(t0h) + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if cutlass.const_expr(HB_REPL == 1): + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + else: + if lane < HB_REPL: # copy `lane` + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4) + + lane * cutlass.Int32(HB_STRIDE), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + t0h = cutlass.Int64(0) + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + t0h = _hang_tick(t0h) + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py new file mode 100644 index 000000000000..067538849ac8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py @@ -0,0 +1,808 @@ +# 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. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP2 rank (routed and shared intermediate 1536 per rank), +over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +INTER = 1536 # capacity: fixed: the routed experts' width per tp2 rank (the routed intermediate per TP rank (TP2)) +NUM_EXPERTS = 128 +SHARED_INTER = 1536 # capacity: fixed: the shared expert's width per tp2 rank +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size (m3_moe's shared-expert CTAs (beside num_ctas routed CTAs)) +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +U1 = 2 * INTER // 32 # FC1 units per expert +# The router GEMV takes at most 2 experts per CTA (router weights held in registers). +MIN_CTAS = 64 # capacity: fixed: an SM split, not a size + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # capacity: fixed: pipeline stages (FC1 stages (items) per unit) +_ITEM_BYTES = _UNIT_BYTES // _S1 # capacity: fixed: a ring item's bytes +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters (4, 7 and 8 tokens; at TP2 4 +# tokens with 6 stages measured 0.67 us/layer faster in the ablation chain than the tagged records). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 1, 0, 0, 1, 1) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit; at TP2 9 / 8 / 5 for 1-2 / 3-4 / 8 tokens). 6 at 4 +# tokens measured 1.9 us/layer faster in the ablation chain than the default 8 (same MMAs, bit-identical outputs). +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 6, 0, 0, 0, 0) +# ``early0``: warp 4 streams group 0 right after its own top-k. Off at 8 tokens (0.6 us/layer faster in the chain: +# there it is on every routing warp's critical path and the ring holds only 5 stages). +EARLY0_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 0) +# ``xw``: FC2 B builder warps per token count (0 = the kernel default). 6 at 3 / 4 / 5 / 8 tokens: -1.04 / -0.63 / +# -0.26 / -0.61 us/call (795995, 795401, bit-identical); at 1 / 2 / 6 / 7 tokens +0.87 / -0.07 / +2.79 / +0.05, so not +# there. +XW_BY_M: Tuple[int, ...] = (0, 0, 6, 6, 6, 0, 0, 6) +# ``sh_take``: the routed CTAs owning a second FC2 output block (H / 32 mod num_ctas of them) hand it to as many +# shared-expert CTAs, which run it FC2-only after their shared work; every block keeps its tiles, scales and finalize, +# so the output is bit-identical. On at 6 / 7 / 8 tokens (-1.34 / -4.47 / -1.24 us/call); off at 1 .. 5, where the +# shared work ends too late for the helper (+0.5 .. +8.1). +# mega1.8: OFF by default. Engine-level gates are bitwise vs the kill switch, but the in-server same-mount check +# (797173) read TEXT-SAME 0/7 against a 4/7 base, cause unknown. The table that was gated is (0, 0, 0, 0, 0, 1, 1, 1). +SH_TAKE_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 0) +# The plan's ``heavy_fc2_only``: 192 FC2 output blocks over 188 routed CTAs leave 4 CTAs with two blocks, whose +# trailing FC2 after the last hand-off otherwise ends the kernel 3-7 us after the rest; they stream no FC1 instead +# (1.7 / 3.4 us/layer faster at 4 / 8 tokens in the ablation chain; bit-identical outputs). +HEAVY_FC2_ONLY_BY_M: Tuple[int, ...] = (0, 0, 0, 1, 0, 0, 0, 1) +# ``hb_repl``: copies of the tagged FC1 -> FC2 records (applied when > 1); CTA bx polls copy bx % hb_repl, which +# spreads the polls of the few hot record lines over L2. 7 and 8 tokens hand off through the B images instead. +HB_REPL_BY_M: Tuple[int, ...] = (32, 32, 8, 8, 1, 1, 1, 1) +# ``lg_repl``: replicas of the router logit exchange at up to LG_REPL_MAX_TOKENS tokens (fewer stores before the +# exchange; bit-identical outputs). Above that the kernel's default applies. +LG_REPL = 32 +LG_REPL_MAX_TOKENS = 4 # capacity: fixed: measured crossover +# The TP2 FLAT stack (MoE lane, 2026-10-17): at every token count, FC2 B through per-(slot, token) records (no B2 +# images), 16 router-logit replicas, the AR push parities loaded in the finalize (no preload) and the PDL release right +# after the grid wait. Bitwise-identical output (gate a 791026, mega1 image, m 1..8 x 4 routings); in-model TP2 pair +# -114 us/step. +# The values are part of the engine's module key. +FLAT_STACK = {"b2img": 0, "lg_repl": 16, "ar_pre": 0, "late_trigger": 0} +# FLAT only up to 4 tokens (MoE lane 10-18, 795150): at 8 tokens the per-m tables (B images, preloaded AR push) are +# 2.96 us/call +# faster than FLAT (47.99 vs 50.95; 4 tokens: FLAT -0.51); batch-2 GEN steps run 8 tokens. Bit-identical either way. +FLAT_MAX_TOKENS = 4 # capacity: fixed: measured FLAT crossover + + +def flat_applies(num_tokens: int) -> bool: + """Whether an engine state for ``num_tokens`` tokens runs the FLAT stack (up to FLAT_MAX_TOKENS).""" + return num_tokens <= FLAT_MAX_TOKENS + + +# Byte-balanced FC1 ranges below one unit per CTA (:func:`_fine_ranges`) for groups too small to give every CTA a +# whole unit; False gives such groups whole units on the CTAs with the least FC2 work. +FINE_PLAN = True +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 + = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row + R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then + down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per + layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + P = I_ // n_cta + split = 128 // (2 * P) + gt = H_ // 128 // split + dkt = I_ // 128 + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + octant = r // P + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(2 * dkt): + jt, kt = gd // dkt, gd % dkt + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, gt + gd, :16384] = img.reshape(-1) + out[s, gt + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def _fine_ranges(weights, total: int): + """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item + counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary + strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit + runs at least to that unit's end. Returns the item counts.""" + n = len(weights) + items = [0] * n + s = 0 + acc = 0.0 + for c in range(n): + acc += weights[c] + e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) + if s % _S1 and e < total: # started inside a unit: run to (at least) its end + e = max(e, min(total, (s // _S1 + 1) * _S1)) + items[c] = e - s + s = e + assert s == total + return items + + +def make_phase_plan( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = FINE_PLAN, + heavy_fc2_only: bool = False, +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + # (heavy_fc2_only) CTAs with more FC2 blocks than the minimum stream no FC1: they run each group's FC2 as soon as + # its hand-off lands, so their trailing FC2 after the last hand-off matches the others' two lagged groups + min_nb = min(nbs) + no_fc1 = [heavy_fc2_only and nbs[c] > min_nb for c in range(num_ctas)] + BIG_W = 1 << 40 + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + for c in range(num_ctas): + if no_fc1[c]: + fc2[c] += BIG_W + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [ + 0 if no_fc1[c] else max(_S1, (limit - f) // _ITEM_BYTES) + for c, f in enumerate(fc2) + ] + + lo, hi = 0, max(f for c, f in enumerate(fc2) if not no_fc1[c]) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 + lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] + sw = sum(wts) + items = _fine_ranges([w * total / sw for w in wts], total) + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +def _shared_takes_fc2_plan(plan: torch.Tensor, num_ctas: int, m_max: int) -> torch.Tensor: + """The plan with SHARED_CTAS more columns for SH_TAKE: routed CTA c < H / 32 mod num_ctas keeps only its first FC2 + block (its FC1 rows unchanged); helper shared CTA SHARED_CTAS - nheavy + c owns the second, with no FC1.""" + nheavy = (HIDDEN // 32) % num_ctas + ext = torch.zeros(plan.shape[0], num_ctas + SHARED_CTAS, plan.shape[2], dtype=plan.dtype) + ext[:, :num_ctas] = plan + h0 = num_ctas + SHARED_CTAS - nheavy + for c in range(nheavy): + b0 = int(plan[TOP_K, c, m_max]) & 0xFFFF + assert (int(plan[TOP_K, c, m_max]) >> 16) == 2 + ext[:, c, m_max] = b0 | (1 << 16) + ext[:, h0 + c, m_max] = (b0 + 1) | (1 << 16) + return ext + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + inter: int = INTER, + early_release: bool = True, + ): + """``inter``: the routed intermediate per rank; this engine is built for INTER (1536) only. + ``early_release``: on a call that pushes into the boundary all-reduce (``arp``), every CTA releases the PDL + dependents right after its grid dependency wait (the kernel's EARLY_AR); False releases them after the push.""" + if inter != INTER: + raise ValueError( + f"m3_moe (TP2): routed intermediate {inter} per rank is not supported (only {INTER})" + ) + self.inter = inter + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + if self.num_ctas > HIDDEN // 32: + # a routed CTA that owns no 32-row FC2 output block deadlocks the FC2 hand-off on its second call + raise ValueError( + f"m3_moe runs at most {HIDDEN // 32} routed CTAs (one FC2 output block each), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, + num_tokens, + M3_MOE_LAG, + 0.0 if num_tokens == 1 else 1.0, + True, + heavy_fc2_only=bool(HEAVY_FC2_ONLY_BY_M[num_tokens - 1]), + ) + sh_take = ( + shared + and SH_TAKE_BY_M[num_tokens - 1] > 0 + and self.num_ctas < HIDDEN // 32 + and 0 < (HIDDEN // 32) % self.num_ctas <= SHARED_CTAS + ) + if sh_take: + plan = _shared_takes_fc2_plan(plan, self.num_ctas, num_tokens) + from tensorrt_llm.logger import logger as _shlog + + _shlog.info( + f"MiniMax-M3 MoE engine SH_TAKE (TP2): {(HIDDEN // 32) % self.num_ctas} shared CTAs take the extra " + f"FC2 blocks, tokens={num_tokens}" + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + **({"sh_take": 1} if sh_take else {}), + "timing": 0, + "i_tp": INTER, + } + if HB_REPL_BY_M[num_tokens - 1] > 1: + self.config["hb_repl"] = HB_REPL_BY_M[num_tokens - 1] + if num_tokens <= LG_REPL_MAX_TOKENS: + self.config["lg_repl"] = LG_REPL + if shared: + self.config["sh_i"] = SHARED_INTER + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if not early_release: + self.config["early_ar"] = 0 + if XW_BY_M[num_tokens - 1] > 0: + self.config["xw"] = XW_BY_M[num_tokens - 1] + from tensorrt_llm.logger import logger as _xlog + + _xlog.info( + f"MiniMax-M3 MoE engine XW (TP2): xw={XW_BY_M[num_tokens - 1]} tokens={num_tokens}" + ) + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + if not EARLY0_BY_M[num_tokens - 1]: + self.config["early0"] = 0 + if flat_applies(num_tokens): + self.config.update(FLAT_STACK) + # engagement line (every rank; once per engine state = once per token count) + from tensorrt_llm.logger import logger as _flog + + _flog.info( + "MiniMax-M3 MoE engine FLAT: " + + " ".join(f"{k}={self.config[k]}" for k in FLAT_STACK) + + f" tokens={num_tokens}" + ) + else: + from tensorrt_llm.logger import logger as _flog + + _flog.info( + f"MiniMax-M3 MoE engine FLAT off above {FLAT_MAX_TOKENS} tokens (per-m tables): tokens={num_tokens}" + ) + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + # Every CTA waits on others (cross-CTA spins on partials, hand-off records, the shared expert's h and y): the + # grid must be resident at once, one CTA per SM. The module's shared-memory / TMEM budget assertion makes one + # CTA fit an SM; the grid must not exceed the SMs. (The kernel also traps after HANG_NS in any such spin.) + if self.num_ctas + mod.NSH > sms: + raise ValueError( + f"m3_moe: a grid of {self.num_ctas} + {mod.NSH} CTAs exceeds the {sms} SMs; its CTAs wait on each " + "other and must all be resident at once" + ) + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full( + (mod.G_CAP * mod.U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw + ) + # the FC1 -> FC2 records; with HB_REPL > 1 copies, HB_STRIDE int32 words apart + hb_words = mod.G_CAP * num_tokens * mod.U1 * 4 + if mod.HB_REPL > 1: + hb_words = mod.HB_REPL * mod.HB_STRIDE + self.hbuf32 = torch.zeros(hb_words, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [3072, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 1536], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + # keyed on the config the kernel module was built with (not self.config), so a module built with any other + # trace-time config can never reuse another build's compiled kernel + key = (_ENGINE_PATH,) + tuple(sorted(self.mod.M3_CONFIG.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe_finalize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe_finalize.py new file mode 100644 index 000000000000..714eafc4661b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe_finalize.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exact MoE finalize for the routed experts' unfinalized output, optionally with the shared expert added. + +The TRTLLM-Gen MoE runner with ``do_finalize=False`` returns the FC2 output in permuted order, the per-token top-k +routing weights and ``expanded_idx_to_permuted_idx``. Its own finalize (``finalizeKernelVecLoad``) computes, per token +and hidden element, ``acc = 0; acc += weight_k * fc2[permuted_idx_k]`` for k in order (skipping ``-1``) in float32 +and rounds to BF16 once, but issues the k gathers one after another, so at prefill sizes it runs ~2.6x its byte +floor (153 us per layer at 16K rows on SM107). This kernel issues the gathers together and keeps the per-element +order, so the result is bitwise the runner's: a BF16 x BF16 product is exact in float32, so fused and unfused +multiply-add agree as well. ``shared`` adds the shared expert as ``shared_output.add_(routed_output)`` does +(``bf16(float(shared) + float(routed))``), which removes that separate pass over the hidden states. +""" + +from typing import Optional + +import torch +import triton +import triton.language as tl + +#: (BLOCK_H, num_warps): the tile only splits the hidden dimension, so it never changes the bits. Tuned on SM107 +#: (2K / 4K / 16K rows: 11.2 / 20.5 / 74.9 us, 1.52 / 1.39 / 1.27x the byte floor; the runner's finalize +#: 23 / 40 / 153 us). +CONFIGS = ((2048, 4),) + + +@triton.jit +def _moe_finalize_kernel(fc2, idx, weights, shared, out, ld_fc2, H: tl.constexpr, TOPK: tl.constexpr, + BLOCK_H: tl.constexpr, SHARED: tl.constexpr): # fmt: skip + t = tl.program_id(0).to(tl.int64) + cols = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H) + cmask = cols < H + acc = tl.zeros((BLOCK_H,), dtype=tl.float32) + for k in tl.static_range(TOPK): + p = tl.load(idx + t * TOPK + k) + w = tl.load(weights + t * TOPK + k).to(tl.float32) + valid = p >= 0 + x = tl.load( + fc2 + tl.maximum(p, 0).to(tl.int64) * ld_fc2 + cols, mask=cmask & valid, other=0.0 + ) + acc = tl.where(valid, tl.fma(w, x.to(tl.float32), acc), acc) + r = acc.to(tl.bfloat16) + if SHARED: + s = tl.load(shared + t * H + cols, mask=cmask, other=0.0) + r = (s.to(tl.float32) + r.to(tl.float32)).to(tl.bfloat16) + tl.store(out + t * H + cols, r, mask=cmask) + + +def m3_moe_finalize( + fc2: torch.Tensor, + weights: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + shared: Optional[torch.Tensor] = None, +) -> Optional[torch.Tensor]: + """BF16 [tokens, H] = sum_k weights[:, k] * fc2[expanded_idx_to_permuted_idx[:, k]] (+ ``shared``), bitwise to the + TRTLLM-Gen finalize (and to ``shared.add_(that)``); None when the operands are not the runner's layout.""" + if ( + fc2.dim() != 2 + or fc2.dtype != torch.bfloat16 + or fc2.stride(1) != 1 + or weights.dim() != 2 + or expanded_idx_to_permuted_idx.shape != weights.shape + or expanded_idx_to_permuted_idx.dtype != torch.int32 + or not weights.is_contiguous() + or not expanded_idx_to_permuted_idx.is_contiguous() + or ( + shared is not None + and ( + shared.dtype != torch.bfloat16 + or not shared.is_contiguous() + or shared.shape != (weights.shape[0], fc2.shape[1]) + ) + ) + ): + return None + tokens, topk = weights.shape + hidden = fc2.shape[1] + out = torch.empty((tokens, hidden), dtype=torch.bfloat16, device=fc2.device) + if tokens == 0: + return out + block_h, warps = CONFIGS[0] + _moe_finalize_kernel[(tokens, triton.cdiv(hidden, block_h))]( + fc2, + expanded_idx_to_permuted_idx, + weights, + shared if shared is not None else out, + out, + fc2.stride(0), + hidden, + topk, + block_h, + shared is not None, + num_warps=warps, + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py new file mode 100644 index 000000000000..48a7827f59b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hand-off of an MXFP8-quantized activation from the kernel that produced it to the next MXFP8 GEMM. + +A producer that already holds a BF16 row in registers (the multicast all-reduce + RMSNorm, the split-K attention +merge) also writes its MXFP8 form (E4M3 rows + UE8M0 scales in the 128x4-swizzled layout of +``trtllm::mxfp8_quantize(x, True)``) and publishes it here keyed by the BF16 tensor. The MXFP8 linear whose input +is that tensor takes it and skips its own quantize kernel. Every take clears the slot, so a hand-off is only ever +consumed by the first MXFP8 linear that runs after the producer. + +Module state shared across files (one instance each per process): + +* ``SLOT`` (PrequantSlot): published by ``m3_mc_allreduce`` (``quant=True``), the sparse-decode merge and the fused + decode attention (MXFP8 output); cleared by the sparse decode before it runs; taken by the MXFP8 linear dispatch + and by the one-launch QKV + attention + o_proj path (which also peeks ``pending``). +* ``AR_PUSH`` (ArPushSlot): armed by the fast forward around the attention call; the o_proj GEMV takes the target + and marks ``pushed``; the fast forward checks ``take_pushed`` before the attention boundary all-reduce. +* ``MOE_PUSH`` (MoePushSlot): armed by the fast forward before the MoE; the routed-expert engine takes ``engine`` + and the shared expert's down projection takes ``shared``, each marks its pushed output; the fast forward + checks ``take_pushed`` before the MoE boundary all-reduce. +""" + + +class PrequantSlot: + def __init__(self): + self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) + + def publish(self, x, xq, xsf): + self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) + + def clear(self): + self.pending = None + + def take(self, x): + """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" + pending, self.pending = self.pending, None + if ( + pending is None + or x.dim() != 2 + or x.data_ptr() != pending[0] + or x.shape[0] != pending[1] + or x.shape[1] != pending[2] + ): + return None + return pending[3], pending[4] + + +SLOT = PrequantSlot() + + +class ArPushSlot: + """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. + + The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output + it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" + + def __init__(self): + self.target = None + self.pushed = None # data_ptr of the pushed output + + def arm(self, target): + self.target = target + self.pushed = None + + def take_target(self): + target, self.target = self.target, None + return target + + def take_pushed(self, x): + pushed, self.pushed = self.pushed, None + return pushed is not None and x.data_ptr() == pushed + + +AR_PUSH = ArPushSlot() + + +class MoePushSlot: + """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, + ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv + ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" + + def __init__(self): + self.engine = None + self.shared = None + self.engine_pushed = None # data_ptr of the pushed routed output + self.shared_pushed = None # data_ptr of the pushed shared output + + def arm(self, engine, shared): + self.engine, self.shared = engine, shared + self.engine_pushed = self.shared_pushed = None + + def disarm(self): + self.engine = self.shared = None + + def take_engine(self): + engine, self.engine = self.engine, None + return engine + + def take_shared(self): + shared, self.shared = self.shared, None + return shared + + def take_pushed(self, shared, routed): + """(routed pushed, shared pushed) for these outputs; clears the marks.""" + e = routed is not None and self.engine_pushed == routed.data_ptr() + s = shared is not None and self.shared_pushed == shared.data_ptr() + self.engine_pushed = self.shared_pushed = None + return e, s + + +MOE_PUSH = MoePushSlot() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_report.py new file mode 100644 index 000000000000..7f864bd240c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_report.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast-path reports a kernel module makes from inside a call: ``taken`` (an engagement line), ``missed`` (a covered +shape that cannot run: raises under ``modeling_v2='require'``), ``uncovered`` and ``capture_missed``. + +The policy (logging once, raising under require) is the target's: ``modeling.py`` installs its functions here at import, +so a kernel reports without importing the target module. Until then every report is a no-op. +""" + +from __future__ import annotations + +from typing import Callable, Optional + + +def _noop(*args, **kwargs) -> None: + return None + + +taken: Callable[..., None] = _noop +missed: Callable[..., None] = _noop +uncovered: Callable[..., None] = _noop +capture_missed: Callable[..., None] = _noop + + +def install( + *, + on_taken: Callable[[str, Optional[int], str], None], + on_missed: Callable[[str, Optional[int], str], None], + on_uncovered: Callable[[str, str], None], + on_capture_missed: Callable[[str, Optional[int], str], None], +) -> None: + """Route this module's reports to the target's policy.""" + global taken, missed, uncovered, capture_missed + taken, missed, uncovered, capture_missed = on_taken, on_missed, on_uncovered, on_capture_missed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py new file mode 100644 index 000000000000..30855eb7f05d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M fp32 router projection: ``logits[M, E] = float(x[M, K]) @ W[E, K]^T`` in one kernel. + +Replaces the bf16->fp32 cast, the TF32 GEMM and its split-K reduce on decode +steps. One CTA per ``ROWS`` experts: the fp32 weight rows are bulk-copied into +shared memory before ``griddepcontrol.wait`` (overlapping the producer kernel +under PDL), the bf16 activations after it. Products and sums are exact fp32 FMA. + +The kernels are ``torch.ops.trtllm.minimax_m3_router_gemv`` / ``_router_gemv_fp4`` / ``_router_tc``, reached +through the ``ops/moe`` entries. +""" + +from typing import Optional + +import torch + +from ..ops.moe.m3_router_gemv import m3_router_gemv as _router_gemv_op +from ..ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 as _router_gemv_fp4_op +from ..ops.moe.m3_router_tc import m3_router_tc as _router_tc_op + +# Opt-in dynamic shared memory left after the kernel's static reduction buffer (SM100: 227 KiB). +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit + + +def m3_router_gemv( + x: torch.Tensor, weight: torch.Tensor, rows: int = 1, use_pdl: bool = False +) -> Optional[torch.Tensor]: + """fp32 ``x @ weight.T`` for bf16 ``x`` with at most 16 rows; None when unsupported (only ``rows=1`` is built).""" + if ( + rows != 1 + or not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or weight.shape[0] % rows != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or rows * x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return _router_gemv_op(x, weight, rows, use_pdl) + + +def m3_router_gemv_fp4( + x: torch.Tensor, weight: torch.Tensor, sf_scale: torch.Tensor, use_pdl: bool = False +) -> Optional[tuple]: + """Router logits plus the NVFP4 activation the routed experts consume (bit-identical to + ``fp4_quantize(x, sf_scale, 16, False, False)``): (logits [M, E] fp32, x_fp4 [M, K / 2] uint8, + x_sf [M * K / 16] uint8 linear). None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or sf_scale is None + or sf_scale.dtype != torch.float32 + or sf_scale.numel() != 1 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return tuple(_router_gemv_fp4_op(x, weight, sf_scale.reshape(1), use_pdl)) + + +def m3_router_tc( + x: torch.Tensor, + weight: torch.Tensor, + sf_scale: Optional[torch.Tensor] = None, + use_pdl: bool = False, +) -> Optional[tuple]: + """Router logits (TF32 tensor cores, like the native TF32 GEMM) for 1..48 tokens with 128 experts, plus the + NVFP4 activation when ``sf_scale`` is given. Returns (logits,) or (logits, x_fp4, x_sf); None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 48 + or weight.dtype != torch.float32 + or weight.shape[0] != 128 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or (x.stride(0) * 2) % 16 != 0 + or x.data_ptr() % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or 16 * (x.shape[1] // 8 + 4) * 4 + x.shape[0] * (x.shape[1] // 8 + 8) * 2 + > _MAX_DYNAMIC_SMEM + or (sf_scale is not None and (sf_scale.dtype != torch.float32 or sf_scale.numel() != 1)) + ): + return None + res = _router_tc_op(x, weight, None if sf_scale is None else sf_scale.reshape(1), use_pdl) + return tuple(res) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_sampled_accept.py new file mode 100644 index 000000000000..93ffda80ecea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_sampled_accept.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). + +Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +active must stay on the generic path; the caller decides. + +Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +(the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens heavier +than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the nucleus +with probability >= top_p, so most rows decide on their second exchange. + +Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based Philox +keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one 4-word +record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and reads +the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the same +order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing and +two call halves keep a rank one call ahead from overwriting a slot still being read. +""" + +from typing import Optional + +import torch + +from . import m3_report + +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu) and called as +# torch.ops.trtllm.minimax_m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) +# The TP2 vocab shard limit of trtllm::minimax_m3_sm107tp2_sampled_accept (kMaxVs): 200064 / 2 = 100032 columns. +MAX_VS = 100352 # capacity: fixed: the vocabulary shard bound at tp2 + + +class SampledAcceptWorkspace: + """Exchange slots behind one multicast mapping (two call halves), the (finished rows, call) counters, the per-row + token scratch and an extra-attempt counter. Created collectively, outside CUDA-graph capture.""" + + def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): + self.uc, self.mc_ptr, self.rank, self.world = uc, mc_ptr, rank, world + self.handle, self.comm = handle, comm + dev = uc.device + self.counters = torch.zeros(2, dtype=torch.int32, device=dev) + self.tokens = torch.zeros(MAX_ROWS, dtype=torch.int32, device=dev) + self.extra_attempts = torch.zeros(1, dtype=torch.int32, device=dev) + + @classmethod + def create(cls, mapping, device) -> "SampledAcceptWorkspace": + words = int( + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept_ws_words(max(mapping.tp_size, 1)) + ) + if mapping.tp_size == 1: + uc = torch.zeros(words, dtype=torch.int64, device=device) + return cls(uc, uc.data_ptr(), 0, 1) + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 8, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("SampledAcceptWorkspace.create failed on another rank") + return cls(uc, mc.data_ptr(), mapping.tp_rank, mapping.tp_size, handle, comm) + + +_WS: Optional[SampledAcceptWorkspace] = None + + +def workspace(mapping, device) -> Optional[SampledAcceptWorkspace]: + global _WS + if _WS is None: + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("spec_tail", None, "the sampled-accept workspace") + return None + _WS = SampledAcceptWorkspace.create(mapping, device) + return _WS + + +def sampled_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: Optional[torch.Tensor], + seeds: torch.Tensor, + offsets: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, + ws: Optional[SampledAcceptWorkspace] = None, + use_pdl: bool = True, +): + """Sampled verification (temperature + top-p) from this rank's bf16 logits shard [nreq * (draft_len + 1), V / tp]: + outputs as shard_accept (accepted [nreq, D+1], num_accepted [nreq], gather_ids [nreq]); None + when unavailable. Per-row sampling parameters and Philox (seed, offset) are the spec metadata's, one per row. + top_ks / min_ps are not applied: rows with top-k or min-p active must be routed to the generic path by the + caller.""" + del top_ks, min_ps + nreq, draft_len = draft_tokens.shape + rows = nreq * (draft_len + 1) + if logits.dim() != 2 or logits.shape[0] != rows or rows > MAX_ROWS or logits.shape[1] > MAX_VS: + return None + if logits.dtype != torch.bfloat16 or logits.stride(1) != 1: + return None + ws = ws if ws is not None else workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + temperatures = temperatures[:rows].float().contiguous() + top_ps = top_ps[:rows].float().contiguous() + seeds = seeds[:rows].long().contiguous() + offsets = offsets[:rows].long().contiguous() + _launch( + logits, vocab_offset, temperatures, top_ps, seeds, offsets, ws, draft_tokens, batch_idx, + accepted, num_accepted, gather_ids, force_total, next_new, force_vec, use_pdl, rows, + ) # fmt: skip + return accepted, num_accepted, gather_ids + + +def _launch( + logits, vocab_offset, temperatures, top_ps, seeds, offsets, ws, draft_tokens, batch_idx, + accepted, num_accepted, gather_ids, force_total, next_new, force_vec, use_pdl, rows, +) -> None: # fmt: skip + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept( + logits, + int(vocab_offset), + temperatures, + top_ps, + seeds, + offsets, + ws.uc, + ws.mc_ptr, + ws.counters, + ws.tokens, + ws.rank, + ws.world, + draft_tokens.int().contiguous(), + batch_idx.int().contiguous(), + accepted, + num_accepted, + gather_ids, + int(force_total), + next_new, + force_vec, + ws.extra_attempts, + bool(use_pdl), + PREWAIT_PARAMS, + EARLY_TRIGGER, + _ilp_walk(rows), + ) + + +# The sampled verify loads each row's sampling parameters before its PDL dependency wait (overlapping the head GEMM's +# tail; the same values and tokens), and releases its dependents right after that wait (the dependent parks at its +# own wait; the same values). +PREWAIT_PARAMS = True +EARLY_TRIGGER = True +# SA-PW: the owner walk takes its tokens 8 at a time (the same draw, tokens and Philox draws as the one-token walk, +# bit for bit; -2.2..-2.5 us per call on Rubin TP2). Off in the TP2 deployment (frozen deploy value). +ILP_WALK = False +ILP_WALK_PATH = "m3_sampled_accept_ilp_walk" + + +def _ilp_walk(rows: int) -> bool: + """ILP_WALK, logged once per row count as taken; off, one 'not taken' line (logged once).""" + if not ILP_WALK: + m3_report.uncovered(ILP_WALK_PATH, "the one-token owner walk (ILP_WALK off)") + return False + m3_report.taken( + ILP_WALK_PATH, + rows, + "the sampled verify's owner walk takes 8 tokens at a time", + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py new file mode 100644 index 000000000000..6e6e23e0c0b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). + +Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long``, +one 1024-thread CTA per row), from a CTA of 384 threads: + +* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and + mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's + valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned + key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped + block and masks the result, so a thread's loads issue back to back. +* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of + the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. +* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by + that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate + against the whole row instead. +* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. + +The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait +(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are +compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and +loads them before the executor's warmup, so a decode step never compiles one. +""" + +from typing import Dict, Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.utils import ( + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +# The 384-thread select reads its row's valid-block count before its grid wait: off (its TP2 in-model pairs did not +# pass); the selection is the same either way. +SELECT_N_VALID_PREWAIT = False + +SELECT_THREADS = 384 +SELECT_WARPS = SELECT_THREADS // 32 +# 1M tokens + draft tokens at 128 tokens per block. +MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row +MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging +TOPK = 16 # capacity: fixed: MSA's top-k +NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) +# The key of -inf (the reference's 0xff800000): never emitted. +NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 +# (keys per thread, rows up to this many valid blocks) +ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) + +# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare +# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group +# minima [16], per-warp candidate counts [SELECT_WARPS]. +_CAND = 0 +_PICK_KEY = _CAND + 2 * MAX_CAND +_PICK_BLOCK = _PICK_KEY + TOPK +_WARP_MIN = _PICK_BLOCK + TOPK +_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 +_WARP_CNT = _GROUP_MIN + 16 +SMEM_WORDS = _WARP_CNT + SELECT_WARPS +SMEM_BYTES = 4 * SMEM_WORDS + + +def _key(score): + """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" + bits = cutlass.Float32(score).bitcast(cutlass.Int32) + return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) + + +def _imin(a, b): + return cutlass.select_(b < a, b, a) + + +def _imax(a, b): + return cutlass.select_(b > a, b, a) + + +def _forced_key(scores, head, block, query, init_blocks, local_start): + raw = scores[head, block, query] + s = cutlass.select_( + block >= local_start, + cutlass.Float32(LOCAL_SCORE), + cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), + ) + return _key(s) + + +@cute.jit +def _select_row( + row, + scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides + n_valid: cute.Tensor, # int32 [>= total_q] + out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] + init_blocks, + local_blocks, + num_kv_heads, + smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory + bar_id: cutlass.Constexpr, + loaded_trigger: cutlass.Constexpr, + raw_prewait, + n_valid_prewait: cutlass.Constexpr, +): + """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier + ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the + row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by + the caller before its grid wait; otherwise it is read here.""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % cutlass.Int32(32) + warp = tid // cutlass.Int32(32) + query = row // num_kv_heads + head = row % num_kv_heads + if cutlass.const_expr(n_valid_prewait): + raw = raw_prewait + else: + raw = n_valid[query] + valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) + local_start = _imax(raw - local_blocks, cutlass.Int32(0)) + cand = cute.make_tensor( + cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) + ) + + if tid < TOPK: + smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) + smem[_PICK_BLOCK + tid] = cutlass.Int32(0) + + if valid <= TOPK: + # Every valid block is picked (thread t holds block t). + if tid < valid: + smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) + smem[_PICK_BLOCK + tid] = tid + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): + items = ITEM_CLASSES[ci][0] + lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] + hi = ITEM_CLASSES[ci][1] + if (valid > lo) & (valid <= hi): + key = cute.make_rmem_tensor(items, cutlass.Int32) + local_min = cutlass.Int32(NO_KEY) + for s in cutlass.range_constexpr(items): + # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. + block = tid + cutlass.Int32(s * SELECT_THREADS) + k = _forced_key( + scores, head, _imin(block, valid - 1), query, init_blocks, local_start + ) + key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) + local_min = _imin(local_min, key[s]) + local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) + if lane < 16: + smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + if tid < 16: + g = cutlass.Int32(NO_KEY) + for w in cutlass.range_constexpr(SELECT_WARPS): + g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) + smem[_GROUP_MIN + tid] = g + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + bound = cutlass.Int32( + -2147483648 + ) # the smallest key: signed keys, unlike the reference's unsigned ones + for g in cutlass.range_constexpr(16): + bound = _imax(bound, smem[_GROUP_MIN + g]) + + # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(items): + cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) + incl = cnt + for d in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync( + incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) + ) + incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) + if lane == 31: + smem[_WARP_CNT + warp] = incl + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + before = cutlass.Int32(0) + total = cutlass.Int32(0) + for w in cutlass.range_constexpr(SELECT_WARPS): + c = smem[_WARP_CNT + w] + before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) + total += c + slot = before + incl - cnt + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + if slot < MAX_CAND: + cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( + tid + cutlass.Int32(s * SELECT_THREADS) + ) + slot += 1 + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if total <= MAX_CAND: + for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): + mine = cand[i] + rank = cutlass.Int32(0) + for j in cutlass.range(total, unroll=8): + rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) + smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) + else: + # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + block = tid + cutlass.Int32(s * SELECT_THREADS) + rank = cutlass.Int32(0) + for j in cutlass.range(valid, unroll=1): + kj = _forced_key(scores, head, j, query, init_blocks, local_start) + less = (kj < key[s]) | ((kj == key[s]) & (j < block)) + rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = key[s] + smem[_PICK_BLOCK + rank] = block + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if warp == 0: + # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. + v = cutlass.Int32(NO_KEY) + if lane < TOPK: + pk = smem[_PICK_KEY + lane] + if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): + v = smem[_PICK_BLOCK + lane] + for size_log in cutlass.range_constexpr(1, 6): + for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): + other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) + ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) + lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) + v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) + if lane < TOPK: + out[query, head, lane] = cutlass.select_( + v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v + ) + + +class _SelectLongDslKernel: + """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count + before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because + ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's + first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it + (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" + + def __init__( + self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool + ): + self.use_pdl = use_pdl + self.early_trigger = early_trigger + self.loaded_trigger = (not early_trigger) and trigger_phase == 1 + self.n_valid_prewait = n_valid_prewait + + @cute.kernel + def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): + row, _, _ = cute.arch.block_idx() + smem = SmemAllocator() + words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) + raw_prewait = cutlass.Int32(0) + if cutlass.const_expr(self.n_valid_prewait): + raw_prewait = n_valid[row // num_kv_heads] + griddepcontrol_wait() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + _select_row( + row, + scores, + n_valid, + out, + init_blocks, + local_blocks, + num_kv_heads, + words, + bar_id=1, + loaded_trigger=self.loaded_trigger, + raw_prewait=raw_prewait, + n_valid_prewait=self.n_valid_prewait, + ) + + @cute.jit + def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): + num_kv_heads = scores.shape[0] + self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( + grid=(scores.shape[2] * num_kv_heads, 1, 1), + block=(SELECT_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + +#: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). +kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} + + +def kernel_key( + use_pdl: bool, early_trigger: bool, trigger_phase: int +) -> Tuple[bool, bool, int, bool]: + """The compile key of a launch (the trigger phase does not matter with the early trigger), with this target's + ``SELECT_N_VALID_PREWAIT``.""" + return ( + bool(use_pdl), + bool(early_trigger), + 0 if early_trigger else int(trigger_phase), + bool(SELECT_N_VALID_PREWAIT), + ) + + +def _compile(key: Tuple[bool, bool, int, bool]): + if key not in kernel_cache: + heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() + scores = make_fake_tensor( + cutlass.Float32, + (heads, blocks, rows), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) + out = make_fake_tensor( + cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) + ) + kernel_cache[key] = cute.compile( + _SelectLongDslKernel(*key), + scores, + n_valid, + out, + cutlass.Int32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + return kernel_cache[key] + + +def select_long_dsl( + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool, + early_trigger: bool, + trigger_phase: int, +) -> None: + """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" + _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( + scores, n_valid_blocks, out, init_blocks, local_blocks + ) + + +def prebuild(keys, device) -> list: + """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" + scores = torch.zeros( + 1, 2048, 4, dtype=torch.float32, device=device + ) # capacity: fixed: a dummy prebuild row, not a size + n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) + out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) + built = [] + for key in keys: + key = kernel_key(*key) + _compile(key)(scores, n_valid, out, 0, 1) + built.append(key) + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py new file mode 100644 index 000000000000..8382d02fb031 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 verification for the MiniMax-M3 target head (vocab-sharded TP head, verify rows <= 48). + +Replaces the all-gather of the target logits, the argmax and the strict-acceptance ops of a greedy verify step with +one kernel over this rank's BF16 logits shard (the target lm_head's own local GEMM output): + +* Every logit becomes a 64-bit key ``ordered(value) << 32 | ~global_index`` (ties pick the lowest index, as + ``torch.argmax``), max-reduced per row: warp shuffles, then ``atomicMax``. +* The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping + (Lamport: 0 = empty, the key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the + slots and runs the strict acceptance. Call n uses slot half n % 2. + +The argmax is taken over the BF16 logits without arithmetic, so the tokens are those of the in-tree path's argmax +over the gathered logits. The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/ +shardAccept.cu) and called as ``torch.ops.trtllm.minimax_m3_shard_accept``. +""" + +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from . import m3_report + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 # capacity: fixed: the kernel's built verify rows +# A verify step of more than SHARD_ROWS rows (13..48 requests at the built draft length) runs in chunks of whole +# requests, each chunk a call of the certified geometry on its own rows; greedy chunks are bitwise the gathered argmax +# (lowest index on ties). Rubin TP2 W1b 797141. Kill: False (past SHARD_ROWS the gathered logits + argmax run). +WIDE_SPEC_TAIL = True + + +@dataclass(eq=False) +class AcceptWorkspace: + """Per-rank exchange slots ([2 halves][SHARD_ROWS rows][world] uint64) behind one multicast mapping, the local + key accumulators and the (finished CTAs, call) counters.""" + + uc: torch.Tensor + mc_ptr: int + keys: torch.Tensor + counters: torch.Tensor + rank: int + world_size: int + handle: Any = None + comm: Any = None + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + if not _mnnvl_workspace_all_succeeded(comm, True): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = mc.data_ptr(), mapping.tp_rank + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None + return AcceptWorkspace( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) + + +def _workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace, created collectively outside CUDA-graph capture on first use.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("m3_shard_accept", None, "the shard-accept workspace") + return None + _WS = create_workspace(mapping, device, SHARD_ROWS) + return _WS + + +def shard_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, +): + """Greedy verification from this rank's BF16 logits shard. + + Args: + logits: [nreq * (draft_len + 1) <= SHARD_ROWS, V / tp] BF16, unit column stride: this rank's columns + (global ids ``vocab_offset ..``) of the verify rows, request-major. + draft_tokens: [nreq, draft_len] int32, the drafts being verified. + batch_idx: [nreq] int32, each request's row in the verify layout. + force_total / force_vec: forced acceptance (see ``m3_shard_accept``). + next_new: optional int32 [nreq, >= 1]; column 0 receives the last accepted token. + + Returns: + (accepted_tokens [nreq, draft_len + 1] int32, num_accepted [nreq] int32, first-pass gather ids [nreq] + int64), or None when the shape is not served or the workspace is missing during a capture. Collective over + the TP group. + """ + nreq, draft_len = draft_tokens.shape + if ( + logits.dim() != 2 + or logits.shape[0] != nreq * (draft_len + 1) + or logits.dtype != torch.bfloat16 + or logits.stride(1) != 1 + ): + return None + if logits.shape[0] > SHARD_ROWS and not (WIDE_SPEC_TAIL and SHARD_ROWS // (draft_len + 1) >= 1): + return None + ws = _workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + if logits.shape[0] > SHARD_ROWS: + # WIDE_SPEC_TAIL: whole requests per launch (SHARD_ROWS // (D + 1) of them), each chunk a call of the + # certified geometry on its own rows; batch_idx is the identity layout (Eagle3 metadata), so a chunk passes + # its local rows and its gather ids are shifted back to the step's verify layout. + per = SHARD_ROWS // (draft_len + 1) + local = _local_rows(per, dev) + for c0 in range(0, nreq, per): + c1 = min(nreq, c0 + per) + r0, r1 = c0 * (draft_len + 1), c1 * (draft_len + 1) + _launch( + logits[r0:r1], vocab_offset, ws, draft_tokens[c0:c1], local[: c1 - c0], force_total, + None if next_new is None else next_new[c0:c1], + None if force_vec is None else force_vec[c0:c1], + accepted[c0:c1], num_accepted[c0:c1], gather_ids[c0:c1], + ) # fmt: skip + gather_ids[c0:c1] += r0 + m3_report.taken( + "wide_spec_tail", logits.shape[0], f"shard accept in {-(-nreq // per)} chunks" + ) + return accepted, num_accepted, gather_ids + _launch( + logits, vocab_offset, ws, draft_tokens, batch_idx, force_total, next_new, force_vec, + accepted, num_accepted, gather_ids, + ) # fmt: skip + return accepted, num_accepted, gather_ids + + +_LOCAL_ROWS: dict = {} + + +def _local_rows(n: int, device) -> torch.Tensor: + """A cached int32 arange(n): a chunk's batch_idx in its own (identity) verify layout.""" + key = (n, device) + t = _LOCAL_ROWS.get(key) + if t is None: + t = _LOCAL_ROWS[key] = torch.arange(n, dtype=torch.int32, device=device) + return t + + +def _launch( + logits, vocab_offset, ws, draft_tokens, batch_idx, force_total, next_new, force_vec, + accepted, num_accepted, gather_ids, +) -> None: # fmt: skip + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_tokens.contiguous(), + batch_idx, + accepted, + num_accepted, + gather_ids, + force_total, + next_new, + _num_sms(), + force_vec, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_spec_tail.py new file mode 100644 index 000000000000..996371370fde --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_spec_tail.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small fused kernels of the MiniMax-M3 Eagle3 step (GB300, decode). + +``chunk_norm`` / ``capture_norm`` apply the drafter's per-layer fc_norm RMSNorms to the captured target hidden states +(the in-tree ``apply_eagle3_fc`` norms, same fp32 arithmetic); ``draft_input`` is a fused draft pass's input (the +embedding lookup, the drafter layer's input_layernorm / hidden_norm and their concat, same arithmetic); +``step_gather`` / ``slot_scatter`` move the step's integer outputs between the spec sampler's store and the engine's +input buffers. +""" + +from typing import Optional, Tuple + +import torch + +# The kernels are compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu) and called as +# torch.ops.trtllm.minimax_m3_{draft_input, chunk_norm, capture_norm, step_gather, slot_scatter}. + + +def draft_input( + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + eps_embed: float, + w_hidden: torch.Tensor, + eps_hidden: float, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(cat([rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden], -1) [M, 2H], hidden[rows] [M, H]) for + int32 ids [M] and int64 rows [M] (rows None: hidden[:M]); ids outside the table embed to zero.""" + m, h = ids.shape[0], table.shape[1] + out = torch.empty((m, 2 * h), dtype=table.dtype, device=table.device) + residual = torch.empty((m, h), dtype=table.dtype, device=table.device) + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, use_pdl + ) + return out, residual + + +def chunk_norm( + src: torch.Tensor, weights: torch.Tensor, eps: float, use_pdl: bool = True +) -> torch.Tensor: + """Per-chunk RMSNorm of src [M, >= C*H] (row stride may exceed C*H) with weights [C, H] -> [M, C*H].""" + c, h = weights.shape + out = torch.empty((src.shape[0], c * h), dtype=src.dtype, device=src.device) + torch.ops.trtllm.minimax_m3_chunk_norm(src, weights, eps, out, use_pdl) + return out + + +# The overlap scheduler's per-step gathers and the spec sampler's slot stores as one kernel each (step_gather / +# slot_scatter) instead of ~10 and 4 small torch ops; called by the target's SpecStepFold. + + +def step_gather( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 +) -> bool: + """See step_gather_kernel; False when the inputs are outside what the kernel covers (caller falls back).""" + ts = (new_tok, next_draft, lens, slots, ppi, ids, draft, pos_off, kv_off) + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or next_draft.dim() != 2 + or new_tok.dim() < 2 + ): + return False + if (new_tok.dim() == 3 and new_tok.shape[2] != 1) or not ( + lens.is_contiguous() + and ids.is_contiguous() + and draft.is_contiguous() + and pos_off.is_contiguous() + and kv_off.is_contiguous() + ): + return False + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 + ) + return True + + +def slot_scatter( + slots, src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft +) -> bool: + """See slot_scatter_kernel (2-D int32 views, zero padding); False to fall back.""" + ts = (src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft) + if any(t.dtype != torch.int32 or not t.is_cuda for t in ts) or slots.dtype != torch.long: + return False + two = (src_new, dst_new, src_next_new, dst_next_new, src_draft, dst_draft) + if any(t.dim() != 2 for t in two) or src_lens.dim() != 1 or dst_lens.dim() != 1: + return False + torch.ops.trtllm.minimax_m3_slot_scatter( + slots, + src_new, + dst_new, + src_next_new, + dst_next_new, + src_lens, + dst_lens, + src_draft, + dst_draft, + ) + return True + + +def capture_norm( + h: torch.Tensor, + residual: Optional[torch.Tensor], + weight: torch.Tensor, + eps: float, + dst: torch.Tensor, +) -> None: + """dst [M, H] (row stride may exceed H) = rmsnorm(h + residual) * weight with chunk_norm's rounding.""" + torch.ops.trtllm.minimax_m3_capture_norm(h, residual, weight, eps, dst) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py new file mode 100644 index 000000000000..f5a1b229c777 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard of a verify step of at most 8 rows, streamed from a lossless 12-bit coding of its BF16 +weight (75% of the bytes; ``ops.gemm.m3_target_head`` has the format). + +``prepare_head`` codes this rank's shard once, before the executor's warmup, and keeps it only when it decodes to the +BF16 weight bit for bit (``round_trip_ok``); a shard that does not keeps the in-tree BF16 GEMM. ``logits`` runs the +kernel: BF16 x BF16 MMAs with fp32 accumulation over the full K and one BF16 rounding per logit, the in-tree head's +arithmetic in another summation order. The kernel is compiled into TensorRT-LLM +(cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu) and called through the catalog entry ``ops.gemm.m3_target_head``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from ..ops.gemm.m3_target_head import CTA_WARPS, MAX_ROWS, encode_chunked, round_trip_ok +from ..ops.gemm.m3_target_head import m3_target_head as _target_head_op + + +@dataclass(eq=False) +class TargetHead: + """One rank's target lm_head shard in the 12-bit coding, and the kernel's split-tile scratch. + + ``coded`` is False when the coding did not reproduce the weight bit for bit (the in-tree GEMM then runs).""" + + weight: torch.Tensor + coded: bool + w: Optional[torch.Tensor] = None + base: int = 0 + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + part: Optional[torch.Tensor] = None + counter: Optional[torch.Tensor] = None + grid: int = 0 + + +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def head_supported(lm_head) -> bool: + """A bias-free 2-D BF16 head of whole 16-row tiles whose K the coding covers (a multiple of 64).""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + return ( + getattr(lm_head, "bias", None) is None + and weight.shape[0] % 16 == 0 + and weight.shape[1] % 64 == 0 + ) + + +def prepare_head(lm_head, grid: Optional[int] = None) -> TargetHead: + """Code ``lm_head``'s local shard and check it decodes to the BF16 weight bit for bit (fails closed: a shard that + does not is returned with ``coded=False`` and nothing kept). Built once, eagerly, before the executor's warmup + (``M3Eagle3Worker.prebuild_target_head``).""" + weight = lm_head.weight + with torch.no_grad(): + w = weight.detach() + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if not round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return TargetHead(weight=weight, coded=False) + if grid is None: + grid = _num_sms() + n = weight.shape[0] + return TargetHead( + weight=weight, + coded=True, + w=packed, + base=base, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + part=torch.zeros(grid * CTA_WARPS * 2 * 4 * 32, dtype=torch.float32, device=weight.device), + counter=torch.zeros((n + 15) // 16, dtype=torch.int32, device=weight.device), + grid=grid, + ) + + +def logits(head: TargetHead, x: torch.Tensor, use_pdl: bool = True) -> torch.Tensor: + """This rank's BF16 logits shard [rows, n] of the rows x [rows <= MAX_ROWS, K] (BF16) from the coded shard.""" + assert head.coded and 1 <= x.shape[0] <= MAX_ROWS, (head.coded, x.shape) + if x.stride(1) != 1 or x.stride(0) % 8 or x.data_ptr() % 16: + x = x.contiguous() + n = head.weight.shape[0] + out = torch.empty((x.shape[0], n), dtype=torch.bfloat16, device=x.device) + return _target_head_op( + head.w, + x, + out, + head.part, + head.counter, + head.exc_off, + head.exc_pos, + head.exc_w, + head.base, + n, + head.grid, + use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py new file mode 100644 index 000000000000..db0077117fe3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The token map of ``thop.trtllm_gen_generation_preprocess`` for one request with several tokens. + +For multi-token generation the op's QKV preprocessing places every row by the workspace token map (``tokensInfo``, +(request, token) per row), but builds that map only when ``isBuildDecoderInfoKernelNeeded()`` holds, which at one +request with the default RoPE it does not. The op then applies whatever map the last multi-request call left in the +workspace: a row can land on another request's length and pages, and its RoPE position and K/V slot are wrong. This +module wraps the op so that, for one request and several tokens, the map the build would write, (0, t) for each row t, +is written into the op's own map slice of the workspace (the offset from the op's layout function) just before the +call. Every other call passes through unchanged. The rows' positions still come from the request's generation length, +exactly as with a freshly built map. + +The write is a device copy from a map prebuilt at install (no allocation, no host sync), so it is captured into the +decode graphs with the call it precedes. +""" + +from __future__ import annotations + +from typing import Callable, Dict, Optional, Tuple + +import torch + +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.logger import logger + +#: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). +_ARG_NAMES = ( + "qkv_input", + "workspace", + "sequence_lengths", + "spec_decoding_generation_lengths", + "spec_decoding_position_offsets", + "kv_cache_block_offsets", + "host_kv_cache_pool_pointers", + "host_kv_cache_pool_mapping", + "kv_scale_orig_quant", + "kv_scale_quant_orig", + "attention_output_orig_quant", + "rotary_inv_freq", + "rotary_cos_sin", + "mrope_position_deltas", + "layer_idx", + "seq_offset", + "num_heads", + "num_kv_heads", + "head_size", + "tokens_per_block", + "kv_cache_quant_mode", + "max_attention_window_size", + "cyclic_attention_window_size", + "num_tokens", + "batch_beam", + "input_seq_length", + "max_past_kv_length", + "rotary_embedding_dim", + "rotary_embedding_base", + "rotary_embedding_scale_type", + "rotary_embedding_scale", + "rotary_embedding_max_positions", + "position_embedding_type", + "bmm1_scale", + "bmm2_scale", + "fp8_context_fmha", + "predicted_tokens_per_seq", + "attention_chunk_size", + "multi_processor_count", + "total_num_blocks", + "kv_factor", + "need_build_kv_cache_metadata", + "cross_attention", + "skip_fmha_workspace", +) +_INT2_BYTES = 8 # capacity: fixed: an element size + +_original: Optional[Callable] = None +_max_rows = 0 +_maps: Dict[torch.device, torch.Tensor] = {} +_offsets: Dict[Tuple, int] = {} + + +def _no_taken(rows: int) -> None: + pass + + +def _no_missed(rows: int, reason: str) -> None: + pass + + +_on_taken: Callable[[int], None] = _no_taken +_on_missed: Callable[[int, str], None] = _no_missed + + +def _indexed(device: torch.device) -> torch.device: + """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: + """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" + device = _indexed(device) + m = _maps.get(device) + if m is None or m.shape[0] < rows: + reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" + _on_missed(rows, reason) + raise RuntimeError(f"m3_tokmap_guard: {reason}") + return m[:rows] + + +def _map_offset(args: dict, rows: int) -> int: + x = args["qkv_input"] + skip = bool(args.get("skip_fmha_workspace", False)) + key = ( + x.dtype, + rows, + int(args["num_heads"]), + int(args["head_size"]), + int(args["rotary_embedding_dim"]), + int(args["num_kv_heads"]), + skip, + ) + off = _offsets.get(key) + if off is None: + layout = thop.get_trtllm_gen_generation_workspace_layout( + x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip + ) + off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) + if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: + raise RuntimeError( + f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " + f"(offset {off}, {size} bytes): the layout changed under this guard" + ) + _offsets[key] = off + return off + + +def _guarded(*args, **kwargs): + a = dict(zip(_ARG_NAMES, args)) + a.update(kwargs) + rows = int(a["num_tokens"]) + if ( + a.get("spec_decoding_generation_lengths") is not None + and int(a["predicted_tokens_per_seq"]) > 1 + and int(a["batch_beam"]) == 1 + and rows > 1 + and not bool(a.get("cross_attention", False)) + ): + ws = a["workspace"] + src = _one_request_map(ws.device, rows) + off = _map_offset(a, rows) + # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. + dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] + dst.view(torch.int32).view(rows, 2).copy_(src) + _on_taken(rows) + return _original(*args, **kwargs) + + +def install( + device: torch.device, + max_rows: int, + on_taken: Callable[[int], None], + on_missed: Callable[[int, str], None] = _no_missed, +) -> bool: + """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to + ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this + call installed the wrapper.""" + global _original, _max_rows, _on_taken, _on_missed + device = _indexed(device) + if device not in _maps or _maps[device].shape[0] < max_rows: + t = torch.arange(max_rows, dtype=torch.int32, device=device) + _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() + _max_rows = max(_max_rows, max_rows) + _on_taken = on_taken + _on_missed = on_missed + if _original is not None: + return False + _original = thop.trtllm_gen_generation_preprocess + thop.trtllm_gen_generation_preprocess = _guarded + logger.info( + "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " + f"(up to {max_rows} rows)" + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/__init__.py new file mode 100644 index 000000000000..77d7b84d9bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..6e3feee8882a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 1..8 decode tokens in one CuTe DSL launch (``m3_moe``, the sm_107 TP2 build: routed +and shared expert width 1536 per rank): the router, the routed NVFP4 experts and the shared MXFP8 expert, on +caller-owned state (an ``M3MoeState`` and one ``M3MoeLayer`` per MoE layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ..m3_moe.op import ( + HIDDEN, + INTER, + SHARED_INTER, + TOP_K, + M3MoeLayer, + M3MoeState, + ar_push_desc, + is_supported, +) + +__all__ = [ + "M3MoeLayer", + "M3MoeState", + "ar_push_desc", + "is_supported", + "m3_moe_engine", + "moe_reference", + "moe_route", + "mxfp8_dequant", + "mxfp8_quant_dequant", + "unswizzle_sf", +] + +# The model's constants the kernels are built with: routed_scaling_factor, the shared expert's SwiGLU-OAI +# (swiglu_alpha, swiglu_limit, the (up + 1) offset). +ROUTED_SCALING = 2.0 +SHARED_ALPHA = 1.702 +SHARED_LIMIT = 7.0 + +_LOG2E = 1.4426950216293334961 +_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) +_F32_MAX = 3.40282346638528859812e38 + + +def _f32(x: float) -> torch.Tensor: + return torch.tensor(x, dtype=torch.float32) + + +def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: + """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" + lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) + lo = lut[(packed & 0xF).long()] + hi = lut[(packed >> 4).long()] + return torch.stack((lo, hi), dim=-1).flatten(-2) + + +def _e2m1_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" + a = x.abs() + q = torch.full_like(a, 6.0) + for hi, value, inclusive in ( + (5.0, 4.0, True), + (3.5, 3.0, False), + (2.5, 2.0, True), + (1.75, 1.5, False), + (1.25, 1.0, True), + (0.75, 0.5, False), + (0.25, 0.0, True), + ): + q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) + return torch.where(x < 0, -q, q) + + +def _e4m3_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" + return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + + +def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: + """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" + bits = x.contiguous().view(torch.int32) + code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + return code.clamp(0, 254) + + +def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: + """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" + flat = sf.reshape(-1).view(torch.uint8) + r = torch.arange(rows, device=flat.device).view(-1, 1) + b = torch.arange(blocks, device=flat.device).view(1, -1) + padded_blocks = (blocks + 3) // 4 * 4 + idx = ( + (r // 128) * (128 * padded_blocks) + + (b // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (b % 4) + ) + return flat[idx] + + +def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" + rows, k = q.shape + codes = unswizzle_sf(sf, rows, k // 32).float() + return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( + 32, dim=1 + ) + + +def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: + """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to + fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" + rows, k = x.shape + blocks = x.float().view(rows, k // 32, 32) + code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) + scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) + return (_e4m3_round(blocks / scale) * scale).view(rows, k) + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16).float() + + +def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" + return (a.double() @ b.double().t()).float() + + +def _row_perm(rows: int, gated: bool) -> torch.Tensor: + """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle + of epilogue tile 128).""" + probe = torch.empty(rows, 1) + perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) + if gated: + perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] + return perm + + +def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: + """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row + order (the global scale is not applied: it is the per-expert alpha).""" + rows, half = w.shape + codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() + vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) + return vals[perm_inv] + + +def moe_route(x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor): + """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" + logits = _gemm(x, router_weight) + sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() + keys = sig + bias + ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] + sk = keys.gather(1, ids) - bias[ids] + total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) + w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) + return ids, _bf16(w) + + +def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: + """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" + rows, k = x.shape + blocks = x.view(rows, k // 16, 16) + vmax = blocks.abs().amax(dim=-1, keepdim=True) + sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) + # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element + # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). + osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) + return (_e2m1_round(blocks * osc) * sff).view(rows, k) + + +def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: + """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 + input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with + act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" + a, c, alpha, beta, lim = (esc[i] for i in range(5)) + gc = torch.where(gate > lim, lim, gate) + gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) + uc = torch.minimum(torch.maximum(up, -lim), lim) + c1 = a * (alpha * _f32(_LOG2E).to(up.device)) + sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) + h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) + blocks = h.view(-1, 16) + amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) + sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) + osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) + return (_e2m1_round((blocks * osc) * c) * sff).view(-1) + + +def _shared_expert(x: torch.Tensor, shared_mxfp8) -> torch.Tensor: + """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the + model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" + gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 + gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) + gate, up = gu[:, :SHARED_INTER], gu[:, SHARED_INTER:] + a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) + ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 + sg = 1.0 / ( + torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 + ) + h = _bf16(ub * (sg * a)) + return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) + + +def moe_reference(hidden: torch.Tensor, layer: M3MoeLayer, shared: bool) -> torch.Tensor: + """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" + x = hidden.float() + m = x.shape[0] + ids, weights = moe_route(x, layer.router_weight, layer.bias) + xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) + inv31 = torch.argsort(_row_perm(2 * INTER, gated=True)).to(x.device) + inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) + y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) + for e in torch.unique(ids).tolist(): + w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) + w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) + esc = layer.escal[:, e] + for t, k in (ids == e).nonzero().tolist(): + acc1 = _gemm(xr[t : t + 1], w31)[0] + # The buffer holds [w3 (up); w1 (gate)]. + h = _fc1_epilogue(acc1[:INTER], acc1[INTER:], esc) + y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) + acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) + for k in range(TOP_K): + acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() + routed = acc.float() + if not shared: + return routed.to(torch.bfloat16) + return (_shared_expert(x, layer.shared_mxfp8) + _bf16(routed)).to(torch.bfloat16) + + +_ENGAGED: set = set() + + +class M3MoeEngine(OpWrapper): + """One MiniMax-M3 MoE layer for M (1..8) decode tokens: ``layer.state.run(layer, hidden, out, arp)``. + + ``hidden`` bf16 [M, 6144] contiguous, M = the state's token count. ``layer`` (from ``M3MoeState.layer``) carries the + TRTLLM-Gen NVFP4 expert buffers read in place, the per-expert scalars, the router weight / bias, the activation's + NVFP4 global scale and, on a ``shared`` state, the shared expert's MXFP8 weights; its state carries the compiled + kernel and the workspace every layer on it shares. Returns bf16 [M, 6144]: bf16(shared + routed) on a ``shared`` + state, else the routed partial. ``out``: write there instead of a new tensor. ``arp`` (``ar_push_desc``): also + push the rows into a multicast all-reduce exchange. + + The routed part is TRTLLM-Gen's NVFP4 MoE on the fp32 router logits: sigmoid (tanh form) + correction bias top-4 + of 128 (ties to the lower id), weights sigmoid / sum * 2.0 rounded to bf16; the activation quantized to NVFP4 with + the layer's global scale; per expert the raw FC1 sums, SwiGLU-OAI (alpha / beta / clamp), the NVFP4 requant with + fc31_scale_c, FC2, bf16(fc2_alpha x sum); the fma chain over the top-4 in order. The shared expert is the target's + MXFP8 SwiGLU-OAI MLP. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # One cell per token count: m_max is a trace-time constant, so each is its own build, and the target builds all + # eight (decode batches of 1..8 requests, spec-decode verify rows up to 8). The whys name what the build differs in. + CELLS: tuple[Cell, ...] = ( + Cell( + why="1 token: FC1 -> FC2 records replicated 32x, logits 32x, last phase without trailing FC2", + spec=dict(tokens=1), + ), + Cell( + why="2 tokens: records 32x, logits 32x, the one build without AR_PRE", + spec=dict(tokens=2), + ), + Cell(why="3 tokens: records 8x, logits 32x, 4 FC2-B builder warps", spec=dict(tokens=3)), + Cell( + why="4 tokens: one request verifying 3 drafts (the TP2 b1 point): 6-stage ring, B2IMG hand-off, " + "the 4 two-block FC2 CTAs stream no FC1", + spec=dict(tokens=4), + ), + Cell( + why="5 tokens: single records, a routing warp takes two tokens, logits 32x", + spec=dict(tokens=5), + ), + Cell( + why="6 tokens: single records, routing warps 8.. take tokens 4..5", spec=dict(tokens=6) + ), + Cell(why="7 tokens: the B2IMG global FC1 -> FC2 hand-off", spec=dict(tokens=7)), + Cell( + why="8 tokens: two requests verifying 3 drafts (the TP2 b2 point): 5-stage ring, x read from global " + "(no XSTAGE), no early group 0, the 4 two-block FC2 CTAs stream no FC1", + spec=dict(tokens=8), + ), + ) + + note = """ + The kernel's elementwise math uses the PTX approximations (rcp.approx, ex2.approx, tanh) where `reference` uses + torch's correctly rounded ones, and its GEMMs sum in tensor-core order where `reference` sums in fp64: a + quantization boundary (NVFP4 activation / intermediate code, routing near-tie) can fall on the other side, so the + two agree within the band the test derives, not bit for bit. The routed part is bit for bit the stock TRTLLM-Gen + NVFP4 MoE (fp4_quantize + fp4_block_scale_moe_runner) on the same routing; the test checks that too. + + The all-reduce push (`arp`) needs a multicast exchange: its 2-rank certification (pushed boundary bit for bit the + unpushed one) is comm/_m3_moe_push_op_matrix.py. Not covered: the non-shared state (`shared=False`), which this + target never builds. Every layer on one state + must run in one stream order: the workspace is shared and tagged per call, never reset. + + No `is_valid`: the compiled call checks every buffer's dtype, shape and alignment against its signature and + raises on a mismatch (a misaligned `hidden`, an `out` that is not bf16 [M, 6144], a malformed `arp`). + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + state = layer.state + if (state.device, state.m) not in _ENGAGED: + # The engagement line (once per device and token count, on every rank: the first call, eager or captured). + _ENGAGED.add((state.device, state.m)) + logger.info( + f"[M3 sm_107 tp2] MoE engine m3_moe engaged: {state.m} tokens, {state.num_ctas} routed + " + f"{state.mod.NSH} shared CTAs, {state.mod.STAGES}-stage ring, on {state.device}" + ) + return state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Plain torch over the layer's buffers (see the class docstring); `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=layer.state.shared) + + +m3_moe_engine = M3MoeEngine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py new file mode 100644 index 000000000000..636dfcd0b614 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV of the decode MoE: ``logits[M, E] = float(x) @ W^T`` at one token.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) +#: The kernel builds each block's output scale through two ``rcp.approx.ftz.f32`` reciprocals (~2^-23 relative): +#: a value this close to an e2m1 midpoint may round to either neighbour, everything further out must match. +TIE_WINDOW = 2.0**-20 + + +def nvfp4_quantize( + x: torch.Tensor, sf_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], + the pre-rounding scaled values [M, K] that ``near_tie`` reads). + + Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. + """ + gs = sf_scale.float().reshape(()) + m, k = x.shape + blocks = x.float().view(m, k // 16, 16) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +def near_tie(scaled: torch.Tensor) -> torch.Tensor: + """Elements whose e2m1 rounding the kernel's approximate reciprocals decide (see ``TIE_WINDOW``).""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=scaled.device, dtype=torch.float32) + a = scaled.abs().unsqueeze(-1) + return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) + + +def fp32_logits(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """``float(x) @ w^T`` as a true fp32 product (TF32 off).""" + with true_fp32_matmul(): + return x.float() @ w.t() + + +class M3RouterGemv(OpWrapper): + """fp32 router logits ``float(x[M, K]) @ w[E, K]^T`` with exact fp32 FMA, one CTA per expert. + + ``x`` is BF16 with 1..16 rows, unit column stride and 16-byte rows; ``w`` is the fp32 router weight; K a + multiple of 1024. ``rows`` experts per CTA; only 1 is built. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target takes the GEMV at exactly one token (2..48 go to the tensor-core router): 128 experts, hidden 6144. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token, 128 experts x hidden 6144", + spec=dict(m=1, e=128, k=6144), + ), + ) + + note = """ + The op accepts up to 16 rows; the target sends one, so wider calls are not cells. The logits differ from a true + fp32 product only in accumulation order (per-thread FMA chains over K / 1024 strided column groups, a warp tree + and an 8-warp sum). + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_router_gemv(x, w, rows, pdl) + + def reference(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + """True fp32 product; ``rows`` and ``pdl`` do not enter the arithmetic.""" + return fp32_logits(x, w) + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> None: + # The activation rows are bulk-copied from x's base address; the op checks the row stride but not the base, + # and a misaligned bulk copy faults the context rather than raising. + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + + +m3_router_gemv = M3RouterGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py new file mode 100644 index 000000000000..0306210adaab --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV plus the routed experts' NVFP4 activation, one token.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_router_gemv import fp32_logits, nvfp4_quantize + + +class M3RouterGemvFp4(OpWrapper): + """``[logits fp32 [M, E], x_fp4 [M, K / 2] uint8, x_sf [M * K / 16] uint8]`` from one kernel. + + The logits are ``m3_router_gemv``'s; the NVFP4 activation is ``fp4_quantize(x, sf_scale, 16, False, False)`` + (16-element blocks, E4M3 scales in the linear layout, per-tensor ``sf_scale``) for the routed experts' FC1. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The fused-router decode step at one token: 128 experts, hidden 6144, the experts' input scale. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token: logits for 128 experts and the NVFP4 activation of hidden 6144", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 8.0), + ), + Cell( + why="a global scale that puts block scales near the E4M3 top, where the scale cast saturates", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 2.0), + ), + ) + + note = """ + The e2m1 codes are exact except within `TIE_WINDOW` of a rounding midpoint, where the kernel's approximate + reciprocals decide the direction. Up to 16 rows are accepted; the target sends one. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_gemv_fp4(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """True fp32 logits and the NVFP4 recipe; ``pdl`` does not enter the arithmetic.""" + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return fp32_logits(x, w), packed, sf + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + # A zero or non-finite global scale makes every block scale 0 or NaN and the data codes garbage; the op + # reads the scalar on the device and cannot refuse it. + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_gemv_fp4 = M3RouterGemvFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..9e21113d8483 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TF32 tensor-core router of the decode MoE for 2..48 tokens, optionally with the NVFP4 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from .m3_router_gemv import nvfp4_quantize + + +def tf32_rna(w: torch.Tensor) -> torch.Tensor: + """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" + bits = w.contiguous().view(torch.int32) + rounded = (bits + 0x1000) & ~0x1FFF + finite = (bits & 0x7F800000) != 0x7F800000 + return torch.where(finite, rounded, bits).view(torch.float32) + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target routes 1..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one token, logits only (the wide target's router_path(1) with no fp32 GEMV rows)", + spec=dict(m=1, k=6144, fp4=False), + ), + Cell( + why="one token with the NVFP4 activation (the fused-router path at one row)", + spec=dict(m=1, k=6144, fp4=True), + ), + Cell( + why="two tokens, logits only (the smallest call the target makes here)", + spec=dict(m=2, k=6144, fp4=False), + ), + Cell( + why="8 tokens with the NVFP4 activation: one 8-token MMA tile, every expert tile quantizes one row", + spec=dict(m=8, k=6144, fp4=True), + ), + Cell( + why="13 tokens: a partial second token tile", + spec=dict(m=13, k=6144, fp4=True), + ), + Cell( + why="the 48-token maximum, six token tiles, logits only", + spec=dict(m=48, k=6144, fp4=False), + ), + Cell( + why="the 48-token maximum with the NVFP4 activation", + spec=dict(m=48, k=6144, fp4=True), + ), + # Every other row count the router takes (modeling's ROUTER_TC), with and without the NVFP4 activation. + *( + Cell( + why=f"{m} rows: a router row count between the ones above", + spec=dict(m=m, k=6144, fp4=m % 2 == 1), + ) + for m in range(2, 49) + if m not in (2, 8, 13, 48) + ), + ) + + note = """ + The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both + sum the same 16 slice partials in the same order. Only the one the test device picks is driven. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> tuple[torch.Tensor, ...]: + """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" + with true_fp32_matmul(): + logits = x.float() @ tf32_rna(w).t() + if sf_scale is None: + return (logits,) + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return logits, packed, sf + + def is_valid( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + if sf_scale is not None: + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py new file mode 100644 index 000000000000..8c2a67356d3f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One kernel for MSA's device-side KV-length correction (MsaAttentionMetadata.on_update_kv_lens). + +The torch formulation is ~23 one-CTA kernels (min, cumsum, gathers, div / clamp / remainder, the block-table lookup, +the valid-block counts) and runs at the head of every captured decode step, twice. This kernel computes the same +int32 results in one launch: + + kv_true[b] = min(kv_lens[b], staged[b]) -> seq_lens_out[b], cu_kv[b + 1] = prefix sum (cu_kv[0] = 0) + pos[t] = kv_true[r] - qo[r] + intra[t] (r = q_batch_row[t]) + out_loc[t] = table[r, min(max(pos, 0) // page, ncols - 1)] * page + max(pos, 0) % page + n_valid[t] = (max(pos + 1, 0) + page - 1) // page + +``msa_kv_update`` is the target's ``MsaKvUpdate`` (``tensorrt_llm._torch.pyexecutor.runtime_hooks``): the MSA +metadata calls it with its own buffers and falls back to the torch formulation when it returns False. The kernel is +``torch.ops.trtllm.minimax_m3_msa_kv_update``, called through this target's ``msa_kv_update`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.msa_kv_update import msa_kv_update as _kv_update_op + + +def msa_kv_update( + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], +) -> bool: + """The fused correction (see module doc); False, with nothing written, when an operand is outside what the + kernel covers. Arguments in ``MsaKvUpdate`` order; ``seq_lens`` are the per-request query lengths (``qo``).""" + ts = [ + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + msa_seq_lens, + out_cache_loc, + ] + ts += [t for t in (cu_kv_lens, n_valid_blocks) if t is not None] + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or block_table.dim() != 2 + or block_table.stride(1) != 1 + ): + return False + if any(not t.is_contiguous() for t in ts if t is not block_table): + return False + if ( + kv_lens.numel() < batch + or kv_lens_staged.numel() < batch + or seq_lens.numel() < batch + or msa_seq_lens.numel() < batch + or q_batch_row.numel() < total_q + or q_intra.numel() < total_q + or out_cache_loc.numel() < total_q + or (cu_kv_lens is not None and cu_kv_lens.numel() < batch + 1) + or (n_valid_blocks is not None and n_valid_blocks.numel() < total_q) + ): + return False + _kv_update_op( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/__init__.py new file mode 100644 index 000000000000..e6c3d806a4a3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) as a CuTe DSL kernel, configured for the TP2 +projection shapes on SM107 (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..3f74be245ee0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1237 @@ +# 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. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # capacity: fixed: a kernel tile constant (weight rows per CTA = the MMA's M) +MMA_N = 8 # token columns +CTA_K = 128 # capacity: fixed: a kernel tile constant (one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row) +MMA_INST_K = 32 +# 4 MMAs per k-tile, one per scale byte of the atom. +K_BLOCKS = CTA_K // MMA_INST_K # capacity: fixed: a kernel tile constant +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], +) -> None: + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py new file mode 100644 index 000000000000..3124827082df --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py @@ -0,0 +1,496 @@ +# 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. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) at the TP2 per-rank projection shapes (SM107). + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring[, contig, nacc, a_tmem, st_async]) of the TP2 per-rank projections. The optional fields default +# to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. SM107 graph us at M 1 / 4 / 8 over weights rotated beyond L2, against the +# stock MXFP8 quantize + FlashInfer CuTe DSL GEMM: QKV 4.11 / 4.12 / 4.10 vs 8.6-8.8, o_proj 3.90 / 3.96 / 3.70 vs +# 6.9-7.2, shared gate_up 3.11 / 3.09 / 3.10 vs 8.1-8.4, shared down 2.22 vs 4.7-4.9, dense gate_up 8.15 / 8.56 / 8.55 +# vs 10.2-10.5, dense down 5.89 / 5.72 / 5.46 vs 8.7-9.0. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + # sparse-layer QKV + index projection (index values disabled); timed at N 5120: 4992 is 39 of its 40 N tiles + (4992, 6144): (4, 12), + # dense-layer QKV + (4608, 6144): (4, 12), + # o_proj + (6144, 4096): (2, 12), + # shared-expert gate_up + (3072, 6144): (4, 12), + # shared-expert down + (6144, 1536): (2, 12), + # dense-layer gate_up + (12288, 6144): (2, 12), + # dense-layer down + (6144, 6144): (2, 12), +} +# Per-token-count configs (N, K, M); none at TP2. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = {} +# Above 8 tokens (wider than the TP2 decode step); none at TP2. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = {} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +# (rule 1b, deterministic) The decode projections at up to STRIDED_MAX_ROWS rows split K over strided k-tiles per +# split rank (contig 0: rank r owns k-tiles r, r + split, ...) instead of contiguous ranges: another fixed fp32 order. +# SM107 TP2 4 rows (783556), us/call contiguous -> strided: sparse QKV 4.100 / 4.102 -> 4.074 / 4.072, dense QKV +# 3.820 / 3.816 -> 3.812 / 3.805, o_proj 3.970 / 4.017 -> 3.945 / 4.001, dense gate_up 8.694 / 8.665 -> 8.627 / 8.606 +# (dense down stays contiguous: strided is slower); deterministic (3 runs bitwise) and max error vs the fp64 product +# equal to the contiguous order's. Kill: False. +STRIDED_K_TILES = True +STRIDED_MAX_ROWS = 8 # capacity: fixed: the strided-input kernels' built rows +_STRIDED = frozenset({(4992, 6144), (4608, 6144), (6144, 4096), (12288, 6144)}) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + if ( + STRIDED_K_TILES + and m is not None + and m <= STRIDED_MAX_ROWS + and (n, k) in _STRIDED + and len(cfg) == 2 + ): + cfg = cfg + (0,) + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # 17..32 rows at K < 1024 decline (an unsplit K below 1024 is slower than the FlashInfer GEMM there); no TP2 + # projection has K < 1024. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + if nacc != 1: + # More than one TMEM accumulator is broken on SM107 (783556 / 783557 / 783579: nondeterministic or garbage + # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + ) + return fn + + +def _views( + x: torch.Tensor, xsf: Optional[torch.Tensor], weight, weight_scale, y, x_bf16: bool +) -> tuple: + k = x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + n = weight.shape[0] + if weight.shape[1] != k: + return None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar) + if resolved is None: + return None + key, runtime_ar = resolved + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros(m, k, dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, device=dev) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/__init__.py new file mode 100644 index 000000000000..acec8fdbb26f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalization entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..94827619dad2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step embedding gather and the first layer's input RMSNorm in one CuTe DSL launch, into caller buffers.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm + +__all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] + +VOCAB = 200064 # MiniMax-M3's vocabulary; the decode table is the full (replicated) embedding + + +class M3EmbeddingFirstnorm(OpWrapper): + """Write ``table[ids]`` into ``raw`` (bf16 [M, 6144], a zero row for an id outside ``[0, V)``) and + ``raw * rsqrt(mean(raw^2) + eps) * weight`` into ``out``; returns ``(out, raw)`` -- the first layer's normed input + and its residual. ``ids`` int32 or int64 [M], M = 1..16; ``table`` bf16 [V, 6144] contiguous; ``weight`` bf16 + [6144]. The first call per (M, id dtype, table shape, PDL) compiles (``DecodeEmbedding`` does it at build).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target builds every (rows 1..16, id dtype) owner; the kernel is one CTA per row, so the rows differ only in + # grid size: the ends of the range and both id dtypes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, int32 ids", spec=dict(rows=1, id_dtype=torch.int32)), + Cell(why="eight requests, int32 ids", spec=dict(rows=8, id_dtype=torch.int32)), + Cell(why="one request, int64 ids", spec=dict(rows=1, id_dtype=torch.int64)), + Cell( + why="six rows (one request verifying a draft), int64 ids", + spec=dict(rows=6, id_dtype=torch.int64), + ), + # Every other row count the decode dispatch takes (modeling's EMBED_FIRSTNORM), both id dtypes. + *( + Cell( + why=f"{rows} rows: a decode row count between the ones above", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in (2, 3, 4, 5, 7) + ), + # B5, the batch-4 decode steps (modeling's EMBED_FIRSTNORM past 8 rows): every row count 9..16, both id dtypes. + *( + Cell( + why=f"{rows} rows (batch-4 decode step)", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in range(9, 17) + ), + ) + + note = """ + The gathered rows are exact. The norm is flashinfer's RMSNorm arithmetic in its reduction order with a fast + rsqrt, where `reference` sums in fp64 and takes torch's rsqrt: within the band the test derives, not bit for bit. + SM107 against the stock embedding + RMSNorm module (flashinfer's kernel, 768 threads): the gathered rows are + bitwise equal; the normed rows differ in one bf16 ulp on ~5e-6 of the elements (~1.6% of rows, 1-2 elements each: + the fp32 sum-of-squares tree differs, so the rsqrt input can differ in its last bits). + An id outside [0, V) gives a zero row and a zero norm (torch's embedding would raise). No accepted input is known + to compute wrongly: the host side rejects any other dtype, shape or a strided buffer. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return m3_embed_norm(ids, table, weight, eps, raw, out, use_pdl) + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x * rstd * weight.double()).to(torch.bfloat16), rows + + +m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/__init__.py new file mode 100644 index 000000000000..d25c41dc9557 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries this target alone calls: the Eagle3 verify ops, the capture / fc_norm RMSNorms and +the step-input staging.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..a1993924a112 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_chunk_norm import rmsnorm_rows + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{4 * n} verify rows of {n} requests (b4 / b8), layer slot {n % 3}, residual added", + spec=dict(m=4 * n, slot=n % 3, residual=True), + ) + for n in range(3, 8) + ), + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ) + + note = """ + fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional + one-ulp move against torch's statistic. + """ + + def __call__( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) + return dst + + def reference( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" + x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h + return rmsnorm_rows(x, w, eps) + + def is_valid( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> None: + # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read + # past its elements without an error. + assert w.is_contiguous(), "w must be dense [H]" + + +m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..453edec27764 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states +(``m3_spec_tail.chunk_norm``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{n} requests' {4 * n} verify rows (b4 / b8), dense [{4 * n}, 3 x 6144]", + spec=dict(m=4 * n, c=3, h=6144, extra_cols=0), + ) + for n in range(3, 8) + ), + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="two requests' 8 verify rows (b2) from the capture buffer, dense [8, 3 x 6144]", + spec=dict(m=8, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ) + + note = """ + fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then + against torch's statistic; the cells' band states it. + """ + + def __call__( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) + return dst + + def reference( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" + c, h = w.shape + x = src[:, : c * h].reshape(src.shape[0], c, h) + return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) + + def is_valid( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ) -> None: + c, h = w.shape + assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + + +m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..4fe14d538a5f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat +(``m3_spec_tail.draft_input``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3DraftInput(OpWrapper): + """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = + hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of + a vocab-local embedding). Returns ``(out, residual)``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling's DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the + # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests + # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", + spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", + spec=dict(m=n, h=6144, gather=True, invalid_id=False), + ) + for n in range(1, 9) + ), + Cell( + why="third pass of 3 requests: one row each, hidden read in place", + spec=dict(m=3, h=6144, gather=False, invalid_id=False), + ), + Cell( + why="one row with an out-of-table id: the embedding half is zero", + spec=dict(m=1, h=6144, gather=False, invalid_id=True), + ), + ) + + note = """ + The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's + summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, + also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the + drafter's embedding is the whole table, so its ids are always in range. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl + ) + return out, residual + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / + ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" + valid = (ids >= 0) & (ids < table.shape[0]) + emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) + emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) + hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] + ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) + return ref, (hid.clone() if residual is not None else None) + + def is_valid( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> None: + h = table.shape[1] + # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. + assert table.is_contiguous() and hidden.is_contiguous(), ( + "table / hidden must be dense [*, H]" + ) + if rows is not None: + assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" + else: + assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" + # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + assert residual is None or ( + residual.is_contiguous() and residual.shape == (ids.shape[0], h) + ), "residual must be contiguous [M, H]" + + +m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py new file mode 100644 index 000000000000..e322e1787230 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter MLP's plain SwiGLU as one C++ launch (modeling's ``DRAFT_SILU_CPP``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3DraftSilu(OpWrapper): + """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: + trtllm::minimax_m3_sm107_tp2_silu_and_mul (minimaxM3Sm107Tp2/siluAndMul.cu), bit for bit the Triton + ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP + runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs + the Triton op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP2 drafter MLP shard: intermediate 18432 / 2 = 9216, gate | up 18432 wide, at the drafter's decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP2 drafter MLP at 1 row", spec=dict(rows=1, d=9216)), + Cell(why="the TP2 drafter MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=9216)), + Cell(why="the TP2 drafter MLP at 8 rows", spec=dict(rows=8, d=9216)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (the same PTX instruction sequence), checked over every + bf16 gate and every bf16 up pattern. + """ + + def __call__(self, x: torch.Tensor) -> Optional[torch.Tensor]: + if ( + x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 16 != 0 + or x.data_ptr() % 16 != 0 + or x.shape[0] > 65535 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp2_silu_and_mul(x, out) + return out + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul(x) + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_draft_silu = M3DraftSilu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..708e6c5caa71 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py @@ -0,0 +1,432 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One pass of the fused greedy Eagle3 draft loop: the drafter's vocab-sharded BF16 head (streamed from its lossless +12-bit coding, ``ops.gemm.m3_target_head``, or as BF16 fragments), the cross-rank argmax and the loop's per-pass +bookkeeping in one kernel (``m3_draft_step.draft_step``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ..gemm.m3_target_head import decode as decode_coded +from .m3_shard_accept import exchange_tokens, local_argmax, make_keys + +#: Requests (head rows) per pass, and exchange rows per call half of the workspace. +MAX_TOKENS = 8 +#: 16-row weight tiles one CTA can own (the kernel's shared-memory accumulator). +MAX_TILES_PER_CTA = 32 +#: Draft length of the fused loop, and the drafter's TP2 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 100032 +HIDDEN = 6144 + + +def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the + 12-bit coded stream with its exception lists decoded.""" + if exc_off is None: + return from_fragments(wf) + n = (exc_off.numel() - 1) * 16 + return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + + +def shard_logits( + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() + + +class M3DraftStep(OpWrapper): + """One fused greedy draft pass over ``rows <= 8`` requests (one head row each). + + Head input ``x[x_rows[b]]`` (pass 0: the last accepted verify row of each request) or ``x[b]``. The head is this + rank's BF16 lm_head shard in fragment order (``m3_draft_step.bf16_fragments``, global ids ``vocab_offset ..``), + or with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base`` its lossless 12-bit coding (``wf`` = the uint8 stream + of ``m3_target_head.encode_chunked``, which must round-trip to the BF16 shard); each logit is the fp32 dot product + rounded to BF16 and the row's token is the global argmax over all ranks' + shards (ties, -0 / +0 included, to the lowest global id), exchanged through the multicast slots + ``[2][8][world]``. The last CTA then writes ``draft_out[b, step]`` and ``ids_next[b]`` (the token), ``pos_next[b] + = pos_src[x_rows[b] or b] + 1``, and advances the draft KV lengths as the in-tree draft loop does between passes + (pass 0 with ``num_acc``: ``kv_lens -= draft_len - num_acc`` and ``seq_lens = 1``; later passes: ``kv_lens += + 1``). ``out_keys`` (optional) receives each row's winning key. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling's DRAFT_STEP): draft_len 3 at the TP2 GEN worker's 1..2 requests, + # the drafter's full-vocabulary BF16 head (100032 rows per rank of 2, hidden 6144; 6252 tiles, 30 per CTA on 212 + # SMs). Pass 0 reads the verify rows' hidden (4 per request) through x_rows and rolls the KV back; the later passes + # read one row per request. BF16 fragments only: on sm_107 the 12-bit coded head is slower than the BF16 stream + # (modeling's DRAFT_HEAD_CODED is False). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"pass 0 at {n} request(s): head rows gathered from {4 * n} verify rows, KV rollback, seq_lens = 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=0, + rank=n % 2, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"pass 1 at {n} requests (b4 / b8): one row each, kv_lens + 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=1, + rank=n % 2, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ) + for n in (4, 8) + ), + Cell( + why="pass 2 at 8 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=8, + draft_len=DRAFT_LEN, + step=2, + rank=0, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + coded=False, + ), + ), + Cell( + why="pass 1 at two requests: one row each, kv_lens + 1, rank 1's ids", + spec=dict( + nreq=2, + draft_len=DRAFT_LEN, + step=1, + rank=1, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ), + Cell( + why="last pass (step 2) at one request, rank 0", + spec=dict( + nreq=1, + draft_len=DRAFT_LEN, + step=2, + rank=0, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ), + Cell( + why="pass 2 at 2 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=2, + draft_len=DRAFT_LEN, + step=2, + rank=1, + world=2, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + coded=False, + ), + ), + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit; + the coded head's weight is the BF16 shard bit for bit; its exceptions decode to +0 in the MMAs and their exact + products are added at the tile's flush) + and the argmax is the in-tree one over BF16 logits, so a token differs from the in-tree path's only where the fp32 + sum, accumulated in another order, rounds to the other side of a BF16 boundary and that decides the row's maximum. + The cells' rows are decidable and the peers' keys sit 0.25 away (or tie exactly); the edge tests (forced BF16 + ties, -0 / +0, top-2 one fp32 step apart that round to one BF16 value, adjacent BF16 values) use integer data + whose fp32 sums are exact in any order and hold the kernel to the in-tree path bit for bit. The cross-rank + exchange is certified on one device with the peers' keys written into the slots (see m3_shard_accept). + """ + + def __call__( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + bnd: Optional[torch.Tensor] = None, + bnd_cnt: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_sm107tp2_draft_step( + x, + x_rows, + rows, + wf, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + num_ctas, + pdl, + out_keys, + exc_off, + exc_pos, + exc_w, + exp_base, + tile_lo, + tile_hi, + total_ctas, + bnd, + bnd_cnt, + ) + + def reference( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + bnd: Optional[torch.Tensor] = None, + bnd_cnt: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Every buffer the pass writes, computed on copies: "logits" (this rank's BF16 logits as fp32), "local_keys" + (this rank's key per row), "keys" (the winning key), "tokens", "draft_out", "ids_next", "pos_next", "kv_lens" + and "seq_lens". + + Reads ``ws_uc``, ``counters`` and ``kv_lens`` before the call. ``keys``, ``ws_mc_ptr``, ``num_ctas``, ``pdl`` + and ``out_keys`` do not change the result. + """ + h = x[x_rows.long()] if x_rows is not None else x[:rows] + logits = shard_logits(h, wf, exc_off, exc_pos, exc_w, exp_base) + best, idx = local_argmax(logits, vocab_offset) + local_keys = make_keys(best, idx) + tokens = exchange_tokens(local_keys, ws_uc, counters, rank, world, MAX_TOKENS) + half = int(counters[1].item()) & 1 + slots = ws_uc.view(torch.int64)[ + half * MAX_TOKENS * world : (half * MAX_TOKENS + rows) * world + ] + slots = slots.view(rows, world).clone() + slots[:, rank] = local_keys + sign = torch.iinfo(torch.int64).min + win = (slots ^ sign).max(dim=1).values ^ sign + d_out = draft_out.clone() + d_out[:rows, step] = tokens + src = x_rows.long() if x_rows is not None else torch.arange(rows, device=x.device) + lens = kv_lens.clone() + sl = seq_lens.clone() if seq_lens is not None else None + if step == 0 and num_acc is not None: + lens[:rows] -= draft_len - num_acc[:rows] + sl[:rows] = 1 + elif step > 0: + lens[:rows] += 1 + out = { + "logits": logits, + "local_keys": local_keys, + "keys": win, + "tokens": tokens, + "draft_out": d_out, + "ids_next": tokens.clone(), + "pos_next": pos_src[src] + 1, + "kv_lens": lens, + } + if sl is not None: + out["seq_lens"] = sl + return out + + def is_valid( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + bnd: Optional[torch.Tensor] = None, + bnd_cnt: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[:rows].any()), "keys must start zeroed" + # Equal-chunk split: the shared-tile partials and arrival counts start (and are left) zero. + assert (bnd is None) == (bnd_cnt is None), "bnd and bnd_cnt come together" + if bnd is not None: + assert not bool(bnd.any()) and not bool(bnd_cnt.any()), ( + "bnd / bnd_cnt must start zeroed" + ) + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + # The kernel reads the head row x + x_rows[b] * K and the position pos_src[x_rows[b]] with no bound. + if x_rows is not None: + assert bool(((x_rows[:rows] >= 0) & (x_rows[:rows] < x.shape[0])).all()), ( + "x_rows must index x" + ) + assert bool((x_rows[:rows] < pos_src.shape[0]).all()), ( + "pos_src must cover the head rows" + ) + else: + assert pos_src.shape[0] >= rows, "pos_src must cover the head rows" + # Each CTA accumulates at most MAX_TILES_PER_CTA tiles in shared memory. + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + assert -(-tiles // num_ctas) <= MAX_TILES_PER_CTA, "too many weight tiles per CTA" + if exc_off is not None: + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), ( + "exc_off must end at the exception count" + ) + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + # A NaN logit maps to a key above every finite one and wins the row silently. + h = x[x_rows.long()] if x_rows is not None else x[:rows] + assert not bool(torch.isnan(h).any()), "NaN head inputs win the argmax" + if num_acc is not None: + na = num_acc[:rows] + assert bool(((na >= 1) & (na <= draft_len + 1)).all()), ( + "num_acc counts must lie in [1, draft_len + 1]" + ) + + +m3_draft_step = M3DraftStep() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..a192e89c5c76 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..m3_force_vec import force_vec_kernel, force_vec_req_kernel + + +class M3ForceVec(OpWrapper): + """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` + (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 + ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + force_vec_kernel[(1,)]( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + """The counts (a new tensor; ``counter`` and ``out`` are not written).""" + c = int(counter.item()) + 1 + slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) + idx = (c * counter_stride + slots * slot_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec = M3ForceVec() + + +class M3ForceVecReq(OpWrapper): + """For each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(keys[b] * key_stride) & pool_mask] < frac), max_total)`` into ``out`` (int32): + one draw per request at its own key (its KV length at the verify), no shared counter. ``block`` is the launch + width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="two requests, fraction 0.3", spec=dict(batch=2, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + force_vec_req_kernel[(1,)]( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + """The counts (a new tensor; ``out`` is not written).""" + k = keys[:batch_size].to(torch.int64) + r = pool[(k * key_stride) & pool_mask] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec_req = M3ForceVecReq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..34fb3387b996 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py @@ -0,0 +1,308 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampled (temperature + top-p) Eagle3 verification over the TP-sharded vocabulary in one kernel +(``m3_sampled_accept.sampled_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_shard_accept import strict_accept + +#: Verify rows per call. +MAX_ROWS = 48 +#: The fused sampling kernel's 1 / T guard. +TEMP_EPSILON = 1e-6 + + +def nucleus_probs( + logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +class M3SampledAccept(OpWrapper): + """Sampled verification of ``rows = nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard. + + Each row draws one token from its top-p nucleus of softmax(logits / T) over the full (all-rank) vocabulary by + distributed rejection sampling (Philox keyed by the row's seed / offset, the row, the attempt and the CTA), then + the strict acceptance runs against ``draft`` as in ``m3_shard_accept``: ``acc_tokens``, ``num_acc``, + ``gather_ids``, optional ``next_new[:, 0]``. ``tokens`` receives the sampled token per row; ``extra_attempts`` + accumulates exchange rounds beyond the second. The exchange goes through the multicast workspace (two call + halves of ``minimax_m3_sampled_accept_ws_words(world)`` words) with the (finished rows, call) ``counters``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3 (T = 4 rows per request) at the TP2 GEN worker's batch: 1 / 2 requests (4 / 8 rows, 8 CTAs per + # row), and the 48-row maximum (4 CTAs per row, 50 KB of staged logits per CTA); the 100032-column target shard + # (vocab 200064 / 2). Single rank (see note). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"b4 / b8 ({4 * n} rows, {n} requests): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=n, + vs=100032, + vocab_offset=100032 * (n % 2), + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ) + for n in range(3, 9) + ), + Cell( + why="b8 (32 rows, 8 CTAs per row = 256 CTAs): AgentX sampling, the draws follow the nucleus", + spec=dict( + nreq=8, + vs=100032, + vocab_offset=0, + check="distribution", + top_p=0.95, + temperature=1.0, + force="total", + ), + ), + Cell( + why="b1 (4 rows), top_p 0.8 over a 12-token head: the draws follow the nucleus distribution", + spec=dict( + nreq=1, vs=100032, vocab_offset=0, check="distribution", top_p=0.8, temperature=0.7 + ), + ), + Cell( + why="b2 (8 rows): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=2, + vs=100032, + vocab_offset=100032, + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ), + Cell( + why="b2 (8 rows) with top_p = 1 (no cut: the first draw stands), next_new and forced counts", + spec=dict( + nreq=2, + vs=100032, + vocab_offset=100032, + check="distribution", + top_p=1.0, + temperature=1.0, + force="vec", + ), + ), + Cell( + why="AgentX sampling (T 1.0, top_p 0.95) at b1 with forced counts: the draws follow the nucleus", + spec=dict( + nreq=1, + vs=100032, + vocab_offset=0, + check="distribution", + top_p=0.95, + temperature=1.0, + force="total", + ), + ), + Cell( + why="48 rows (4 CTAs per row): the second launch geometry, nucleus distribution at top_p 0.5", + spec=dict( + nreq=12, + vs=100032, + vocab_offset=0, + check="distribution", + top_p=0.5, + temperature=1.3, + force="total", + ), + ), + ) + + note = """ + What is sampled is gated statistically, the rest exactly. Every drawn token must lie in the reference nucleus; + the empirical distribution over many calls (distinct Philox offsets) must sit within a stated total-variation + band of the reference's; the acceptance outputs must equal the strict acceptance of the drawn tokens exactly. + The Philox stream itself is not reproduced in torch: the kernel's draw uses the hardware exp / log + approximations (__expf, __logf) and a CTA-order inverse CDF, which torch does not reproduce bit for bit. + + Certified at world 1 only. The decision is replicated from every rank's per-round records, which another rank + computes from its own shard on the fly, so peers cannot be stood in for on one device; the cross-rank path + needs the multi-device job and is not covered here. Rows with top-k or min-p must not reach this op (it applies + neither); is_valid states it. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept( + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + rows = logits.shape[0] + # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the + # wrong rows. + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + # Per-row parameters are read at the row index, unchecked. + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), ( + "per-row parameters must cover every row" + ) + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + # The finished-row counter elects the CTA that runs the acceptance. + assert int(counters[0].item()) == 0, ( + "counters[0] (finished rows) must be zero between calls" + ) + + +m3_sampled_accept = M3SampledAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..13ec40116191 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept_ws_words.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Size of the sampled-verify exchange workspace (``m3_sampled_accept.SampledAcceptWorkspace.create``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Exchange rounds per row, verify rows per call, CTAs per row (at most), uint64 words per record. +MAX_ATTEMPTS = 16 +MAX_ROWS = 48 +MAX_SPLITS = 8 +RECORD_WORDS = 4 + + +class M3SampledAcceptWsWords(OpWrapper): + """uint64 words of the multicast workspace ``minimax_m3_sm107tp2_sampled_accept`` exchanges through for + ``world`` ranks: two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="the tp2 target's workspace", spec=dict(world=2)), + Cell(why="a single rank (local workspace, no multicast)", spec=dict(world=1)), + ) + + note = """ + A host computation (no device work); it is an op so the workspace size and the kernel's indexing come from one + definition. + """ + + def __call__(self, world: int) -> int: + return torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept_ws_words(world) + + def reference(self, world: int) -> int: + """Two halves of attempts x rows x (world x splits) records x record words.""" + return 2 * MAX_ATTEMPTS * MAX_ROWS * world * MAX_SPLITS * RECORD_WORDS + + +m3_sampled_accept_ws_words = M3SampledAcceptWsWords() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..408759009472 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py @@ -0,0 +1,342 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Greedy Eagle3 verification from this rank's logits shard: per-row argmax, the cross-rank (value, index) exchange +and the strict acceptance in one kernel (``m3_shard_accept.shard_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Verify rows per call, and exchange rows per call half of its workspace. +SHARD_ROWS = 48 + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out + + +class M3ShardAccept(OpWrapper): + """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` + (global ids ``vocab_offset ..``). + + Each row's token is its global argmax over all ranks' shards: keys ``ordered(value) << 32 | ~index`` (ties take + the lowest index, as ``torch.argmax``) are max-reduced over the shard, pushed into every rank's exchange slot + through the multicast slots ``[2][48][world]`` (Lamport: 0 = empty) and max-reduced over the ranks by the last + CTA, which then runs the strict acceptance: ``acc_tokens [nreq, D + 1]``, ``num_acc [nreq]`` (``force_vec`` / + ``force_total`` override the counts), ``gather_ids [nreq]`` (the first draft pass's row per request) and, when + given, ``next_new[:, 0]`` (the last accepted token). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every greedy verify step of the target (modeling's SHARD_ACCEPT): draft_len 3 at the TP2 GEN worker's 1..2 + # requests (4..8 rows), plus the 48-row maximum; the target lm_head's BF16 shard logits (100032 columns per rank + # at TP2). The logits come from the head GEMM, so the op only reduces. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"b4 / b8: {4 * n} verify rows ({n} requests), rank {n % 2} of 2", + spec=dict( + nreq=n, draft_len=3, vs=100032, world=2, rank=n % 2, force="none", strided=False + ), + ) + for n in range(3, 9) + ), + Cell( + why="b1: 4 verify rows (1 request x T = 4), rank 0 of 2, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=100032, world=2, rank=0, force="none", strided=False), + ), + Cell( + why="b2: 8 verify rows, rank 1, global ids offset by 100032", + spec=dict(nreq=2, draft_len=3, vs=100032, world=2, rank=1, force="none", strided=False), + ), + Cell( + why="8 rows (2 requests) read from a row-strided view of a wider logits buffer, forced counts", + spec=dict(nreq=2, draft_len=3, vs=100032, world=2, rank=0, force="vec", strided=True), + ), + # W2 (WIDE_SPEC_TAIL): 13..48 requests (52..192 verify rows) as whole-request chunks of 12 (the 48-row + # geometry), each chunk a launch on its own rows; world 1 (the chunked exchange is the per-launch one above). + *( + Cell( + why=f"wide: {4 * n} verify rows ({n} requests) in {-(-n // 12)} chunks, world 1", + spec=dict( + nreq=n, + draft_len=3, + vs=100032, + world=1, + rank=0, + force="vec" if n % 3 == 0 else "none", + strided=n % 5 == 0, + ties=n % 2 == 0, + ), + ) + for n in range(13, 49) + ), + # W2b: the drafter's greedy top-1 (draft length 0: one row per request, token = the global argmax), 1..48 + # requests in one launch and 49..56 in two (spec_worker.greedy_sample_draft_with_tp_gather). + *( + Cell( + why=f"draft top-1: {n} rows (draft length 0), world 1", + spec=dict( + nreq=n, draft_len=0, vs=100032, world=1, rank=0, force="none", strided=n == 20 + ), + ) + for n in (1, 8, 20, 46, 48, 56) + ), + # 9..11 requests (36..44 rows): one launch each (the 12-request cell below is the 48-row one). + *( + Cell( + why=f"wide: {4 * n} verify rows ({n} requests), one launch, world 1", + spec=dict( + nreq=n, draft_len=3, vs=100032, world=1, rank=0, force="none", strided=False + ), + ) + for n in range(9, 12) + ), + Cell( + why="the 48-row maximum (12 requests at draft 3), force_total, rows with exact cross-rank ties", + spec=dict( + nreq=12, + draft_len=3, + vs=100032, + world=2, + rank=1, + force="total", + strided=False, + ties=True, + ), + ), + ) + + note = """ + The op reduces BF16 logits exactly, so the reference is bit-exact, including cross-rank ties (lowest global + index wins). The cross-rank exchange is certified on one device: the other ranks' keys are written into this + rank's slots before the call, which is what their multicast pushes would have put there; the multicast mapping + itself (ws_mc_ptr = this rank's slots, as at world 1) is not exercised across devices here. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + rows = logits.shape[0] + if rows > SHARD_ROWS: + # W2: whole requests per launch; batch_idx is the identity layout (is_valid), so each chunk passes its + # local rows and its gather ids are shifted back to the full verify layout. + k1 = draft.shape[1] + 1 + per = SHARD_ROWS // k1 + nreq = draft.shape[0] + local = torch.arange(per, dtype=torch.int32, device=logits.device) + for c0 in range(0, nreq, per): + c1 = min(nreq, c0 + per) + self( + logits[c0 * k1 : c1 * k1], vocab_offset, keys, counters, ws_uc, ws_mc_ptr, rank, world, + draft[c0:c1].contiguous(), local[: c1 - c0], acc_tokens[c0:c1], num_acc[c0:c1], + gather_ids[c0:c1], force_total, None if next_new is None else next_new[c0:c1], num_ctas, + None if force_vec is None else force_vec[c0:c1], + ) # fmt: skip + gather_ids[c0:c1] += c0 * k1 + return + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + num_ctas, + force_vec, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Lowest-index maximum of each row's BF16 logits, the largest key among it and the peers' keys in + ``ws_uc``, then the strict acceptance. Reads ``ws_uc`` / ``counters`` before the call; the scratch and + output buffers and ``num_ctas`` do not change the result.""" + best, idx = local_argmax(logits.float(), vocab_offset) + if world == 1: + tokens = idx.to( + torch.int32 + ) # no peers: the exchange returns this rank's own pick (any row count) + else: + tokens = exchange_tokens(make_keys(best, idx), ws_uc, counters, rank, world, SHARD_ROWS) + out = strict_accept(tokens, draft, batch_idx, force_total, next_new, force_vec) + out["tokens"] = tokens + return out + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[: logits.shape[0]].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + assert ws_uc.numel() >= 2 * SHARD_ROWS * world, "ws_uc must hold [2][48][world] slots" + # A NaN logit maps to a key above every finite one and wins the row silently. + assert not bool(torch.isnan(logits).any()), "NaN logits win the argmax" + nreq = draft.shape[0] + assert bool(((batch_idx[:nreq] >= 0) & (batch_idx[:nreq] < nreq)).all()), ( + "batch_idx must index the nreq requests" + ) + if logits.shape[0] > SHARD_ROWS: + # The chunked (W2) call relocates each chunk's rows: only the identity layout maps back. + assert torch.equal(batch_idx[:nreq].cpu(), torch.arange(nreq, dtype=batch_idx.dtype)), ( + "a chunked verify needs the identity batch_idx layout" + ) + if force_vec is not None: + fv = force_vec[:nreq] + assert bool(((fv >= 1) & (fv <= draft.shape[1] + 1)).all()), ( + "force_vec counts must lie in [1, D + 1]" + ) + + +m3_shard_accept = M3ShardAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..ff07c54de1f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> +``m3_spec_tail.slot_scatter``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: + """src rows zero-padded (or cut) to ``width`` columns.""" + out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) + w = min(width, src.shape[1]) + out[:, :w] = src[:, :w] + return out + + +class M3SlotScatter(OpWrapper): + """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, + ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The + destinations are the store's transposed [slots, width] views.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of + # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) + # is zero-padded. + CELLS: tuple[Cell, ...] = ( + Cell( + why="8 rows into shuffled slots, widths equal to the store's", + spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), + ), + Cell( + why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", + spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), + ), + ) + + note = """ + Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. + """ + + def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): + torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) + + def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): + """The four destinations after the call, on copies.""" + s = slots.long() + out = {} + for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): + r = dst.clone() + r[s] = padded_rows(src[: s.numel()], dst.shape[1]) + out[name] = r + lr = dl.clone() + lr[s] = sl[: s.numel()] + out["dl"] = lr + return out + + def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): + assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( + "slots must index the store" + ) + # Two rows into one slot race (which one lands is unordered). + assert slots.unique().numel() == slots.numel(), "slots must be distinct" + assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + + +m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py new file mode 100644 index 000000000000..8925fc8ab5b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's speculative step correction as one launch (``RuntimeHooks.spec_offsets_op``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger + +# The engagement line is logged once per process (module state: __call__ reads no instance state). +_ENGAGED: list = [] + + +class M3SpecOffsets(OpWrapper): + """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp2_spec_offsets (minimaxM3Sm107Tp2/specOffsets.cu), the + same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other + layouts (the caller runs the ATen adds).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="1 request, 4 tokens (draft length 3)", spec=dict(tokens=4, seqs=1)), + Cell(why="8 requests, 32 tokens", spec=dict(tokens=32, seqs=8)), + Cell(why="wrap-around values", spec=dict(tokens=8, seqs=2, wrap=True)), + ) + + note = "Integer adds: the reference (the two in-place ATen adds) is exact." + + def __call__(self, pos, pos_off, kv, kv_off) -> bool: + ts = (pos, pos_off, kv, kv_off) + if ( + any( + t.dtype != torch.int32 or t.dim() != 1 or t.stride(0) != 1 or not t.is_cuda + for t in ts + ) + or pos.numel() != pos_off.numel() + or kv.numel() != kv_off.numel() + ): + return False + torch.ops.trtllm.minimax_m3_sm107_tp2_spec_offsets(pos, pos_off, kv, kv_off) + if not _ENGAGED: + _ENGAGED.append(True) + logger.info("[M3 host] spec step correction: one launch taken (spec_offsets_op)") + return True + + def reference(self, pos, pos_off, kv, kv_off) -> None: + pos += pos_off + kv += kv_off + + def is_valid(self, pos, pos_off, kv, kv_off) -> None: + assert pos.numel() == pos_off.numel() and kv.numel() == kv_off.numel(), ( + "equal lengths per pair" + ) + + +m3_spec_offsets = M3SpecOffsets() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..a0da8b337574 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch +(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS: tuple[Cell, ...] = ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + Cell( + why="an unchanged batch of 3 previous requests (c3) in shuffled slots", + spec=dict(p=3, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="an unchanged batch of 4 previous requests (c4) in shuffled slots", + spec=dict(p=4, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + ) + + note = """ + Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid + states them. + """ + + def __call__( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ) + + def reference( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" + nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) + s = slots[:P].long() + ids_r, draft_r = ids.clone(), draft.clone() + ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) + draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) + pos_r = torch.zeros_like(pos_off) + pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] + kv_r = torch.zeros_like(kv_off) + kv_r[b0 : b0 + P] = lens[s] - T + return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} + + def is_valid( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or + # writes another buffer. + s = slots[:P] + assert bool( + ( + (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) + ).all() + ), "slots must index the store" + assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( + "ppi must index lens" + ) + assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" + assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" + assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( + "offset buffers too short" + ) + + +m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py new file mode 100644 index 000000000000..080c3018e4b6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode graph's step-input prologue in one launch (this target's ``RuntimeHooks.pack_prologue_op``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepPrologue(OpWrapper): + """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch + (trtllm::minimax_m3_sm107_tp2_step_prologue, minimaxM3Sm107Tp2/stepPrologue.cu). + + ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one + ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the + window, a device copy's a device address). The launch writes up to ``max_desc`` copies of each kind, every host + copy before any device copy, then zeroes the header; with ``epoch`` it increments ``epoch[0]`` once. The default + prologue does the same with two parallel-scatter launches (``step_inputs._graph_scatter_kernel``) and a header + fill.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP2 decode step's packs: a few dozen host copies of 4 to a few hundred bytes (positions, KV lengths, slots, + # the spec sampler's uploads), a handful of device copies (the graph runner's static inputs), some reading what a + # host copy of the same window wrote. + CELLS: tuple[Cell, ...] = ( + Cell(why="an empty window (a replay with no new upload)", spec=dict(h2d=0, d2d=0)), + Cell(why="host copies only, odd sizes and alignments", spec=dict(h2d=40, d2d=0)), + Cell( + why="host copies feeding device copies (a device copy reads a host copy's destination)", + spec=dict(h2d=24, d2d=8, chained=True), + ), + Cell( + why="every descriptor slot (128 + 128 of the pack's 256)", spec=dict(h2d=128, d2d=128) + ), + ) + + note = """ + Byte copies: the reference (the default prologue's Triton launches and fill on the same window) is exact. + Destinations are raw addresses; is_valid cannot check them. + """ + + def __call__(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + torch.ops.trtllm.minimax_m3_sm107_tp2_step_prologue(window, epoch, max_desc) + + def reference(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + """The default prologue on the same arguments: it writes the same destinations.""" + from tensorrt_llm._torch.pyexecutor.step_inputs import _graph_scatter_kernel + + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + bump = epoch is not None and not device_src + _graph_scatter_kernel[(max_desc,)]( + window, + stage_i64, + epoch if bump else stage_i64, + DEVICE_SRC=device_src, + BLOCK=256, + BUMP_EPOCH=bump, + num_warps=1, + ) + window[:16].zero_() + + def is_valid(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + assert window.dtype == torch.uint8 and window.is_contiguous(), "window: contiguous uint8" + assert window.numel() >= 16 + 24 * max_desc, "window: header and max_desc descriptors" + counts = window[:16].view(torch.int64).tolist() + assert counts[0] + counts[1] <= max_desc, "at most max_desc copies in all" + + +m3_step_prologue = M3StepPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py new file mode 100644 index 000000000000..b383cd4036a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py @@ -0,0 +1,1747 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +GB300/SM103 tuning targets GQA group 16, top-k 16 and four-token Eagle3 +verification. Split-K, query pairing and launch tables require four KV heads; +the linear-softmax policy also covers one and two KV heads. Other hardware +and shapes retain the default policies, including supported SM100 NVFP4. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* A packed NVFP4 cache is read in place, dequantized in-register from its E2M1 + nibbles and E4M3 block scales. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* A single top-k chunk writes the final output directly, avoiding scratch and + the otherwise redundant LSE merge. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. + +This is the MiniMax-M3 NVFP4 / sm_107 / tp2 target's copy. On top of the +in-tree decode it routes long-context TP4 steps to the one-launch +``m3_attn_lc`` kernel, short-sequence verify steps of at least +``MQ_MIN_TOKENS`` rows to the multi-query producer, releases PDL dependents +early, and has the merge emit the MXFP8 o_proj hand-off (``m3_prequant.SLOT``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + check_decode_span_shape, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +from . import m3_attn_lc, m3_prequant, m3_report + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +_TL_DTYPES = { + torch.bfloat16: tl.bfloat16, + torch.float16: tl.float16, + torch.float32: tl.float32, +} + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size + +# NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale +# byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. +NVFP4_SF_VEC_SIZE = 16 + + +# Fewest query rows for which the merge writes the MXFP8 o_proj hand-off: every decode step, so the MXFP8 o_proj +# never runs its own quantize kernel after attention. +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# The sparse decode (LC2, the Triton merge) also writes the MXFP8 o_proj hand-off (bit for bit +# torch.ops.trtllm.mxfp8_quantize of its BF16 rows) and publishes it to m3_prequant.SLOT. +MERGE_MXFP8_HANDOFF = True +# The long-context decode (m3_attn_lc) stages the request's block-table row before its PDL grid wait: off (measured +# +1.6..+1.8 us at 32 rows; flat to -0.2 us below), up to LC_PREWAIT_MAX_ROWS query rows when on. +LC_PREWAIT = False +LC_PREWAIT_MAX_ROWS = 16 # capacity: fixed: measured occupancy crossover +# The long-context decode's shapes (modeling's SPARSE_DECODE_LC keys: "mxfp8" = a step of at most MergeMxfp8.MAX_ROWS +# rows) and the identity-selection producer's (modeling's SPARSE_DECODE_PRODUCER keys), with their report paths. +LC_KEYS = frozenset({"mxfp8"}) +LC_PATH = "m3_sparse_decode_lc" +PRODUCER_KEYS = frozenset({"per_row"}) +PRODUCER_PATH = "sparse_decode_producer" + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and + swizzled UE8M0 scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips + its own quantization. One buffer of ``MAX_ROWS`` rows (one 128-row scale tile) per width serves every layer: each + o_proj consumes it before the next layer's merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + @staticmethod + def handoff(mxfp8, rows: int): + """The hand-off of a ``rows``-row step: its E4M3 rows and the 128-row scale tiles that cover them. The o_proj + GEMM checks the swizzled scales against the step's padded rows, so the step gets the tiles that cover its rows, + never the whole buffer.""" + q8, sf = mxfp8 + width = q8.shape[-1] + return q8[:rows], sf[: (rows + 127) // 128 * 128 * (width // 32)] + + +MERGE_MXFP8 = MergeMxfp8() + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + if IDENTITY: + # Slot `it` of the identity selection, bounded by this row's own kv extent above. + blk = it + else: + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, + # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + if IDENTITY: + blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) + blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) + else: + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for cand in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + if IDENTITY: + blk_0 = tl.where(valid_0, idx_0, sentinel) + blk_1 = tl.where(valid_1, idx_1, sentinel) + else: + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the + # pool, reads nothing. + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. Because this trigger precedes this grid's own + # wait, it says nothing about what upstream has finished: a dependent may read only step-invariant data + # (weights) before its griddepcontrol.wait, never this step's attention output, q, top-k table or lengths. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) +# Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) + + +# Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query +# tokens; below it the per-token producer is as fast. +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover +# Warps per multi-query producer program. +MQ_NUM_WARPS = 8 + + +def _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, +) -> None: + """minimax_m3_sparse_attn_decode through m3_attn_lc: same table, output and MXFP8 hand-off contract.""" + use_pdl = get_env_enable_pdl() + total_q, num_heads, head_dim = q.shape + prewait = bool(LC_PREWAIT) and total_q <= LC_PREWAIT_MAX_ROWS + mxfp8 = None + m3_prequant.SLOT.clear() + if MERGE_MXFP8_HANDOFF and MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(num_heads * head_dim, output.device) + # The early trigger lets the o_proj launch and stream its weights while attention runs. + m3_attn_lc.m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + linear_softmax=False, + mxfp8=mxfp8, + use_pdl=use_pdl, + prewait=prewait, + n1_ids=m3_attn_lc.N1_IDS.take(), + ) + if prewait: + m3_report.taken(LC_PREWAIT_PATH, total_q) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), *MergeMxfp8.handoff(mxfp8, total_q) + ) + + +# Engagement path of the identity-selection producers reading their blocks from the kv lengths (IDENTITY). +DEVICE_IDENTITY_PATH = "sparse_decode_device_identity" +# Engagement path of LC2's pre-wait block-table staging (LC_PREWAIT, up to its row cap). +LC_PREWAIT_PATH = "sparse_decode_lc_prewait" + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + q may be FP8, as a fused QK-norm/RoPE producer emits it. The kernel widens + it to output.dtype in-register rather than making the caller materialize a + widened copy, which would cost a standalone kernel on every sparse layer. + The attention math is unchanged either way, since E4M3 widens exactly. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. + + Passing the four NVFP4 scale tensors selects the packed-NVFP4 cache instead: + k_paged/v_paged then hold two E2M1 elements per byte, the block scales are + paged E4M3 bytes in the layouts msa_scatter writes, and the two per-tensor + dequant scales are single-element FP32 device tensors. + + num_topk_chunks overrides the split-K factor and exists for tests: the + merged result must not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) + nvfp4_args = (k_block_scale, v_block_scale, k_global_scale, v_global_scale) + kv_nvfp4 = any(arg is not None for arg in nvfp4_args) + scale_cols = head_dim // NVFP4_SF_VEC_SIZE + if kv_nvfp4: + capability = torch.cuda.get_device_capability(q.device) + if capability not in _SM100F_CAPABILITIES: + raise NotImplementedError( + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; " + f"got SM{capability[0] * 10 + capability[1]}." + ) + _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) + # Two E2M1 elements per byte; Triton must see them unsigned so that the + # high-nibble shift does not sign-extend. + k_paged, v_paged = k_paged.view(torch.uint8), v_paged.view(torch.uint8) + k_block_scale, v_block_scale = ( + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused arguments. + scale_arg = kv_scale if use_scale else output + k_sc_arg, v_sc_arg = (k_block_scale, v_block_scale) if kv_nvfp4 else (output, output) + k_gs_arg, v_gs_arg = (k_global_scale, v_global_scale) if kv_nvfp4 else (output, output) + + widen_q = q.dtype in _FP8_DTYPES + if widen_q and output.dtype not in _TL_DTYPES: + raise ValueError( + f"MiniMax-M3 sparse decode cannot widen FP8 q into {output.dtype}; " + f"supported compute dtypes are {sorted(d.__str__() for d in _TL_DTYPES)}." + ) + # Packed NVFP4 dequantizes to FP16. The production producer supplies FP8 + # q, which widens exactly to FP16, so keep both dot operands there instead + # of converting every K/V tile to BF16. Accumulation remains FP32. + if widen_q: + compute_dtype = tl.float16 if kv_nvfp4 else _TL_DTYPES[output.dtype] + else: + compute_dtype = tl.float32 + + # Long-context TP2 steps: attention, split-K merge and MXFP8 hand-off in one launch. A selected table on an NVFP4 + # cache with FP8 q at the TP2 heads is the kernel's (modeling's SPARSE_DECODE_LC); operands it declines are a miss. + lc_key = "mxfp8" if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS else "plain" + lc_covered = ( + kv_nvfp4 + and widen_q + and not identity_blocks + and num_topk_chunks is None + and num_heads == 32 + and num_kv_heads == 2 + and lc_key in LC_KEYS + ) + if lc_covered and not m3_attn_lc.lc_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ): + m3_report.missed( + LC_PATH, + total_q, + "the step's q / KV pages / scales / top-k table / output are not the layout the kernel is built for", + ) + lc_covered = False + if lc_covered: + _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + ) + return + + query_group_size = ( + _sm100f_nvfp4_query_group_size( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else 1 + ) + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = ( + _sm100f_nvfp4_num_topk_chunks( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else None + ) + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one + # program per (request, block) covers all of them and dequantizes each block once. + use_mq = ( + identity_blocks + and total_q >= MQ_MIN_TOKENS + and kv_nvfp4 + and num_kv_heads == 1 + and decode_query_len > 1 + and head_dim == 128 + and gqa_group_size * decode_query_len <= 128 + ) + if identity_blocks and kv_nvfp4 and widen_q: + # The identity-selection route is modeling's SPARSE_DECODE_PRODUCER ("per_row"; "multi_query" needs one KV + # head, so it never runs at the TP2 heads); a variant the range does not certify is a miss (raises under + # require). + producer_key = "multi_query" if use_mq else "per_row" + if producer_key not in PRODUCER_KEYS: + m3_report.missed( + PRODUCER_PATH, + total_q, + f"the identity-selection producer variant {producer_key!r} is not certified at {num_kv_heads} KV heads", + ) + if use_mq: + num_topk_chunks = max_topk + direct_output = num_topk_chunks == 1 + if direct_output: + # With one chunk, acc_o is already normalized and final. Point the + # producer at the caller's output and compile away the LSE store. + # Keeping a real dummy pointer preserves Triton's pointer contract. + o_arg = output + lse_arg = output + o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) + lse_strides = (0, 0, 0) + else: + # Persistent arena rather than torch.empty, so the partials keep one + # address across CUDA graph replays. FP32 keeps the split-K factor from + # perturbing the merged result. + reserve = torch.cuda.is_current_stream_capturing() + o_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + o_strides = tuple(o_arg.stride()) + lse_strides = tuple(lse_arg.stride()) + + use_pdl = get_env_enable_pdl() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = ( + _sm100f_nvfp4_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + if total_q % query_group_size or decode_query_len % query_group_size: + raise ValueError( + f"query group {query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." + ) + + producer_kernel = ( + _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel + ) + device_identity = identity_blocks + if device_identity and not use_mq: + m3_report.taken(DEVICE_IDENTITY_PATH, total_q) + single_kernel_args = ( + { + # No split-K arrival counters: the separate merge kernel combines the chunks. + "merge": (output, None), + "tail": (output.stride(0), output.stride(1), output.stride(2), num_kv_heads), + # An early release pays off when the dependent is the consumer of the + # final output (it can prefetch), or a merge grid that itself releases + # its dependent at entry: the o_proj then launches during this kernel and + # streams its weights while attention runs. Above 8 query rows the + # concurrent weight stream slows the larger attention grid more than the + # prefetch saves (GB300 TP4: -1.1% at 48 rows). + "options": {"FUSED_MERGE": False, "EARLY_TRIGGER": direct_output or total_q <= 8}, + } + if query_group_size == 1 + else {"merge": (), "tail": (), "options": {}} + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + else {} + ) + if use_mq: + from .triton_sparse_decode_mq import launch_mq_producer + + launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + sm_scale=sm_scale, + decode_query_len=decode_query_len, + use_pdl=use_pdl, + early_trigger=False, + num_warps=MQ_NUM_WARPS, + ) + else: + producer_kernel[ + (total_q // query_group_size) * num_topk_chunks, + num_kv_heads, + ]( + q, + k_paged, + v_paged, + scale_arg, + k_sc_arg, + v_sc_arg, + k_gs_arg, + v_gs_arg, + topk_idx, + o_arg, + lse_arg, + *single_kernel_args["merge"], + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + k_sc_arg.stride(0), + k_sc_arg.stride(1), + v_sc_arg.stride(0), + v_sc_arg.stride(1), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + *o_strides, + *lse_strides, + block_table.stride(0), + block_table.shape[1], + k_paged.shape[0], + *single_kernel_args["tail"], + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + WIDEN_Q=widen_q, + COMPUTE_DTYPE=compute_dtype, + USE_PDL=use_pdl, + KV_NVFP4=kv_nvfp4, + USE_LINEAR_SOFTMAX=use_linear_softmax, + SCALE_COLS=scale_cols, + DIRECT_OUTPUT=direct_output, + # Identity selection from the lengths this kernel reads, not the staged table. + IDENTITY=device_identity, + **single_kernel_args["options"], + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + m3_prequant.SLOT.clear() + if not direct_output: + width = num_heads * head_dim + mxfp8 = ( + MERGE_MXFP8.buffers(width, output.device) + if MERGE_MXFP8_HANDOFF + and MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS + and head_dim == 128 + and output.dim() == 3 + and output.is_contiguous() + and output.dtype == torch.bfloat16 + else None + ) + mxfp8_args = ( + { + "q8_ptr": mxfp8[0], + "sf_ptr": mxfp8[1], + "num_sf_k_tiles": width // 128, + "MXFP8_OUT": True, + } + if mxfp8 is not None + else {} + ) + merge_launch_options = ( + _sm100f_nvfp4_merge_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_arg, + lse_arg, + output, + head_dim, + *o_strides, + *lse_strides, + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + # Release the o_proj at the merge's entry so it streams its weights during the merge. + EARLY_TRIGGER=True, + **mxfp8_args, + **pdl_launch, + **merge_launch_options, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, width), *MergeMxfp8.handoff(mxfp8, total_q) + ) + + +__all__ = [ + "MERGE_MXFP8", + "MERGE_MXFP8_MIN_TOKENS", + "MergeMxfp8", + "NVFP4_SF_VEC_SIZE", + "SPARSE_BLOCK_SIZE", + "_sm100f_nvfp4_merge_launch_options", + "_sm100f_nvfp4_num_topk_chunks", + "_sm100f_nvfp4_query_group_size", + "_sm100f_nvfp4_use_linear_softmax", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py new file mode 100644 index 000000000000..048d31b82d68 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multi-query block-sparse decode producer for MiniMax-M3 short-sequence verify steps. + +In short-sequence mode every request attends at most ``topk`` KV blocks and block selection is the identity, so +all ``decode_query_len`` query tokens of a request read the same blocks. The per-token producer +(``_gqa_sparse_decode_kernel``) then dequantizes every block once per query token. This producer runs one program +per (request, block) for all of the request's query tokens: rows are (token, head) pairs, each block is +dequantized once, and the causal limit is a per-row mask. + +It writes the same per-block partials (output scaled by 1/l and the V global scale, log2 LSE) as the per-token +producer with one block per chunk, so ``_merge_topk_attn_out_kernel`` merges them unchanged. +""" + +import triton +import triton.language as tl + +from .triton_sparse_decode import _dequant_nvfp4_rows + + +@triton.jit( + do_not_specialize=["num_reqs"], do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"] +) +def _gqa_sparse_decode_mq_kernel( + q_ptr, # [num_reqs * QLEN, num_heads, head_dim], FP8 or FP16/BF16 + k_ptr, # packed NVFP4 K: [num_pages, 1, page_size, head_dim / 2] uint8 + v_ptr, # packed NVFP4 V: same layout + k_sc_ptr, # K block scales [num_pages, 1, page_size, head_dim / 16], token-linear + v_sc_ptr, # V block scales, 4x4 token-quad order + k_gs_ptr, + v_gs_ptr, + t_ptr, # topk table [1, total_q, topk]; the request's last token's row is read + o_ptr, # partials [topk, total_q, num_heads, head_dim] fp32 + lse_ptr, # partials [topk, total_q, num_heads] fp32 (log2) + block_table_ptr, + seq_lens, + num_reqs, + sm_scale, + stride_qn, + stride_qh, + stride_k_blk, + stride_k_pos, + stride_v_blk, + stride_v_pos, + stride_ksc_blk, + stride_vsc_blk, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + QLEN: tl.constexpr, + GROUP: tl.constexpr, # query heads per KV head (16) + BLOCK_R: tl.constexpr, # next power of two of QLEN * GROUP + HEAD_DIM: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SCALE_COLS: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr, +): + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) # top-k position == block index under identity selection + sm_scale_log2e = sm_scale * 1.4426950409 * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + + off_r = tl.arange(0, BLOCK_R) + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, HEAD_DIM) + off_dp = tl.arange(0, HEAD_DIM // 2) + off_s = tl.arange(0, SCALE_COLS) + row_t = off_r // GROUP + row_h = off_r % GROUP + r_mask = off_r < QLEN * GROUP + tok = pid_r * QLEN + row_t + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + pid_r) + # Token t of the request sits at seq_len - QLEN + t and attends positions [0, kv_len_t). + kv_len = tl.maximum(seq_len - QLEN + row_t + 1, 0) + blk = tl.load(t_ptr + (pid_r * QLEN + QLEN - 1) * stride_tn + pid_c * stride_tk) + blk_ok = blk >= 0 + blk = tl.where(blk_ok, blk, 0) + page = tl.load(block_table_ptr + pid_r * stride_bt_b + blk, mask=blk_ok, other=0).to(tl.int64) + # Rows past the last token's position are never read: their cache bytes may be stale (0 * NaN in the PV dot). + n_valid = tl.where( + blk_ok, tl.minimum(tl.maximum(seq_len - blk * BLOCK_SIZE_K, 0), BLOCK_SIZE_K), 0 + ) + n_mask = off_n < n_valid + + q = tl.load( + q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :], + mask=r_mask[:, None], + other=0.0, + ).to(tl.float16) + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + page * stride_k_blk + off_n[:, None] * stride_k_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + page * stride_ksc_blk + off_n[:, None] * SCALE_COLS + off_s[None, :], + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ) + valid = ( + (blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None]) & n_mask[None, :] & r_mask[:, None] + ) + qk = tl.dot(q, tl.trans(k_deq).to(tl.float16)) * sm_scale_log2e + qk = tl.where(valid, qk, float("-inf")) + m = tl.max(qk, axis=1) + m_safe = tl.where(m == float("-inf"), 0.0, m) + p = tl.exp2(qk - m_safe[:, None]) + p_sum = tl.sum(p, axis=1) + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + page * stride_v_blk + off_n[:, None] * stride_v_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ).to(tl.float16) + acc = tl.dot(p.to(tl.float16), v) + # An empty partial stores zero with LSE -inf, so the merge gives it weight 0. + scale = tl.where(p_sum > 0.0, 1.0 / p_sum, 0.0) * v_global_scale + lse = tl.where(p_sum > 0.0, m + tl.log2(p_sum), float("-inf")) + acc = acc * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + tok[:, None] * stride_o_b + + row_h[:, None] * stride_o_h + + off_d[None, :], + acc, + mask=r_mask[:, None], + ) + tl.store(lse_ptr + pid_c * stride_l_c + tok * stride_l_b + row_h * stride_l_h, lse, mask=r_mask) + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + *, + sm_scale, + decode_query_len, + use_pdl, + early_trigger, + num_warps=8, + num_stages=1, +): + """Launch the multi-query producer; arguments as minimax_m3_sparse_attn_decode after its NVFP4 views.""" + total_q, num_heads, head_dim = q.shape + num_reqs = total_q // decode_query_len + max_topk = int(topk_idx.shape[-1]) + rows = decode_query_len * num_heads + pdl_launch = {"launch_pdl": True} if use_pdl else {} + _gqa_sparse_decode_mq_kernel[(num_reqs, max_topk)]( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + o_arg, + lse_arg, + block_table, + seq_lens, + num_reqs, + sm_scale, + q.stride(0), + q.stride(1), + k_paged.stride(0), + k_paged.stride(2), + v_paged.stride(0), + v_paged.stride(2), + k_block_scale.stride(0), + v_block_scale.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + o_arg.stride(0), + o_arg.stride(1), + o_arg.stride(2), + lse_arg.stride(0), + lse_arg.stride(1), + lse_arg.stride(2), + block_table.stride(0), + QLEN=decode_query_len, + GROUP=num_heads, + BLOCK_R=triton.next_power_of_2(rows), + HEAD_DIM=head_dim, + BLOCK_SIZE_K=int(k_paged.shape[2]), + SCALE_COLS=head_dim // 16, + USE_PDL=use_pdl, + EARLY_TRIGGER=early_trigger, + num_warps=num_warps, + num_stages=num_stages, + **pdl_launch, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/weights.py new file mode 100644 index 000000000000..fa0bc1ba7208 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/weights.py @@ -0,0 +1,865 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_107 / tp2. + +Two tables, one per checkpoint the target reads: + +* ``MODEL_WEIGHTS`` (``MiniMaxM3Tp2Weights``): the vision-language checkpoint's language model, + ``language_model.*``. Every rank is handed the whole checkpoint, so every tensor-parallel split is this + table's: the Q heads, the K/V/index-Q heads (two each per rank), the dense MLPs and the shared expert's + intermediate, the o_proj / down-projection input columns, every routed expert's intermediate (all 128 + experts live on every rank), and the vocabulary of the shell's ``lm_head``. The index-K head is replicated. +* ``DRAFT_WEIGHTS`` (``MiniMaxM3Tp2DraftWeights``): the one-layer BF16 Eagle3 drafter's own checkpoint + (``speculative_model``), loaded through the shell's ``load_draft_weights`` when speculation is on. + +Storage is HF ``[out, in]`` row-major. Five families need a transform, each a layout a kernel demands: + +* **MXFP8 block scales** (every attention projection, the dense MLPs and the shared expert): the checkpoint + stores UE8M0 bytes ``[N, K / 32]``; the GEMMs read the 128x4-swizzled order + (``trtllm::block_scale_interleave``). A fused projection concatenates its shards' scales first and swizzles + once. +* **The five-way sparse-layer projection** ``[Q | K | V | index-Q | index-K]`` and the three-way dense one + ``[Q | K | V]``: the per-rank shards concatenated along N. +* **Gated MLPs** (dense and shared expert): ``[gate ; up]`` rows. +* **NVFP4 routed experts** (layers 3..59): ``[up (w3) ; gate (w1)]`` rows, then the trtllm-gen row shuffle of + the weight and of its E4M3 scales (gated-activation reorder + 128-row epilogue tile shuffle, + ``trtllmgen_maybe_get_cached_*_permute_indices``) and the 128x4 swizzle of the scales. +* **Gemma boundary norms**: ``input_layernorm``, ``post_attention_layernorm`` and ``model.norm`` are stored as + Gemma norms, ``(1 + w) * x``; the fused all-reduce + RMSNorm kernels apply ``w * x``, so ``1`` is added + (in BF16) at load. The per-head Q/K/index norms stay Gemma: their kernels apply the offset themselves. + +The routed experts' per-tensor scalars (input scales and second-level weight scales, every expert) load raw; +``modeling`` derives the kernels' global scales, alphas and activation constants from them once the weights +are real. + +Loading contract: ``load(model, weights)`` consumes the engine-provided checkpoint dict, fills every declared +parameter exactly once and asserts coverage in both directions: every declared parameter written, and every +checkpoint key either consumed or in ``expected_unconsumed`` (the vision tower and its projector, which this +text-only target does not build, and the language model's multi-token-prediction head). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) + +#: ``layers=ALL_LAYERS`` is the common case: a weight every decoder layer carries. +ALL_LAYERS = object() + +_LM = "language_model." +#: Checkpoint key families this target does not build: the vision tower and its two projectors (text only), and the +#: language model's multi-token-prediction head (speculation, when on, is the separate Eagle3 drafter). +_UNBUILT_PREFIXES = ( + "vision_tower.", + "multi_modal_projector.", + "patch_merge_mlp.", + _LM + "model.mtp.", +) + +MX_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 elements along K +NV_BLOCK = 16 # NVFP4: one E4M3 scale per 16 elements along K +EPILOGUE_TILE_M = ( + 128 # the trtllm-gen MoE runner's row tile: fixes its weight and scale row shuffles +) + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + ``shape`` takes the dimension bundle and returns the declared shape. ``src`` is a callable + ``(d, i) -> [(key or key tuple, index into the param | None), ...]`` with keys relative to the checkpoint + root; a key tuple hands the transform those tensors together. ``layers`` is ``ALL_LAYERS``, ``None`` (not + per-layer) or ``d -> layer indices``. ``transform(src, d)`` runs on the destination's device before the copy. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + +# ------------------------------------------------------------------------------------------------ shards + + +def _col( + t: torch.Tensor, d: SimpleNamespace, tp: int | None = None, rank: int | None = None +) -> torch.Tensor: + """This rank's rows (column-parallel output split), even split, sliced on the host and moved to the + destination device (only this rank's bytes are read).""" + tp = d.tp if tp is None else tp + rank = d.rank if rank is None else rank + n = t.shape[0] // tp + return _dev(t[rank * n : (rank + 1) * n], d) + + +def _row(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """This rank's input columns (row-parallel input split), even split, on the destination device.""" + n = t.shape[1] // d.tp + return _dev(t[:, d.rank * n : (d.rank + 1) * n], d) + + +def _dev(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """``t`` on the destination device (``d.dev``, set per fill).""" + return t.to(d.dev, non_blocking=True) + + +def _swizzle(scale: torch.Tensor) -> torch.Tensor: + """128x4 swizzle of a row-major ``[N, K / 32]`` UE8M0 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(scale.contiguous()) + + +# --------------------------------------------------------------------------------------------- transforms + + +def _gemma_fold(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A Gemma boundary norm as the plain norm the fused all-reduce epilogue applies: ``w + 1`` in BF16.""" + return t + 1.0 + + +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) + + +def _qkv_sparse_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[Q | K | V | index-Q | index-K]``: Q by attention heads, K/V/index-Q by KV heads, index-K replicated.""" + q, k, v, iq, ik = t + return torch.cat( + [_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0 + ).contiguous() + + +def _qkv_sparse_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v, iq, ik = t + return _swizzle( + torch.cat([_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0) + ) + + +def _qkv_dense_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +def _qkv_dense_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return _swizzle(torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0)) + + +def _index_qk_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[index-Q | index-K]`` in BF16: this rank's index-Q heads and the replicated index-K, the E4M3 values cast + (the scale tensors, the last two sources, are consumed only).""" + iq, ik = t[0], t[1] + return torch.cat([_col(iq, d), _dev(ik, d)], dim=0).to(d.dtype).contiguous() + + +def _gate_up_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return torch.cat([_col(gate, d), _col(up, d)], dim=0).contiguous() + + +def _gate_up_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return _swizzle(torch.cat([_col(gate, d), _col(up, d)], dim=0)) + + +def _row_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _row(t, d).contiguous() + + +def _row_s(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _swizzle(_row(t, d)) + + +def _col_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _col(t, d).contiguous() + + +#: Permute indices of the trtllm-gen expert relayouts, keyed by (shape, family, scale group): one expert's are +#: every expert's, so each is computed once per load. +_PERMUTE_CACHE: dict = {} + + +def _expert_w31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 operand: ``[up ; gate]`` packed-FP4 rows of this rank, then the gated-activation row + reorder and the epilogue-tile shuffle.""" + up, gate = t + w = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + w, _PERMUTE_CACHE, EPILOGUE_TILE_M, is_gated_act_gemm=True + ) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 E4M3 scales: ``[up ; gate]`` rows of this rank, the same row reorder at scale + granularity, then the 128x4 swizzle.""" + up, gate = t + s = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK, is_gated_act_gemm=True + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled.view(torch.uint8).reshape(shape)) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + + +def _expert_w2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 operand: this rank's input columns, then the epilogue-tile row shuffle.""" + w = _row(t, d).contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(w, _PERMUTE_CACHE, EPILOGUE_TILE_M) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 E4M3 scales: this rank's columns, the row shuffle at scale granularity, the swizzle.""" + s = _row(t, d).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w2_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled).view(torch.float8_e4m3fn).reshape(shape) + ) + + +# ------------------------------------------------------------------------------------------------ sources + + +def _one(template: str): + return lambda d, i, t=template: [(_LM + t.format(i=i), None)] + + +def _many(*templates: str): + return lambda d, i, ts=templates: [(tuple(_LM + t.format(i=i) for t in ts), None)] + + +def _per_expert(*suffixes: str): + """One source per routed expert, written at its expert index (all experts are local at tp2).""" + + def build(d, i, sx=suffixes): + out = [] + for e in range(d.num_experts): + keys = tuple(_LM + f"model.layers.{i}.block_sparse_moe.experts.{e}.{s}" for s in sx) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _dense_layers(d): + return range(d.first_sparse) + + +def _sparse_layers(d): + return range(d.first_sparse, d.num_layers) + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense layers, and every layer without the five-way packed + projection (``fuse_qkv_index_projection=False``).""" + return range(d.first_sparse) if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return range(d.first_sparse, d.num_layers) if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else range(d.first_sparse, d.num_layers) + + +_ATTN = "model.layers.{i}.self_attn" +_MOE = "model.layers.{i}.block_sparse_moe" + + +class MiniMaxM3Tp2Weights: + """MiniMax-M3 NVFP4 / sm_107 / tp2: the language model's table and the passes over it. + + ``declare`` allocates every weight it names, ``manifest`` maps checkpoint keys onto them, ``load`` fills + them; ``_targets`` is the one place that decides which parameters exist. + """ + + #: Free the per-expert relayout scratch after each expert stack. + RELEASE_AFTER = "_e_s2" + + WEIGHTS: tuple[W, ...] = ( + # --- the two boundary norms of every layer (Gemma, folded) --- + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + transform=_gemma_fold, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + transform=_gemma_fold, + ), + # --- attention: per-head Q/K norms (Gemma, unfolded), the packed QKV and the output projection --- + W("q_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".q_norm.weight")), + W("k_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".k_norm.weight")), + W( + "qkv_w", + shape=lambda d: (d.qkv_dense_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in "qkv")), + transform=_qkv_dense_w, + layers=_three_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_dense_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many(*(f"{_ATTN}.{p}_proj.weight_scale_inv" for p in "qkv")), + transform=_qkv_dense_s, + layers=_three_way_qkv_layers, + ), + W( + "qkv_w", + shape=lambda d: (d.qkv_sparse_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in ("q", "k", "v", "index_q", "index_k"))), + transform=_qkv_sparse_w, + layers=_five_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_sparse_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + *( + f"{_ATTN}.{p}_proj.weight_scale_inv" + for p in ("q", "k", "v", "index_q", "index_k") + ) + ), + transform=_qkv_sparse_s, + layers=_five_way_qkv_layers, + ), + # The separate BF16 index projection [index-Q | index-K] of the unfused sparse layers. The built-in model + # builds it as an unquantized Linear and fills it through the fused gate/up loader, which casts the + # checkpoint's E4M3 values to BF16 and does not apply their block scales: so does this table (the scale + # tensors are consumed, not applied). + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + src=_many( + _ATTN + ".index_q_proj.weight", + _ATTN + ".index_k_proj.weight", + _ATTN + ".index_q_proj.weight_scale_inv", + _ATTN + ".index_k_proj.weight_scale_inv", + ), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_layers, + ), + W( + "o_w", + shape=lambda d: (d.hidden, d.q_width), + dtype=torch.float8_e4m3fn, + src=_one(_ATTN + ".o_proj.weight"), + transform=_row_w, + ), + W( + "o_s", + shape=lambda d: (d.hidden * (d.q_width // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_ATTN + ".o_proj.weight_scale_inv"), + transform=_row_s, + ), + # --- dense SwiGLU-OAI MLP (layers 0..2) --- + W( + "mlp_gu_w", + shape=lambda d: (2 * d.dense_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight", "model.layers.{i}.mlp.up_proj.weight" + ), + transform=_gate_up_w, + layers=_dense_layers, + ), + W( + "mlp_gu_s", + shape=lambda d: (2 * d.dense_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight_scale_inv", + "model.layers.{i}.mlp.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_dense_layers, + ), + W( + "mlp_dn_w", + shape=lambda d: (d.hidden, d.dense_inter), + dtype=torch.float8_e4m3fn, + src=_one("model.layers.{i}.mlp.down_proj.weight"), + transform=_row_w, + layers=_dense_layers, + ), + W( + "mlp_dn_s", + shape=lambda d: (d.hidden * (d.dense_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one("model.layers.{i}.mlp.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_dense_layers, + ), + # --- MoE (layers 3..59): the fp32 router and its bias, the shared expert --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + dtype=torch.float32, + src=_one(_MOE + ".gate.weight"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "sh_gu_w", + shape=lambda d: (2 * d.shared_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + _MOE + ".shared_experts.gate_proj.weight", _MOE + ".shared_experts.up_proj.weight" + ), + transform=_gate_up_w, + layers=_sparse_layers, + ), + W( + "sh_gu_s", + shape=lambda d: (2 * d.shared_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _MOE + ".shared_experts.gate_proj.weight_scale_inv", + _MOE + ".shared_experts.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_sparse_layers, + ), + W( + "sh_dn_w", + shape=lambda d: (d.hidden, d.shared_inter), + dtype=torch.float8_e4m3fn, + src=_one(_MOE + ".shared_experts.down_proj.weight"), + transform=_row_w, + layers=_sparse_layers, + ), + W( + "sh_dn_s", + shape=lambda d: (d.hidden * (d.shared_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_MOE + ".shared_experts.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_sparse_layers, + ), + # --- the routed experts, NVFP4 (trtllm-gen layouts), all 128 on every rank --- + W( + "e_w31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_per_expert("w3.weight", "w1.weight"), + transform=_expert_w31, + layers=_sparse_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight"), + transform=_expert_w2, + layers=_sparse_layers, + ), + W( + "e_s31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w3.weight_scale", "w1.weight_scale"), + transform=_expert_s31, + layers=_sparse_layers, + ), + W( + "e_s2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w2.weight_scale"), + transform=_expert_s2, + layers=_sparse_layers, + ), + # The per-tensor scalars of every expert, raw; modeling derives the kernels' scalars from them. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + layers=_sparse_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_scale"), + ("isc3", "w3.input_scale"), + ("isc2", "w2.input_scale"), + ("ws2_1", "w1.weight_scale_2"), + ("ws2_3", "w3.weight_scale_2"), + ("ws2_2", "w2.weight_scale_2"), + ) + ), + # --- not per-layer --- + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=_one("model.embed_tokens.weight"), + layers=None, + ), + W( + "final_norm", + shape=lambda d: (d.hidden,), + src=_one("model.norm.weight"), + transform=_gemma_fold, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + sparse = cfg.sparse_attention_config + freq = [int(f) for f in sparse["sparse_attention_freq"]] + tp = mapping.tp_size + head_dim = int(cfg.head_dim) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + index_heads = max(int(sparse["sparse_num_index_heads"]) // tp, 1) + index_dim = int(sparse["sparse_index_dim"]) + return SimpleNamespace( + fused=bool( + getattr( + core.model_config.sparse_attention_config, "fuse_qkv_index_projection", False + ) + ), + num_layers=cfg.num_hidden_layers, + first_sparse=freq.index(1), + hidden=cfg.hidden_size, + head_dim=head_dim, + heads=heads, + kv_heads=kv_heads, + index_heads=index_heads, + index_dim=index_dim, + q_width=heads * head_dim, + qkv_dense_rows=(heads + 2 * kv_heads) * head_dim, + qkv_sparse_rows=(heads + 2 * kv_heads) * head_dim + index_heads * index_dim + index_dim, + dense_inter=int(cfg.dense_intermediate_size) // tp, + shared_inter=int(cfg.shared_intermediate_size) * int(cfg.n_shared_experts) // tp, + moe_inter=int(cfg.intermediate_size) // tp, + num_experts=int(cfg.num_local_experts), + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None)] + return [(f"l{i}_{entry.name}", i) for i in layers] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares (meta under the engine's meta init).""" + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key or key tuple, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, core, weights) -> set: + """The vision tower and its projectors, and the multi-token-prediction head: in the checkpoint, never built + by this target.""" + return {k for k in weights.keys() if k.startswith(_UNBUILT_PREFIXES)} + + def load(self, model, weights) -> None: + """Fill every declared weight and the shell's ``lm_head`` shard; assert coverage both ways.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + assert set(manifest) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest) ^ set(core.w.keys()), + ) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(core, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(core.w[param_key], ckpt, index, transform, weights, d, consumed) + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + # The shell owns the vocab-parallel head: this rank's rows of the full table. + _fill(model.lm_head.weight, _LM + "lm_head.weight", None, _col_w, weights, d, consumed) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + } + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the unbuilt families (vision, MTP head); " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _materialize(entry) -> torch.Tensor: + """``[:]`` on a lazy safetensors slice; a 0-dim tensor (the per-tensor scalars) rejects that index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +def _fill(param, ckpt, index, transform, weights, d, consumed) -> None: + keys = ckpt if isinstance(ckpt, tuple) else (ckpt,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + d.dev = dst.device + # Host tensors: a transform slices this rank's part on the host and moves only that to the device. + src_all = tuple(_materialize(weights[key]) for key in keys) + if transform is not None: + src = transform(src_all if isinstance(ckpt, tuple) else src_all[0], d) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + if src.dim() == 0 and dst.dim() == 0: + pass + elif src.numel() == 1 and dst.numel() == 1: + src = src.reshape(dst.shape) + src = src.to(dst.device, non_blocking=True) + assert tuple(dst.shape) == tuple(src.shape), (ckpt, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + _release(weights, keys) + + +def _release(weights, keys) -> None: + """Drop consumed checkpoint tensors from the engine's dict as the load goes (bounded host memory).""" + if hasattr(weights, "mark_consumed_keys"): + weights.mark_consumed_keys(keys) + else: + for key in keys: + weights.pop(key, None) + + +MODEL_WEIGHTS = MiniMaxM3Tp2Weights() + + +# -------------------------------------------------------------------------------------------- the drafter + + +def _draft_one(template: str): + return lambda d, i, t=template: [(t, None)] + + +def _draft_many(*templates: str): + return lambda d, i, ts=templates: [(tuple(ts), None)] + + +def _draft_fc_norms(d): + return range(d.capture_layers) + + +def _draft_qkv(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +class MiniMaxM3Tp2DraftWeights: + """The one-layer BF16 Eagle3 drafter (``speculative_model``), tp2. + + The fc over the three captured target layers and its per-layer fc_norms, the layer's three norms, its QKV + (input ``[embeds | hidden]``, so twice the hidden width; Q by heads, K/V by KV heads) and o_proj, its gated + MLP, the final norm and the full-vocabulary head (vocab-parallel). The checkpoint's ``embed_tokens`` is the + target's table (the drafter reads the target's); it is a predicted non-load. + """ + + WEIGHTS: tuple[W, ...] = ( + W( + "fc", + shape=lambda d: (d.hidden, d.capture_layers * d.hidden), + src=_draft_one("fc.weight"), + layers=None, + ), + W( + "fc_norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"fc_norm.{i}.weight", None)], + layers=_draft_fc_norms, + ), + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.input_layernorm.weight"), + layers=None, + ), + W( + "hidden_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.hidden_norm.weight"), + layers=None, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.post_attention_layernorm.weight"), + layers=None, + ), + W( + "qkv", + shape=lambda d: (d.qkv_rows, 2 * d.hidden), + src=_draft_many(*(f"layers.0.self_attn.{p}_proj.weight" for p in "qkv")), + transform=_draft_qkv, + layers=None, + ), + W( + "o", + shape=lambda d: (d.hidden, d.q_width), + src=_draft_one("layers.0.self_attn.o_proj.weight"), + transform=_row_w, + layers=None, + ), + W( + "gu", + shape=lambda d: (2 * d.inter, d.hidden), + src=_draft_many("layers.0.mlp.gate_proj.weight", "layers.0.mlp.up_proj.weight"), + transform=_gate_up_w, + layers=None, + ), + W( + "dn", + shape=lambda d: (d.hidden, d.inter), + src=_draft_one("layers.0.mlp.down_proj.weight"), + transform=_row_w, + layers=None, + ), + W("final_norm", shape=lambda d: (d.hidden,), src=_draft_one("norm.weight"), layers=None), + ) + + def dims(self, draft) -> SimpleNamespace: + cfg = draft.draft_config.pretrained_config + mapping = draft.draft_config.mapping + tp = mapping.tp_size + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + return SimpleNamespace( + hidden=cfg.hidden_size, + capture_layers=draft.capture_layers, + head_dim=head_dim, + q_width=heads * head_dim, + qkv_rows=(heads + 2 * kv_heads) * head_dim, + inter=cfg.intermediate_size // tp, + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + if entry.layers is None: + return [(entry.name, None)] + return [(f"{entry.name}{i}", i) for i in entry.layers(d)] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, draft) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + d = self.dims(draft) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, draft, weights) -> set: + """The drafter checkpoint's own embedding: the drafter reads the target's table.""" + return {"embed_tokens.weight"} & set(weights.keys()) + + def load(self, draft, weights) -> None: + d = self.dims(draft) + manifest = self.manifest(draft) + assert set(manifest) == set(draft.w.keys()), set(manifest) ^ set(draft.w.keys()) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(draft, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(draft.w[param_key], ckpt, index, transform, weights, d, consumed) + _fill(draft.lm_head.weight, "lm_head.weight", None, _col_w, weights, d, consumed) + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + f"drafter checkpoint coverage: unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = MiniMaxM3Tp2DraftWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py new file mode 100644 index 000000000000..302bb6b15c31 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/__init__.py @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 NVFP4 / sm_107 / tp4.""" + +from ..kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py new file mode 100644 index 000000000000..550316d84740 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py @@ -0,0 +1,5936 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_107 / tp4 -- self-contained modeling code. + +MiniMax-M3's language model: 60 layers, hidden 6144, 64 query heads and 4 KV heads of dim 128. Layers 0-2 are dense +attention with a dense SwiGLU-OAI MLP; layers 3-59 are MSA block-sparse attention (an index branch scores 128-token +blocks and the top 16 are attended) with a MoE of 128 routed experts (top 4, sigmoid scores with a correction bias, +renormalized and scaled by 2) plus one shared expert. The checkpoint is a modelopt mixed-precision export: NVFP4 routed +experts, MXFP8 projections, fp32 router, bf16 embedding and head; the KV cache is NVFP4 for the sparse layers and FP8 +for the dense ones. + +``tp4`` is four Rubin GPUs (SM 10.7): head-split attention (16 query heads, one KV head and one index head per rank) and +tensor-parallel experts (every expert on every rank, intermediate 768 per rank), no attention DP. Speculative +decoding, when configured, is one-model Eagle3 (draft length 3): the BF16 drafter (``M3Eagle3Draft``) and the BF16 +target head, both at checkpoint precision; decode steps of up to 8 requests verify and draft through fused kernels +at that precision (``M3Eagle3Worker``). + +**The core** (``MiniMaxM3Tp4ModelingV2``) holds the weights (``weights.MODEL_WEIGHTS``) and the runtime resources the +steps use -- one MSA attention backend per layer, the TP all-reduce communicator, the RoPE table, the MXFP8 GEMM +dispatch state the engine tunes -- and dispatches each step, on the host, from its context count alone +(``_select_target``): + +* ``DecodeTarget`` takes every pure generation step. Up to ``DECODE_MAX_TOKENS`` tokens (up to 8 requests, each one + token or a draft verify of 1 + 3 rows) with 1-D input ids it runs the decode-step kernels and picks each layer + boundary's collective by the step's token bucket (``DECODE_PLANS``); it runs its own generic layer for the first + layer when the fused embedding does not take the step, for every layer whose output Eagle3 captures, and for every + layer of a wider step. +* ``PrefillTarget`` takes every step with context rows: prefill chunks and mixed steps, every layer in its generic + form over the MSA backend. + +The two targets share no body: each owns its layer, attention and MoE compositions and its ``step_args``, the one +projection of the step metadata its forward reads. What both call are the module-level op adapters, each picking the +kernel for one call from shapes and layouts (``_linear``, ``_qk_norm_rope``, ``_sparse_qkv_producer``, the router and +MoE-runner calls, the all-reduce + norm calls), which read no step metadata themselves. + +Every fast path the decode step can take, the rows each covers (``FAST_PATHS``, each paired with the entry whose +``CELLS`` certify it) and what a step that cannot take one does are below: **uncovered** (outside every covered +range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise +``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). + +**Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the +next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the +long-context decode for the o_proj. It publishes it in ``ops/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +MXFP8 projection on a decode graph takes it (``_linear``), and every take clears the slot. The decode boundaries' +push targets (the o_proj / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's +exchange) are passed explicitly from the layer to the producer and back. + +**What this target asserts rather than adapts**: SM 10.3, the topology and checkpoint (``routing.py``); the +MiniMax-M3 MSA sparse attention runtime with the fused five-way QKV/index projection, an FP8 indexer and an FP8 or +NVFP4 KV cache, no torch.compile and no piecewise prefill graphs (``_validate_runtime``); text-only input (the vision +tower is not built and a step carrying image input raises). +""" + +from __future__ import annotations + +import atexit +import copy +import dataclasses +import os +import threading +import time +from dataclasses import dataclass +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + resolved_mode, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionInputType, + PositionalEmbeddingParams, + RopeParams, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, + write_msa_phase_kv, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + _ensure_draft_vocab_size, + _get_requested_draft_moe_backend, + _set_draft_kv_cache_quant_algo, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + TensorParallelMode, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen, is_torch_compiling +from tensorrt_llm._utils import mpi_rank, prefer_pinned +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import weights as _weights +from .ops import m3_attn_lc, m3_fold_tail, m3_mc_allreduce, m3_prequant, m3_report, m3_spec_tail +from .ops import m3_draft_lowm as _draft_lowm_kernel +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import m3_drafter_attn as _drafter_attn_kernel +from .ops import m3_fullctx_decode as _fullctx_decode_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_index_score_t2sc as _index_score_t2sc_kernel +from .ops import m3_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_select_long_dsl as _select_long_dsl_kernel +from .ops import m3_shard_accept as m3_shard_accept_kernel +from .ops import m3_target_head as m3_target_head_kernel +from .ops import m3_tokmap_guard as _tokmap_guard +from .ops import msa_kv_update as _msa_kv_update_kernel +from .ops.activation.m3_silu_and_mul import m3_silu_and_mul +from .ops.attention.m3_dense_decode import m3_dense_decode +from .ops.attention.m3_dense_qk_norm_rope_fp8 import m3_dense_qk_norm_rope_fp8 +from .ops.attention.m3_draft_attention import m3_draft_attention +from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope +from .ops.attention.m3_fullctx_decode import m3_fullctx_decode +from .ops.attention.m3_index_score import SPLIT_K as INDEX_SCORE_SPLIT_K +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_index_score_t2sc import m3_index_score_t2sc +from .ops.attention.m3_msa_attention import m3_msa_attention +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 +from .ops.attention.m3_qkv_insert import m3_qkv_insert +from .ops.attention.m3_select_blocks import m3_select_blocks +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_select_long_dsl import m3_select_long_dsl +from .ops.attention.m3_sparse_decode import m3_sparse_decode +from .ops.comm.m3_allreduce import m3_allreduce +from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm +from .ops.embedding.m3_embedding import m3_embedding +from .ops.gemm.m3_bf16_linear import m3_bf16_linear +from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu +from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS +from .ops.gemm.m3_draft_lowm import m3_draft_lowm +from .ops.gemm.m3_fp32_linear import m3_fp32_linear +from .ops.gemm.m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + m3_mx_gemv, +) +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.gemm.m3_mxfp8_pinned_gemm import m3_mxfp8_pinned_gemm +from .ops.gemm.m3_mxfp8_pinned_quantize import m3_mxfp8_pinned_quantize +from .ops.gemm.m3_target_head import MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS +from .ops.gemm.m3_target_head import m3_target_head +from .ops.long_select import long_select_blocks +from .ops.moe.m3_moe16_engine import M3Moe16State, m3_moe16_engine +from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .ops.spec.m3_draft_input import m3_draft_input +from .ops.spec.m3_draft_silu import m3_draft_silu +from .ops.spec.m3_draft_step import MAX_TOKENS as DRAFT_STEP_KERNEL_MAX_TOKENS +from .ops.spec.m3_draft_step import m3_draft_step +from .ops.spec.m3_force_vec import m3_force_vec +from .ops.spec.m3_force_vec_req import m3_force_vec_req +from .ops.spec.m3_shard_accept import m3_shard_accept +from .ops.spec.m3_spec_offsets import m3_spec_offsets +from .ops.spec.m3_step_prologue import m3_step_prologue +from .ops.triton_sparse_decode import MERGE_MXFP8, MERGE_MXFP8_MIN_TOKENS + +_SM = (10, 7) # capacity: fixed: the device capability, not a size +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The largest generation step the decode target takes: 8 requests verifying 1 + 3 drafts and their padding rows. +DECODE_MAX_TOKENS = 48 # capacity: fixed: the decode kernels' step bound + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 tp4 measurements in the comments).""" + + name: str + max_tokens: int + # The attention boundary's multicast all-reduce + residual + RMSNorm (5.05 vs 5.27 us native at 1 token). + attn_mc_ar: bool + # The same kernel at the MoE boundary with the shared + routed add folded in: wins up to 48 tokens. + moe_mc_ar: bool + # The MoE-boundary all-reduce also emits the next QKV activation in MXFP8 for the decode GEMV / GEMM. + qkv_prequant: bool + + +# tiny: every step up to 8 tokens (c1..c8 without speculation, the c1 verify, every draft pass); small: the c2 +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (ops/m3_mc_allreduce.BUCKETS). +# SM 10.7 TP4: one bucket. MC_AR_TP4 puts every decode boundary on the multicast all-reduce + residual + RMSNorm +# (trtllm::minimax_m3_sm107_tp4_mc_allreduce_norm, its row kernel: bitwise to the stock TP4 oneshot all-reduce fusion +# on the residual and the normed rows) with the MXFP8 copy for the next QKV; False = the built-in all-reduce. +MC_AR_TP4 = True +# Up to this many rows the decode all-reduces (not the MoE finalize) take the 8-CTA cluster kernel instead of the row +# kernel: the same stock oneshot order (opts bit 32: per-column sum of squares, warp butterflies, 6-warp groups, 4 group +# sums from 0.f), so bitwise the row kernel's output, with each row's exchange spread over 8 CTAs. Rubin TP4 synth, both +# orders: c1 (4 rows) -12, c2 (8 rows) -38 us/step, but c4 / c8 (16 / 32 rows) +5 / +45, where +# 8x the CTAs per row cost more than they save. 0 = the row kernel at every row count. +MC_AR_CLUSTER_MAX_ROWS = 8 # capacity: fixed: measured cluster-kernel crossover +# Every decode / drafter boundary all-reduce launches with the max-smem carveout preference (the cluster kernel's +# option 8 at every row count, the row kernel's carveout instance), so the next kernel's CTAs (the QKV GEMV's staged +# tiles) can co-reside early. Launch attribute only: bitwise unchanged. Kill: False (the GB option table). +MC_AR_CARVEOUT = True +# The fused draft loop's boundary all-reduces (M3Eagle3Worker._m3_draft_ar_norm) launch WITHOUT that carveout: the +# kernel after the drafter's last boundary is the draft step (one CTA per SM, ~145 KB of shared memory) rather than a +# staged GEMV, and TP4 alone of the sm_107 targets ran its draft pass ~74-80 us in-model vs ~57 us standalone at world 4 +# in the step census. Launch attribute only: bitwise unchanged. Kill: MC_AR_CARVEOUT (the shared setting). +DRAFT_MC_AR_CARVEOUT = False +# The fused draft loop's two all-reduce boundaries per pass (o_proj -> post-attention norm, down -> final norm) of up +# to this many rows run the multicast all-reduce + residual + RMSNorm's row kernel in its drafter form (opts bit 16: +# FlashInfer fused_add_rmsnorm's reduction order on the unrounded FP32 sum, so bitwise the drafter's stock plain +# all-reduce + RMSNorm) instead of the built-in all-reduce and the separate norm. Kill: 0. +DRAFT_MC_AR_MAX_ROWS = ( + 8 if MC_AR_TP4 else 0 +) # capacity: fixed: the drafter boundary kernel's certified rows +# Those drafter boundaries run the 8-CTA cluster kernel's drafter form (opts bit 128; TP8's drafter boundary kernel) +# instead of the 1-CTA row kernel. Exact-class: deterministic, the updated residual bitwise, the RMS sum of +# squares in the cluster's order instead of FlashInfer's, so a normed row can differ by one BF16 rounding step (7/64 +# rows at 1 row, 14/256 at 4). Rubin world 4, both orders: 4.39 -> 3.75 us per boundary at 1 row, +# 4.79 -> 4.14 at 4 rows, ~ -2..-3.5 us/step. Kill: False (the row kernel, bitwise the stock boundary). +# OFF: in-model timing was flat and natural acceptance length was grey (lower bound -1.16% vs the -1% line). +DRAFT_MC_AR_CLUSTER = False +# The drafter's o_proj and down projection, on the copied split-K kernel (ops/m3_lowm_splitk), also push their rows +# into the drafter boundary all-reduce's exchange (the row kernel's slots: the canonical words its own push would store, +# at the (token, 768-column slice) call parity), which then runs prepushed 2 instead of pushing x itself; the copy +# releases its dependents only after its own grid wait and the pushes (the prepushed >= 2 contract). Bitwise. The TP8 +# port: ~ -4 us/step at world 8. Kill: False. +DRAFT_AR_PUSH = True +# The drafter boundary all-reduce fed by its projection's push (DRAFT_AR_PUSH): its own taken / not-taken lines. +DRAFT_AR_PUSH_PATH = "m3_draft_ar_push" +DECODE_PLANS = ( + DecodePlan( + "tiny", DECODE_MAX_TOKENS, attn_mc_ar=MC_AR_TP4, moe_mc_ar=MC_AR_TP4, qkv_prequant=MC_AR_TP4 + ), +) + +# The dense layers' decode attention (trtllm-gen, msa_decode.dense_decode_max_seq_len): the max_seq_len it is handed +# per request count. It only sizes the multi-CTA KV split (every CTA still covers its share of the actual length; the +# result changes only by the split-merge order), and the model maximum over-splits long-context caches at one +# request: GB300 TP4, 16 heads, FP8 KV, 64K context: 15.1 us at 16K vs 28.6 us at 1M. Two requests are neutral (the +# grid is unchanged) and keep the step's bound. +# The generic layer's boundaries (prefill, mixed steps, Eagle3 capture layers): the most rows that take the multicast +# all-reduce + residual + RMSNorm, alone and with the shared + routed add folded in. TP4 skeleton: none (the built-in +# all-reduce) until the multicast all-reduce is certified on Rubin. +GENERIC_MC_AR_MAX_TOKENS = 0 # capacity: fixed: off at TP4 (the built-in all-reduce) +GENERIC_MC_AR_FOLDED_MAX_TOKENS = 0 # capacity: fixed: off at TP4 (the built-in all-reduce) +# An Eagle3 capture layer whose capture norm is armed (CAPTURE_NORM) has its fused feed-forward boundary all-reduce +# (the 8-CTA cluster kernel up to MC_AR_CLUSTER_MAX_ROWS, the row kernel above) also write the capture: +# rmsnorm(updated residual) * fc_norm[i], bit for bit the capture_norm kernel, which then does not run. The MoE +# finalize runs the shared GB300 op and keeps the capture_norm kernel. Wired at the generic layer's fused boundary +# (layers.M3DecoderLayer._feed_forward_boundary) and through M3DecodeTarget._layer's boundaries (``capture``) for a +# capture layer on the fast decode layer. Kill: False (the capture_norm kernel after the boundary). +CAPTURE_NORM_IN_BOUNDARY = True + +# TP4 skeleton: no cap (the built-in split) until a TP4 SM 10.7 cap is measured and certified. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) + +# The dense layers' decode attention on PrimTS split-KV decode (ops/m3_fullctx_decode) instead of the trtllm-gen +# call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the +# BF16 merge. At TP4 (16 query heads, 1 KV head per rank) the stock split is capped at 16 cluster CTAs per (request, +# KV head), i.e. 16 of 212 CTAs at one request: stock dense fmhaSm107a 33.4-33.8 us per call and the stock drafter +# fmhaSm100f 45.4 us per pass at 165k. Rubin TP4, FP8 KV, 165k, chained us per call +# (microbench), every recipe passing the fp32 oracle: (4, 1) 13.70 (stock 33.6), (1, 1) 11.52. At one +# request the K/V stream is latency-bound, so both deepen the K/V load pipeline (the load pipeline only; the same +# bits): 12 stages at one query row, 8 at four (the long-context sweep: (4, 1) 12.95 / 17.00 / 20.98 us +# per call at 165k / 323k / 512k with 8 stages vs 13.03 / 17.15 / 21.30 with 12; (1, 1) is fastest with 12). Two +# requests (microbench; stock at c2: dense fmhaSm107a 26.5 us, drafter fmhaSm100f 44.0-44.6 us per pass, +# census): (4, 2) 17.34, (1, 2) 14.16. +# K/V stages in flight of the (4, 1) recipe (12 was the previous depth; the same bits at any depth). +FULLCTX_DECODE_Q4_KV_STAGES = 8 # capacity: fixed: a pipeline depth, not a size +FULLCTX_DECODE_RECIPES: Mapping[Tuple[int, int], Mapping[str, Any]] = MappingProxyType( + { + (4, 1): MappingProxyType( + { + "config": { + "tile_size_q": 32, + "groups_tokens_heads_q": True, + "kv_stages": FULLCTX_DECODE_Q4_KV_STAGES, + }, + "splits": 96, + } + ), + (1, 1): MappingProxyType({"config": {"kv_stages": 12}, "splits": 108}), + (4, 2): MappingProxyType( + {"config": {"tile_size_q": 32, "groups_tokens_heads_q": True}, "splits": 48} + ), + (1, 2): MappingProxyType({"config": {}, "splits": 80}), + } +) # capacity: fixed: pinned recipes per certified (rows, requests) +# Kill constant: False puts the dense layers back on the trtllm-gen call at every shape. +FULLCTX_DECODE_ENABLED = True +# The dense layers' FP8 q/k/v (decode steps of the hybrid cache): the op minimax_m3_sm107tp4_dense_qk_norm_rope_fp8 +# (head-count generic) rounds Q/K to BF16 and then to E4M3 in its store, the bits of the BF16 fused_qk_norm_rope output +# cast to FP8. False (or a build without the op) keeps the single-rounding fused_qk_norm_rope FP8 output. +DENSE_QKV_FP8_VIA_BF16 = True +# On a full-context decode step that op also stores the new K/V rows into the layer's paged pool (the bytes +# write_msa_phase_kv's indexed copies store, one launch instead of about ten). False writes them with +# write_msa_phase_kv. +DENSE_KV_INSERT_FUSED = True +# The engagement line of that fused store (taken, one per row count and rank). +DENSE_KV_INSERT_PATH = "m3_dense_kv_insert" +# The drafter MLP's C++ SwiGLU (DRAFT_SILU_CPP, M3Eagle3Draft.mlp): its own taken / not-taken lines. +DRAFT_SILU_PATH = "m3_draft_silu" +# The split-KV merge on the BF16 op (minimax_m3_sm107tp4_fullctx_reduce, head-count generic) instead of PrimTS's FP16 +# reducer and a cast. +FULLCTX_DECODE_OWN_REDUCER = True +# The merge also writes the MXFP8 copy of its BF16 rows (E4M3 + 128x4-swizzled UE8M0, bit for bit +# trtllm::mxfp8_quantize) into the MERGE_MXFP8 buffers and publishes it to m3_prequant.SLOT, so the dense layer's +# MXFP8 o_proj skips its own quantization. Kill: False leaves the o_proj to quantize its BF16 input. +FULLCTX_DECODE_MXFP8 = True +FULLCTX_DECODE_MXFP8_PATH = "m3_fullctx_merge_mxfp8" +# The K/V length bound the full-context plans are compiled for; a step whose cache length bound is longer is a miss. It +# covers the long-context serving config's max_seq_len (1 << 20) plus the speculative rows the cache reserves past it. +FULLCTX_DECODE_MAX_KV_LEN = ( + 1 << 21 +) # capacity: fixed: the plans' KV-length bound (1M context + spec rows) +# The Eagle3 drafter layer's generation phase on the same full-context decode (ops/m3_drafter_attn: the in-tree +# RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its +# split at 8 CTAs per (request, KV head) on SM 10.7: 45.4 us per draft pass at 165k. Kill constant. +DRAFTER_FULLCTX_ENABLED = True +# On those passes the drafter's RoPE, FP8 query and FP8 K/V append run as one launch (the op +# minimax_m3_sm107tp4_drafter_rope_fp8_kv: the in-tree preprocess's GPT-NeoX rotation and unit-scale E4M3 store, and +# the torch E4M3 cast of its BF16 query) instead of thop.trtllm_gen_generation_preprocess and the cast. Kill constant. +DRAFTER_ROPE_FUSED = True +# The engagement line of that launch (taken, one per row count and rank). +DRAFTER_ROPE_PATH = "m3_drafter_rope" +# thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request +# multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by +# that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, +# token-0 drafts). ops/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# calls: what the build writes (padded rows included), so bitwise the op with a fresh map. Kill: False. +PREPROCESS_TOKMAP_FIX = True +PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" +# How the rotation's products are summed, as the in-tree preprocess build sums them (0: two products then the add, +# 1: fma(cos, x, sin * pair), 2: fma(sin, pair, cos * x)); the op's bitwise test pins it. +DRAFTER_ROPE_MODE = 1 # capacity: fixed: a launch mode, not a size + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the index scorer launches and stages its pre-wait loads while the producer runs (its +# outputs are read only after the scorer's own grid wait). +PRODUCER_EARLY_TRIGGER = True +# Steps of up to this many rows keep the producer's trigger at its exit instead. On Rubin TP4 the early-launched scorer +# (its grid doubles with the requests, its pre-wait index-K prefetch running) is resident through the producer's whole +# post-QKV segment and slows the producer's dependent loads; from 16 rows the overlap pays. In-model, locked memory +# clock, tactic- and FlashInfer-pinned, A and B in one job: -75 / -176 us per decode step at 1 / 2 requests, +# +67 at 4 with the trigger at the exit. Timing only: the scorer reads the producer's outputs after its own grid wait. +# Kill: 0 (the early trigger at every row count). Set to 0 on the shipped stack: with the long-row select's late release +# and post-load trigger, the scorer's pre-wait overlap with the producer is worth more than the producer's own +# segment (both orders, double-pinned: -12 / -42 us at 1 / 2 requests with 0). +PRODUCER_LATE_TRIGGER_MAX_ROWS = 0 # capacity: fixed: measured trigger crossover (0: off) +# The producer reads its per-step inputs (position ids, cache slots, RoPE rows) before its grid wait, leaving only the +# packed QKV row after it: this target's copy of the producer (trtllm::minimax_m3_sm107_tp4_producer_pw, +# minimaxM3Sm107Tp4/producerPw4.cu, the TP2 producerPw2 at 16 query heads). Their writers are host-to-device copies +# before the replay and plain-launched kernels earlier in the step (the input preprocessing, msa_kv_update or its torch +# fallback), so they are final before the producer can start; a PDL-launched writer of any of them would have to turn +# this off. The call site's packed.contiguous() and position-id copies are plain launches too. Bit-identical outputs. +# Kill: False (the shared minimax_m3_producer_pw at 16 query heads, which reads them after its wait). +PRODUCER_PREWAIT_STEP_INPUTS = True + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. +EMBED_FIRSTNORM_MAX_TOKENS = 8 # capacity: fixed: built rows; wider steps embed generically +# The o_proj pushes its rows into the attention boundary's multicast all-reduce (prepushed mode 2) up to this many +# tokens; wider steps run the GEMV and a plain multicast all-reduce. +ATTN_PUSH_MAX_TOKENS = 8 if MC_AR_TP4 else 0 # capacity: fixed: measured push crossover +# The dense layers' MLP down projection pushes its rows into the post-MLP boundary's multicast all-reduce (prepushed +# mode 2), as the o_proj does into the attention boundary (TP2: those hops end 5.5-5.8 us after the down +# GEMV unpushed vs 2.4-2.5 us pushed). Bitwise equal to the unpushed boundary (the same bf16 words). Kill: False. +DENSE_PUSH_TP4 = True +DENSE_PUSH_MAX_TOKENS = ( + 8 if (MC_AR_TP4 and DENSE_PUSH_TP4) else 0 +) # capacity: fixed: the GEMV push rows +AR_PUSH_MODE = 2 # capacity: fixed: an exchange mode, not a size +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch). It pushes its sum into the +# MoE boundary's multicast all-reduce (prepushed mode 2) on the rows MOE_PUSH covers (with MC_AR_TP4). +MOE_ENGINE_MAX_TOKENS = 8 # capacity: fixed: the engine's built rows +MOE_PUSH_MAX_TOKENS = ( + MOE_ENGINE_MAX_TOKENS if MC_AR_TP4 else 0 +) # capacity: fixed: the engine's rows +# The 9..13-row engine (router, routed and shared experts in one kernel); its bf16 sum goes to the boundary. +MOE16_MIN_TOKENS = 9 # capacity: fixed: the 9..13-row engine's lower bound +MOE16_MAX_TOKENS = 13 # capacity: fixed: the 9..13-row engine's built rows +# From this many tokens the MoE runs unfinalized through the TRTLLM-Gen runner and its finalize (top-k gather, +# routing weights, shared expert add) runs inside the boundary all-reduce. 9..13 take it only when their engine is +# missing (outside require mode). TP4 skeleton: no rows (the built-in MoE finalize and all-reduce). +MC_FINALIZE_MIN_TOKENS = 9 # capacity: fixed: the rows above the engines +MC_FINALIZE_MAX_TOKENS = 0 # capacity: fixed: off at TP4 +# A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel +# over the target lm_head's BF16 shard logits (ops/m3_shard_accept), at the configured draft length. +SHARD_ACCEPT_DRAFT_LEN = 3 # capacity: fixed: the built Eagle3 draft length +SHARD_ACCEPT_MAX_REQUESTS = DECODE_MAX_TOKENS // ( + SHARD_ACCEPT_DRAFT_LEN + 1 +) # capacity: fixed: decode rows / (draft length + 1) +# A sampled (temperature + top-p, advanced_sampling_mode=no_topk) Eagle3 verify step of up to this many requests +# verifies with the fused sampled acceptance (ops/m3_sampled_accept) over this rank's logits shard: every request +# count a decode step holds. +SAMPLED_VERIFY_MAX_REQUESTS = SHARD_ACCEPT_MAX_REQUESTS # capacity: fixed: a decode step's requests +# The sampled verify runs the sm_107 tp4 copy of the TP2 sampled-accept kernel +# (trtllm::minimax_m3_sm107tp4_sampled_accept: a row's CTAs launch as one thread-block cluster, bounded record polls). +# Same records and tokens. Kill: False (the shared kernel, trtllm::minimax_m3_sampled_accept: every CTA of the grid +# co-resident, unbounded polls). +SAMPLED_ACCEPT_SM107TP4 = True +# The sm_107 tp4 sampled verify loads each row's sampling parameters (temperature, top-p, Philox seed and offset: step +# inputs, not the head's output) before its PDL dependency wait, overlapping the head GEMM's tail; the same values and +# tokens. Kill: False (they load after the wait; the shared kernel always loads them after). Rubin TP2 microbench: +# verify end 20.48 -> 19.97 us at top-p 0.95. +SAMPLED_ACCEPT_PREWAIT_PARAMS = True +# The sm_107 tp4 sampled verify releases its dependents (the first draft pass's input) right after its dependency wait +# instead of at the acceptance epilogue: the dependent launches early and parks at its own griddepcontrol.wait, which +# returns only once the verify has completed, so the values are the same; it saves the launch latency after the verify +# (TP2: draft_input started 2.0-2.4 us after the verify / fc join). Kill: False. Off until its sampled-column +# in-model pairs read. +SAMPLED_ACCEPT_EARLY_TRIGGER = False +# SA-1R: the sm_107 tp4 sampled verify also multicasts each CTA's largest logits with its round-0 record (82: each +# thread's top 2 pooled, the CTA's top 8 plus the bound of everything it does not list; 9: an exact top-9 list), so a +# row whose first candidate is provably in the nucleus decides in round 0 instead of exchanging a second round. The +# same tokens, accepted counts and Philox draws as the two-round kernel, bit for bit (a fp32 margin m * 2^-21 * S +# covers the summation order; any doubt runs round 1). The prototype was exact on 39,968 rows; real +# long-context serving logits decide all 4 rows in round 0 on 79.5% (82) / 80.8% (9) of steps. Off (0, the two-round +# kernel): exact, but the untimed decode step reads +47 / +65 us slower with it on at TP4 (the earlier dependent +# launch it enables costs more than the saved round). +SAMPLED_ACCEPT_1R_K = 0 # capacity: fixed: a kernel variant, not a size +# The sm_107 tp4 sampled verify's owner walk (the one thread that finds its row's drawn token by a serial +# prefix over its strided tokens) takes those tokens 8 at a time: their weights are computed independently, then +# summed in the same order with the same early exit, so the draw, tokens and Philox draws are the one-token walk's, +# bit for bit. Rubin TP4 bench (one request, top_p 0.95): -1.06 / -1.07 us per call at SA-1R K 0 / 82. Kill: False +# (the one-token walk, the stock instantiation; the shared kernel has no 8-token walk). +# Off until its in-model TP4 pair passes. +SAMPLED_ACCEPT_ILP_WALK = False +# A step verified by the in-tree path takes the fused first drafter inputs (the capture norms and the first pass's +# gather) up to this many requests: every request count a decode step holds. +FIRST_DRAFTER_INPUTS_MAX_REQUESTS = SHARD_ACCEPT_MAX_REQUESTS # capacity: fixed: step requests +# A fused-verify step of up to this many verify rows (one or two requests at draft length 3) computes the target +# logits shard from the 12-bit coded copy of the BF16 head (ops/m3_target_head; the same BF16 x BF16 -> fp32 +# arithmetic, 75% of the bytes) instead of the in-tree head GEMM; wider verifies keep the in-tree GEMM. +# Off at TP4 (as at TP2): at the TP4 shard (50,016 x 6144) the in-tree BF16 cuBLAS head is faster than the coded kernel, +# 47.9 vs 56.8 us at 4 and 8 rows (TP2: 93 vs 109). True re-enables the coded head. +TARGET_HEAD_CODED = False +TARGET_HEAD_MAX_ROWS = ( + TARGET_HEAD_KERNEL_MAX_ROWS if TARGET_HEAD_CODED else 0 +) # capacity: fixed: the coded head's built rows (off at TP4) +# On a fused-verify step of up to this many verify rows the drafter's first fc is forked onto the aux stream right +# after the last captured target layer, overlapping the remaining target layers and the head (M3Eagle3Worker). +FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover (0: off) +# A sampled fused verify of up to this many verify rows on the in-tree head GEMM forks the drafter's first fc after the +# head instead of beside it (M3Eagle3Worker), as the coded head's verify already does: the fc's split-K grid otherwise +# holds the SMs first and the head GEMM starts late (TP2: 17.4 us late). Kill: 0. +FC_AFTER_HEAD_MAX_ROWS = 8 # capacity: fixed: measured fork-after-head crossover +# A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (M3Eagle3Worker): per pass the +# draft input in one launch (ops/m3_spec_tail.draft_input), the drafter layer's in-tree attention, MLP and +# all-reduces, and one kernel for the draft head (its 12-bit coded copy, or the BF16 fragments when the coding does +# not decode bit for bit), the cross-rank argmax and the pass's bookkeeping (ops/m3_draft_step), at the +# configured draft length. The head kernel takes one row per request, at most 8. +# The fused draft loop streams the 12-bit coded copy of the draft head (when it decodes bit for bit); False keeps the +# BF16 fragments. +DRAFT_HEAD_CODED = False +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +DRAFT_STEP_MAX_REQUESTS = DRAFT_STEP_KERNEL_MAX_TOKENS +# The Eagle3 capture layers (decoder layers whose output the drafter reads) run the fast decode layer, with its +# o_proj / MoE / dense-down pushes into the boundary all-reduces, and capture the MoE / dense boundary's reduced +# residual sum (the same BF16 words the generic layer's fused boundary captures). Kill: False (the generic layer, +# no boundary pushes). Rubin TP2: the capture layers' hops end 4.4-5.3 us after their producing +# GEMV unpushed vs 2.4-2.5 us pushed. +CAPTURE_LAYERS_FAST = True +# An Eagle3 capture layer (one with a next norm) ends at the fused boundary (all-reduce + residual + next RMSNorm in +# one kernel) and captures that kernel's reduced residual sum. Kill: False keeps the feed-forward module's own +# all-reduce, captures hidden + residual and runs the next norm separately, as the built-in model does; it also +# turns CAPTURE_LAYERS_FAST off, which captures the fused boundary's sum. +CAPTURE_FUSED_BOUNDARY = True +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) run FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm, the TP2 target's port with the TP4 table) at the rows its table is certified for, up to this +# many rows; other calls keep the in-tree cuBLAS GEMM (split-K + reduce). Kill: 0. Rubin TP4 sweep (locked memory +# clock) at 1 row: qkv 8.15 -> 5.90, o 5.36 -> 3.40, gate_up 12.93 -> 9.93, down 7.68 -> 6.16 us. +DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the low-M tactic table's certified rows +# The drafter MLP's activation (the plain SwiGLU of its in-tree GatedMLP) runs this target's C++ op +# (ops/spec/m3_draft_silu, minimaxM3Sm107Tp4/siluAndMul.cu) instead of the Triton silu_and_mul: bit for bit, one launch +# either way (no Triton on the generation decode path). The fused draft loop's MLP (M3Eagle3Draft.mlp) routes the +# call; the in-tree draft path keeps the Triton op. Kill: False. +DRAFT_SILU_CPP = True +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(m for m in rows if m <= DRAFT_LOWM_MAX_ROWS) + for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() +} +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (ops/triton_sparse_decode.MQ_MIN_TOKENS). +SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover +# The long-row top-16 selector takes engines of 129..LONG_SELECT_MAX_BLOCKS KV blocks of 128 tokens; it holds 2, 4 or +# 9 keys per thread by the row's valid-block count. At most 128 blocks the built-in selector runs. +LONG_SELECT_MIN_BLOCKS = 129 # capacity: fixed: past the built-in selector's 128 +LONG_SELECT_MAX_BLOCKS = 0 # capacity: fixed: off at TP4 (the DSL select takes the long rows) +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) # capacity: fixed: compiled keys per thread +# The long-row select releases the sparse decode after it right after its own grid wait (the decode stages q and the +# request length before its wait) up to this many query rows. On Rubin TP4 the early-launched attention CTAs crowd the +# select and the o_proj GEMV already at 2 requests (8 rows), so the early release is off at every row count: in-model, +# locked memory clock, pinned tactics, long-context draws, 2 instances vs the stored baseline, the replay period with +# the trigger phase below is -209 us per step at 2 requests (mean of draws), -39 at 4 and -12 at 8. +SELECT_EARLY_MAX_ROWS = 0 # capacity: fixed: measured early-release crossover (0: off) +# The long-context sparse decode runs this target's copy of the op (trtllm::minimax_m3_sm107_tp4_sparse_decode_lc, +# minimaxM3Sm107Tp4/sparseDecodeLc4.cu): its split-K merge forms the 16 weights once per CTA (16 threads, through shared +# memory) instead of in each of the 128 threads. Same values in the same order: bitwise the image op's output. The +# SM107 TP2 sparse decode's -0.17 / -0.23 us per layer at 1 / 2 requests. Kill: False (the image op). +SPARSE_DECODE_LC_MERGE_ONCE = True +# With SPARSE_DECODE_LC_MERGE_ONCE, that op's CTAs also pull the layer's o_proj weight and scales into L2 before +# their grid dependency (cp.async.bulk.prefetch.L2, a hint: no value changes), so the o_proj GEMV after it finds them +# L2-warm. The o_proj's own pre-wait staging streams part of them already; the cold-vs-hot GEMV gap left is ~0.18 us +# per layer. Kill: False. +SPARSE_DECODE_LC_PREFETCH_OPROJ = True +# With SPARSE_DECODE_LC_MERGE_ONCE, that op also records each step's selected block ids per layer and pulls into L2 +# the K/V pages + scales of the blocks the NEXT sparse layer selected at the previous step (65-68% of a layer's +# distinct selected blocks recur one step later, measured on all 4 ranks), so that layer's sparse decode reads +# them L2-warm. Hints only; up to N1_MAX_ROWS query rows. Kill: False. +SPARSE_DECODE_LC_PREFETCH_NEXT_KV = True +N1_MAX_ROWS = 128 # capacity: fixed: the next-KV id buffer's rows +# With SPARSE_DECODE_LC_MERGE_ONCE, short-sequence identity steps (every request within 16 blocks, +# MSA_SHORT_SEQUENCE_IDENTITY) also run that op instead of the Triton producer: CTA r attends block r of its request +# when r < min(16, blocks), taken from the device sequence length (the table is not read, so a host-staged identity +# table that misses the in-flight accepted tokens cannot drop a block). Exact-class like every LC row (band vs the +# fp64 reference); removes the Triton decode from the GEN path at TP4. Kill: False (the Triton identity producer). +# Identity steps skip the select, so the producer's entry trigger used to release LC4's pre-wait q read before the +# producer wrote q (NUL tokens in GSM8K outputs). LC4 reads q after the grid wait, which clears that race (0 NUL). +# Matched ON-vs-OFF GSM8K pairs, greedy and sampled, show no accuracy change; acceptance length +0.024. +SPARSE_DECODE_LC_IDENTITY = True +# Up to SELECT_LONG_TRIGGER_MAX_ROWS query rows and without the early release, the long-row select releases the sparse +# decode once every score of its row is loaded (SELECT_LONG_TRIGGER_PHASE 1, before the exact rank of the candidates) +# instead of at its exit (0, the kill value); wider steps keep the exit. The release still follows the select's own +# grid wait, which the sparse decode's pre-wait q / length / block-table reads rely on. +SELECT_LONG_TRIGGER_PHASE = 1 # capacity: fixed: a trigger phase, not a size +SELECT_LONG_TRIGGER_MAX_ROWS = 8 # capacity: fixed: measured trigger crossover +# The long-row select ranks a row's candidates on warp 0 when they number at most 64 (a 64-key bitonic network; the +# other warps exit) and reads the row's valid count before its grid wait; the picks are the same. Kill: False (the +# all-thread rank). In-model, locked memory clock, double-pinned, A and B in one job, both orders: flat within +# noise (mean of draws -5 / -60 us at 1 / 2 requests; replay period at the 323K request +14 / +54); a FLAT-but-exact +# stack item (post-wait select time 1.76 -> 1.54 us per layer at 1 request, 165K tokens). OFF by default: at TP2 the +# warp-0 tail runs ~0.95 us/layer slower in-model next to the early-launched sparse decode (+53 us/step); +# on at TP4 only after an in-model A,B shows a gain. +SELECT_LONG_SL3 = False +# The long-row select runs as ops/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per row +# instead of one 1024-thread CTA, the same selection bit for bit, with 4, 11 or 24 keys per thread +# (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices (SELECT_LONG_SL3 does not apply). TP4 has the TP8 target's +# per-rank select geometry (1 KV head, 4 rows per request), where on the Rubin decode chain (57 layers behind the +# producer stream, 1 / 2 requests) the select is -0.15 / -0.12 us per layer at 165K tokens, 0.00 / -0.01 at 323K and +# -0.04 / -0.04 at 500K against the 1024-thread kernel. Kill: False (the 1024-thread kernel). +SELECT_LONG_DSL = True +LONG_SELECT_DSL_KEY_WIDTHS = (4, 11, 24) # capacity: fixed: compiled keys per thread +# The MSA index scorer (ops/attention/m3_index_score) prefetches the older index-K pages into L2 before its grid wait +# at every batch; up to this many requests it also issues Q and the first K tile right after the wait and caps that +# prefetch at one block per CTA (larger batches are bandwidth bound, where the extra early tiles only crowd the +# prefetch). On Rubin TP4 the 2-request step keeps the early schedule: with the producer's trigger at its exit +# (PRODUCER_LATE_TRIGGER_MAX_ROWS) the scorer no longer overlaps the producer, and the capped prefetch shortens its +# exposed pre-wait stream. In-model, locked memory clock, tactic- and FlashInfer-pinned, A and B in one job, both +# orders: -260 us per decode step at 2 requests (replay period at the 323K P90 request; -82 mean of draws), flat at +# 1 / 4 / 8. Timing only: the tiles are read after the wait. Kill: 1. +INDEX_SCORE_EARLY_MAX_REQUESTS = 2 # capacity: fixed: measured early-issue crossover +# T2SC: up to 2 requests the sparse decode scores with this target's copy of the sm_107 TP2 / TP8 scorer +# (ops/m3_index_score_t2sc; ops/attention/m3_index_score_t2sc), whose split is chosen per request count and which +# TMA-loads an all-old first block before its grid wait; the scores are bit-identical +# to the in-tree scorer. TP4 has the TP8 target's per-rank indexer geometry (1 index head over the replicated index-K +# of its KV group), where this schedule is ~100 us per decode step faster at 1 request than split 512 without the +# prestage (TP8 in-model, greedy early column, both orders). Wider steps keep the sm_103-derived scorer above. +# Kill: False (the sm_103-derived scorer at every request count). +INDEX_SCORE_T2SC = True +# At 1 request the split is one wave of CTAs (1272 = 212 SMs x 6 resident at this CTA's shared memory): at the +# 1-index-head geometry (57-layer chain behind the producer, scores and the selection bitwise) -0.03 / -0.27 / +0.06 us +# per layer at 165K / 323K / 500K tokens against 1536. Kill: 1536. +INDEX_SCORE_T2SC_SPLIT_K: Mapping[int, int] = MappingProxyType( + {1: 1272, 2: 768} +) # capacity: fixed: measured splits per request count + +# The decode MXFP8 GEMVs, per projection (N, K): the rows each serves. On an activation its producer quantized +# (the MoE boundary for the QKV, the attention merge for the o_proj) the GEMV serves every row up to 48 but the +# unsplit K=768 shared down projection at 17..32 rows, where FlashInfer's GEMM is faster. +# The QKV MXFP8 GEMVs' (split, ring) at decode rows up to GEMV_QKV_CFG_MAX_ROWS: a TP4-owned override of the shared +# GB300-tuned table (ops/mx_gemv/op.py). SM107 TP4 sweep +# (4 rows): sparse-layer QKV (8, 8) 2.70 us vs (4, 12) 3.01, dense-attention QKV (8, 8) 2.60 vs 2.83. Split 8 +# sums K in a different fp32 split-K order than split 4: exact-class, not bitwise to the shared config. Kill: False. +GEMV_QKV_SPLIT8_TP4 = True +GEMV_QKV_CFG_MAX_ROWS = 8 # capacity: fixed: the measured override's rows +MX_GEMV_CFG: Dict[Tuple[int, int], tuple] = { + QKV_SPARSE: (8, 8), + QKV: (8, 8), +} # capacity: fixed: (split, ring) per projection +# (exact-class, deterministic) The o_proj GEMV at decode rows up to GEMV_QKV_CFG_MAX_ROWS with strided k-tiles per split +# rank (contig 0: rank r owns k-tiles r, r + split, ...) instead of contiguous ranges, the shared (2, 8): another +# fixed fp32 order. SM107 TP4 4 rows: 2.513 / 2.533 vs 2.550 / 2.563 us; deterministic (3 runs bitwise) and +# max error vs the fp64 product equal to the contiguous order's. Kill: False. +GEMV_OPROJ_STRIDED_TP4 = True +MX_GEMV_CFG_STRIDED: Dict[Tuple[int, int], tuple] = { + O_PROJ: (2, 8, 0) +} # capacity: fixed: (split, ring, contig) per projection + + +def mx_gemv_cfg(n: int, k: int, rows: int) -> Optional[tuple]: + """The TP4 GEMV config override for an [rows, k] x [n, k]^T MXFP8 GEMV call, or None for the shared table's.""" + if rows > GEMV_QKV_CFG_MAX_ROWS: + return None + if GEMV_QKV_SPLIT8_TP4 and (n, k) in MX_GEMV_CFG: + return MX_GEMV_CFG[(n, k)] + if GEMV_OPROJ_STRIDED_TP4: + return MX_GEMV_CFG_STRIDED.get((n, k)) + return None + + +# Synth only (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS, fractional): the fused verify's forced accepted counts are +# drawn per request, keyed by the request's KV length at the verify (its prompt plus what it has accepted), instead of +# from the server-wide counter, so a request's forced counts do not depend on what the server ran before (A/B draws +# become history-independent). Production sampling never forces counts. Kill: False (the shared counter, the in-tree +# sequence). +FORCE_AL_PER_REQUEST = True + +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) # capacity: decode rows +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + QKV_SPARSE: _ALL_ROWS, + QKV: _ALL_ROWS, + O_PROJ: _ALL_ROWS, + SHARED_GATE_UP: _ALL_ROWS, + SHARED_DOWN: _ALL_ROWS - frozenset(range(17, 33)), + DENSE_GATE_UP: _ALL_ROWS, + DENSE_DOWN: _ALL_ROWS, +} # capacity: decode rows per projection +# On a bf16 activation the GEMV quantizes it in its prologue, which pays only up to 8 32-element blocks per thread: +# the most rows each projection takes. +MX_GEMV_BF16_MAX_TOKENS: Dict[Tuple[int, int], int] = { + QKV_SPARSE: 16, + QKV: 16, + O_PROJ: 32, + SHARED_GATE_UP: 16, + SHARED_DOWN: 16, + DENSE_GATE_UP: 8, + DENSE_DOWN: 16, +} # capacity: fixed: measured bf16-GEMV crossover per projection +# MX-TAC (the TP2 target's, ported): the decode-graph MXFP8 projections the GEMVs do not serve (SHARED_DOWN at 17..32 +# rows, rows over DECODE_MAX_TOKENS or a bf16 activation over MX_GEMV_BF16_MAX_TOKENS, or a GEMV decline) take one fixed +# decision instead of the AutoTuner's per-start, per-rank choice: FlashInfer's CuTe DSL quantize and CuTe DSL GEMM. All +# of stock's candidates (quantize {TRT-LLM, CuTe DSL} x GEMM {CUTLASS, CuTe DSL}) give the same bits at the TP4 +# projections (gemm/test_modeling_v2_m3_pinned_gemm.py), so the pin is stock's output on every rank and start. Kill +# constant: False returns these projections to the AutoTuner. +MX_TAC_TP4 = True +# The tactics below are MXFP8QuantizeRunner.CUTE_DSL / FlashInferMXFP8GemmRunner.CUTE_DSL (0). +MXFP8_DECODE_QUANT_TACTIC = 0 # capacity: fixed: a tactic id, not a size +MXFP8_DECODE_GEMM_TACTIC = 0 # capacity: fixed: a tactic id, not a size + +Key = Hashable + + +@dataclass(frozen=True) +class RowRange: + """The shapes one fast path takes, and the catalog entry whose cells certify each of them. + + ``keys`` are the covered shapes: row counts for most paths, (query rows per request, requests) for the attention + kernels, a variant name for thresholds. ``certifies`` maps a cell's spec of ``entry`` to the keys that cell + drives (empty for a cell of another variant). + """ + + path: str + keys: FrozenSet[Key] + entry: OpWrapper + certifies: Callable[[Mapping[str, Any]], Iterable[Key]] + + def covers(self, key: Key) -> bool: + return key in self.keys + + def certified(self) -> FrozenSet[Key]: + return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + + +def _rows(lo: int, hi: int) -> FrozenSet[int]: + return frozenset(range(lo, hi + 1)) + + +def _spec_rows(field: str, **match: Any) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + """A cell's ``spec[field]`` when the cell's other fields match ``match``.""" + + def certifies(spec: Mapping[str, Any]) -> Iterable[Key]: + if any(spec.get(name) != value for name, value in match.items()): + return () + return (spec[field],) + + return certifies + + +def _const(*keys: Key) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + return lambda spec: keys + + +def _sparse_decode_producer(spec: Mapping[str, Any]) -> Iterable[Key]: + if not spec["identity"]: + return () + rows = spec["decode_query_len"] * len(spec["seq_lens"]) + return ("multi_query" if rows >= SPARSE_DECODE_MQ_MIN_TOKENS else "per_row",) + + +def _long_select_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The key widths a cell's rows reach: every width up to its blocks' when rows of mixed lengths are drawn.""" + blocks = spec["num_blocks"] + widest = 2 if blocks <= 2048 else 4 if blocks <= 4096 else 9 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_KEY_WIDTHS if w <= widest) + return (widest,) + + +EMBED_FIRSTNORM = RowRange( + "embed_firstnorm", + _rows(1, EMBED_FIRSTNORM_MAX_TOKENS), + m3_embedding_firstnorm, + _spec_rows("rows"), +) +SPARSE_DECODE_PRODUCER = RowRange( + "sparse_decode_producer", + frozenset({"per_row", "multi_query"}), + m3_sparse_decode, + _sparse_decode_producer, +) +# Every token count of a sparse layer on an NVFP4 KV cache at the TP4 heads (prefill and decode): the producer takes the +# geometry, not a row range. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp4"}), m3_producer, _const("nvfp4_tp4")) +INDEX_SCORE = RowRange( + "m3_index_score", + frozenset({"early", "prefetch"}), + m3_index_score, + lambda spec: ("early" if spec["early"] else "prefetch",), +) + + +def index_score_t2sc_schedule(requests: int) -> Tuple[int, bool]: + """T2SC's (CTAs per request, early issue) at this many requests (a request count in INDEX_SCORE_T2SC_SPLIT_K).""" + return INDEX_SCORE_T2SC_SPLIT_K[requests], requests <= INDEX_SCORE_EARLY_MAX_REQUESTS + + +def _index_score_t2sc_requests(spec: Mapping[str, Any]) -> Iterable[Key]: + """A cell certifies its request count when it drives this target's 1 index head with that count's schedule.""" + requests = len(spec["seq_lens"]) + if spec.get("heads") != 1 or requests not in INDEX_SCORE_T2SC_SPLIT_K: + return () + if (spec["split_k"], spec["early"]) != index_score_t2sc_schedule(requests): + return () + return (requests,) + + +INDEX_SCORE_T2SC_PATH = RowRange( + "m3_index_score_t2sc", + frozenset(INDEX_SCORE_T2SC_SPLIT_K) if INDEX_SCORE_T2SC else frozenset(), + m3_index_score_t2sc, + _index_score_t2sc_requests, +) +LONG_SELECT = RowRange("long_select", frozenset(), m3_select_long, _long_select_widths) + + +def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The 384-thread select's key widths a cell's rows reach (every width up to its blocks' for mixed rows).""" + blocks = spec["num_blocks"] + widest = 4 if blocks <= 4 * 384 else 11 if blocks <= 11 * 384 else 24 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_DSL_KEY_WIDTHS if w <= widest) + return (widest,) + + +LONG_SELECT_DSL = RowRange( + "long_select_dsl", + frozenset(LONG_SELECT_DSL_KEY_WIDTHS) if SELECT_LONG_DSL else frozenset(), + m3_select_long_dsl, + _long_select_dsl_widths, +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS) if MC_AR_TP4 else frozenset(), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=0, drafter=None), +) +# Rows of a fused-draft-loop all-reduce boundary (drafter hidden 6144) on the multicast all-reduce's drafter form. +DRAFT_MC_AR = RowRange( + "m3_draft_mc_allreduce_norm", + _rows(1, DRAFT_MC_AR_MAX_ROWS), + m3_mc_allreduce_norm, + _spec_rows("m", drafter=True), +) +DENSE_PUSH = RowRange( + "dense_push", + _rows(1, DENSE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +ATTN_PUSH = RowRange( + "attn_push", + _rows(1, ATTN_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MOE_PUSH = RowRange( + "moe_push", + _rows(1, MOE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MC_FINALIZE = RowRange( + "mc_finalize_allreduce_norm", + _rows(MC_FINALIZE_MIN_TOKENS, MC_FINALIZE_MAX_TOKENS), + m3_mc_finalize_allreduce_norm, + _spec_rows("m"), +) +MOE_ENGINE = RowRange( + "m3_moe_engine", _rows(1, MOE_ENGINE_MAX_TOKENS), m3_moe_engine, _spec_rows("tokens") +) +MOE16_ENGINE = RowRange( + "m3_moe16_engine", + _rows(MOE16_MIN_TOKENS, MOE16_MAX_TOKENS), + m3_moe16_engine, + _spec_rows("tokens"), +) +# The router of a pure generation step of 9..16 rows when the MoE runs unfused (above the MoE engine's 8; at 9..13 an +# MoE16 engine miss, 14..16 always): the TF32 tensor-core router (minimax_m3_router_tc, 16-CTA clusters, one reduction +# order), the arithmetic the built-in gate takes on SM107 at these rows. Devices that cannot co-schedule 16-CTA +# clusters keep the fp32 router. Context steps never take it. +ROUTER_TC_GEN_MIN_ROWS = 9 # capacity: fixed: past the MoE engine's 8 rows +ROUTER_TC_GEN_MAX_ROWS = 16 # capacity: fixed: the built-in gate's router_tc range +ROUTER_TC_GEN = RowRange( + "m3_router_tc", + _rows(ROUTER_TC_GEN_MIN_ROWS, ROUTER_TC_GEN_MAX_ROWS), + m3_router_tc, + _spec_rows("m"), # every cell checks the logits, with or without the NVFP4 activation +) +# The GEMM paths are keyed (projection (N, K), rows): each projection runs its own split-K / ring / MMA +# configuration, so a row count is certified per projection. +MX_GEMV = RowRange( + "m3_mx_gemv", + frozenset((nk, m) for nk, rows in MX_GEMV_ROWS.items() for m in rows), + m3_mx_gemv, + lambda spec: ((spec["nk"], spec["m"]),), +) +MX_GEMV_BF16 = RowRange( + "m3_mx_gemv_bf16", + frozenset((nk, m) for nk, hi in MX_GEMV_BF16_MAX_TOKENS.items() for m in _rows(1, hi)), + m3_mx_gemv_bf16, + lambda spec: ((spec["nk"], spec["m"]),), +) +# (draft length, requests) of a greedy verify step. +SHARD_ACCEPT = RowRange( + "m3_shard_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SHARD_ACCEPT_MAX_REQUESTS + 1)), + m3_shard_accept, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) + +# Requests of a decode step whose dense-layer attention runs at a capped KV split. +DENSE_DECODE_SPLIT = RowRange( + "m3_dense_decode", + frozenset(DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), + m3_dense_decode, + lambda spec: (len(spec["kv_lens"]),) if spec["capped"] else (), +) + +# (query rows per request, requests) of a decode step whose dense-layer attention runs on PrimTS split-KV decode. +FULLCTX_DECODE = RowRange( + "m3_fullctx_decode", + frozenset(FULLCTX_DECODE_RECIPES) if FULLCTX_DECODE_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + +# (query rows per request, requests) of a generation-only drafter pass that runs on the full-context decode. +DRAFTER_FULLCTX = RowRange( + "m3_drafter_fullctx", + frozenset(FULLCTX_DECODE_RECIPES) if DRAFTER_FULLCTX_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + +# Verify rows of a fused-verify step whose target logits come from the coded head. +TARGET_HEAD = RowRange( + "m3_target_head", + _rows(1, TARGET_HEAD_MAX_ROWS), + m3_target_head, + _spec_rows("m"), +) + +# (draft length, requests) of a fused draft loop, and the rows of its passes' draft input: the first pass's +# requests x (draft length + 1), one per request on the later passes. +DRAFT_STEP = RowRange( + "m3_draft_step", + frozenset((DRAFT_STEP_DRAFT_LEN, n) for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1)), + m3_draft_step, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) +# ((N, K), rows) of a drafter BF16 linear on its low-M kernel. +DRAFT_LOWM = RowRange( + "m3_draft_lowm", + frozenset((nk, m) for nk, rows in DRAFT_LOWM_ROWS.items() for m in rows), + m3_draft_lowm, + lambda spec: ((spec["nk"], spec["m"]),), +) +DRAFT_INPUT = RowRange( + "m3_draft_input", + frozenset( + rows + for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1) + for rows in (n, n * (DRAFT_STEP_DRAFT_LEN + 1)) + ), + m3_draft_input, + _spec_rows("m"), +) + +# The step-loop hooks (RuntimeHooks), one switch each. They move a decode step's input preparation +# off the GPU timeline between step graphs. +STEP_PACK_INPUTS = True +# The pack's host copies go out in one parallel scatter (a program per copy) when their destinations do not overlap: the +# serial one-CTA scatter pays one dependent global round trip per copy (SM107: 12.5 / 15.6 / 31.7 us at 16 / 24 / 40 +# copies vs 6.5 / 6.5 / 6.9 us parallel). +STEP_PACK_PARALLEL_SCATTER = True +# The pack is issued by the decode graph's own prologue (a copy of a pinned window + the parallel scatter), so no +# operation is left between the step graphs (needs STEP_PACK_PARALLEL_SCATTER). +STEP_PACK_IN_GRAPH = True +# The window reaches the prologue's device buffer on a side copy stream while the previous graph runs; the prologue +# waits on that copy and scatters device to device, so the replay no longer starts with a host copy. Needs +# STEP_PACK_IN_GRAPH. +STEP_PACK_EARLY_COPY = True +# The step epoch (step_inputs.step_epoch): one int64 device counter the decode graph's prologue increments once per +# replay (program 0 of its first scatter launch, no extra kernel) and never resets. The per-layer megakernel pieces +# tag their per-step device flags with it (contract v1: tag = epoch << 8 | layer). Needs STEP_PACK_IN_GRAPH. +STEP_PACK_EPOCH = True +# The prologue is one launch of this target's op (ops/spec/m3_step_prologue, minimaxM3Sm107Tp4/stepPrologue.cu): the +# host copies, a CTA barrier, the device copies, the header clear and the epoch, where the default prologue runs two +# Triton scatter launches and an ATen fill (3 graph kernels -> 1, no Triton at the head of the decode graph). Same bytes +# written. Needs STEP_PACK_EARLY_COPY. Kill: False. +STEP_PACK_CPP_PROLOGUE = True +# PRE-WAIT WRITER CONTRACT (rubin-20261015/PREWAIT_AUDIT.md item b): the kernels that read step inputs (seq_lens, +# block tables, positions, cache slots, the sampled-accept pre-wait parameters) before their grid-dependency wait rely +# on every writer of them being a host-to-device copy before the replay or a PLAIN launch earlier in the step: +# msa_kv_update, the step-input prologue (Triton scatter or the C++ op), the spec fold tail (Triton or C++), the pack +# scatter and any step-head op (the spec step correction, STEP_SPEC_OFFSETS_ONE_LAUNCH). Never give one of them a PDL +# attribute with an early trigger. +STEP_SKIP_UNCHANGED_H2D = True +# A changed packed 2-D metadata upload stages only each row's changed column span. On a step where a request crosses +# into a new page the MSA block table changed by one entry, but the whole padded table (32 KB a request) was staged, +# overflowing the 16 KB in-graph window into an eager upload and scatter between the step graphs (57-69 us at the +# ~1 in 19 block-crossing steps; TTFT inter-graph scan). Kill: False. +STEP_H2D_CHANGED_SPAN = True +# A speculative overlap step re-derives the MSA KV-length metadata once, after the forward's KV-length correction, +# instead of also before it (RuntimeHooks.kv_lens_update_once): the first pass (a fill, the out_cache_loc tail fill and +# an msa_kv_update at the head of the decode graph, 2.13 us at TP4) is rewritten in full by the second. Kill: False. +STEP_KV_LENS_UPDATE_ONCE = True +# The overlap scheduler's speculative step correction (position_ids += offsets, kv_lens += offsets) runs as this +# target's one plain launch (RuntimeHooks.spec_offsets_op, ops/spec/m3_spec_offsets) instead of two ATen adds at the +# head of the decode graph: the same int32 adds. TP8 graph bench (the work is per token, not per rank; 64 steps +# per graph, a predecessor kernel each): -1.27 us/step at 4 tokens / 1 request, -0.72..-0.78 at 8-32 tokens. +# Kill: False. +STEP_SPEC_OFFSETS_ONE_LAUNCH = True +STEP_LOOP_ON_EXECUTION_STREAM = True +STEP_MERGE_TUNED_TACTICS = True +STEP_SKIP_GREEDY_RNG_UPLOAD = True +# Sampled speculative batches stage their four per-step Philox seed/offset uploads into the step-input pack +# (STEP_PACK_INPUTS) instead of issuing four copies between the step graphs. +STEP_PACK_SPEC_RNG_UPLOADS = True +# The spec step's small uploads: the accepted-draft counts ride the step-input pack, and the sampler's slot ids stay in +# a device buffer re-uploaded only when they change (a pageable upload every step otherwise). +STEP_SPEC_SMALL_H2D = True +STEP_SPEC_FOLD = True +# The forward of a decode-only step stores its outputs into the spec sampler's store and stages the next step's inputs +# at its tail (one launch inside the step graph), so neither the sampler's store nor the next step's gathers run +# between the step graphs (M3SpecStepFold.in_forward). +STEP_SPEC_FOLD_IN_FORWARD = True +# The forward's tail (STEP_SPEC_FOLD_IN_FORWARD) runs as this target's C++ op (ops/m3_fold_tail cpp=, +# minimaxM3Sm107Tp4/foldTail.cu) instead of the Triton kernel: the same integer copies, one launch either way (no +# Triton on the generation decode path). Kill: False. +STEP_FOLD_TAIL_CPP = True +# With STEP_SPEC_FOLD_IN_FORWARD the spec sampler's host copies of its store (new tokens, lengths, draft tokens) still +# ran on the main stream between the step graphs (the forward, not the sampler, wrote the store, and it rewrites it at +# the next replay's tail). They run on the D2H side stream instead, and the forward's tail waits on an event the +# sampler records after them before it rewrites the store (M3SpecStepFold.store_copies_done). +STEP_SPEC_STORE_COPY_SIDE = True +# The spec sampler copies its per-slot store to the host on its D2H side stream, so the copies overlap the next +# step's input staging instead of sitting between the step graphs. +STEP_SAMPLER_D2H_SIDE_STREAM = True +# The iteration log times the device step with one event per iteration boundary on a side stream, so no timing +# record sits between two step graphs on the main stream (the reported step times are the same). +STEP_ONE_TIMING_EVENT = True +MSA_SHORT_SEQUENCE_IDENTITY = True +# The sparse decode (the LC, the Triton merge) also writes the MXFP8 o_proj hand-off (bit +# for bit torch.ops.trtllm.mxfp8_quantize of its BF16 rows) and publishes it to m3_prequant.SLOT. Kill: False publishes +# nothing, so the o_proj quantizes its BF16 input itself as the built-in path does. +MERGE_MXFP8_HANDOFF = True +MSA_KV_UPDATE = True +MSA_SKIP_UNCHANGED_BLOCK_OFFSETS = True +MSA_REUSE_POOL_VIEWS = True +# Step-input preparation fills the MSA block-id table by slice-copying each request's page-index view (MSA +# RuntimeHooks.block_ids_from_views) instead of building a Python list of the whole padded row; the table is the same. +MSA_BLOCK_IDS_FROM_VIEWS = True + +FAST_PATHS: Tuple[RowRange, ...] = ( + EMBED_FIRSTNORM, + DENSE_DECODE_SPLIT, + FULLCTX_DECODE, + DRAFTER_FULLCTX, + SPARSE_DECODE_PRODUCER, + PRODUCER, + INDEX_SCORE, + INDEX_SCORE_T2SC_PATH, + LONG_SELECT, + LONG_SELECT_DSL, + MC_ALLREDUCE, + DRAFT_MC_AR, + ATTN_PUSH, + DENSE_PUSH, + MOE_PUSH, + MC_FINALIZE, + MOE_ENGINE, + MOE16_ENGINE, + ROUTER_TC_GEN, + MX_GEMV, + MX_GEMV_BF16, + SHARD_ACCEPT, + TARGET_HEAD, + DRAFT_STEP, + DRAFT_LOWM, + DRAFT_INPUT, +) + +# -------------------------------------------------------------------------------------------------------------------- +# Selection: pure functions of the step shape +# -------------------------------------------------------------------------------------------------------------------- + + +def plan_for(num_tokens: int) -> DecodePlan: + for plan in DECODE_PLANS: + if num_tokens <= plan.max_tokens: + return plan + raise AssertionError(f"no decode plan covers {num_tokens} tokens") + + +def is_decode_step(num_tokens: int) -> bool: + """Whether a pure generation step of ``num_tokens`` rows runs the decode kernels (else generic layers).""" + return num_tokens <= DECODE_MAX_TOKENS + + +def moe_path(num_tokens: int) -> Optional[str]: + """The MoE kernel of a decode step: an engine, the unfinalized MoE + finalize all-reduce, or None (generic).""" + if MOE_ENGINE.covers(num_tokens): + return MOE_ENGINE.path + if MOE16_ENGINE.covers(num_tokens): + return MOE16_ENGINE.path + if MC_FINALIZE.covers(num_tokens): + return MC_FINALIZE.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# The dense layers' swigluoai activation (where the fold above would apply) as this target's C++ op +# (ops/gemm/m3_dense_swiglu, minimaxM3Sm107Tp8/swigluOai.cu: bit for bit the Triton silu_and_mul), PDL-launched with +# an early release so the down GEMV stages its weights during it; the down GEMV reads the activation only after its own +# grid wait. TP2 chain bench: -0.69 us per dense layer vs the Triton kernel. Kill constant: False +# runs the Triton kernel. +DENSE_SWIGLU_CPP = True +DENSE_SWIGLU_CPP_PATH = "m3_dense_swiglu" +# FlashInfer's CuTe DSL MXFP8 GEMM on the producer's MXFP8 activation (rows the GEMV does not serve). +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +# The MXFP8 quantize (tuned backend) and FlashInfer's CuTe DSL GEMM: every row the decode kernels do not serve. +GENERIC_GEMM_PATH = "mxfp8_gemm" + + +def linear_path(n: int, k: int, num_tokens: int, prequant: bool, low_m_gemm: bool) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. ``low_m_gemm`` is accepted for the call sites' + symmetry: no low-M GEMM configuration is certified at TP4, so it selects nothing. + """ + if prequant: + if num_tokens in MX_GEMV_ROWS.get((n, k), ()): + return MX_GEMV_PATH + return PREQUANT_GEMM_PATH + if num_tokens <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0): + return MX_GEMV_BF16_PATH + return GENERIC_GEMM_PATH + + +# -------------------------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A covered decode shape could not take its fast path under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() + + +def set_require(require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: + """The (PDL launch, early trigger, trigger phase) keys of the 384-thread select a decode step can launch, compiled + before warmup; none when ``SELECT_LONG_DSL`` is off.""" + if not LONG_SELECT_DSL.keys: + return () + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + pdl = get_env_enable_pdl() + keys = [(pdl, False, SELECT_LONG_TRIGGER_PHASE), (pdl, False, 0)] + if SELECT_EARLY_MAX_ROWS: + keys.append((pdl, True, 0)) + return tuple(dict.fromkeys(keys)) + + +def missed(path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if _require[0]: + raise FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " + f"{reason}" + ) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def taken(path: str, rows: Optional[int], what: str = "") -> None: + """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line + every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" + at = "" if rows is None else f" at {rows} rows" + detail = f" ({what})" if what else "" + _log_once( + logger.info, + path, + f"taken: {rows}" + (f": {what}" if what else ""), + f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", + ) + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +# Kernel modules report from inside their calls through ops/m3_report: the same policy. +m3_report.install( + on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed +) + + +# -------------------------------------------------------------------------------------------------------------------- +# Target constants +# -------------------------------------------------------------------------------------------------------------------- + +# The shared expert folds into the routed engine from one token (dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# SMs left to the concurrent kernels while a persistent routed engine runs. +MOE_ENGINE_RESERVED_SMS = 24 # capacity: fixed: an SM split, not a size +# Id of the target's final hidden state in Eagle3's ``layers_to_capture``. +FINAL_HIDDEN_LAYER_ID = -1 # capacity: fixed: an Eagle3 layer id, not a size +# The trtllm-gen MoE runner's routing kind for MiniMax-M3 (sigmoid + correction bias, renormalized, scaled) and its +# gated activation (SwiGLU with the per-expert alpha / beta / clamp operands). +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) + +# -------------------------------------------------------------------------------------------------------------------- +# Streaming: also emit a response on decoding iterations 2..FIRST_CONTENT_EMIT_ITERS (RuntimeHooks), so a request's +# first content leaves one step after its first token. 1: the stock schedule. +FIRST_CONTENT_EMIT_ITERS = 2 + +# The runtime-value dump (TRTLLM_MODELING_V2_VALUE_DUMP=1): the metadata of every fast-path and entry decision a step +# takes, recorded on eager steps and captures (replay runs no Python), printed once at exit. +# -------------------------------------------------------------------------------------------------------------------- + +_VALUE_DUMP = os.environ.get("TRTLLM_MODELING_V2_VALUE_DUMP") == "1" +_VALUES: set = set() + + +def _record(*meta: Hashable) -> None: + if _VALUE_DUMP: + _VALUES.add(meta) + + +def _print_values() -> None: + if _VALUES: + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "0")) + for meta in sorted(_VALUES, key=repr): + print(f"[modeling_v2 value dump] rank {rank}: {meta}", flush=True) + + +if _VALUE_DUMP: + atexit.register(_print_values) + +# -------------------------------------------------------------------------------------------------------------------- +# Configuration: the text tower and the runtime the target assumes + +# -------------------------------------------------------------------------------------------------------------------- + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing + sees only its KV cache algorithm; this check at construction refuses every other runtime. + + The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index + projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. + """ + sparse = model_config.sparse_attention_config + if ( + sparse is None + or getattr(sparse, "algorithm", None) != "minimax_m3" + or sparse.implementation != "msa" + ): + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs +# -------------------------------------------------------------------------------------------------------------------- + + +def _mxfp8_linear( + core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + _record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +def _gemv_variants(n: int, k: int, push: Optional[str], world: int): + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch routes to an [N, K] + projection's GEMV: an MXFP8 activation on the rows ``MX_GEMV_ROWS`` gives it, a bf16 one up to + ``MX_GEMV_BF16_MAX_TOKENS``; each pushing when the projection's boundary push covers the rows.""" + push_rows = {"attn": ATTN_PUSH.keys, "dense": DENSE_PUSH.keys, "moe": MOE_PUSH.keys}.get( + push, frozenset() + ) + variants = [] + for rows in range(1, DECODE_MAX_TOKENS + 1): + for x_bf16, served in ( + (False, rows in MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + bucket = next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + variants.append((rows, x_bf16, (world, bucket, AR_PUSH_MODE))) + return variants + + +def _precompile_cfg( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + cfg: tuple, + use_pdl: bool, + ar: Optional[Tuple[int, int, int]] = None, +) -> Optional[tuple]: + """``mx_gemv.precompile`` for a call with this target's (split, ring) override ``cfg`` (``mx_gemv_cfg``): the + shared precompile resolves the shared table's config only. Compiles without launching, outside capture, also while + compiles are frozen; returns the kernel's key, or None when such a call is declined.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = mx_gemv._resolve(m, n, k, False, True, cfg, use_pdl, push) + if resolved is None: + return None + key, runtime_ar = resolved + if mx_gemv.is_compiled(key): + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("_precompile_cfg allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros(m, k, dtype=torch.float8_e4m3fn, device=dev) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = mx_gemv._views(x, xsf, weight, weight_scale, y, False) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + mx_gemv._compile(key, args, m, runtime_ar, stream, mx_gemv._extra_args(dev)) + return key + + +_MX_TAC_ENGAGED: set = set() + + +def _note_mx_tac(kind: str, n: int, k: int, rows: int) -> None: + """MX-TAC's engagement line, once per (activation kind, projection, rows).""" + key = (kind, n, k, rows) + if key not in _MX_TAC_ENGAGED: + _MX_TAC_ENGAGED.add(key) + logger.info( + f"[M3 sm_107 tp4] MX-TAC engaged: pinned MXFP8 decision (quantize tactic " + f"{MXFP8_DECODE_QUANT_TACTIC}, GEMM tactic {MXFP8_DECODE_GEMM_TACTIC}) " + f"for the [{n}, {k}] projection at {rows} rows, {kind} activation" + ) + + +def _uncompiled_note() -> str: + key = mx_gemv.last_uncompiled + return "" if key is None else f" (its kernel {key} was not compiled before warmup)" + + +def _linear( + core, x: torch.Tensor, wk: str, sk: str, low_m_gemm: bool, push: Optional[tuple] = None +) -> Tuple[torch.Tensor, bool]: + """One MXFP8 projection ``x @ w[wk]^T``; returns (output, whether its GEMV pushed into ``push``). + + Outside decode CUDA graphs: the engine-tuned MXFP8 quantize + GEMM. Inside them the projection's (N, K) and the + step's rows pick one path (``linear_path``): + + 1. an activation its producer already quantized to MXFP8 (``m3_prequant.SLOT``, taken here): the tcgen05 decode + GEMV on the rows it serves, else the decode-graph MXFP8 GEMM on that activation (AutoTuner-selected backend + per token bucket); + 2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; + 3. otherwise the engine-tuned MXFP8 quantize + GEMM. + + A kernel declining rows its path covers is a miss; the chain then continues down the list. ``push`` (an + ``m3_mx_gemv`` ``ar`` tuple) is the boundary all-reduce exchange the GEMV writes its rows into. ``low_m_gemm`` + is accepted for the call sites' symmetry (``linear_path``): no TP4 low-M GEMM configuration is built. + """ + w, s = core.w[wk], core.w[sk] + x = x.contiguous() + mode = _mx_mode(core.mxfp8) + if not mode.decode_graph: + return _mxfp8_linear(core, x, w, s, mode), False + use_pdl = get_env_enable_pdl() + rows = x.shape[0] + n, k = w.shape + prequant = m3_prequant.SLOT.take(x) + if prequant is None and x.dtype != torch.bfloat16: + path = GENERIC_GEMM_PATH + else: + path = linear_path(n, k, rows, prequant is not None, low_m_gemm) + _record("linear", (n, k), rows, path, push is not None) + output = None + # MX-TAC: the decode-graph MXFP8 quantize and GEMM at the target's fixed tactics, no AutoTuner. + pinned = MX_TAC_TP4 and core.mxfp8.quant_method.tune_decode_graph_backends + if prequant is not None: + act, act_scale = prequant + if path == MX_GEMV_PATH: + cfg = mx_gemv_cfg(n, k, rows) + output = m3_mx_gemv(act, act_scale, w, s, True, cfg, use_pdl, push) + if output is not None: + split = f", split {cfg[0]}" if cfg is not None else "" + taken(path, rows, f"[{n}, {k}] mxfp8 in{split}") + return output, push is not None and mx_gemv.last_pushed + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if path == PREQUANT_GEMM_PATH: + uncovered(MX_GEMV_PATH, f"the [{n}, {k}] projection at rows the GEMV does not serve") + if pinned: + output = m3_mxfp8_pinned_gemm( + act, act_scale, w, s, torch.bfloat16, MXFP8_DECODE_GEMM_TACTIC + ) + _note_mx_tac("mxfp8", n, k, rows) + return output, False + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16(x, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] bf16 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, rows, f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note() + ) + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + if pinned and x.dtype == torch.bfloat16: + act, act_scale = m3_mxfp8_pinned_quantize(x, MXFP8_DECODE_QUANT_TACTIC) + output = m3_mxfp8_pinned_gemm( + act, act_scale, w, s, torch.bfloat16, MXFP8_DECODE_GEMM_TACTIC + ) + _note_mx_tac("bf16", n, k, rows) + return output, False + # Stock's decode-graph dispatch for this shape (the tuned quantize + GEMM). + return _mxfp8_linear(core, x, w, s, mode), False + + +def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: + """A dense MLP's SwiGLU-OAI activation: this target's C++ op (``DENSE_SWIGLU_CPP``, bit for bit the Triton kernel, + PDL-launched with an early release) up to its certified rows, in any graph or eager step; else the Triton + kernel.""" + rows = int(gu.shape[0]) + path = DENSE_SWIGLU_CPP_PATH + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + if DENSE_SWIGLU_CPP and not is_torch_compiling(): + if rows > max_rows: + uncovered(path, f"the dense swigluoai op: {rows} rows, past its certified {max_rows}") + else: + act = m3_dense_swiglu( + gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + ) + if act is not None: + taken(path, rows) + return act + uncovered(path, "the gate_up layout") + return m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer boundaries: all-reduce + residual + RMSNorm +# -------------------------------------------------------------------------------------------------------------------- + + +def _can_mc(x, folded: bool = False) -> bool: + """Whether the generic boundary takes the multicast all-reduce + residual + RMSNorm (``GENERIC_MC_AR_MAX_TOKENS``, + with the shared + routed add folded in ``GENERIC_MC_AR_FOLDED_MAX_TOKENS``).""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (GENERIC_MC_AR_FOLDED_MAX_TOKENS if folded else GENERIC_MC_AR_MAX_TOKENS) + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + +def _mc_allreduce_norm(core, x, x2, residual, norm_w, *, quant=False, prepushed=0, capture=None): + """The multicast all-reduce + optional second addend + residual + RMSNorm (``m3_mc_allreduce_norm``), or None + without a workspace or for operands the kernel does not take (the same on every rank). With ``quant`` the normed + rows are also published, MXFP8-quantized, to ``m3_prequant.SLOT`` for the next MXFP8 projection; ``capture`` = + (out, fc_norm weight, eps): the boundary also writes the Eagle3 capture norm of the updated residual. Up to + ``MC_AR_CLUSTER_MAX_ROWS`` rows the one-shot 8-CTA cluster form runs; ``MC_AR_CARVEOUT`` sets its shared-memory + carveout.""" + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + if workspace is None: + return None + _record("m3_mc_allreduce_norm", x.shape[0], x2 is not None, quant, prepushed) + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm_w, + core.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + oneshot_cluster=x.shape[0] <= MC_AR_CLUSTER_MAX_ROWS, + carveout=MC_AR_CARVEOUT, + capture=capture, + ) + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm.""" + _record("m3_allreduce", x.shape[0]) + return m3_allreduce( + x, + core.allreduce, + AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ), + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The MoE: router, routed experts (trtllm-gen NVFP4), shared expert +# -------------------------------------------------------------------------------------------------------------------- + + +@dataclass +class _MoeScalars: + """The routed experts' kernel scalars of one layer, derived from the checkpoint's per-tensor scales once the + weights are real (``MiniMaxM3Tp4ModelingV2._derive_moe``).""" + + fc31_input_scale: torch.Tensor + fc2_input_scale: torch.Tensor + fc31_alpha: torch.Tensor + fc2_alpha: torch.Tensor + fc31_scale_c: torch.Tensor + act_alpha: torch.Tensor + act_beta: torch.Tensor + act_clamp: torch.Tensor + + +def _router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The fp32 router logits of MoE layer ``i`` when the MoE runs unfused (the decode steps' engines route + in-kernel): ``F.linear(float(h), w)``, the built-in gate's arithmetic outside its tensor-core router's rows. The + experts quantize their bf16 input themselves.""" + _record("m3_fp32_linear", h.shape[0] if h.dim() == 2 else 0) + return m3_fp32_linear(h, core.w[f"l{i}_router"]) + + +def _decode_router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The router logits of a pure generation step: ``ROUTER_TC_GEN``'s rows take the TF32 tensor-core router when + the device co-schedules 16-CTA clusters (decided at prebuild) and the operands are the built-in gate's (2-D bf16, + unit column stride, 16-byte aligned rows); every other step the fp32 router (``_router_logits``).""" + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + ROUTER_TC_GEN.covers(rows) + and core.router_tc_cluster16 + and h.dtype == torch.bfloat16 + and h.stride(1) == 1 + and (h.stride(0) * 2) % 16 == 0 + ): + taken(ROUTER_TC_GEN.path, rows) + _record("m3_router_tc", rows) + return m3_router_tc(h, core.w[f"l{i}_router"], None, get_env_enable_pdl())[0] + return _router_logits(core, i, h) + + +def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): + """The trtllm-gen NVFP4 MoE runner over this rank's tensor-parallel slice of every expert, routing in-kernel. + + ``x_sf`` None means ``x`` is the bf16 activation, quantized here with the experts' global input scale. Returns + the finalized rows, or the unfinalized ``[gemm2 out, expert scales, expanded->permuted index]``. + """ + moe = core.moe[i] + if x_sf is None: + x, x_sf = fp4_quantize(x, moe.fc31_input_scale, 16, False, False) + rows = x.shape[0] + x_sf = x_sf.view(rows, -1).flatten() + _record("fp4_block_scale_moe_runner", rows, do_finalize) + outs = fp4_block_scale_moe_runner( + router_logits, + core.w[f"l{i}_router_bias"], + x, + x_sf.view(torch.float8_e4m3fn), + core.w[f"l{i}_e_w31"], + core.w[f"l{i}_e_s31"].view(torch.float8_e4m3fn), + None, + moe.act_alpha, + moe.act_beta, + moe.act_clamp, + core.w[f"l{i}_e_w2"], + core.w[f"l{i}_e_s2"].view(torch.float8_e4m3fn), + None, + moe.fc31_scale_c, + moe.fc31_alpha, + moe.fc2_alpha, + core.num_experts, + core.top_k, + None, + None, + core.moe_inter, + 0, + core.num_experts, + core.routed_scaling_factor, + ROUTING_MINIMAX2, + do_finalize, + act_type=ACT_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=core.max_num_tokens, + use_dp=False, + ) + return outs[0] if do_finalize else outs + + +def _split_chunk(rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +# -------------------------------------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------------------------------------- + + +def _nvfp4_layer(manager, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _producer_covered(core) -> bool: + """Whether ``PRODUCER`` covers this target's sparse layers: an NVFP4 KV cache at the TP4 heads (16 query, one KV + and one index head), with the range on. A covered layer that cannot take ``m3_producer`` is a miss.""" + return ( + core.main_kv_is_nvfp4 + and (core.heads, core.kv_heads, core.index_heads) == (16, 1, 1) + and PRODUCER.covers("nvfp4_tp4") + ) + + +def _producer_layout(kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 pools are the layout ``m3_producer`` is built for: [pages, 2, 1, 128, 64] with + its scale pool [pages, 2, 1, 128, 8].""" + return ( + kv_data is not None + and kv_scale is not None + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride()[2:] == (128 * 64, 64, 1) + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride()[2:] == (128 * 8, 8, 1) + and kv_scale.shape[0] == kv_data.shape[0] + ) + + +def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or data.shape[0] != scales.shape[0] + or tuple(data.shape[1:]) != (2, num_heads, 128, 64) + or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) + or data.dtype not in (torch.int8, torch.uint8) + or scales.dtype != torch.uint8 + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or cache.stride(1) < num_heads * cache.stride(2) + or cache.stride(0) < 2 * cache.stride(1) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + + +def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: + if covered: + missed(PRODUCER.path, rows, reason) + return None + + +def _sparse_qkv_producer( + core, i: int, packed, position_ids, kv_cache_manager, msa_idx_k_cache, msa_out_cache_loc +): + """The horizontal producer of a sparse layer: Q / index-Q (FP8) out, K / V / index-K written into the paged + caches (NVFP4 or FP8 main cache), for every step type (token-major). An NVFP4 cache at the TP4 heads takes + ``m3_producer`` (``PRODUCER``); a covered layer that cannot is a miss. Returns (q, index_q), or None for caches + neither kernel takes. + + ``kv_cache_manager``, ``msa_idx_k_cache`` (the per-layer index-K cache getter) and ``msa_out_cache_loc`` are the + calling target's ``step_args`` entries.""" + covered = _producer_covered(core) + rows = int(packed.shape[0]) + if packed.dtype != torch.bfloat16 or position_ids is None: + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + manager = kv_cache_manager + if manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = manager.get_buffers(i, kv_layout="HND") + index_k_cache = msa_idx_k_cache(i) + out_cache_loc = msa_out_cache_loc + layer_uses_nvfp4 = _nvfp4_layer(manager, i) + scale_buffers = ( + manager.get_block_scale_buffers(i, kv_layout="HND") if layer_uses_nvfp4 else None + ) + kv_quant_scale = core.inv_kv_scales[i] if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supported_main = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_fused_write_layout(buffers, scale_buffers, core.kv_heads) + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale.is_cuda + and kv_quant_scale.dtype == torch.float32 + and kv_quant_scale.numel() >= 3 + and kv_quant_scale.is_contiguous() + ) + else: + supported_main = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, core.kv_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= core.kv_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + cos_sin = core.rotary_cos_sin + supported_rope = ( + cos_sin.is_cuda + and cos_sin.dtype == torch.float32 + and cos_sin.is_contiguous() + and cos_sin.dim() == 3 + and tuple(cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main + or not supported_index + or not supported_rope + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < rows + ): + return _producer_miss( + covered and layer_uses_nvfp4, + rows, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the layout the " + "producer is built for", + ) + w = core.w + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + layer_uses_nvfp4 + and PRODUCER.covers("nvfp4_tp4") + and (core.heads, core.kv_heads, core.index_heads) == (16, 1, 1) + and _producer_layout(buffers, scale_buffers) + ): + early_trigger = PRODUCER_EARLY_TRIGGER and rows > PRODUCER_LATE_TRIGGER_MAX_ROWS + _record("m3_producer", rows, early_trigger) + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + core.heads, + get_env_enable_pdl(), + early_trigger, + PRODUCER_PREWAIT_STEP_INPUTS, + ) + # Every token count takes the producer: rows past a decode step log once as one line. + decode_rows = is_decode_step(rows) + taken( + PRODUCER.path, + rows if decode_rows else None, + f"nvfp4_tp4, {'entry' if early_trigger else 'exit'} trigger, " + f"prewait_step_inputs={PRODUCER_PREWAIT_STEP_INPUTS}" + if decode_rows + else "nvfp4_tp4, past the decode rows", + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool {tuple(scale_buffers.shape)} is not the " + "producer's layout", + ) + _record("m3_qkv_insert", rows, layer_uses_nvfp4) + q, index_q = m3_qkv_insert( + packed.contiguous(), + buffers.view(torch.uint8) if layer_uses_nvfp4 else buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.heads, + core.kv_heads, + core.index_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + ) + return q.flatten(1), index_q.flatten(1) + + +def _qk_norm_rope( + core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool, kv_dest: Optional[tuple] = None +) -> torch.Tensor: + """The per-head Gemma RMSNorm of Q and K and the partial NeoX RoPE of a layer's packed [Q | K | V], in one kernel: + in place, or an FP8 copy (Q / K normed and roped, V cast) when ``out_fp8``. A dense layer's FP8 copy carries the + BF16 output's bits cast to E4M3 (``m3_dense_qk_norm_rope_fp8``, ``DENSE_QKV_FP8_VIA_BF16``); with ``kv_dest`` = + (the layer's HND pool, the step's cache slots, live tokens, a [tokens, q heads, 128] FP8 buffer) that kernel also + stores the live tokens' K / V rows into the pool and every row's Q into the buffer.""" + qkv = qkv.contiguous() + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + out_fp8 + and DENSE_QKV_FP8_VIA_BF16 + and i not in core.sparse_layers + and qkv.dtype == torch.bfloat16 + and core.head_dim == 128 + ): + _record("m3_dense_qk_norm_rope_fp8", qkv.shape[0], kv_dest is not None) + return m3_dense_qk_norm_rope_fp8( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + True, + *(kv_dest or ()), + ) + assert kv_dest is None, "only the dense FP8 kernel stores K / V" + args = ( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + 1.0, # factor: no YARN + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + _record("qk_norm_rope", qkv.shape[0], out_fp8) + if out_fp8: + return m3_qk_norm_rope_fp8(*args) + return fused_qk_norm_rope(*args) + + +def _fp8_index_qk_norm_rope( + core, i: int, idx_qk: torch.Tensor, position_ids, msa_idx_k_cache, msa_out_cache_loc +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + _record("m3_fp8_indexer_qk_norm_rope", rows) + return m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer helpers both targets' bodies call: the Eagle3 capture hook and the norm after a layer +# -------------------------------------------------------------------------------------------------------------------- + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + +def _next_norm(core, i: int) -> torch.Tensor: + """The norm the boundary after layer ``i`` applies: the next layer's input norm, the final norm after the last.""" + return core.w[f"l{i + 1}_in_norm"] if i + 1 < core.num_layers else core.w["final_norm"] + + +# -------------------------------------------------------------------------------------------------------------------- +# The core +# -------------------------------------------------------------------------------------------------------------------- + +# Engine states (compiled kernel + workspace) per (device, token count[, shared fold]), shared by every layer: a step +# runs the layers one after another on one stream. +_ENGINE_STATES: Dict[tuple, object] = {} +_MOE16_STATES: Dict[tuple, object] = {} + + +class MiniMaxM3Tp4ModelingV2(DecoderModel): + """The MiniMax-M3 text model of this target: its weights, its runtime resources, and the per-step dispatch.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_runtime(model_config) + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + self.mapping = mapping + self.vocab_size = cfg.vocab_size + d = _weights.MODEL_WEIGHTS.dims(self) + self.num_layers = d.num_layers + self.hidden = d.hidden + self.heads = d.heads + self.kv_heads = d.kv_heads + self.head_dim = d.head_dim + self.index_heads = d.index_heads + self.index_dim = d.index_dim + self.q_size = d.heads * d.head_dim + self.kv_size = d.kv_heads * d.head_dim + self.moe_inter = d.moe_inter + self.num_experts = d.num_experts + self.top_k = int(cfg.num_experts_per_tok) + self.routed_scaling_factor = float(getattr(cfg, "routed_scaling_factor", 1.0)) + self.rms_eps = float(cfg.rms_norm_eps) + self.swiglu_alpha = float(getattr(cfg, "swiglu_alpha", 1.702)) + self.swiglu_beta = 1.0 # the (up + 1) offset of SwiGLU-OAI + self.swiglu_limit = float(getattr(cfg, "swiglu_limit", 7.0)) + self.max_num_tokens = model_config.max_num_tokens + self.moe_max_num_tokens = model_config.moe_max_num_tokens + self.sparse_layers = frozenset(_sparse_layer_ids(cfg)) + self.moe_layers = frozenset(_moe_layer_ids(cfg)) + mode = model_config.quant_config.quant_mode + self.main_kv_is_fp8 = bool(mode is not None and mode.has_fp8_kv_cache()) + self.main_kv_is_nvfp4 = bool(mode is not None and mode.has_fp4_kv_cache()) + # The sparse layers' QKV and index projections: one five-way packed projection, or the QKV projection and a + # separate BF16 index projection (fuse_qkv_index_projection=False) run on two streams. + self.fused_projection = bool( + getattr(model_config.sparse_attention_config, "fuse_qkv_index_projection", False) + ) + + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # The engine reads these two off the module tree: the MXFP8 GEMM state it tunes, and the opt-in to the + # automatic FlashInfer MXFP8 path on decode graphs. + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + + # The built-in fused all-reduce + residual + RMSNorm of the layer boundaries the multicast kernel does not take. + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + + # Streams: the shared / routed expert overlap and, when the MoE chunks its input, the chunking overlap. + self.aux_stream_moe = torch.cuda.Stream() + chunked = self.moe_max_num_tokens < model_config.max_num_tokens * mapping.dp_size + self.aux_stream_moe_chunk = torch.cuda.Stream() if chunked else None + self.moe_events = {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers} + self.moe_chunk_events = { + i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers + } + # The unfused sparse layers' index branch overlaps the QKV branch on its own stream. + self.aux_stream_attn = None if self.fused_projection else torch.cuda.Stream() + self.attn_events = ( + {} + if self.fused_projection + else {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.sparse_layers} + ) + + # Partial NeoX RoPE (rotary_dim of head_dim): the cos/sin table every attention kernel reads. + rope = RopeParams.from_config(cfg) + self.rope_theta = rope.theta + self.rotary_dim = int(rope.dim) + self.rotary_cos_sin = rope.create_rope_const_params(interleave=False)[1].reshape( + rope.max_positions, 2, -1 + ) + + # One MSA attention backend per layer (dense and sparse; the step metadata is the backend's). + quant_config = model_config.get_quant_config() + self.attn: List[MiniMaxM3MsaSparseAttention] = [] + for i in range(self.num_layers): + backend = create_attention( + model_config.attn_backend, + i, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=None, + quant_config=quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + q_scaling=1.0, + sparse_params=model_config.sparse_attention_config.to_sparse_params( + pretrained_config=cfg, layer_idx=i + ), + ) + backend.update_quant_config(quant_config) + if not isinstance(backend, MiniMaxM3MsaSparseAttention): + raise ValueError( + f"layer {i}: expected the MiniMax-M3 MSA backend, got {type(backend).__name__}" + ) + if i not in self.sparse_layers: + # The KV split trtllm-gen dense decode runs at, per request count. + backend.dense_decode_max_seq_len_by_batch = DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + self.attn.append(backend) + + # Built once the weights are real (``post_load_weights``): the KV (Q, K, V) scales, the MoE scalars, the + # decode engines and workspaces, the two targets. + self.kv_scales: Dict[int, torch.Tensor] = {} + self.inv_kv_scales: Dict[int, torch.Tensor] = {} + self.moe: Dict[int, _MoeScalars] = {} + self.mx_alpha: Optional[torch.Tensor] = None + self.moe_engines: Dict[int, Dict[Tuple[int, bool, bool], Optional[object]]] = {} + self.moe16_engines: Dict[int, Dict[int, Optional[object]]] = {} + self.embed_owners: Dict[tuple, DecodeEmbedding] = {} + # Whether the 9..16-row generation router (ROUTER_TC_GEN) runs; decided at prebuild. + self.router_tc_cluster16 = False + # Each sparse layer's selected block ids per query row at its last long-context decode, read by the previous + # sparse layer's decode to prefetch this layer's KV (``SPARSE_DECODE_LC_PREFETCH_NEXT_KV``). + self.n1_ids: Dict[int, torch.Tensor] = ( + { + i: torch.full( + (N1_MAX_ROWS, 16), + -1, + dtype=torch.int32, + device=torch.device("cuda", torch.cuda.current_device()), + ) + for i in sorted(self.sparse_layers) + } + if SPARSE_DECODE_LC_PREFETCH_NEXT_KV + else {} + ) + self._built_for: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def sm_scale(self, i: int) -> float: + return (self.head_dim**-0.5) / float(self.attn[i].q_scaling) + + # ---------------------------------------------------------------------------------------------- after the load + + def post_load_weights(self) -> None: + """Derive the kernels' state from the loaded weights, then build both targets and every decode kernel. + + Idempotent per set of weight storages: the loader's post-load walk reaches this module directly and through + the shell, and a second call on the same storages does nothing. + """ + identity = ( + self.w["embed"].data_ptr(), + self.w["l0_in_norm"].data_ptr(), + self.w["final_norm"].data_ptr(), + ) + if self._built_for == identity: + return + device = self.w["embed"].device + derived = nn.Module() + derived.register_buffer( + "mx_alpha", torch.ones(1, dtype=torch.float32, device=device), persistent=False + ) + for i in range(self.num_layers): + # No KV calibration in the checkpoint: unit (Q, K, V) scales and their reciprocals. + derived.register_buffer( + f"l{i}_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + derived.register_buffer( + f"l{i}_inv_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + for i in sorted(self.moe_layers): + for name, value in dataclasses.asdict(self._derive_moe(i)).items(): + derived.register_buffer(f"l{i}_{name}", value, persistent=False) + self.derived = derived + self.mx_alpha = derived.mx_alpha + self.kv_scales = {i: getattr(derived, f"l{i}_kv_scales") for i in range(self.num_layers)} + self.inv_kv_scales = { + i: getattr(derived, f"l{i}_inv_kv_scales") for i in range(self.num_layers) + } + self.moe = { + i: _MoeScalars( + **{ + f.name: getattr(derived, f"l{i}_{f.name}") + for f in dataclasses.fields(_MoeScalars) + } + ) + for i in sorted(self.moe_layers) + } + self._build_embed_owners() + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self._built_for = identity + + def cache_derived_state(self) -> None: + """The staged loaders' entry to the same build (``post_load_weights``).""" + self.post_load_weights() + + def _derive_moe(self, i: int) -> _MoeScalars: + """The routed experts' kernel scalars from the checkpoint's per-tensor scales (trtllm-gen NVFP4 layout). + + Global input scales: the reciprocal of the largest per-expert input scale (gate and up must agree). Per expert: + ``alpha = 1 / (input_scale * (1 / weight_scale_2))`` for the gate (fc31_alpha), the up half (its GEMM2-input + rescale ``fc31_scale_c = fc2_input_scale * up_alpha``) and the down projection; SwiGLU-OAI's activation + constants per expert, with beta and the clamp in the pre-dequant GEMM output domain (divided by fc31_alpha). + """ + w = self.w + isc1, isc3, isc2 = w[f"l{i}_e_isc1"].data, w[f"l{i}_e_isc3"].data, w[f"l{i}_e_isc2"].data + if not torch.allclose(isc1, isc3): + raise ValueError(f"layer {i}: the routed experts' w1 and w3 input scales differ") + fc31_input_scale = isc1.max().reciprocal() + fc2_input_scale = isc2.max().reciprocal() + ws2_1, ws2_3, ws2_2 = ( + w[f"l{i}_e_ws2_1"].data, + w[f"l{i}_e_ws2_3"].data, + w[f"l{i}_e_ws2_2"].data, + ) + fc31_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_1)) + up_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_3)) + fc2_alpha = 1.0 / (fc2_input_scale * (1.0 / ws2_2)) + fc31_scale_c = fc2_input_scale * up_alpha + n = self.num_experts + device = fc31_alpha.device + act_alpha = torch.full((n,), self.swiglu_alpha, dtype=torch.float32, device=device) + act_beta = torch.full((n,), self.swiglu_beta, dtype=torch.float32, device=device) + act_clamp = torch.full((n,), self.swiglu_limit, dtype=torch.float32, device=device) + act_beta.div_(fc31_alpha) + act_clamp.div_(fc31_alpha) + return _MoeScalars( + fc31_input_scale=fc31_input_scale, + fc2_input_scale=fc2_input_scale, + fc31_alpha=fc31_alpha, + fc2_alpha=fc2_alpha, + fc31_scale_c=fc31_scale_c, + act_alpha=act_alpha, + act_beta=act_beta, + act_clamp=act_clamp, + ) + + def _build_embed_owners(self) -> None: + """The fused embedding + first RMSNorm's caller-owned buffers, one per (rows, id dtype) of a decode step.""" + table, weight = self.w["embed"], self.w["l0_in_norm"] + self.embed_owners = { + (rows, dtype): DecodeEmbedding(table, weight, self.rms_eps, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + + def projections(self) -> List[Tuple[str, str, Optional[str]]]: + """Every MXFP8 projection in layer order: (weight key, scale key, the boundary its GEMV pushes into).""" + out = [] + for i in range(self.num_layers): + out.append((f"l{i}_qkv_w", f"l{i}_qkv_s", None)) + out.append((f"l{i}_o_w", f"l{i}_o_s", "attn")) + if i in self.moe_layers: + out.append((f"l{i}_sh_gu_w", f"l{i}_sh_gu_s", None)) + out.append((f"l{i}_sh_dn_w", f"l{i}_sh_dn_s", "moe")) + else: + out.append((f"l{i}_mlp_gu_w", f"l{i}_mlp_gu_s", None)) + out.append((f"l{i}_mlp_dn_w", f"l{i}_mlp_dn_s", "dense")) + return out + + def decode_query_lens(self) -> List[int]: + """The query rows per request a decode step can have: 1, and 1 + each draft length speculation runs.""" + lens = {1} + spec = getattr(self.model_config, "spec_config", None) + if spec is not None: + for drafts in ( + getattr(spec, "max_draft_len", None), + getattr(spec, "max_total_draft_tokens", None), + *(getattr(spec, "draft_len_schedule", None) or {}).values(), + ): + if drafts: + lens.add(1 + int(drafts)) + return sorted(lens) + + def prebuild_kernels(self) -> Dict[str, object]: + """Compile (and load) every kernel a decode forward dispatches to that would otherwise compile on its first + call: each projection's decode GEMVs (every row count, activation and boundary push), the MoE engines of + ``MOE_ENGINE`` / ``MOE16_ENGINE``'s token counts, and the index scorer at the decode query lengths. Runs on the + current (main) stream before the executor's warmup, in the same order on every rank, then synchronizes; after + it, GEMV compiles are frozen and the forward only looks kernels up (a missing one is a miss). It also builds + the T2SC scorer (``INDEX_SCORE_T2SC_PATH``), the 384-thread long-row select (``LONG_SELECT_DSL``) and the dense + layers' full-context decode plans (``FULLCTX_DECODE``). Returns the seconds each part took and the scorer + keys.""" + times: Dict[str, object] = {} + # The 9..16-row generation router: decided once, from the device, never inside a CUDA-graph capture. + # The built-in gate's conditions, mirrored exactly: 128 experts, a hidden size in 1024s, 16-CTA clusters. + self.router_tc_cluster16 = ( + self.num_experts == 128 + and self.hidden % 1024 == 0 + and bool(torch.ops.trtllm.minimax_m3_router_tc_cluster16(self.hidden)) + ) + logger.info( + f"[M3 prebuild] rank {mpi_rank()}: TF32 tensor-core router for {ROUTER_TC_GEN_MIN_ROWS}.." + f"{ROUTER_TC_GEN_MAX_ROWS}-row generation steps: " + + ( + "engaged (16-CTA clusters)" + if self.router_tc_cluster16 + else "not used (no 16-CTA clusters)" + ) + ) + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + if PREPROCESS_TOKMAP_FIX: + # Every path that can run the stock generation preprocess (the drafter's full-context decode or its + # fallback, any trtllm-gen fallback), so it is installed whatever else is enabled. + _tokmap_guard.install( + torch.device("cuda", torch.cuda.current_device()), + DECODE_MAX_TOKENS, + lambda rows: taken(PREPROCESS_TOKMAP_PATH, rows, "one-request token map"), + lambda rows, reason: missed(PREPROCESS_TOKMAP_PATH, rows, reason), + ) + mx_gemv.freeze(False) + use_pdl = get_env_enable_pdl() + world = int(self.mapping.tp_size) + start = time.perf_counter() + for wk, sk, push in self.projections(): + weight, scale = self.w[wk], self.w[sk] + mx_gemv.ensure_compiled(weight, scale, use_pdl) + n, k = weight.shape + for rows, x_bf16, ar in _gemv_variants(n, k, push, world): + # An MXFP8 call this target gives a (split, ring) override (``mx_gemv_cfg``, as ``_linear`` passes + # it) compiles that config. + cfg = None if x_bf16 else mx_gemv_cfg(n, k, rows) + if cfg is not None: + _precompile_cfg(weight, scale, rows, cfg, use_pdl, ar=ar) + else: + mx_gemv.precompile(weight, scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + torch.cuda.synchronize() + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for i in sorted(self.moe_layers): + self._prebuild_engines(i) + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = self._prebuild_index_scorer( + self.decode_query_lens(), self.w["embed"].device + ) + times["scorer_t2sc_keys"] = self._prebuild_index_scorer_t2sc( + self.decode_query_lens(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["scorer_s"] = time.perf_counter() - start + start = time.perf_counter() + times["select_keys"] = [] + if long_select_dsl_keys(): + # The 384-thread long-row select (LONG_SELECT_DSL): one kernel per launch key. + times["select_keys"] = _select_long_dsl_kernel.prebuild( + long_select_dsl_keys(), self.w["embed"].device + ) + torch.cuda.synchronize() + times["select_s"] = time.perf_counter() - start + start = time.perf_counter() + times["fullctx_plans"] = 0 + if FULLCTX_DECODE.keys: + # The dense layers' full-context decode: one compiled PrimTS plan per covered (rows, requests). + times["fullctx_plans"] = _fullctx_decode_kernel.prebuild( + self.w["embed"].device, + {key: FULLCTX_DECODE_RECIPES[key] for key in FULLCTX_DECODE.keys}, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), + get_env_enable_pdl(), + ) + torch.cuda.synchronize() + times["fullctx_s"] = time.perf_counter() - start + mx_gemv.freeze(True) + logger.info( + "[M3 prebuild] decode kernels built before warmup: GEMVs %.1f s, MoE engines %.1f s, index scorer " + "%.1f s (%d variants), long-row select %.1f s (%d variants), full-context decode %.1f s (%d plans)" + % ( + times["gemv_s"], + times["moe_s"], + times["scorer_s"], + len(times["scorer_keys"]), + times["select_s"], + len(times["select_keys"]), + times["fullctx_s"], + times["fullctx_plans"], + ) + ) + return times + + # ------------------------------------------------------------------------------------------ the MoE engines + + def _expert_operands(self, i: int): + """The routed experts' operands the engines read (trtllm-gen NVFP4 SwiGLU-OAI layout), and the input scale.""" + w, m = self.w, self.moe[i] + tensors = [ + w[f"l{i}_e_w31"].data, + w[f"l{i}_e_s31"].data, + w[f"l{i}_e_w2"].data, + w[f"l{i}_e_s2"].data, + m.fc31_alpha, + m.fc31_scale_c, + m.act_alpha, + m.act_beta, + m.act_clamp, + m.fc2_alpha, + ] + return tensors, m.fc31_input_scale + + def _shared_weights(self, i: int): + w = self.w + return ( + w[f"l{i}_sh_gu_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_gu_s"].data, + w[f"l{i}_sh_dn_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_dn_s"].data, + ) + + def engine(self, i: int, num_tokens: int, shared: bool) -> Optional[object]: + """Layer ``i``'s prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in (and the shared + expert with ``shared``). None after a miss: the build failed or was never made (a forward never builds one).""" + return self._lookup( + i, self.moe_engines[i], (num_tokens, True, shared), MOE_ENGINE.path, num_tokens + ) + + def moe16_engine(self, i: int, num_tokens: int) -> Optional[object]: + """Layer ``i``'s prebuilt 9..13-row engine (router, routed and shared experts in one kernel), or None.""" + if not MOE16_MIN_TOKENS <= num_tokens <= MOE16_MAX_TOKENS: + return None + return self._lookup(i, self.moe16_engines[i], num_tokens, MOE16_ENGINE.path, num_tokens) + + def _lookup(self, i: int, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {i}'s engine" + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, what) + else: + missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + missed(path, num_tokens, "its build failed (logged)") + return engine + + def _prebuild_engines(self, i: int) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of layer ``i`` can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in, and ``MOE16_ENGINE``'s. Whether an + engine is kept is agreed across the TP group (one host collective per engine): a build that fails on one rank + drops that engine on every rank, so the ranks never mix an engine (and its boundary push) with the generic + MoE.""" + engines = self.moe_engines.setdefault(i, {}) + engines16 = self.moe16_engines.setdefault(i, {}) + for num_tokens in sorted(MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in engines: + engines[key] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build( + i, MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + f"[M3 dispatch] {MOE_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + for num_tokens in sorted(MOE16_ENGINE.keys): + if num_tokens not in engines16: + engines16[num_tokens] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build(i, MOE16_ENGINE.path, num_tokens, self._build_moe16_engine), + f"[M3 dispatch] {MOE16_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + + def _try_build(self, i: int, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(i, num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {i}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, i: int, num_tokens: int, shared: bool) -> object: + tensors, scale = self._expert_operands(i) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + kwargs["shared_weights"] = self._shared_weights(i) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = M3MoeState( + device, num_tokens=num_tokens, num_ctas=sms - MOE_ENGINE_RESERVED_SMS, shared=shared + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + def _build_moe16_engine(self, i: int, num_tokens: int) -> object: + tensors, scale = self._expert_operands(i) + device = tensors[0].device + state = _MOE16_STATES.get((str(device), num_tokens)) + if state is None: + state = _MOE16_STATES[(str(device), num_tokens)] = M3Moe16State( + device, num_tokens=num_tokens + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + shared_weights=self._shared_weights(i), + ) + # Compile (first use of the state) and warm. + m3_moe16_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + # ------------------------------------------------------------------------------------------ the index scorer + + def _scorer_variants(self, query_lens) -> list: + """The scorer keys a decode step can take: every sparse layer's (index dtype, index heads per rank, index dim) + at each decode query length, with the early-issue schedule and without it, where the kernel serves it.""" + runner = _index_score_kernel.IndexScoreRunner + keys = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=config.num_index_heads, + head_dim=config.sparse_index_dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for early in (True, False): + key = ( + dtype, + int(config.num_index_heads), + int(qlen), + int(config.sparse_index_dim), + int(INDEX_SCORE_SPLIT_K), + bool(early), + ) + if key not in keys: + keys.append(key) + return keys + + def _prebuild_index_scorer(self, query_lens, device) -> list: + """Compile every scorer key and launch each once on real-shaped dummy operands (loading its module), on the + current stream. The kernel's batch, token, page and block counts are symbolic: its compile needs no step + metadata. Returns the keys.""" + runner = _index_score_kernel.IndexScoreRunner + keys = self._scorer_variants(query_lens) + page = runner.SUPPORTED_PAGE_SIZE + for key in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + rows = 1 if early else INDEX_SCORE_EARLY_MAX_REQUESTS + 1 + blocks = 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros( + rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device + ).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + runner.kernel_cache[key]( + idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + ) + return keys + + def _prebuild_index_scorer_t2sc(self, query_lens, device) -> list: + """T2SC (``INDEX_SCORE_T2SC_PATH``): compile this target's copy of the TP2 / TP8 scorer at every covered request + count's schedule (``index_score_t2sc_schedule``) for every sparse layer geometry and decode query length, and + launch each once on real-shaped dummy operands (loading its module). Returns the keys.""" + runner = _index_score_t2sc_kernel.IndexScoreRunner + keys = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=config.num_index_heads, + head_dim=config.sparse_index_dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for requests in sorted(INDEX_SCORE_T2SC_PATH.keys): + split_k, early = index_score_t2sc_schedule(requests) + key = ( + dtype, + int(config.num_index_heads), + int(qlen), + int(config.sparse_index_dim), + int(split_k), + bool(early), + ) + if all(key != seen for seen, _ in keys): + keys.append((key, requests)) + page = runner.SUPPORTED_PAGE_SIZE + for key, rows in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + blocks = 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros( + rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device + ).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + runner.kernel_cache[key]( + idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + ) + return [key for key, _ in keys] + + # --------------------------------------------------------------------------------------------- the dispatch + + def _select_target(self, attn_metadata: AttentionMetadata): + """Which target runs this step: one host-side integer, fixed per CUDA-graph capture. A step with context rows + (a mixed step included) is prefill's; every pure generation step is decode's, whatever its width.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + capture: Callable[..., None] = _plain_capture, + capture_target: Optional[Callable[..., Any]] = None, + **kwargs, + ) -> torch.Tensor: + """``capture_target``: the worker's capture-norm target (``CAPTURE_NORM_IN_BOUNDARY``) for a capture layer + on the decode fast path; None: the capture runs its own norm.""" + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if kwargs.get("lora_params") is not None: + raise ValueError("this MiniMax-M3 target serves no LoRA adapters") + return self._select_target(attn_metadata).forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The targets +# -------------------------------------------------------------------------------------------------------------------- + + +class PrefillTarget: + """A step with context rows: prefill chunks, and the mixed steps that carry generation rows beside them. Every + layer runs in its generic form; the attention core is the MSA backend's for every row.""" + + def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + Counts and per-step host scalars are read here; device tensors are the metadata's persistent buffers, held + by reference; ``msa_idx_k_cache`` is the metadata's per-layer index-K cache getter. ``metadata`` is the + prepared object itself: the MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_tokens=md.num_tokens, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + """``capture_target`` is unused: the generic boundary takes no multicast all-reduce at this TP + (``GENERIC_MC_AR_MAX_TOKENS`` 0), so a capture layer always runs its own capture norm.""" + core = self.core + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) prefill CUDA graphs" + ) + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + _record("prefill", int(step["num_tokens"]), int(step["num_contexts"])) + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + hidden, residual = self._layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + return hidden + + def _swiglu_mlp( + self, h: torch.Tensor, prefix: str, low_m_gemm: bool, dense: bool = False + ) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP with ``dense``, the shared expert): gate_up, SwiGLU-OAI, down, no + reduction.""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + out, _ = _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm) + return out + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: + the MoE scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: the fp32 router logits and the experts quantizing their bf16 input + (chunked).""" + return self._routed_chunks(i, h, _router_logits(self.core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool): + """The MoE block of layer ``i`` without its reduction. Returns ``(out, routed)``: with ``return_parts`` the + shared and routed parts (routed None when the engine folded the shared expert in), else their sum in + ``out``.""" + core = self.core + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, False) + folded = m3_moe_engine(h.contiguous(), engine, None, None) + return folded, None + routed_out, shared_out = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the + # same collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step): + """Attention of layer ``i`` up to its o_proj (no reduction). Returns (o_proj output, pushed): a prefill step + arms no boundary push, so pushed is always False.""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v for an FP8 cache, else in place in bf16), the + # MSA backend (it writes the layer's cache), o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection + and its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and + the FP8 indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the + o_proj. A layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): the MSA + backend (``m3_msa_attention``: the KV write unless a producer wrote the caches, the sparse layers' top-k block + selection, the FMHA over the context rows and any generation rows beside them).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + sparse = i in core.sparse_layers + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + +class DecodeTarget: + """A pure generation step (no context rows), of any width. + + Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids, per layer: attention (o_proj all-reduce deferred, its rows + pushed into the boundary all-reduce up to 16 tokens) -> attention boundary (all-reduce + residual + post-attention + RMSNorm) -> MoE (the 1..8 / 9..13-row engines, the first pushing its sum into the boundary all-reduce; the + unfinalized MoE with the finalize in the all-reduce from 14 rows) or dense MLP -> MoE boundary (all-reduce + + shared/routed add + residual + the next layer's RMSNorm, emitting the next QKV's MXFP8 activation). A layer whose + output Eagle3 captures runs the generic layer, as does the first layer when the fused embedding does not take the + step. A wider step (or one carrying embeddings) runs every layer in its generic form, with this target's attention + (the decode selection and sparse decode where the producer wrote the caches). + """ + + def __init__(self, core: MiniMaxM3Tp4ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # Set by a fused feed-forward boundary that wrote the layer's Eagle3 capture (see _layer's ``capture``). + self._capture_written = False + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + States the two context counts instead of reading them back: this target runs only when ``num_contexts`` is + zero, so reading them would ask the metadata to confirm what the routing already decided. + + Graph classes: counts, the generation span, the identity-selection flag and the worst-case score extent are + host values fixed per capture (prepare() stages them before the forward); the block table, lengths, cache + slots, identity table and valid-block counts are the metadata's persistent buffers, held by reference and + refreshed in place every step; ``msa_idx_k_cache`` / ``msa_proxy_max_score_view`` are the metadata's getters + of the per-layer index-K cache and of the score scratch view. ``metadata`` is the prepared object itself: the + MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=0, + num_ctx_tokens=0, + num_tokens=md.num_tokens, + num_seqs=md.num_seqs, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + msa_decode_span=md.msa_decode_span, + msa_block_table=md.msa_block_table, + msa_seq_lens_cuda=md.msa_seq_lens_cuda, + msa_short_ok=md._msa_short_ok, + msa_short_table=md.msa_short_table, + msa_proxy_max_score_view=md.msa_proxy_max_score_view, + msa_worst_case_max_k_tiles=md.msa_worst_case_max_k_tiles, + msa_n_valid_blocks=md.msa_n_valid_blocks, + max_seq_len=md.max_seq_len, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + # The decode kernels' rows (a per-capture constant): up to DECODE_MAX_TOKENS tokens of 1-D input ids. + kernels = ( + inputs_embeds is None and input_ids.ndim == 1 and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + _record("decode", num_tokens, plan.name) + hidden, residual = self._embed(input_ids, num_tokens) + else: + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + _record("decode", int(step["num_tokens"]), "generic") + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + is_capture = spec_metadata is not None and spec_metadata.is_layer_capture(i) + if ( + kernels + and residual is not None + and is_capture + and CAPTURE_LAYERS_FAST + and CAPTURE_FUSED_BOUNDARY + ): + # The capture layer on the fast path (its boundary pushes); the capture reads the MoE / dense + # boundary's reduced residual sum. With CAPTURE_NORM_IN_BOUNDARY the boundary all-reduce writes the + # capture norm itself. + capture_out = ( + capture_target(spec_metadata, i, num_tokens) + if capture_target is not None and CAPTURE_NORM_IN_BOUNDARY + else None + ) + self._capture_written = False + hidden, residual = self._layer( + i, plan, position_ids, hidden, residual, step, capture=capture_out + ) + taken("capture_layer_fast", num_tokens, f"layer {i}") + if capture_out is not None and self._capture_written: + taken("capture_norm_in_boundary", num_tokens, f"layer {i}") + capture(spec_metadata, i, residual, None, written=True) + else: + capture(spec_metadata, i, residual, None) + elif not kernels or residual is None or is_capture: + hidden, residual = self._generic_layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + else: + hidden, residual = self._layer(i, plan, position_ids, hidden, residual, step) + return hidden + + def _embed(self, input_ids: torch.Tensor, num_tokens: int): + """(hidden, residual) of the fused embedding + first RMSNorm; (embedding, None) when it does not run, so the + first layer runs its own input norm.""" + core = self.core + if not EMBED_FIRSTNORM.covers(num_tokens): + uncovered(EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens") + return m3_embedding(input_ids, core.w["embed"]), None + owner = core.embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + missed( + EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 after the load)" + if owner is None + else "the input ids are not contiguous", + ) + return m3_embedding(input_ids, core.w["embed"]), None + _record("m3_embedding_firstnorm", num_tokens, input_ids.dtype) + taken(EMBED_FIRSTNORM.path, num_tokens, str(input_ids.dtype)) + return m3_embedding_firstnorm( + input_ids, owner.table, owner.weight, owner.eps, owner.raw, owner.output, owner.use_pdl + ) + + def _bucket_flags(self, num_tokens: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(self.core.mapping, self.core.hidden) + workspace = None if workspace is None else m3_mc_allreduce._bucket(workspace, num_tokens) + flags = None if workspace is None else workspace.flags_for(8) + return None if flags is None else (workspace, flags) + + @staticmethod + def _no_bucket(path: str, num_tokens: int) -> None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, num_tokens: int): + """The o_proj's push tuple into the attention boundary's all-reduce, or None (after a miss). + + None comes only from the token count or from workspace state created collectively in the same eager call on + every rank, so all ranks agree. (A lone non-pushing rank would still be safe: its all-reduce pushes the same + bf16 rows into the same set A.)""" + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(ATTN_PUSH.path, num_tokens) + return None + workspace, flags = found + return ( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + AR_PUSH_MODE, + ) + + def _moe_push_targets(self, num_tokens: int): + """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None + (after a miss). + + This decision must match across ranks: prepushed 3 reads set B of every rank, which a non-pushing rank never + fills. None comes only from the token count or from state created in the same eager call on every rank + (workspace, flags, the descriptor cache), and the engine it is armed for is agreed at build time + (``_prebuild_engines``), so it does.""" + hidden = self.core.hidden + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(MOE_PUSH.path, num_tokens) + return None + workspace, flags = found + key = (workspace.mc.data_ptr(), flags.data_ptr()) + desc = self._moe_push_descs.get(key) + if desc is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor") + return None + desc = self._moe_push_descs[key] = ar_push_desc( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + flags.device, + ) + set_b = workspace.mc.data_ptr() + 2 * 4 * workspace.max_tokens * workspace.world_size * ( + hidden // 2 + ) + return desc, ( + set_b, + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + 2, + ) + + def _boundary(self, use_mc: bool, x, residual, norm_w, x2=None, quant=False, capture=None): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce. + ``capture``: the Eagle3 capture the multicast kernel also writes (``_capture_written`` records it).""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant, capture=capture) + if fused is not None: + if capture is not None: + self._capture_written = True + rows = x.shape[0] + taken( + MC_ALLREDUCE.path, + rows, + ("x + x2" if x2 is not None else "x") + + (", MXFP8 copy" if quant else "") + + (", cluster" if rows <= MC_AR_CLUSTER_MAX_ROWS else ""), + ) + return fused + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + missed( + MC_ALLREDUCE.path, + x.shape[0], + m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + ) + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _prepushed(self, i: int, x, residual, norm_w, quant, prepushed, x2=None, capture=None): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed, capture=capture + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + if capture is not None: + self._capture_written = True + rows = x.shape[0] + taken( + MC_ALLREDUCE.path, + rows, + f"prepushed {prepushed}" + (", cluster" if rows <= MC_AR_CLUSTER_MAX_ROWS else ""), + ) + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step, capture=None): + """``capture``: (out, fc_norm weight, eps) for the feed-forward boundary to write the layer's Eagle3 capture + norm (CAPTURE_NORM_IN_BOUNDARY, for a capture layer on this path); a fused boundary that writes it sets + ``_capture_written``.""" + core = self.core + num_tokens = hidden.shape[0] + # Padded rows: the boundary all-reduce polls and re-arms rows t < M = residual.shape[0], CUDA-graph padding + # rows included. Every producer pushes all rows of the same tensor: the o_proj / shared down GEMV t < m = + # x.shape[0] (its row count comes from the host shape, not a live count), the MoE engine every row of the + # num_tokens it was built for (= num_tokens here). The all-reduce checks x.sizes() == residual.sizes(), so + # no padded row is left with a stale flag. + # The GEMV dispatch (and with it every push into a boundary all-reduce) is a decode-graph one. + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + uncovered( + "decode_graph", + "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + ) + push = plan.attn_mc_ar and decode_graph and ATTN_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if push else None + hidden, attn_pushed = self._attention(i, position_ids, hidden, step, armed) + post_norm = core.w[f"l{i}_post_norm"] + if armed is not None and attn_pushed: + taken(ATTN_PUSH.path, num_tokens) + hidden, residual = self._prepushed(i, hidden, residual, post_norm, False, AR_PUSH_MODE) + else: + if armed is not None: + missed( + ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden, residual = self._boundary(plan.attn_mc_ar, hidden, residual, post_norm) + + norm_w = _next_norm(core, i) + if i not in core.moe_layers: + # A dense layer: its MLP's down projection pushes into the post-MLP boundary's all-reduce, as the o_proj + # does into the attention one (the same slot; the gate_up projection never pushes). + dense_push = plan.attn_mc_ar and decode_graph and DENSE_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if dense_push else None + out, dense_pushed = self._swiglu_mlp(hidden, f"l{i}_mlp", True, armed, dense=True) + if armed is not None and dense_pushed: + taken(DENSE_PUSH.path, num_tokens) + return self._prepushed( + i, out, residual, norm_w, False, AR_PUSH_MODE, capture=capture + ) + if armed is not None: + missed( + DENSE_PUSH.path, + num_tokens, + "the dense MLP's down projection did not push its rows into the armed boundary all-reduce", + ) + return self._boundary(plan.attn_mc_ar, out, residual, norm_w, capture=capture) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + path = moe_path(num_tokens) + moe_push = None + if path == MOE_ENGINE.path: + # The engine runs inside the MoE; looked up (or found missing) here first so a push is armed only for an + # engine that will take it. + engine = core.engine(i, num_tokens, shared=True) + if ( + engine is not None + and plan.moe_mc_ar + and decode_graph + and MOE_PUSH.covers(num_tokens) + ): + moe_push = self._moe_push_targets(num_tokens) + elif path == MOE16_ENGINE.path and decode_graph: + # Router, routed and shared experts in one kernel; its bf16 sum goes to the boundary all-reduce. + engine = core.moe16_engine(i, num_tokens) + if engine is not None: + _record("m3_moe16_engine", num_tokens) + out = m3_moe16_engine(hidden.contiguous(), engine) + taken(MOE16_ENGINE.path, num_tokens) + return self._boundary( + plan.moe_mc_ar, out, residual, norm_w, quant=quant, capture=capture + ) + if path in (MOE16_ENGINE.path, MC_FINALIZE.path) and plan.moe_mc_ar: + if self._bucket_flags(num_tokens) is None: + self._no_bucket(MC_FINALIZE.path, num_tokens) + else: + (shared, fc2, scales, perm), _, _, _ = self._moe( + i, hidden, return_parts=False, do_finalize=False + ) + _record("m3_mc_finalize_allreduce_norm", num_tokens, quant) + fused = m3_mc_allreduce.m3_mc_finalize_allreduce_norm( + fc2, + perm.view(num_tokens, -1).to(torch.int32), + scales.reshape(num_tokens, -1), + shared, + residual, + norm_w, + core.rms_eps, + m3_mc_allreduce.get_workspace(core.mapping, core.hidden), + use_pdl=get_env_enable_pdl(), + quant=quant, + ) + if fused is None: + raise RuntimeError( + "the MC finalize all-reduce rejected the unfinalized MoE outputs" + ) + taken(MC_FINALIZE.path, num_tokens) + return fused + shared, routed, engine_pushed, shared_pushed = self._moe( + i, hidden, return_parts=True, push=moe_push + ) + if moe_push is not None: + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + if engine_pushed: + taken(MOE_PUSH.path, num_tokens, "engine, routed + shared") + return self._prepushed(i, shared, residual, norm_w, quant, 2, capture=capture) + else: + if engine_pushed and shared_pushed: + taken(MOE_PUSH.path, num_tokens, "routed, shared") + return self._prepushed( + i, shared, residual, norm_w, quant, 3, x2=routed, capture=capture + ) + if engine_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + missed( + MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary( + plan.moe_mc_ar, shared, residual, norm_w, x2=routed, quant=quant, capture=capture + ) + + def _swiglu_mlp( + self, + h: torch.Tensor, + prefix: str, + low_m_gemm: bool, + push: Optional[tuple] = None, + dense: bool = False, + ): + """A gated MXFP8 MLP (the dense MLP with ``dense``, the shared expert): gate_up, SwiGLU-OAI, down, no + reduction. Returns (output, whether the down projection's GEMV pushed into ``push``).""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm, push) + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: the MoE + scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor, do_finalize: bool): + """Router + routed experts: the router logits of a generation step (``_decode_router_logits``) and the experts + quantizing their bf16 input.""" + core = self.core + router_logits = _decode_router_logits(core, i, h) + if not do_finalize: + return _routed_experts(core, i, h, None, router_logits, False) + return self._routed_chunks(i, h, router_logits) + + def _moe( + self, i: int, h: torch.Tensor, *, return_parts: bool, do_finalize: bool = True, push=None + ): + """The MoE block of layer ``i`` without its reduction. + + ``push`` = (the engine's ``arp`` descriptor, the shared down projection's ``ar`` tuple) the decode layer armed + for the MoE boundary's exchange sets, or None. Returns ``(out, routed, engine_pushed, shared_pushed)``: with + ``return_parts`` the shared and routed parts (routed None when the engine folded the shared expert in), else + their sum in ``out``; with ``do_finalize=False`` ``out`` is ``[shared, fc2 out, expert scales, permuted + index]``. + """ + core = self.core + engine_arp, shared_ar = push if push is not None else (None, None) + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + do_finalize + and h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, engine_arp is not None) + folded = m3_moe_engine(h.contiguous(), engine, None, engine_arp) + taken(MOE_ENGINE.path, rows) + return folded, None, engine_arp is not None, False + routed_out, (shared_out, shared_pushed) = maybe_execute_in_parallel( + lambda: self._routed(i, h, do_finalize), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False, shared_ar), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the same + # collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if not do_finalize: + return [shared_out, *routed_out], None, False, shared_pushed + if return_parts: + return shared_out, routed_out, False, shared_pushed + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None, False, shared_pushed + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _generic_layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed, _, _ = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = None): + """Attention of layer ``i`` up to its o_proj (no reduction). ``push`` is the attention boundary's exchange the + o_proj may write into. Returns (o_proj output, pushed).""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v: an FP8 cache, or the FP8 pool the dense layers + # sit in within a hybrid NVFP4 cache on a decode step), the attention, o_proj. A covered (query rows, + # requests) shape attends on the full-context PrimTS decode (``FULLCTX_DECODE``), the norm/RoPE kernel then + # storing the new K / V rows itself (``DENSE_KV_INSERT_FUSED``); every other step the MSA backend. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + out_fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + plan = self._fullctx_plan(i, qkv.device, step) + kv_dest = None + if ( + plan is not None + and out_fp8 + and DENSE_KV_INSERT_FUSED + and DENSE_QKV_FP8_VIA_BF16 + and qkv.dtype == torch.bfloat16 + and core.head_dim == 128 + ): + kv_dest = ( + step["kv_cache_manager"].get_buffers(i, kv_layout="HND"), + step["msa_out_cache_loc"], + plan[1], + torch.empty( + (qkv.shape[0], core.heads, core.head_dim), + dtype=torch.float8_e4m3fn, + device=qkv.device, + ), + ) + fused = _qk_norm_rope(core, i, qkv, position_ids, out_fp8, kv_dest) + if kv_dest is not None: + taken(DENSE_KV_INSERT_PATH, kv_dest[2]) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + if kv_dest is not None: + q = kv_dest[3].view(q.shape[0], -1) + o = None + if plan is not None: + o = self._fullctx_decode(i, q, k, v, step, plan, kv_written=kv_dest is not None) + if o is None: + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _fullctx_plan(self, i: int, device: torch.device, step): + """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs on + ops/m3_fullctx_decode: a covered (query rows per request, requests) shape over the layer's FP8 pool. None + otherwise (the MSA backend runs).""" + path = FULLCTX_DECODE.path + if is_torch_compiling(): + uncovered(path, "not a pure generation step on the MSA backend") + return None + span = step["msa_decode_span"] + if span is None or span.row_first != 0: + uncovered(path, "no decode span starting at row 0") + return None + nreq = int(step["num_seqs"]) + qlen = int(span.query_len) + if not FULLCTX_DECODE.covers((qlen, nreq)): + uncovered(path, f"{qlen} query rows x {nreq} requests is not a covered shape") + return None + if int(step["max_seq_len"]) > FULLCTX_DECODE_MAX_KV_LEN: + uncovered(path, "the cache length bound is beyond the plans' K/V capacity") + return None + num_tokens = nreq * qlen + pool, _ = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + if pool.dtype != torch.float8_e4m3fn: + missed(path, num_tokens, "the dense layer's KV pool is not FP8") + return None + own_reducer = FULLCTX_DECODE_OWN_REDUCER + if own_reducer and not _fullctx_decode_kernel.has_own_reducer(): + missed(path, num_tokens, "the BF16 split-KV merge op is not built") + return None + # Compiled before warmup (``prebuild_kernels``); the forward never compiles one. + plan = _fullctx_decode_kernel.plan_for(device, qlen, nreq, own_reducer) + if plan is None: + missed(path, num_tokens, "the PrimTS decode plan is not prebuilt") + return None + return plan, num_tokens + + def _fullctx_decode( + self, i: int, q, k, v, step, plan_tokens, *, kv_written: bool + ) -> torch.Tensor: + """Dense layer ``i``'s decode attention on the PrimTS plan of ``_fullctx_plan``: the step's new K / V go into + the paged cache as the MSA backend writes them (``write_msa_phase_kv``; already there when ``kv_written``), + then every query row attends its request's full causal context. The merge also emits the o_proj's MXFP8 + input (``FULLCTX_DECODE_MXFP8``) when the plan supports it.""" + core = self.core + backend = core.attn[i] + plan, num_tokens = plan_tokens + nreq = int(step["num_seqs"]) + if not kv_written: + write_msa_phase_kv( + backend, + k[:num_tokens], + v[:num_tokens], + step["metadata"], + AttentionInputType.mixed, + token_offset=0, + ) + pool, subpages = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + k_pages, v_pages = _fullctx_decode_kernel.kv_page_views(pool, subpages) + mxfp8 = None + if ( + FULLCTX_DECODE_MXFP8 + and output.shape[0] == num_tokens + and _fullctx_decode_kernel.emits_mxfp8() + and MERGE_MXFP8_MIN_TOKENS <= num_tokens <= MERGE_MXFP8.MAX_ROWS + ): + mxfp8 = MERGE_MXFP8.buffers(output.shape[1], output.device) + _record("m3_fullctx_decode", num_tokens, kv_written, mxfp8 is not None) + emitted = plan.run( + fp8_e4m3_saturate(q[:num_tokens].reshape(num_tokens, core.heads, core.head_dim)), + k_pages, + v_pages, + step["msa_block_table"][:nreq], + step["msa_seq_lens_cuda"][:nreq], + core.sm_scale(i), + output[:num_tokens], + mxfp8=mxfp8, + ) + if emitted: + m3_prequant.SLOT.publish(output, mxfp8[0][:num_tokens], mxfp8[1]) + taken(FULLCTX_DECODE_MXFP8_PATH, num_tokens) + taken(FULLCTX_DECODE.path, num_tokens) + return output + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step, push): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection and + its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and the FP8 + indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the o_proj. A + layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): a sparse + layer whose caches the producer wrote runs this target's selection and sparse decode when the generation span + starts at row 0; everything else the MSA backend (``m3_msa_attention``: the KV write unless a producer wrote the + caches, the sparse layers' top-k block selection, the FMHA).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + span = step["msa_decode_span"] + sparse = i in core.sparse_layers + if sparse and k_live is None and span is not None and span.row_first == 0: + self._decode_sparse_core(i, q_live, idx_q_live, step, out_live) + return output + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step of a sparse layer whose caches + the producer already wrote. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise the CuTe + DSL scorer writes the per-index-head block maxima, reduced to the KV head, and the long-row selector takes the + top 16 (the built-in selector beyond its shapes). + """ + core = self.core + backend = core.attn[i] + config = backend.m3_config + num_tokens = int(q.shape[0]) + rows = int(step["num_seqs"]) + qlen = int(step["msa_decode_span"].query_len) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + num_kv_heads = int(backend.indexer.config.num_kv_heads) + identity = step["msa_short_ok"] and step["msa_short_table"] is not None + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + else: + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = step["msa_idx_k_cache"](i) + max_score = step["msa_proxy_max_score_view"]( + config.num_index_heads, step["msa_worst_case_max_k_tiles"], num_tokens + ) + if INDEX_SCORE_T2SC_PATH.covers(rows): + # T2SC: this target's copy of the TP2 / TP8 scorer at the request count's split and prestage. + split_k, early = index_score_t2sc_schedule(rows) + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(split_k), + bool(early), + ) + if key not in _index_score_t2sc_kernel.IndexScoreRunner.kernel_cache: + missed( + INDEX_SCORE_T2SC_PATH.path, + num_tokens, + f"the T2SC index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score_t2sc", key) + m3_index_score_t2sc( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + split_k, + early, + ) + taken(INDEX_SCORE_T2SC_PATH.path, rows) + else: + early = rows <= INDEX_SCORE_EARLY_MAX_REQUESTS + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(INDEX_SCORE_SPLIT_K), + bool(early), + ) + if key not in _index_score_kernel.IndexScoreRunner.kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + missed( + INDEX_SCORE.path, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score", key) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + INDEX_SCORE_SPLIT_K, + early, + ) + taken(INDEX_SCORE.path, rows) + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + group = config.num_index_heads // config.num_kv_heads + max_score_kv = ( + max_score.view( + config.num_kv_heads, group, max_score.shape[1], max_score.shape[2] + ).amax(dim=1) + if group > 1 + else max_score + ) + early_trigger = num_tokens <= SELECT_EARLY_MAX_ROWS + trigger_phase = ( + SELECT_LONG_TRIGGER_PHASE if num_tokens <= SELECT_LONG_TRIGGER_MAX_ROWS else 0 + ) + dsl = False + if LONG_SELECT_DSL.keys and 128 < max_score_kv.shape[1]: + # Compiled before warmup (prebuild_kernels); a missing kernel falls back to the 1024-thread one. + key = _select_long_dsl_kernel.kernel_key( + get_env_enable_pdl(), early_trigger, trigger_phase + ) + dsl = key in _select_long_dsl_kernel.kernel_cache + if not dsl: + missed( + LONG_SELECT_DSL.path, + num_tokens, + f"the 384-thread select {key} was not compiled before warmup", + ) + table = long_select_blocks( + max_score_kv, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + use_pdl=get_env_enable_pdl(), + early_trigger=early_trigger, + trigger_phase=trigger_phase, + fast_tail=SELECT_LONG_SL3, + dsl=dsl, + ) + if table is not None: + taken((LONG_SELECT_DSL if dsl else LONG_SELECT).path, num_tokens) + else: + uncovered( + LONG_SELECT.path, + f"{max_score_kv.shape[1]} KV blocks, outside {LONG_SELECT_MIN_BLOCKS}..{LONG_SELECT_MAX_BLOCKS}: " + "the built-in selector", + ) + table = m3_select_blocks( + max_score_kv, + MSA_REQUIRED_TOPK, + n_valid_blocks, + config.init_blocks, + config.local_blocks, + True, + ) + topk = table.permute(1, 0, 2) + manager = step["kv_cache_manager"] + k_paged, v_paged = msa_paged_kv(manager, i) + nvfp4 = {} + if manager.is_nvfp4_layer(i): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(backend, core.kv_scales[i]) + scales = manager.get_block_scale_buffers(i, "HND") + nvfp4 = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + _record("m3_sparse_decode", num_tokens, qlen, identity) + # The long-context decode (if it runs) prefetches this layer's o_proj weights into L2 and the blocks the next + # sparse layer selected at the previous step; any other route leaves the slots unused, and they are cleared + # after the call either way. + if SPARSE_DECODE_LC_MERGE_ONCE and SPARSE_DECODE_LC_PREFETCH_OPROJ: + m3_attn_lc.PREFETCH.tensors = (core.w[f"l{i}_o_w"], core.w[f"l{i}_o_s"]) + if SPARSE_DECODE_LC_MERGE_ONCE and i in core.n1_ids and num_tokens <= N1_MAX_ROWS: + m3_attn_lc.NEXT_KV.tensors = self._next_kv_prefetch(i, manager) + try: + m3_sparse_decode( + q.view(num_tokens, core.heads, core.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=core.sm_scale(i), + output=output.view(num_tokens, core.heads, core.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4, + ) + finally: + m3_attn_lc.PREFETCH.tensors = None + m3_attn_lc.NEXT_KV.tensors = None + + def _next_kv_prefetch(self, i: int, manager): + """(layer ``i``'s selected-block id buffer, the next sparse layer's, that layer's K / V pools and block scales) + for the long-context decode's next-layer prefetch; the last sparse layer's next is the first one (next step). + None when the next layer's cache is not NVFP4.""" + ids = self.core.n1_ids + order = sorted(ids) + nxt = next((j for j in order if j > i), order[0]) + if not manager.is_nvfp4_layer(nxt): + return None + k_next, v_next = msa_paged_kv(manager, nxt) + scales = manager.get_block_scale_buffers(nxt, "HND") + return (ids[i], ids[nxt], k_next, v_next, scales[:, 0], scales[:, 1]) + + +# -------------------------------------------------------------------------------------------------------------------- +# Speculative decoding: one-model Eagle3 at checkpoint precision +# -------------------------------------------------------------------------------------------------------------------- + + +class M3SpecStepFold: + """``SpecStepFold`` of this target. + + Between two decode graphs of a one-model Eagle3 step, the spec sampler stores the step's outputs into its per-slot + store and the engine gathers them back into its flat input buffers for the next step. This object does both in one + launch each (``m3_spec_tail.slot_scatter`` / ``step_gather``): integer copies only, the same values the engine's and + the sampler's own copies write. + + With ``in_forward`` (``STEP_SPEC_FOLD_IN_FORWARD``) the forward of a decode-only step also does both at its tail + (``store_in_forward``, one ``m3_fold_tail`` launch inside the step graph): it stores the step's outputs for the rows + ``begin_step`` armed and stages the next step's inputs for those rows in that order. The arming is a device-side row + count and slot list staged before the replay, so a captured graph whose step was not armed writes nothing. The + sampler then skips its store for exactly the armed rows (``sampler_rows_stored``), and the next step skips its + gathers when its rows are those rows (``begin_step``). + + One instance per target model; the target passes it to ``RuntimeHooks.spec_step_fold``. + """ + + def __init__( + self, in_forward: bool = False, store_copy_wait: bool = False, cpp_tail: bool = False + ) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + ] = None + self.in_forward = in_forward + # The forward's tail as this target's C++ op instead of the Triton kernel (m3_fold_tail.fold_tail(cpp=)). + self.cpp_tail = in_forward and cpp_tail + # Device arming read by the forward's tail: ctrl[0] = armed rows (0 = write nothing), slots[:rows]. + self._ctrl: Optional[torch.Tensor] = None + self._slots: Optional[torch.Tensor] = None + # The rows (sequence slots) the next forward stores and stages for, as begin_step armed them. + self._armed: Optional[tuple[int, ...]] = None + # The rows the previous forward staged the next step's inputs for (confirmed by sampler_rows_stored). + self._staged_for: Optional[tuple[int, ...]] = None + # Output buffers the tail ran on: the captured graphs' static outputs (by address) and the last eager step's. + self._graph_outputs: set[int] = set() + self._eager_outputs: Optional[torch.Tensor] = None + self._logged = False + # store_copy_wait: the sampler records this (external) event on its D2H side stream after copying the store + # to the host, and the forward's tail waits on it before it rewrites the store; the sampler can then copy + # on the side stream on the steps this forward stored the rows (see SpecSamplerBase._copy_outputs_to_host). + self._store_copy_wait = in_forward and store_copy_wait + self.store_copies_done: Optional[torch.cuda.Event] = None + + def bind_engine_buffers( + self, + input_ids: torch.Tensor, + draft_tokens: torch.Tensor, + position_id_offsets: torch.Tensor, + kv_lens_offsets: torch.Tensor, + max_batch_size: int, + ) -> None: + self.step_inputs = (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) + if self._store_copy_wait: + # Recorded once now so a captured wait has a completed record to wait on before the first copy. + self.store_copies_done = torch.cuda.Event(external=True) + self.store_copies_done.record() + logger.info( + "[M3 host] spec store host copies on the D2H side stream, awaited by the forward's tail" + ) + if self.cpp_tail: + logger.info( + "[M3 host] spec fold tail: one C++ launch (no Triton) at the forward's tail" + ) + if self.in_forward: + self._ctrl = torch.zeros(1, dtype=torch.int32, device=input_ids.device) + self._slots = torch.zeros( + max(max_batch_size, 1), dtype=torch.long, device=input_ids.device + ) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list[int]], num_rows: int) -> bool: + if not self.in_forward or self._ctrl is None: + return False + rows = tuple(seq_slots) if seq_slots is not None else None + staged = rows is not None and rows == self._staged_for + self._staged_for = None + armed = rows if rows is not None and 0 < len(rows) <= self._slots.numel() else None + # Read only by this step's forward, so both uploads may wait for the step-input pack's flush. + if armed is not None: + slots = torch.tensor(armed, dtype=torch.long, pin_memory=prefer_pinned()) + copy_h2d_if_changed(self._slots[: len(armed)], slots, pack=True) + ctrl = torch.tensor( + [len(armed) if armed is not None else 0], dtype=torch.int32, pin_memory=prefer_pinned() + ) + copy_h2d_if_changed(self._ctrl, ctrl, pack=True) + self._armed = armed + return staged + + # The launches the store wait may precede. store_copies_done is an external event: under graph capture its wait + # is a node of its own, so the kernel launched right after it loses its programmatic (PDL) edge and with it the + # overlap of its pre-wait part. Only a launch that waits for its predecessor's completion anyway may follow it: + # the fold tail (a plain launch, no griddepcontrol). + STORE_WAIT_BEFORE = frozenset({"fold_tail"}) + + def _wait_store_copies(self, before: str) -> None: + """Order the store's rewrite after the sampler's host copies of it; ``before`` names the launch that follows. + Under graph capture a launch outside STORE_WAIT_BEFORE fails the capture.""" + if self.store_copies_done is None: + return + if before not in self.STORE_WAIT_BEFORE and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"M3SpecStepFold: the store wait would precede {before!r}, a launch outside " + f"{sorted(self.STORE_WAIT_BEFORE)}; an external-event wait node drops the next kernel's PDL edge" + ) + torch.cuda.current_stream().wait_event(self.store_copies_done) + + def store_in_forward(self, outputs: dict) -> None: + """At the forward's tail: store the armed rows' outputs and stage the next step's inputs (one launch).""" + if self._ctrl is None or self.store is None or self.step_inputs is None: + return + store = self.store + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + new_tokens = outputs.get("new_tokens") + next_new_tokens = outputs.get("next_new_tokens") + new_tokens_lens = outputs.get("new_tokens_lens") + next_draft_tokens = outputs.get("next_draft_tokens") + if any( + t is None or not isinstance(t, torch.Tensor) + for t in (new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens) + ): + return + rows = new_tokens.shape[0] + # The previous step's host copies of the store run on the sampler's side stream. + self._wait_store_copies(before="fold_tail") + if not m3_fold_tail.fold_tail( + self._ctrl, + self._slots, + min(rows, self._slots.numel()), + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + next_draft_tokens, + store.next_draft_tokens, + new_tokens_lens, + store.new_tokens_lens, + input_ids, + draft_tokens, + position_id_offsets, + kv_lens_offsets, + next_new_tokens.shape[1], + cpp=self.cpp_tail, + ): + return + if torch.cuda.is_current_stream_capturing(): + self._graph_outputs.add(new_tokens.data_ptr()) + else: + self._eager_outputs = new_tokens + + def gather( + self, + new_tokens: torch.Tensor, + next_draft_tokens: torch.Tensor, + new_tokens_lens: torch.Tensor, + previous_slots: torch.Tensor, + previous_pos_indices: torch.Tensor, + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + if self.step_inputs is None: + return False + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + return m3_spec_tail.step_gather( + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len, + tokens_per_gen_step, + draft_token_buffer_width, + input_ids, + num_tokens, + draft_tokens, + num_draft_tokens, + position_id_offsets, + kv_lens_offsets, + first_previous_row, + ) + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list[int]) -> bool: + armed, self._armed = self._armed, None + new_tokens = outputs.get("new_tokens") + ran = isinstance(new_tokens, torch.Tensor) and ( + new_tokens is self._eager_outputs or new_tokens.data_ptr() in self._graph_outputs + ) + self._eager_outputs = None + stored = ( + ran + and armed is not None + and num_skip == 0 + and tuple(seq_slots) == armed + and new_tokens.shape[0] == len(armed) + ) + self._staged_for = armed if stored else None + if stored and not self._logged: + self._logged = True + logger.info("[M3 host] spec step stored and next step staged in the forward tail") + return stored + + def scatter( + self, + seq_slots: torch.Tensor, + new_tokens: torch.Tensor, + next_new_tokens: torch.Tensor, + new_tokens_lens: torch.Tensor, + next_draft_tokens: torch.Tensor, + ) -> bool: + store = self.store + if store is None: + return False + return m3_spec_tail.slot_scatter( + seq_slots, + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + new_tokens_lens, + store.new_tokens_lens, + next_draft_tokens, + store.next_draft_tokens, + ) + + +class _DraftForward: + """``draft_model.model`` for the in-tree worker: the drafter layer's forward, called as + ``model(input_ids=, position_ids=, hidden_states=, attn_metadata=, spec_metadata=, all_rank_num_tokens=)`` and + returning (hidden, hidden to carry to the next pass). ``d2t`` None: the draft vocabulary is the target's.""" + + d2t = None + + def __init__(self, draft: "M3Eagle3Draft") -> None: + self._draft = draft + + def __call__(self, **kwargs): + return self._draft.forward_layer(**kwargs) + + +class M3Eagle3Draft(nn.Module): + """The one-layer BF16 Eagle3 drafter (``speculative_model``): its weights (``weights.DRAFT_WEIGHTS``), its + vocab-parallel BF16 head and logits processor, the drafter layer's TRTLLM attention backend (layer index = the + target's layer count, FP8 KV) and its TP all-reduce. + + What the worker reads: ``model`` (the layer forward), ``apply_eagle3_fc`` (the per-layer fc_norms + fc over the + captured target hidden states), ``lm_head`` and ``logits_processor``; the fused draft loop reads the layer's parts + (``draft_input_weights``, ``attention``, ``mlp``, ``ar_norm``). The drafter embeds with the target's table. + """ + + def __init__( + self, + core: MiniMaxM3Tp4ModelingV2, + draft_config: ModelConfig, + start_layer_idx: int, + capture_layers: int, + ) -> None: + super().__init__() + cfg = draft_config.pretrained_config + _ensure_draft_vocab_size(cfg) + eagle = getattr(cfg, "eagle_config", None) or {} + mapping = draft_config.mapping + self.draft_config = draft_config + # The engine's draft weight mapper binds to these two. + self.model_config = draft_config + self.config = cfg + self.capture_layers = int(capture_layers) + self.layer_idx = int(start_layer_idx) + self.dtype = cfg.torch_dtype + self.hidden = int(cfg.hidden_size) + self.hidden_size_in = int(getattr(cfg, "target_hidden_size", cfg.hidden_size)) + self.rms_eps = float(cfg.rms_norm_eps) + self.return_hidden_post_norm = bool( + eagle.get("return_hidden_post_norm", False) or getattr(cfg, "norm_output", False) + ) + if ( + int(cfg.num_hidden_layers) != 1 + or eagle.get("eh_proj_before_attn", False) + or eagle.get("norm_before_fc", False) + or not getattr(cfg, "fc_norm", False) + or getattr(cfg, "kv_lora_rank", None) + or self.capture_layers <= 1 + or self.hidden_size_in != self.hidden + or cfg.draft_vocab_size != cfg.vocab_size + or getattr(cfg, "attention_bias", False) + or getattr(cfg, "mlp_bias", False) + or getattr(cfg, "bias", False) + or getattr(cfg, "tie_word_embeddings", False) + ): + raise ValueError( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle3 drafter with per-layer " + "fc_norms over " + "several captured layers, the target's hidden size and vocabulary and no biases" + ) + tp = mapping.tp_size + self.heads = cfg.num_attention_heads // tp + self.kv_heads = (cfg.num_key_value_heads + tp - 1) // tp + self.head_dim = int( + getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + ) + self.q_size = self.heads * self.head_dim + self.kv_size = self.kv_heads * self.head_dim + self.w = _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + self.lm_head = LMHead( + cfg.draft_vocab_size, + cfg.hidden_size, + dtype=cfg.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=draft_config.lm_head_gather_output, + reduce_output=False, + use_custom_cublas_mm=False, + quant_config=None, + ) + self.logits_processor = LogitsProcessor() + self.allreduce = AllReduce( + mapping=mapping, strategy=draft_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + quant_config = draft_config.get_quant_config() + self.attn = create_attention( + draft_config.attn_backend, + self.layer_idx, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, rope=RopeParams.from_config(cfg) + ), + quant_config=quant_config, + skip_create_weights_in_init=draft_config.skip_create_weights_in_init, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): + raise ValueError( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + # The drafter embeds with the target's table, read off the core on every call (the engine materializes the + # core's weights after construction). + self._core = [core] + self.model = _DraftForward(self) + # The projections (weight keys) on the fixed-tactic low-M kernels (``DRAFT_LOWM``), set by + # ``prebuild_linears``. + self._lowm_keys: frozenset = frozenset() + + def prebuild_linears(self) -> int: + """Put the drafter's BF16 projections ``DRAFT_LOWM`` covers on the fixed-tactic low-M kernels and compile + every covered call. Once the weights are real, before the executor's warmup, on every rank. Returns the + number of low-M calls compiled.""" + keys = ("fc", "qkv", "o", "gu", "dn") + built = _draft_lowm_kernel.prebuild([self.w[k] for k in keys], DRAFT_LOWM.keys) + self._lowm_keys = frozenset(k for k in keys if _draft_lowm_kernel.installed(self.w[k])) + return built + + def _linear(self, x: torch.Tensor, key: str) -> torch.Tensor: + """``x @ w[key]^T``: the low-M kernel at the rows ``DRAFT_LOWM`` covers (``prebuild_linears``), else the plain + GEMM.""" + w = self.w[key] + if key in self._lowm_keys: + out = _draft_lowm_kernel.draft_linear(w, x, None) + if out is not None: + return out + return m3_bf16_linear(x, w) + + @property + def embed_table(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def load_weights_from_target_model(self, target_model: nn.Module) -> None: + """Nothing to alias: the drafter reads the target's embedding table through the core and owns its head.""" + + # ------------------------------------------------------------------------------------------------- the layer + + def fc(self, hs: torch.Tensor) -> torch.Tensor: + return self._linear(hs, "fc") + + def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The captured target hidden states [T, layers x H] to the drafter's input width: each layer's fc_norm, + then the fc (the in-tree ``Eagle3ForCausalLM.apply_eagle3_fc``).""" + hidden_states = hidden_states.to(self.dtype) + if hidden_states.shape[-1] != self.hidden: + chunks = hidden_states.chunk(self.capture_layers, dim=-1) + hidden_states = torch.cat( + [ + flashinfer_rmsnorm(c, self.w[f"fc_norm{i}"], self.rms_eps) + for i, c in enumerate(chunks) + ], + dim=-1, + ) + hidden_states = self.fc(hidden_states) + return hidden_states + + def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: + """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is + False, where the caller fuses the reduction).""" + qkv = self._linear(x, "qkv") + _record("m3_draft_attention", int(md.num_tokens)) + o = m3_draft_attention(self.attn, qkv, md) + out = self._linear(o, "o") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool, silu_cpp: bool = False) -> torch.Tensor: + """gate_up, the plain SwiGLU (with ``silu_cpp``, ``DRAFT_SILU_CPP``: this target's C++ op, bit for bit the + Triton one, in one launch), down.""" + gu = self._linear(h, "gu") + act = None + if silu_cpp and DRAFT_SILU_CPP: + act = m3_draft_silu(gu) + if act is None: + uncovered( + DRAFT_SILU_PATH, "the drafter SwiGLU input layout (the Triton op runs it)" + ) + else: + taken(DRAFT_SILU_PATH, int(gu.shape[0]), "drafter MLP SwiGLU, one C++ launch") + if act is None: + act = m3_silu_and_mul(gu, None, None, None) + out = self._linear(act, "dn") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def ar_norm(self, x: torch.Tensor, residual: torch.Tensor, norm_key: str): + """One draft-layer boundary in the in-tree structure: the projection's plain all-reduce (BF16 out), then the + fused residual add + RMSNorm (FP32 add, norm of the FP32 sum). Returns (normed, residual sum).""" + return flashinfer_fused_add_rmsnorm( + m3_allreduce(x, self.allreduce, None), residual, self.w[norm_key], self.rms_eps + ) + + def forward_layer( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + all_rank_num_tokens: Optional[List[int]] = None, + ): + """The drafter layer (the in-tree ``Eagle3DraftModel.forward`` for one layer): the embedded tokens and the + fc'd hidden states normed and concatenated, attention, the post-attention norm with the residual, the MLP and + the final norm. Returns (normed output, the hidden state the next pass carries).""" + previous = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + embeds = m3_embedding(input_ids, self.embed_table).to(self.dtype) + residual = hidden_states + h = flashinfer_rmsnorm(hidden_states, self.w["hidden_norm"], self.rms_eps) + e = flashinfer_rmsnorm(embeds, self.w["in_norm"], self.rms_eps) + h = torch.cat([e, h], dim=-1) + h = self.attention(h, attn_metadata, reduce=True) + h, residual = flashinfer_fused_add_rmsnorm( + h, residual, self.w["post_norm"], self.rms_eps + ) + h = self.mlp(h, reduce=True) + spec_metadata.maybe_capture_hidden_states(self.layer_idx, h, residual) + h, to_save = flashinfer_fused_add_rmsnorm( + h, residual, self.w["final_norm"], self.rms_eps + ) + return (h, h) if self.return_hidden_post_norm else (h, to_save) + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous + + +def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, capture_norm) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "tp_size", 1) > 1 + ) + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify and fused draft loop, at the in-tree numerics. + + * Verify: on decode-only batches whose verify rows the fused kernels cover (``tail_eligible``), the target's BF16 + lm_head computes this rank's vocabulary shard of the logits (the in-tree head's own local GEMM; at most + ``TARGET_HEAD_MAX_ROWS`` verify rows stream a lossless 12-bit coding of the same BF16 weight) and one kernel + replaces the logits all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. + * Draft: the drafter's first pass may take the per-layer fc_norms in one launch (or folded into the hidden-state + capture) and run its fc beside the target head GEMM (forked right after the last captured target layer at most + ``FC_EARLY_MAX_ROWS`` verify rows). On a decode-only step ``DRAFT_STEP`` covers, the draft loop is fused: each + pass's draft input in one launch, the drafter layer's attention and MLP with their all-reduces fused with the + residual add and the next RMSNorm, and one kernel for the BF16 head, the cross-rank argmax and the pass's + bookkeeping. + + Every other batch takes the in-tree paths. Methods whose docstring says "Mirrors" copy the named in-tree method + with the fused branches inserted. + """ + + def __init__(self, spec_config, mapping, use_separate_draft_kv_cache: bool = False) -> None: + super().__init__( + spec_config, mapping=mapping, use_separate_draft_kv_cache=use_separate_draft_kv_cache + ) + # Per-step state, reset at the top of every _forward_impl: the first draft pass's gather rows when this step + # takes the fused first drafter inputs, and the fc output a fork computed beside the target head. + self._m3_first_gather_ids: Optional[torch.Tensor] = None + self._m3_fc_hidden: Optional[torch.Tensor] = None + self._m3_target_lm_head: Optional[nn.Module] = None + # (fc_norm0 weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple] = None + # Per step: the captured layers written with the capture norm (set by capture_hidden_states). + self._m3_capture_normed: set = set() + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + # Armed per step by arm_capture_norm: the drafter whose first fc is forked at the last captured target layer, + # and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + # The 12-bit coded copy of the target head shard the fused verify streams, built before the warmup. + self._m3_target_head = None + self._m3_force_vec_warm = False + # The streamed copy of the drafter's BF16 head shard the fused draft loop reads, built before the warmup. + self._m3_draft_head = None + # Per step: the next step's input tokens [nreq, D + 1], column 0 written by the fused verify (or the seed) and + # columns 1.. by the fused draft passes; set as _m3_next_new_tokens once the fused draft loop filled it. + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + # Per step: the first draft inputs are the fused ones (accepted tokens as they lie, fc output of every row). + self._m3_fused_first_inputs = False + + def capture_target( + self, spec_metadata, layer_id: int, num_tokens: int + ) -> Optional[Tuple[torch.Tensor, torch.Tensor, float]]: + """Where this step's capture norm writes layer ``layer_id``'s capture, for a boundary all-reduce that writes it + itself (``CAPTURE_NORM_IN_BOUNDARY``): (the capture buffer's [num_tokens, H] column block, the layer's + fc_norm weight, eps), or None (no capture norm armed, not a captured layer, or more rows than the buffer). + The boundary writes every row it reduces; rows past num_capture_tokens are unread.""" + capture_norm = self._m3_capture_norm + if capture_norm is None or spec_metadata.hidden_states is None: + return None + if num_tokens > spec_metadata.hidden_states.shape[0]: + return None + hidden_size = spec_metadata.hidden_size + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + out = spec_metadata.hidden_states[ + :num_tokens, i * hidden_size : (i + 1) * hidden_size + ] + return out, capture_norm[0][i], capture_norm[1] + return None + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None, written: bool = False + ) -> None: + """The Eagle3 capture with this step's capture norm (``written``: the layer's boundary all-reduce already + wrote it, see ``capture_target``); on a step armed for the early fc fork (``arm_capture_norm``), the last + captured layer then forks the drafter's first fc onto the aux stream, where it overlaps the remaining target + layers and the head (joined by ``_m3_target_accept``).""" + if not written: + _eagle3_capture(spec_metadata, layer_id, hidden_states, residual, self._m3_capture_norm) + if self._m3_capture_norm is not None: + # Written with this step's capture norm: by the boundary all-reduce (written) or just above. + self._m3_capture_normed.add(layer_id) + drafter = self._m3_fc_early_drafter + layers = spec_metadata.layers_to_capture + if drafter is not None and layers and layer_id == layers[-1]: + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = self._m3_fc_fork( + spec_metadata, spec_metadata.num_capture_tokens, drafter + ) + + def prebuild(self, lm_head: nn.Module, draft: Optional[M3Eagle3Draft]) -> None: + """The fused verify's coded target head and the fused draft loop's head copy, from the loaded weights, + before the executor's warmup, on every rank.""" + start = time.perf_counter() + if self.prebuild_target_head(lm_head): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] coded target head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + start = time.perf_counter() + if self.prebuild_draft_head(draft): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Mirrors Eagle3OneModelWorker._forward_impl, with the fused verify branch.""" + runtime_draft_len = spec_metadata.runtime_draft_len + # skip the draft forward if the runtime draft length is 0 + if runtime_draft_len == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._m3_first_gather_ids = None + self._m3_next_new = None + self._m3_next_new_tokens = None + self._m3_fused_first_inputs = False + if logits is None: + # tail_eligible() skipped the gathered target logits: verify from this rank's logits shard. + accepted_tokens, num_accepted_tokens, self._m3_first_gather_ids = ( + self._m3_target_accept( + hidden_states, + spec_metadata, + batch_size, + runtime_draft_len, + draft_model, + kv_lens=getattr(attn_metadata, "kv_lens_cuda", None), + ) + ) + else: + self._execute_guided_decoder_if_present(logits) + + # Sample and accept tokens. ``input_ids`` is required by the relaxed-acceptance path (scans for + # thinking-phase tokens); ignored otherwise. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, attn_metadata, spec_metadata + ) + if self.draft_only_eligible(attn_metadata, spec_metadata, draft_model): + self._m3_seed_first_pass(num_accepted_tokens, spec_metadata, batch_size) + + # Mamba hybrid models need state updates after token acceptance (a no-op on this target's cache manager). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None) + if self.sa_enhancer is not None and sa_manager is not None: + self.sa_enhancer.extend_and_prepare( + sa_manager=sa_manager, + request_ids=spec_metadata.request_ids, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_gens=num_gens, + num_contexts=num_contexts, + max_draft_len=runtime_draft_len, + ) + + if self._uses_external_shared_target_kv: + next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + ) + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here so the post-loop restore (below) + # can put attn_metadata back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + # prepare next new tokens to support overlap scheduler + if self._m3_next_new_tokens is not None: + # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). + next_new_tokens = self._m3_next_new_tokens + else: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + # Commit model-owned target state only after draft preparation has succeeded; SpecWorkerBase aborts the + # still-pending snapshots on any exception above this point. + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices[:batch_size], + num_contexts, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """The fused first draft inputs on a step that verified with the fused kernels (or was seeded for them), + else the in-tree ones.""" + fused = self._m3_first_drafter_inputs( + position_ids, + accepted_tokens, + input_ids.shape[0], + attn_metadata.num_contexts, + spec_metadata, + draft_model, + ) + self._m3_fused_first_inputs = fused is not None + if fused is not None: + taken("m3_first_drafter_inputs", attn_metadata.num_seqs) + fused.update(attn_metadata=attn_metadata, spec_metadata=spec_metadata) + return fused + return super().prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + def _m3_structural_ok(self, attn_metadata, spec_metadata) -> bool: + """The batch conditions shared by the fused verify and the fused first drafter inputs: a decode-only plain + Eagle3 step without attention DP, guided decoding, relaxed acceptance or a suffix-automaton enhancer, of at + most DECODE_MAX_TOKENS verify rows. Host state only, so the TP ranks agree.""" + if self.is_mtp_eagle or self.use_dynamic_tree: + return False + if self._uses_external_shared_target_kv or self.sa_enhancer is not None: + return False + if ( + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp + ): + return False + if getattr(self.spec_config, "use_relaxed_acceptance_for_thinking", False): + return False + if attn_metadata.num_contexts != 0 or attn_metadata.num_seqs == 0: + return False + draft_len = spec_metadata.runtime_draft_len + if draft_len == 0 or attn_metadata.num_seqs * (draft_len + 1) > DECODE_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Whether this step verifies with the fused kernels instead of the gathered target logits. + + Decode-only Eagle3 batches (``_m3_structural_ok``) over an unpadded vocab-sharded bias-free target head, all + greedy at a request count ``SHARD_ACCEPT`` covers, or sampled under the temperature + top-p contract (see + _m3_sampled_verify_ok), with nothing else reading the target logits. Every input is host state shared by all + TP ranks, so the ranks agree. + """ + if not self._m3_structural_ok(attn_metadata, spec_metadata): + return False + if getattr(spec_metadata, "enable_penalty", False): + return False + if not _sharded_head(lm_head) or getattr(lm_head, "bias", None) is not None: + return False + draft_len = spec_metadata.runtime_draft_len + if spec_metadata.is_all_greedy_sample: + key = (draft_len, attn_metadata.num_seqs) + if not SHARD_ACCEPT.covers(key): + uncovered( + SHARD_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + elif attn_metadata.num_seqs > SAMPLED_VERIFY_MAX_REQUESTS: + uncovered("m3_sampled_accept", f"sampled verify at {attn_metadata.num_seqs} requests") + return False + elif not self._m3_sampled_verify_ok(spec_metadata): + return False + self._m3_target_lm_head = lm_head + return True + + def draft_only_eligible(self, attn_metadata, spec_metadata, draft_model) -> bool: + """Whether a step verified by the in-tree path can still take the fused first drafter inputs: the structural + conditions of tail_eligible, without the target-sampling ones, up to ``FIRST_DRAFTER_INPUTS_MAX_REQUESTS`` + requests.""" + if attn_metadata.num_seqs > FIRST_DRAFTER_INPUTS_MAX_REQUESTS: + uncovered( + "m3_first_drafter_inputs", f"in-tree verify at {attn_metadata.num_seqs} requests" + ) + return False + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None + + def _m3_seed_first_pass(self, num_accepted_tokens, spec_metadata, batch_size: int) -> None: + """The first draft pass's gather rows (the last accepted row of each request) from the in-tree acceptance, as + _m3_target_accept leaves them.""" + k1 = spec_metadata.runtime_draft_len + 1 + n = num_accepted_tokens[:batch_size].to(torch.int64) + rows = spec_metadata.batch_indices_cuda[:batch_size].to(torch.int64) + self._m3_first_gather_ids = rows * k1 + n - 1 + + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], + ): + """The in-tree first draft inputs of a decode-only step, with the per-layer fc_norm RMSNorms of + ``apply_eagle3_fc`` in one kernel (unless the capture already applied them) before the fc; the input ids are + the accepted tokens as they lie. None to use the in-tree path.""" + if ( + self._m3_first_gather_ids is None + or self.is_mtp_eagle + or num_contexts != 0 + or draft_model is None + ): + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + ): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not draft_model.w["fc_norm0"]: + if torch.cuda.is_current_stream_capturing(): + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), + ) + self._m3_fc_norm_w = cache + hidden = self._m3_fc_hidden + self._m3_fc_hidden = None + if hidden is None or hidden.shape[0] != num_tokens: + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + hs = m3_spec_tail.chunk_norm(hs, cache[1], draft_model.rms_eps) + else: + hs = self._m3_renorm_raw_captures(hs, spec_metadata, draft_model) + hidden = draft_model.fc(hs) + input_ids = accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].reshape(-1) + return {"input_ids": input_ids, "position_ids": position_ids, "hidden_states": hidden} + + def _m3_renorm_raw_captures( + self, hs: torch.Tensor, spec_metadata, draft_model: M3Eagle3Draft + ) -> torch.Tensor: + """The captured rows ``hs`` [rows, layers * H] for a reader that skips the fc_norm (the capture norm is armed): + unchanged when every captured layer was written with this step's norm (capture_hidden_states records them), + else -- not expected on any path -- a copy with each layer captured raw normalized by its own fc_norm here, + logged as a miss, so a raw row never reaches the drafter's fc.""" + missing = [ + i + for i, layer_id in enumerate(spec_metadata.layers_to_capture) + if layer_id not in self._m3_capture_normed + ] + if not missing: + return hs + missed( + "m3_capture_norm", + int(hs.shape[0]), + f"captured layer(s) {[spec_metadata.layers_to_capture[i] for i in missing]} were written without the " + "capture norm; their fc_norm is applied before the drafter's fc", + ) + hs = hs.clone() + hidden = spec_metadata.hidden_size + for i in missing: + block = hs[:, i * hidden : (i + 1) * hidden] + block.copy_( + flashinfer_rmsnorm( + block.contiguous(), draft_model.w[f"fc_norm{i}"], draft_model.rms_eps + ) + ) + return hs + + def _m3_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + after: Optional[torch.cuda.Event] = None, + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (after the event + ``after`` when given, else after everything queued on the current stream); the result is picked up by + _m3_first_drafter_inputs. Returns the stream to join, or None when the fc path does not apply (the first + drafter inputs then compute it).""" + self._m3_fc_hidden = None + if draft_model is None: + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + or self._m3_fc_norm_w is None + ): + return None + aux = self._m3_aux_stream + if aux is None: + aux = self._m3_aux_stream = torch.cuda.Stream() + if after is not None: + aux.wait_event(after) + else: + aux.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(aux): + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + hs = m3_spec_tail.chunk_norm( + hs, self._m3_fc_norm_w[1], draft_model.rms_eps, use_pdl=False + ) + else: + hs = self._m3_renorm_raw_captures(hs, spec_metadata, draft_model) + self._m3_fc_hidden = draft_model.fc(hs) + return aux + + def _m3_fc_join(self, fork: torch.cuda.Stream) -> None: + """Join the fc fork onto the current stream and keep its output's memory alive for that stream. + + _m3_fc_hidden is allocated in fork (aux) stream order and read on the current stream. Without the record, + dropping it returns its block to the aux stream's pool, where later aux work can take it while the current + stream's reads are still queued (eager warm-up / mixed steps). Nothing is recorded under graph capture, as in + maybe_execute_in_parallel: a captured graph relies on the next fork's aux.wait_stream(main) to order the + reuse.""" + main = torch.cuda.current_stream() + main.wait_stream(fork) + hidden = self._m3_fc_hidden + if ( + hidden is not None + and hidden.is_cuda + and not (torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()) + ): + hidden.record_stream(main) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model: Optional[M3Eagle3Draft] + ) -> None: + """Before every target forward: whether this step's Eagle3 captures apply the drafter's fc_norm themselves. + + On a fused-verify step whose fc_norm weights are already stacked (by an earlier fused first pass), the + capture writes rmsnorm(h + r) * fc_norm[i] (one kernel per captured layer instead of add + copy) and the fused + first drafter inputs read it without their own chunk_norm. A fused-verify step of at most + ``FC_EARLY_MAX_ROWS`` verify rows is also armed to fork the drafter's first fc at the last captured layer + (``capture_hidden_states``). + """ + self._m3_capture_norm = None + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = None + # The captured layers this step wrote with its capture norm (capture_hidden_states). + self._m3_capture_normed = set() + # A fork result no first draft pass consumed must not reach a later step's (same-size) first inputs. + self._m3_fc_hidden = None + if draft_model is None: + return + if not self.tail_eligible(attn_metadata, spec_metadata, lm_head): + return + # The fork reads every captured column at the last captured layer, so a final-state capture, which is written + # after the model returns, keeps the fc at its stock place. + if ( + attn_metadata.num_seqs * (spec_metadata.runtime_draft_len + 1) <= FC_EARLY_MAX_ROWS + and FINAL_HIDDEN_LAYER_ID not in spec_metadata.layers_to_capture + ): + self._m3_fc_early_drafter = draft_model + fc_norm_w = self._m3_fc_norm_w + if ( + fc_norm_w is not None + and fc_norm_w[0] is draft_model.w["fc_norm0"] + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.rms_eps) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata) -> bool: + """Whether sampled verify rows can use the fused sampled op: only under advanced_sampling_mode=no_topk, the + deploy-time guarantee that no row uses top-k or min-p (that op applies neither).""" + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + return mode is not None and mode.skips_top_k and not mode.skips_top_p + + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + kv_lens: Optional[torch.Tensor] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path. ``kv_lens``: the + requests' KV lengths at this verify (the per-request forced-acceptance key, ``FORCE_AL_PER_REQUEST``).""" + force_total = 0 + force_vec = None + if self.force_num_accepted_tokens != 0.0: + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # The draw pool is built by an H2D copy and the draw kernel JIT-compiles per block width, so both + # happen here, before the first captured step (a scratch counter keeps the real draw sequence). + self._ensure_force_accept_rng_state(hidden_states.device) + if not self._m3_force_vec_warm: + scratch = torch.zeros(1, dtype=torch.int64, device=hidden_states.device) + dummy = torch.empty(64, dtype=torch.int32, device=hidden_states.device) + for block in (16, 32, 64): + m3_force_vec( + self._force_accept_rng_pool, + scratch, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + m3_force_vec_req( + self._force_accept_rng_pool, + dummy, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + self._m3_force_vec_warm = True + if not is_warmup: + force_total = min(int(self.force_num_accepted_tokens) + 1, draft_len + 1) + force_vec = self._m3_force_vec( + batch_size, force_total, draft_len + 1, hidden_states.device, kv_lens + ) + draft_tokens = spec_metadata.draft_tokens.view(batch_size, draft_len) + if draft_tokens.dtype != torch.int32: + draft_tokens = draft_tokens.to(torch.int32) + # Column 0 receives the last accepted token, the next step's first input (the fused draft loop adds the rest). + self._m3_next_new = torch.empty( + (batch_size, draft_len + 1), dtype=torch.int32, device=hidden_states.device + ) + rows = hidden_states[: batch_size * (draft_len + 1)] + lm_head = self._m3_target_lm_head + batch_idx = spec_metadata.batch_indices_cuda[:batch_size] + head = self._m3_coded_head(lm_head, rows.shape[0]) + # The Eagle3 fc on the aux stream: the fork made at the last captured layer, else one beside the head. A + # sampled verify's coded head holds every SM while it streams: its fc is forked after it instead (overlapping + # the sampled acceptance), where it does not delay the head's start. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + head_first = ( + fork is None + and (head is not None or rows.shape[0] <= FC_AFTER_HEAD_MAX_ROWS) + and not spec_metadata.is_all_greedy_sample + ) + if fork is None and not head_first: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the head's own local GEMM (or its coded copy), whose + # all-gather the fused kernels replace (they exchange per-row records across ranks, not logits). + if head is not None: + taken(TARGET_HEAD.path, rows.shape[0]) + logits = m3_target_head_kernel.logits(head, rows) + else: + logits = lm_head.apply_linear(rows, None) + head_done = None + if head_first: + head_done = torch.cuda.Event() + head_done.record() + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept_kernel.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = SHARD_ACCEPT.path + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept_kernel.sampled_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + spec_metadata.temperatures[:n_rows], + spec_metadata.top_ks[:n_rows], + spec_metadata.top_ps[:n_rows], + None, + seeds, + offsets, + force_total, + self._m3_next_new, + force_vec=force_vec, + one_round_k=SAMPLED_ACCEPT_1R_K if SAMPLED_ACCEPT_SM107TP4 else 0, + prewait_params=SAMPLED_ACCEPT_PREWAIT_PARAMS, + early_trigger=SAMPLED_ACCEPT_EARLY_TRIGGER, + ilp_walk=SAMPLED_ACCEPT_ILP_WALK, + ) + path = "spec_tail" + if head_done is not None: + # Enqueued after the sampled acceptance (which starts first) but waits on the head only. + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model, after=head_done) + if fork is not None: + self._m3_fc_join(fork) + if result is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, None, "the fused verify's exchange workspace") + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + taken( + SHARD_ACCEPT.path if spec_metadata.is_all_greedy_sample else "m3_sampled_accept", + batch_size, + f"draft length {draft_len}" + + ( + "" + if spec_metadata.is_all_greedy_sample + else f", SA-1R one_round_k {SAMPLED_ACCEPT_1R_K if SAMPLED_ACCEPT_SM107TP4 else 0}" + ), + ) + return result + + def prebuild_target_head(self, lm_head: nn.Module) -> bool: + """Build the 12-bit coded copy of the target head shard the fused verify streams at <= TARGET_HEAD_MAX_ROWS + rows (rebuilt from the current weight values on every call), on every rank. True when the coded copy decodes + to the BF16 weight bit for bit; otherwise (or for a head the fused verify does not read) nothing is kept and + the head's own GEMM runs.""" + self._m3_target_head = None + if not _sharded_head(lm_head) or not m3_target_head_kernel.head_supported(lm_head): + return False + head = m3_target_head_kernel.prepare_head(lm_head) + self._m3_target_head = head + if not head.coded: + logger.warning( + f"[M3 target head] the coded shard {tuple(lm_head.weight.shape)} does not decode to the BF16 weight " + "bit for bit; the fused verify keeps the head's BF16 GEMM" + ) + return head.coded + + def _m3_coded_head(self, lm_head: nn.Module, rows: int): + """The coded target head for a fused verify of ``rows`` rows, or None for the head's own GEMM: above + ``TARGET_HEAD_MAX_ROWS`` rows (uncovered), on a shard whose coding failed its round trip (uncovered: the GEMM + is the exact path), and when the copy was not built for this weight (a miss).""" + path = TARGET_HEAD.path + if not TARGET_HEAD.covers(rows): + uncovered(path, f"{rows} verify rows (the head's own GEMM runs)") + return None + head = self._m3_target_head + if head is None or head.weight is not lm_head.weight: + missed( + path, + rows, + "the coded target head is not built for the head's current weight (prebuild)", + ) + return None + if not head.coded: + uncovered( + path, "the coded shard failed its bitwise round trip (the head's own GEMM runs)" + ) + return None + return head + + def _m3_force_vec( + self, + batch_size: int, + base_total: int, + max_total: int, + device, + kv_lens: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + ``FORCE_AL_PER_REQUEST`` (with ``kv_lens``): each request draws from the fixed-seed pool at its own key, its + KV length at this verify (its prompt plus what it has accepted so far), so its draws do not depend on the + server's history or on the other requests; every TP rank holds the same lengths, so the ranks agree. Without + lengths: exactly as _apply_force_accepted_tokens draws for a decode-only batch (one shared counter bump per + step, the (counter, slot) hash), the generic path's sequence. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: + return None + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + if FORCE_AL_PER_REQUEST and kv_lens is not None: + keys = kv_lens[:batch_size] + if keys.dtype != torch.int32: + keys = keys.to(torch.int32) + return m3_force_vec_req( + self._force_accept_rng_pool, + keys, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (batch_size - 1).bit_length()), + ) + return m3_force_vec( + self._force_accept_rng_pool, + self._force_accept_rng_counter, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (batch_size - 1).bit_length()), + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[M3Eagle3Draft]) -> bool: + """Build the streamed copy of the drafter's BF16 head shard that the fused draft loop reads (rebuilt from the + current weight values on every call), on every rank. False when the drafter is not one the loop serves.""" + self._m3_draft_head = None + if draft_model is None or not self._m3_drafter_supported(draft_model): + return False + self._m3_draft_head = m3_draft_step_kernel.prepare_head( + draft_model.lm_head, self.mapping.tp_rank, DRAFT_HEAD_CODED + ) + return True + + def _m3_drafter_supported(self, draft_model: M3Eagle3Draft) -> bool: + """The drafter the fused draft loop runs (``M3Eagle3Draft`` validates its shape at construction) over this + worker's tensor-parallel group with a bias-free vocab-sharded BF16 head without padding.""" + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if self._d2t is not None: + return False + lm_head = draft_model.lm_head + return _sharded_head(lm_head) and m3_draft_step_kernel.head_supported(lm_head) + + def _m3_draft_ar_norm( + self, + draft_model: M3Eagle3Draft, + x: torch.Tensor, + residual: torch.Tensor, + norm_key: str, + pushed: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """One drafter boundary of the fused draft loop: the multicast all-reduce + residual + RMSNorm in its drafter + form (``DRAFT_MC_AR``; bitwise the stock boundary), else the built-in path (``M3Eagle3Draft.ar_norm``). The + choice depends on shapes and a collectively created workspace only, so the TP ranks agree. ``pushed``: the + projection already pushed its rows into the all-reduce (``DRAFT_AR_PUSH``).""" + rows = x.shape[0] + path = DRAFT_MC_AR.path + if DRAFT_MC_AR.covers(rows): + workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) + out = None + if workspace is not None: + out = m3_mc_allreduce.m3_mc_allreduce_norm( + x.contiguous(), + residual.contiguous(), + draft_model.w[norm_key], + draft_model.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + drafter_norm=True, + carveout=DRAFT_MC_AR_CARVEOUT, + prepushed=2 if pushed else 0, + drafter_cluster=DRAFT_MC_AR_CLUSTER, + ) + if out is not None: + taken( + path, + rows, + "drafter boundary (cluster kernel)" + if DRAFT_MC_AR_CLUSTER + else "drafter boundary (row kernel)", + ) + if pushed: + taken(DRAFT_AR_PUSH_PATH, rows, "the projection pushed its rows (prepushed 2)") + elif DRAFT_AR_PUSH: + uncovered( + DRAFT_AR_PUSH_PATH, + "the boundary's projection did not push (the AR pushes x)", + ) + return out + if pushed: + raise RuntimeError( + "the drafter's rows were pushed into a boundary all-reduce that did not run" + ) + missed(path, rows, "the multicast all-reduce workspace or operand layout") + return draft_model.ar_norm(x, residual, norm_key) + + def _m3_arm_draft_push(self, weight: torch.Tensor, rows: int) -> None: + """``DRAFT_AR_PUSH``: let the projection ``weight``'s next low-M split-K call push its rows into the boundary + all-reduce that ``_m3_draft_ar_norm`` then runs (the multicast path only; shapes and a collectively created + workspace decide it, so the TP ranks agree).""" + if not (DRAFT_AR_PUSH and DRAFT_MC_AR.covers(rows)): + return + workspace = m3_mc_allreduce.get_workspace(self.mapping, weight.shape[0]) + if workspace is not None: + _draft_lowm_kernel.arm_push(weight, workspace, rows) + + def _m3_fused_loop_ok( + self, attn_metadata, spec_metadata, batch_size: int, draft_model: Optional[M3Eagle3Draft] + ) -> bool: + """Whether this step drafts with the fused loop: a decode-only step the fused verify or the seed prepared + (``_m3_first_gather_ids``) with the fused first draft inputs, a supported drafter, a (draft length, requests) + ``DRAFT_STEP`` covers, the prebuilt head and the exchange workspace. Host state only, so the TP ranks + agree.""" + if ( + self._m3_first_gather_ids is None + or not self._m3_fused_first_inputs + or draft_model is None + ): + return False + path = DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + uncovered( + path, "the drafter or its attention metadata is not one the fused loop serves" + ) + return False + draft_len = spec_metadata.runtime_draft_len + if not DRAFT_STEP.covers((draft_len, batch_size)): + uncovered(path, f"draft length {draft_len} at {batch_size} requests") + return False + head = self._m3_draft_head + if head is None or head.weight is not draft_model.lm_head.weight: + missed( + path, + batch_size, + "the drafter head's streamed copy is not built for its current weight (prebuild)", + ) + return False + if m3_draft_step_kernel.workspace(self.mapping, head.wf.device) is None: + capture_missed(path, None, "the draft-step exchange workspace") + return False + return True + + def _forward_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ): + """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop.""" + if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): + return self._m3_fused_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + return super()._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + + def _m3_fused_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model: M3Eagle3Draft, + draft_kv_cache_manager, + batch_size: int, + num_accepted_tokens, + ): + """Mirrors Eagle3OneModelWorker._forward_linear_draft_loop for a decode-only Eagle3 step, with each pass's + glue fused (see ``_m3_fused_loop_ok``): + + * the draft input -- the embedding lookup, the layer's input norm and hidden norm and their concat -- in one + launch (``m3_spec_tail.draft_input``; the second pass gathers the last accepted rows there); + * the drafter layer's attention and MLP, each followed by the all-reduce with the residual add and the next + RMSNorm (the post-attention norm, then the drafter's final norm); + * one kernel for the BF16 head over the head rows (pass 0: the last accepted row of each request), the + cross-rank argmax and the bookkeeping the in-tree loop does between passes: the draft token, the next pass's + input id and position, the draft KV lengths (and pass 0's sequence lengths) (``m3_draft_step``). + + The host metadata transitions are the in-tree loop's. The drafts land in columns 1.. of the step's next-input + tokens (column 0 is the last accepted token), which ``_forward_impl`` returns as they lie. + """ + draft_len = spec_metadata.runtime_draft_len + table = draft_model.embed_table + dev = table.device + head = self._m3_draft_head + gather_ids = self._m3_first_gather_ids + ids = inputs["input_ids"] + hidden = inputs["hidden_states"] + pos = inputs["position_ids"] + if ids.dtype != torch.int32: + ids = ids.to(torch.int32) + if pos.dtype != torch.int32: + pos = pos.to(torch.int32) + num_acc = num_accepted_tokens[:batch_size] + if num_acc.dtype != torch.int32: + num_acc = num_acc.to(torch.int32) + next_new = self._m3_next_new + if next_new is None: + # Seeded from the in-tree acceptance: the last accepted token is the first-pass input at the gather row. + next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=dev) + next_new[:, 0] = ids[gather_ids] + draft_tokens = next_new[:, 1:] + w = draft_model.w + eps = draft_model.rms_eps + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) as md: + for i in range(draft_len): + x, residual = m3_spec_tail.draft_input( + ids, + table, + hidden, + gather_ids if i == 1 else None, + w["in_norm"], + eps, + w["hidden_norm"], + eps, + ) + taken(DRAFT_INPUT.path, x.shape[0]) + rows = x.shape[0] + self._m3_arm_draft_push(w["o"], rows) + h = draft_model.attention(x, md, reduce=False) + h, residual = self._m3_draft_ar_norm( + draft_model, + h, + residual, + "post_norm", + pushed=_draft_lowm_kernel.take_pushed(h), + ) + self._m3_arm_draft_push(w["dn"], rows) + h = draft_model.mlp(h, reduce=False, silu_cpp=True) + h, to_save = self._m3_draft_ar_norm( + draft_model, + h, + residual, + "final_norm", + pushed=_draft_lowm_kernel.take_pushed(h), + ) + if draft_model.return_hidden_post_norm: + to_save = h + ids_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + pos_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + if not m3_draft_step_kernel.draft_step( + head, + h, + batch_size, + gather_ids if i == 0 else None, + self.mapping, + draft_tokens, + i, + ids_next, + pos, + pos_next, + md.kv_lens_cuda, + num_acc if i == 0 else None, + md._seq_lens_cuda if i == 0 else None, + draft_len, + ): + raise RuntimeError( + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" + ) + taken(DRAFT_STEP.path, batch_size, f"draft length {draft_len}") + ids, pos, hidden = ids_next, pos_next, to_save + if i == 0: + # The draft step set the device sequence lengths to 1 and rolled the KV lengths back. + md._seq_lens[:batch_size].fill_(1) + md.on_update() + has_kv_cache = md.kv_cache_manager is not None + if has_kv_cache and hasattr(md, "host_request_types"): + md.host_request_types[: md.num_contexts].fill_(1) + if has_kv_cache: + md.num_contexts = 0 + self._prepare_flash_mla_generation_layout(md, 0, batch_size) + md.update_for_spec_dec() + md.use_spec_decoding = False + else: + md.update_for_spec_dec() + self._m3_next_new_tokens = next_new + return draft_tokens + + +def _build_speculation(shell, model_config: ModelConfig): + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise ValueError( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + _set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return draft_config, draft, worker, use_separate + + +# -------------------------------------------------------------------------------------------------------------------- +# Runtime hooks and the shell +# -------------------------------------------------------------------------------------------------------------------- + + +def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: + """The step-loop hooks this target opts into (each behind its constant above). + + * ``pack_step_inputs``: a decode step's small input copies go out as one upload and one scatter at the graph + replay instead of a chain of copies the next graph waits on (~30 us of idle GPU at one request). + * ``skip_unchanged_h2d``: steady decode re-uploads slot indices, query lengths and page tables that rarely change. + * ``loop_on_execution_stream``: the overlap loop's own stream work and the forward on one stream, so step N+1's + inputs wait on graph N in program order rather than through two cross-stream event hops. + * ``merge_tuned_tactics_across_ranks``: tensor-parallel peers capture the same GEMM tactics (rank-local profiling + noise otherwise splits them, and the slowest rank sets every step). + * ``skip_greedy_rng_upload``: all-greedy speculative batches never read the Philox seeds. + * ``spec_step_fold``: the spec sampler's slot stores and the engine's next-step input gathers in one launch each. + * ``pack_parallel_scatter`` / ``pack_in_graph`` / ``pack_early_copy``: the pack's copies scatter in parallel, issued + by the decode graph's own prologue from a window uploaded ahead of the replay on a side copy stream. + * ``pack_spec_rng_uploads`` / ``spec_small_h2d``: the spec step's seeds, accepted counts and slot ids ride the pack + or stay on the device. + * ``sampler_d2h_side_stream`` / ``iteration_timing_one_event``: the sampler's store copies and the iteration timing + leave the main stream. + * MSA: identity selections for short sequences (and short/long graph families), the one-launch KV-length update, + unchanged block-table copies skipped, pool views built once, block-id tables copied from the page-index views. + """ + return RuntimeHooks( + first_content_emit_iters=FIRST_CONTENT_EMIT_ITERS, + mxfp8_pinned_tactics=True, + late_compile_detect=True, + pack_step_inputs=STEP_PACK_INPUTS, + pack_parallel_scatter=STEP_PACK_PARALLEL_SCATTER, + pack_in_graph=STEP_PACK_IN_GRAPH, + pack_early_copy=STEP_PACK_EARLY_COPY, + pack_step_epoch=STEP_PACK_EPOCH, + pack_prologue_op=m3_step_prologue if STEP_PACK_CPP_PROLOGUE else None, + skip_unchanged_h2d=STEP_SKIP_UNCHANGED_H2D, + h2d_changed_span=STEP_H2D_CHANGED_SPAN, + kv_lens_update_once=STEP_KV_LENS_UPDATE_ONCE, + spec_offsets_op=m3_spec_offsets if STEP_SPEC_OFFSETS_ONE_LAUNCH else None, + loop_on_execution_stream=STEP_LOOP_ON_EXECUTION_STREAM, + merge_tuned_tactics_across_ranks=STEP_MERGE_TUNED_TACTICS, + skip_greedy_rng_upload=STEP_SKIP_GREEDY_RNG_UPLOAD, + pack_spec_rng_uploads=STEP_PACK_SPEC_RNG_UPLOADS, + spec_small_h2d=STEP_SPEC_SMALL_H2D, + spec_step_fold=spec_fold, + sampler_d2h_side_stream=STEP_SAMPLER_D2H_SIDE_STREAM, + iteration_timing_one_event=STEP_ONE_TIMING_EVENT, + msa=MsaRuntimePolicy( + short_sequence_identity=MSA_SHORT_SEQUENCE_IDENTITY, + kv_update=_msa_kv_update_kernel.msa_kv_update if MSA_KV_UPDATE else None, + skip_unchanged_block_offsets=MSA_SKIP_UNCHANGED_BLOCK_OFFSETS, + reuse_pool_views=MSA_REUSE_POOL_VIEWS, + block_ids_from_views=MSA_BLOCK_IDS_FROM_VIEWS, + ), + ) + + +class ModelingV2MiniMaxM3Nvfp4Sm107Tp4( + DecoderModelForCausalLM[MiniMaxM3Tp4ModelingV2, PretrainedConfig] +): + """MiniMax-M3 NVFP4 / sm_107 / tp4, text only, with one-model Eagle3 or no speculation. + + The shell owns the vocab-parallel BF16 ``lm_head`` and the logits processor (``DecoderModelForCausalLM``); with + speculation, the drafter (``M3Eagle3Draft``, its weights from the speculative checkpoint through + ``load_draft_weights``) and the worker (``M3Eagle3Worker``) that verifies and drafts each step. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + assert torch.cuda.get_device_capability() == _SM, ( + f"this target is built for sm_{_SM[0]}{_SM[1]}, the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + # Under modeling_v2='require' a covered decode shape that cannot take its fast path raises. + set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + model_config = _text_model_config(model_config) + cfg = model_config.pretrained_config + spec_config = getattr(model_config, "spec_config", None) + if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: + # The one-model Eagle3 layer has no NVFP4 trtllm-gen decode kernel: keep its modules and the shared + # draft cache on FP8 while the target stays NVFP4. + model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 + super().__init__( + MiniMaxM3Tp4ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self.spec_config = spec_config + self.draft_model: Optional[M3Eagle3Draft] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[M3Eagle3Worker] = None + self.use_separate_draft_kv_cache = False + self.layer_idx = FINAL_HIDDEN_LAYER_ID + spec_fold = None + if spec_config is not None: + ( + self.draft_config, + self.draft_model, + self.spec_worker, + self.use_separate_draft_kv_cache, + ) = _build_speculation(self, model_config) + if STEP_SPEC_FOLD: + spec_fold = M3SpecStepFold( + in_forward=STEP_SPEC_FOLD_IN_FORWARD, + store_copy_wait=STEP_SPEC_STORE_COPY_SIDE, + cpp_tail=STEP_FOLD_TAIL_CPP, + ) + self._spec_fold = spec_fold + self.runtime_hooks = _runtime_hooks(spec_fold) + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights: Dict, weight_mapper=None) -> None: + """The drafter's checkpoint (``speculative_model``), loaded by the engine after the target's.""" + _weights.DRAFT_WEIGHTS.load(self.draft_model, weights) + + def set_guided_decoder(self, guided_decoder) -> bool: + """Mirrors ``SpecDecOneEngineForCausalLM.set_guided_decoder``: guided decoding with speculation is installed + on the worker.""" + if hasattr(self.spec_worker, "set_guided_decoder"): + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + def cache_derived_state(self) -> None: + """Build the targets and every decode kernel once the weights are real, before the executor's warmup; with + speculation also the drafter layer's full-context decode attention (``DRAFTER_FULLCTX``), the fused verify's + coded target head, the fused draft loop's head copy and the drafter's low-M GEMMs (``DRAFT_LOWM``).""" + super().cache_derived_state() + self.model.post_load_weights() + draft = self.draft_model + if draft is not None and DRAFTER_FULLCTX.keys: + path = DRAFTER_FULLCTX.path + installed = _drafter_attn_kernel.install( + draft, + {k: FULLCTX_DECODE_RECIPES[k] for k in DRAFTER_FULLCTX.keys}, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), + get_env_enable_pdl(), + fused_rope_mode=(DRAFTER_ROPE_MODE if DRAFTER_ROPE_FUSED else None), + on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), + on_taken=lambda rows: taken(path, rows), + on_uncovered=lambda reason: uncovered(path, reason), + on_missed=lambda rows, reason: missed(path, rows, reason), + ) + logger.info(f"[M3 prebuild] {path} installed on {installed} drafter attention layer(s)") + if self.spec_worker is not None: + self.spec_worker.prebuild(self.lm_head, draft) + if draft is not None: + # The drafter linears' low-M kernels, before the warmup. + start = time.perf_counter() + built = draft.prebuild_linears() + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] drafter low-M GEMMs built before warmup: %d calls, %.1f s" + % (built, time.perf_counter() - start) + ) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + """One engine step: the target forward, then the logits (or, with speculation, the worker's verify and + draft). Text only.""" + multimodal_params = kwargs.pop("multimodal_params", None) or [] + if any(p is not None and getattr(p, "multimodal_data", None) for p in multimodal_params): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + worker = self.spec_worker + capture = _plain_capture + if worker is not None: + capture = worker.capture_hidden_states + if spec_metadata is not None: + worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + capture=capture, + capture_target=worker.capture_target if worker is not None else None, + **kwargs, + ) + if getattr(attn_metadata, "csa2_remote_tail_mode", None) == "source": + # A context worker stopped at the encoder/decoder boundary: a disposable token for the context-only + # completion, and no draft worker. + return self.logits_processor.forward(hidden_states, self.lm_head, attn_metadata, False) + if spec_metadata is not None and spec_metadata.is_layer_capture(FINAL_HIDDEN_LAYER_ID): + capture(spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + if worker is None: + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], self.lm_head, attn_metadata, True + ) + spec_input_ids = input_ids if input_ids is not None else kwargs.get("orig_input_ids") + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if spec_input_ids is not None: + spec_input_ids = spec_input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens) + outputs = worker( + input_ids=spec_input_ids, + position_ids=spec_position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + spec_fold = self._spec_fold + if ( + spec_fold is not None + and spec_fold.in_forward + and attn_metadata.num_contexts == 0 + and isinstance(outputs, dict) + ): + spec_fold.store_in_forward(outputs) + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/__init__.py new file mode 100644 index 000000000000..b8c7426e3933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Activation ops of this target: the gated MLPs' fused SwiGLU.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..5bcdee3fbb82 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SiluAndMul(OpWrapper): + """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with + ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain + SwiGLU ``silu(gate) * up``. + + The MXFP8 dense MLPs (layers 0-2, intermediate 3072 per rank) and the MoE's shared expert (intermediate 768 per + rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 4608 per + rank) runs it plain. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=768, oai=True)), + Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=768, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=3072, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=3072, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=4608, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=4608, oai=False), + ), + ) + + note = """ + The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two + differ only by the fast sigmoid's last bits (the test's band). + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + gate, up = x.float().chunk(2, dim=-1) + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + a = 1.0 if alpha is None else alpha + b = 0.0 if beta is None else beta + return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" + + +m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/__init__.py new file mode 100644 index 000000000000..a7f7fee7a234 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention ops of this target: MSA block selection, the sparse decode and the KV-length update.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..ba008f493b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per +request count (``modeling.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 16 +KV_HEADS = 1 +HEAD_DIM = 128 +#: The split this target hands trtllm-gen at one request (modeling.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model +#: maximum the in-tree step passes otherwise. +CAPPED = 16384 +MODEL_MAX = 1 << 20 + + +def dense_attention_reference( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first + ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's + tokens in order). Returns fp32 [total_q, heads, head_dim].""" + batch, kv_heads, max_len, _ = keys.shape + heads = q.shape[1] + group = heads // kv_heads + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + pos = torch.arange(max_len, device=q.device) + with true_fp32_matmul(): + for b in range(batch): + full = int(seq_lens[b]) + rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] + limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) + valid = pos[None, :] < limit[:, None] # [qlen, max_len] + for h in range(kv_heads): + qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] + logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + probs = torch.softmax(logits, dim=-1) + out[ + b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group + ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) + return out + + +def gather_paged_kv( + kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool + (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + page = pool.shape[2] + pages = (int(seq_lens.max()) + page - 1) // page + slots = block_table[:, :pages].long() + k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] + v = pool[slots * subpages + 1] + batch, _, kv_heads, _, d = k.shape + + def order(t): + return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) + + return order(k), order(v) + + +class M3DenseDecode(OpWrapper): + """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len + ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. + + Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into + ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: + every CTA still covers its share of each request's actual length, which may exceed it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="two requests keep the model maximum (no cap at two requests)", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), + ), + ) + + note = """ + The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped + one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV + band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + minimax_m3_trtllm_gen_dense_decode( + q, + kv_cache_manager, + layer_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + max_seq_len=max_seq_len, + max_num_requests=max_num_requests, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). + + ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + if keys.dtype == torch.float8_e4m3fn: + q = fp8_e4m3_saturate(q) + return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> None: + # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. + assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" + # Each verify row attends a causal prefix of at least one key. + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + + +m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..d9e8ca93fff7 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_qk_norm_rope_fp8.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head RMSNorm + partial NeoX RoPE of the packed [Q | K | V] into FP8, the bits of the BF16 +``fused_qk_norm_rope`` output cast to E4M3; optionally also storing the live tokens' K / V into the paged FP8 cache.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3DenseQkNormRopeFp8(OpWrapper): + """``qkv`` [T, (q + 2 kv) heads x 128] bf16 -> [T, same] E4M3: Q and K heads normed and roped in fp32, rounded to + BF16 and then to E4M3 (saturating), V cast; ``qkv`` untouched + (``trtllm::minimax_m3_sm107tp4_dense_qk_norm_rope_fp8``). + + With ``kv_buffers`` (the layer's HND FP8 pool [pages, 2, kv heads, 128, 128]), ``out_cache_loc`` (the step's int32 + cache slots) and ``num_live_tokens``, the kernel also stores the live tokens' K and V rows into the pool (the bytes + the MSA backend's KV write stores) and every row's Q into ``q_compact`` [T, q heads, 128] + (``..._dense_qk_norm_rope_fp8_kv``). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp4 heads (16 q, 1 kv)", spec=dict(m=1, q=16, kv=1)), + Cell( + why="a 4-request verify (16 rows) storing K / V", + spec=dict(m=16, q=16, kv=1, kv_write=True), + ), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=16, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=16, kv=1)), + ) + + note = """ + Bit-exact: `reference` runs the catalog's in-place BF16 ``fused_qk_norm_rope`` on a copy and casts every element to + E4M3, which is this kernel's rounding sequence (fp32 -> BF16 -> E4M3) for Q and K and the cast for V. + """ + + def __call__( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + use_gemma: bool, + kv_buffers: Optional[torch.Tensor] = None, + out_cache_loc: Optional[torch.Tensor] = None, + num_live_tokens: int = 0, + q_compact: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + args = ( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + use_gemma, + ) + if kv_buffers is not None: + return torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv( + *args, kv_buffers, out_cache_loc, num_live_tokens, q_compact + ) + return torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8(*args) + + def reference( + self, + qkv: torch.Tensor, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + is_neox: bool, + position_ids: torch.Tensor, + use_gemma: bool, + kv_buffers: Optional[torch.Tensor] = None, + out_cache_loc: Optional[torch.Tensor] = None, + num_live_tokens: int = 0, + q_compact: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """The FP8 copy (the K / V store into ``kv_buffers`` is checked by the test against the reference rows).""" + head_dim = qkv.shape[-1] // (num_heads_q + num_heads_k + num_heads_v) + out = qkv.clone() + fused_qk_norm_rope( + out, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + use_gemma, + False, + 0, + 0, + ) + return out.to(torch.float8_e4m3fn) + + +m3_dense_qk_norm_rope_fp8 = M3DenseQkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..ed4e3048d780 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention on the TRTLLM backend: RoPE, KV-cache append and causal attention.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class M3DraftAttention(OpWrapper): + """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the + backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's + NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` + ``num_tokens``); returns [num_tokens, q heads x 128] bf16. + + ``backend`` (``TrtllmAttention``) and ``metadata`` are the runtime's objects for the drafter layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="first draft pass of a 4-request verify (4 x 4 rows)", spec=dict(reqs=4, qlen=4)), + Cell(why="a later draft pass, one row per request", spec=dict(reqs=8, qlen=1)), + Cell(why="a context chunk's draft rows", spec=dict(reqs=1, context=512)), + ) + + note = """ + `reference` is exact (float64) causal softmax attention: Q rotated by the backend's NeoX RoPE at each token's + position, over the drafter layer's K / V rows read back from its FP8 cache (dequantized by the backend's + ``kv_scale_quant_orig``). What remains between the two is the FMHA's own arithmetic (Q's rounding to the + kernel's input precision, its P rounding and accumulation order), which the test bounds. + """ + + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors one KV head per rank and un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + assert backend.num_kv_heads == 1, "mirrors the one-KV-head-per-rank layout only" + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + keys = pool[pages, 0, 0].to(torch.float64).reshape(-1, dim)[:length] * kv_scale + values = pool[pages, 1, 0].to(torch.float64).reshape(-1, dim)[:length] * kv_scale + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] + rows = q[first : first + n] + x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,kd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) + + +m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..48543be75675 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q +out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3Fp8IndexerQkNormRope(OpWrapper): + """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and + writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, + then RoPE on its leading ``rotary_dim`` channels (NeoX). + + The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; + the fused layers' producer writes the index cache itself. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, one index head per rank", spec=dict(m=1, heads=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, heads=1)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=1)), + ) + + note = """ + `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 + before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of + values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. + """ + + def __call__( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + index_k_cache, + out_cache_loc, + num_heads_q, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + position_ids, + ) + + def reference( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, + then the E4M3 cast.""" + out = fused_qk_norm_rope( + qk.clone(), + num_heads_q, + 1, + 0, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) + return ( + q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), + k.to(torch.float8_e4m3fn), + ) + + +m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py new file mode 100644 index 000000000000..2c6c2133b2bb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' full-context decode attention at TP4 on Rubin (``ops/m3_fullctx_decode``): PrimTS split-KV +decode at the recipe ``modeling.FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) + +from .. import m3_fullctx_decode as kernel +from .m3_dense_decode import dense_attention_reference, gather_paged_kv + +__all__ = ["m3_fullctx_decode"] + +HEADS = kernel.HEADS +KV_HEADS = kernel.KV_HEADS +HEAD_DIM = kernel.HEAD_DIM +#: The plan capacity (the per-request K/V length bound PrimTS sizes its policy for): dispatch's +#: FULLCTX_DECODE_MAX_KV_LEN, which the forward's prebuilt plans use. +MAX_KV_LEN = 1 << 21 + + +class M3FullCtxDecode(OpWrapper): + """Decode attention of every query row over its request's full causal context in the paged FP8 pool, written + into ``output`` (BF16 [T, heads, head_dim]) in place. + + ``recipe`` is ``modeling.FULLCTX_DECODE_RECIPES[(qlen, requests)]``: the PrimTS Q tile and the split-KV fanout. + ``own_reducer`` merges the splits on the BF16 op (``modeling.FULLCTX_DECODE_OWN_REDUCER``) rather than PrimTS's + FP16 reducer and a cast. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over an AgentX-length (165k) context", + spec=dict(kv_lens=(165525,), qlen=4), + ), + Cell( + why="one request, one draft-pass row, over a 165k context", + spec=dict(kv_lens=(165522,), qlen=1), + ), + Cell( + why="two requests verifying 1 + 3 rows at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165525, 120011), qlen=4), + ), + Cell( + why="two requests, one row each, at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165522, 120008), qlen=1), + ), + Cell( + why="one request verifying 1 + 3 rows over a short (4.1k) context: most splits are empty", + spec=dict(kv_lens=(4104,), qlen=4), + ), + ) + + note = """ + Same operands and arithmetic as the stock trtllm-gen call (E4M3 q, E4M3 K/V, P quantized to E4M3 per split, + fp32 accumulation): the split differs, so the result differs from stock as a stock split change does (the P + quantization is per split). Deterministic, and the same bits eager and in a CUDA graph. Held to the fp32 + reference within the FP8-KV band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + shape = (decode_query_len, int(seq_lens.shape[0])) + kernel.prebuild(q.device, {shape: recipe}, max_kv_len, own_reducer, get_env_enable_pdl()) + plan = kernel.plan_for(q.device, *shape, own_reducer) + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + k_pages, v_pages = kernel.kv_page_views(pool, subpages) + plan.run(fp8_e4m3_saturate(q), k_pages, v_pages, block_table, seq_lens, sm_scale, output) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V, q rounded to E4M3 first as the kernel runs it.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + return dense_attention_reference( + fp8_e4m3_saturate(q), keys, values, seq_lens, decode_query_len, sm_scale + ) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> None: + pool, _ = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + assert pool.dtype == torch.float8_e4m3fn, "an FP8 E4M3 pool" + assert tuple(pool.shape[1:]) == (KV_HEADS, kernel.PAGE, HEAD_DIM), ( + "the TP4 dense-layer pool geometry" + ) + assert q.shape[1:] == (HEADS, HEAD_DIM), "16 query heads of 128 per rank" + assert int(q.shape[0]) == int(seq_lens.shape[0]) * decode_query_len, ( + "q covers exactly the batch" + ) + if torch.cuda.is_current_stream_capturing(): + return # the length checks read the device: eager calls only + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + assert bool((seq_lens <= max_kv_len).all()), "a request longer than the plan capacity" + + +m3_fullctx_decode = M3FullCtxDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..4fb706182485 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score``.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_index_score"] + +PAGE = 128 +HEAD_DIM = 128 +#: CTAs per request (each scores blocks split, split + SPLIT_K, ...; the scores do not depend on it). +SPLIT_K = 512 + + +def compiled_scorer(key: tuple): + """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss + (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" + from ..m3_index_score import IndexScoreRunner + + if key not in IndexScoreRunner.kernel_cache: + IndexScoreRunner._compile(*key) + return IndexScoreRunner.kernel_cache[key] + + +class M3IndexScore(OpWrapper): + """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the + positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of + ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is + written). FP8 or BF16 q and index-K, fp32 products and sums. + + ``early`` (the target's one-request kernel) issues the Q and first K tiles right after the grid wait from a + pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise every older block of + the CTA is prefetched. Neither changes a score. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over 20000 tokens: the early-issue kernel", + spec=dict(seq_lens=(20000,), qlen=4, early=True), + ), + Cell( + why="one request, one row, 65536 tokens: early issue, 512 blocks on 512 CTAs", + spec=dict(seq_lens=(65536,), qlen=1, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most CTAs lie past the length and are drained", + spec=dict(seq_lens=(129,), qlen=4, early=True), + ), + Cell( + why="one request longer than 512 blocks (the CTAs loop): early issue", + spec=dict(seq_lens=(70000,), qlen=4, early=True), + ), + Cell( + why="two requests: the prefetching kernel", + spec=dict(seq_lens=(20000, 3000), qlen=4, early=False), + ), + Cell( + why="eight requests of mixed lengths: the prefetching kernel", + spec=dict( + seq_lens=(300, 1500, 2049, 4096, 5000, 7777, 8192, 33001), qlen=4, early=False + ), + ), + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py new file mode 100644 index 000000000000..39a27ee64b71 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score_t2sc.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score_t2sc`` (T2SC: the sm_107 +TP2 / TP8 scorer as this target's copy, at its 1 index head per rank; ``modeling.INDEX_SCORE_T2SC``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_index_score_t2sc"] + +PAGE = 128 +HEAD_DIM = 128 + + +def compiled_scorer(key: tuple): + """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss + (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer_t2sc``).""" + from ..m3_index_score_t2sc import IndexScoreRunner + + if key not in IndexScoreRunner.kernel_cache: + IndexScoreRunner._compile(*key) + return IndexScoreRunner.kernel_cache[key] + + +class M3IndexScoreT2sc(OpWrapper): + """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the + positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of + ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is + written). FP8 or BF16 q and index-K, fp32 products and sums. + + ``early`` (the target's kernel at every request count it covers) issues the Q and first K tiles right after the + grid wait from a pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise + every older block of the CTA is prefetched. Neither, nor ``split_k``, changes a score. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request, 1 + 3 rows over 165521 tokens (AgentX ISL), 1 index head: 1536 CTAs, early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most of the 1536 CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request of 400000 tokens: 3125 blocks, so the 1536 CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request decoding one row over 65536 tokens", + spec=dict(seq_lens=(65536,), qlen=1, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request, 1 + 3 rows over 165521 tokens, 1 index head (TP4): 1272 CTAs (one wave), early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="one request of 129 tokens at 1272 CTAs: the early tiles of most CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="one request of 400000 tokens at 1272 CTAs: 3125 blocks, so the CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="two requests verifying 4 rows each over 165521 and 20000 tokens: 768 CTAs each, early issue", + spec=dict(seq_lens=(165521, 20000), qlen=4, heads=1, split_k=768, early=True), + ), + Cell( + why="two requests, one of 300 tokens and one of 262144 (2048 blocks: the 768 CTAs loop)", + spec=dict(seq_lens=(300, 262144), qlen=4, heads=1, split_k=768, early=True), + ), + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score_t2sc = M3IndexScoreT2sc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..0ce3c864d30f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MiniMax-M3 MSA attention backend's step: KV write, sparse top-k block selection, FMHA.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + +from .m3_sparse_decode_lc import dequant_nvfp4_pages + + +class M3MsaAttention(OpWrapper): + """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` + [T, heads * 128] (bf16, the step's live rows). + + * ``k`` / ``v`` given (the dense layers): written into the layer's paged main cache first (NVFP4 with + ``kv_scale_orig_quant``, else FP8 / BF16 as the pool is). None: a producer already wrote the caches. + * ``sparse``: the indexer scores the layer's index-K cache (already written) with ``idx_q`` and publishes the + top-k blocks; the FMHA attends those blocks. Otherwise the FMHA attends the whole page table. + + The FMHA dispatches on the step's metadata (context rows: the MSA prefill FMHA; generation rows: the MSA decode + FMHA, the dense layers' trtllm-gen decode at their capped KV split). ``backend`` and ``metadata`` are the runtime's + objects for this layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dense layer, a prefill chunk, FP8 pool", + spec=dict(layer="dense", phase="context", tokens=1024), + ), + Cell( + why="dense layer, a decode step of 4 requests", + spec=dict(layer="dense", phase="generation", reqs=4), + ), + Cell( + why="sparse layer, a prefill chunk beyond 16 blocks, NVFP4 pool", + spec=dict(layer="sparse", phase="context", tokens=4096), + ), + Cell( + why="sparse layer, a mixed step (context + generation rows)", + spec=dict(layer="sparse", phase="mixed", tokens=512, reqs=2), + ), + ) + + note = """ + `reference` is exact (float64) softmax attention over the cache rows the step attends, read back from the paged + cache the call wrote (or a producer wrote before it) and dequantized as the FMHA dequantizes them: FP8 rows as + E4M3, NVFP4 rows as E2M1 x E4M3 block scale x the per-tensor ``kv_scale_quant_orig``. A sparse layer's top-k blocks + are re-selected from ``idx_q`` and the index-K cache with the backend's rule (per-block max over the causal + keys, the query's own block forced in, ties to the lower block). On an FP8 pool the FMHA takes Q as a plain E4M3 + cast, which the reference applies too. What remains between the two is the FMHA's own arithmetic (its P + rounding and accumulation order), which the test bounds. + """ + + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors one KV head per rank, + the target's layout. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + assert backend.num_kv_heads == 1, "mirrors the one-KV-head-per-rank layout only" + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + if nvfp4: + keys = dequant_nvfp4_pages( + pool[:, 0].view(torch.uint8), scale_pool[:, 0], pages, quad_scales=False + ) + values = dequant_nvfp4_pages( + pool[:, 1].view(torch.uint8), scale_pool[:, 1], pages, quad_scales=True + ) + keys, values = keys * k_global, values * v_global + else: + keys = pool[pages, 0, 0].to(torch.float64) + values = pool[pages, 1, 0].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + qpos = cached[row] + torch.arange(n, device=q.device) + allowed = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + if sparse: + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = _selected_blocks( + idx64[first : first + n], idx_keys, allowed, qpos, page, config + ) + allowed &= chosen.repeat_interleave(page, dim=1)[:, :length] + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = torch.einsum("thd,kd->thk", q64[first + lo : first + hi], keys) * sm_scale + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + (ids >= n_valid - config.local_blocks) & (ids < n_valid) + ) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen + + +m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py new file mode 100644 index 000000000000..5a3d1b64ac58 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer, with its constant loads staged +before the PDL grid-dependency wait and an early dependents trigger (``kernels/minimaxM3/producerPw.cu``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +__all__ = ["m3_producer"] + +HEAD_DIM = 128 +ROTARY_DIM = 64 +PAGE = 128 +#: Query heads per rank at TP4 (MiniMax-M3: 64 query heads, 1 KV head and 1 index head per rank). +HEADS_Q = 16 +#: NVFP4 block: 16 values share one E4M3 scale. +FP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _norm_rope( + x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(FP4_BLOCK) + + +def scale_index(within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return (within_page // 4) * 32 + cols * 4 + within_page % 4 + + +class M3Producer(OpWrapper): + """The NVFP4 QKV + indexer producer of a sparse layer: ``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert`` for one + KV head, one index head, head dim 128 and rotary dim 64, with the same arguments and outputs bit for bit. + + ``packed`` [T, (heads_q + 4) * 128] BF16 holds each token's q heads, k, v, index-q and index-k. Every head but v + gets the Gemma RMSNorm with its weight and NeoX RoPE on its first 64 dims (``rotary_cos_sin[position]``). Returns + (q FP8 [T, heads_q, 128], index-q FP8 [T, 1, 128]); writes the token's slot ``out_cache_loc[t]`` (skipped when + negative or past the pool): k and v as NVFP4 (16-value blocks, E4M3 scales against ``kv_quant_scale[1]`` / + ``[2]``, the V scales in the 4-token swizzle) and index-k as FP8. ``early`` triggers the dependent kernel on + entry (its outputs are read only after the dependent's own grid wait). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token (c1 without speculation)", spec=dict(tokens=1, scale=1.0)), + Cell(why="a c1 verify: 4 tokens", spec=dict(tokens=4, scale=3.0)), + Cell(why="a c4 verify: 16 tokens, small activations", spec=dict(tokens=16, scale=1e-3)), + Cell( + why="a c8 verify: 32 tokens, large activations, a skipped and an out-of-pool slot", + spec=dict(tokens=32, scale=40.0, bad_slots=True), + ), + Cell(why="the widest decode step: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="a prefill chunk of 1000 tokens, early trigger off", + spec=dict(tokens=1000, scale=1.0, early=False), + ), + ) + + note = """ + The kernel is the native producer's arithmetic, compiled with the library's --use_fast_math, so it is held to that + kernel bit for bit on every output and cache byte. The torch reference states the math in fp32: the FP8 outputs are + held to it within one E4M3 step, and the NVFP4 K/V, decoded, within one E2M1 step of their block scale. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if prewait_step_inputs: + # This target's copy (producerPw4): the same outputs bit for bit, with the position ids, slots and RoPE + # rows read before its grid wait (none of them may be written by a PDL predecessor). + return tuple( + torch.ops.trtllm.minimax_m3_sm107_tp4_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + True, + ) + ) + q, index_q = torch.ops.trtllm.minimax_m3_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + ) + return q, index_q + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", + "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change + them.""" + t = packed.shape[0] + x = packed.float().view(t, heads_q + 4, HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = rope[:, 0], rope[:, 1] + q = torch.stack( + [_norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + ) + return { + "q": q, + "k": _norm_rope(x[:, heads_q], k_weight, cos, sin, eps), + "v": x[:, heads_q + 1], + "index_q": _norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), + "index_k": _norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), + } + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> None: + t = packed.shape[0] + # The RoPE row is read at the position with no bound. + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + # Two tokens writing one slot race. + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" + + +m3_producer = M3Producer() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..f92a3e6efaa1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3QkNormRopeFp8(OpWrapper): + """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded + once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place + entry. + + The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps + of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp4 heads (16 q, 1 kv)", spec=dict(m=1, q=16, kv=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=16, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=16, kv=1)), + ) + + note = """ + `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from + fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. + """ + + def __call__( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + + def reference( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" + out = fused_qk_norm_rope( + qkv.clone(), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return out.to(torch.float8_e4m3fn) + + +m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..ef79579c0c21 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into +the paged caches.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3QkvInsert(OpWrapper): + """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the + partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and + index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. + + ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool + [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 + [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp4 + heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), + Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), + ) + + note = """ + The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` + returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + common = ( + num_heads, + num_kv_heads, + num_index_heads, + head_dim, + rotary_dim, + eps, + q_norm, + k_norm, + index_q_norm, + index_k_norm, + rotary_cos_sin, + position_ids, + ) + if kv_scale is not None: + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common + ) + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, index_k_cache, out_cache_loc, *common + ) + + def reference( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of + Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as + packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" + m = packed.shape[0] + sizes = [ + num_heads * head_dim, + num_kv_heads * head_dim, + num_kv_heads * head_dim, + num_index_heads * head_dim, + head_dim, + ] + q, k, v, iq, ik = packed.float().split(sizes, dim=-1) + cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) + half = rotary_dim // 2 + + def norm_rope(x, w, heads): + x = x.reshape(m, heads, head_dim) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) + x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] + c, s = cos.unsqueeze(1), sin.unsqueeze(1) + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) + + return { + "q": norm_rope(q, q_norm, num_heads), + "index_q": norm_rope(iq, index_q_norm, num_index_heads), + "k": norm_rope(k, k_norm, num_kv_heads), + "v": v.reshape(m, num_kv_heads, head_dim), + "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], + } + + +m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..76841cca732d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SelectBlocks(OpWrapper): + """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), + with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the + row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` + backs the result head-major so ``permute(1, 0, 2)`` is contiguous. + + The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="eight decode rows of 17..128 blocks", + spec=dict(total_q=8, n_blocks=128, mixed=True), + ), + Cell( + why="one decode row at the 128-block edge", + spec=dict(total_q=1, n_blocks=128, mixed=False), + ), + Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), + ) + + note = """ + Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. + """ + + def __call__( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + return torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output + ) + + def reference( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + """Plain torch selection; ``head_major_output`` changes the backing only.""" + heads, n_blocks, total_q = max_score_kv.shape + out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) + blocks = torch.arange(n_blocks, device=max_score_kv.device) + for t in range(total_q): + nv = int(n_valid_blocks[t]) + for h in range(heads): + s = max_score_kv[h, :, t].clone().double() + s[blocks >= nv] = float("-inf") + forced = (blocks < min(init_blocks, nv)) | ( + (blocks >= nv - local_blocks) & (blocks < nv) + ) + s[forced] = float("inf") + order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] + ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) + out[t, h, : ids.numel()] = ids + return out + + +m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..3d076c8181d6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). +MAX_BLOCKS = 9216 +#: Blocks the selection keeps per (query, KV head). +TOPK = 16 +#: Scores the forcing substitutes: initial blocks rank above local ones, and both above any finite score. +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 + +# M3's checkpoint: no forced initial blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=1, init_blocks=0, local_blocks=1) + + +class M3SelectLong(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query's top-16 blocks. + + ``scores`` is the per-KV-head block maxima [num_kv_heads, num_blocks, + total_q] the CuTe DSL index scorer wrote (amax over the KV head's index + heads). Per (query, KV head) row the op keeps the ``n_valid_blocks[query]`` + leading blocks (clamped to ``[0, num_blocks]``), forces the last + ``local_blocks`` of them (counted from the *unclamped* valid count) to + ``LOCAL_SCORE`` and the first ``init_blocks`` to ``INIT_SCORE`` -- local wins + where both apply -- and selects the 16 largest, ties to the smaller block. + A pick whose forced score is -inf is emitted as -1. The ids are written + ascending with a -1 tail. + + Same contract as ``torch.ops.trtllm.minimax_m3_select_blocks`` (the in-tree + ``select_blocks_from_maxscore``) and bit-identical to it, except on a row + where more than 1024 valid blocks tie at the 16th-largest score: this op + breaks that tie to the smallest blocks as the contract says, the in-tree op + does not. The test checks both against ``reference``. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="32K-context engine (257 worst-case k-tiles), Eagle3 verify at one request: 4 query rows", + spec=dict(num_blocks=257, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="decode without speculation at batch 8 on a 2048-block engine: the 2-key-per-thread path at its edge", + spec=dict(num_blocks=2048, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="1M-context engine at batch 8 x qlen 4: rows on all three key widths, short rows and padding rows", + spec=dict(num_blocks=8193, total_q=32, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 1024 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + note = """ + Certified at one KV head per rank (TP4) with M3's forcing (`init_blocks=0`, + `local_blocks=1`). The kernel takes any head count and forcing; neither is + driven here. + + `scores` may be strided (the op reads it by stride); the target hands over + the scorer's contiguous view, which is what the cells drive. Rows of at + most 128 blocks are within the kernel's range, but the target keeps those on + the in-tree op, so they are not certified here. + + Ordering is by the descending-score key `orderedKey`: a positive NaN ranks + above +inf, +0 above -0, and a negative NaN below -inf -- which, unlike + -inf, is emitted as a block id rather than -1. The cells drive positive + NaN only. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sm107_tp4_select_long( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + fast_tail, + ) + return out + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and + ``trigger_phase`` and ``fast_tail`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() # [total_q, num_blocks] + forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where( + in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max + ) + k = min(TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) + ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> None: + """The inputs this op takes and answers wrongly. + + Negative forcing counts are accepted and push the forced ranges past the + row instead of failing; the in-tree op this one stands in for refuses + them, so a caller relying on this one would silently diverge from it. + An ``out`` overlapping ``scores`` is overwritten while it is still read. + """ + assert init_blocks >= 0 and local_blocks >= 0, ( + f"init_blocks ({init_blocks}) and local_blocks ({local_blocks}) must be non-negative; the " + "op accepts negative counts and selects as if the forced ranges were elsewhere" + ) + storage = scores.untyped_storage() + assert not (storage.data_ptr() <= out.data_ptr() < storage.data_ptr() + storage.nbytes()), ( + "out must not alias scores: the op writes it while scores are still read" + ) + + +m3_select_long = M3SelectLong() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py new file mode 100644 index 000000000000..dbfb06ea7c2d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long_dsl.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_select_long import _M3, MAX_BLOCKS, M3SelectLong + + +class M3SelectLongDsl(OpWrapper): + """``m3_select_long``'s contract, cells and reference, run by ``ops/m3_select_long_dsl``: the same selection + bit for bit as the 1024-thread ``minimax_m3_sm107_tp4_select_long`` (also on rows where more than 2048 blocks tie + at the 16th-largest score), from one 384-thread CTA per row. ``fast_tail`` (the 1024-thread kernel's warp-0 rank) + does not apply and is ignored. Returns ``out``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="32K-context engine (257 worst-case k-tiles), Eagle3 verify at one request: 4 query rows", + spec=dict(num_blocks=257, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="decode without speculation at batch 8 on a 2048-block engine: the 2-key-per-thread path at its edge", + spec=dict(num_blocks=2048, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="1M-context engine at batch 8 x qlen 4: rows on all three key widths, short rows and padding rows", + spec=dict(num_blocks=8193, total_q=32, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 1024 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> torch.Tensor: + """``m3_select_long``'s reference: the selection is the same bit for bit.""" + return M3SelectLong.reference( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + fast_tail, + ) + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> None: + """``m3_select_long``'s domain.""" + M3SelectLong.is_valid( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + fast_tail, + ) + + note = """ + Certified at one KV head per rank (TP4, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`), and bit for bit against the + 1024-thread `m3_select_long` on every cell. The kernel takes any head count + and forcing; other forcing is not driven here. + + The kernel is compiled per (PDL launch, early trigger, trigger phase) with + symbolic shapes and strides; the target compiles its keys before warmup + (`ops/m3_select_long_dsl.prebuild`), and a call outside them compiles + on first use. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + ) -> torch.Tensor: + from ..m3_select_long_dsl import select_long_dsl + + if scores.shape[1] > MAX_BLOCKS: + raise RuntimeError( + f"m3_select_long_dsl supports at most {MAX_BLOCKS} blocks, got {scores.shape[1]}" + ) + select_long_dsl( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + +m3_select_long_dsl = M3SelectLongDsl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..96171a6a9a29 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This target's block-sparse decode attention (Triton, with the one-launch long-context route).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..triton_sparse_decode import minimax_m3_sparse_attn_decode +from .m3_sparse_decode_lc import PAGE, sparse_attention_reference + + +class M3SparseDecode(OpWrapper): + """``minimax_m3_sparse_attn_decode`` as ``M3Attention._decode_sparse_core`` calls it. + + Attention of each query token over the blocks ``topk_idx`` selected for it + (``[num_kv_heads, T, topk]``, ascending ids with a -1 tail), each block + truncated at the token's causal extent, written into ``output`` in place. + The target passes FP8 q, the packed NVFP4 sparse-layer cache with its block + and per-tensor scales, and BF16 ``output`` [T, 16, 128]. Returns ``output``. + + Which kernel runs is decided inside, per step: a selected (non-identity) + table on the TP4 geometry goes to the one-launch ``m3_sparse_decode_lc`` + kernel; an identity table (every request within 16 blocks, staged by the + metadata) runs the Triton producer, or from 17 query rows the multi-query + producer, then the LSE merge. Every route also publishes the MXFP8 o_proj + hand-off to ``m3_prequant.SLOT``; this entry certifies the attention output. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="identity selection, decode without speculation at batch 8: the Triton producer + merge", + spec=dict( + seq_lens=(1, 128, 129, 700, 2048, 1500, 64, 0), decode_query_len=1, identity=True + ), + ), + Cell( + why="identity selection, Eagle3 verify at one request (T=4): one block per producer program", + spec=dict(seq_lens=(1900,), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 4 (T=16): the last row count below the MQ producer", + spec=dict(seq_lens=(300, 2048, 4, 1000), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 8 (T=32): the multi-query producer, padding rows", + spec=dict( + seq_lens=(2048, 130, 5, 0, 1777, 900, 2000, 0), decode_query_len=4, identity=True + ), + ), + Cell( + why="selected blocks, Eagle3 verify at batch 2: routed to the one-launch long-context kernel", + spec=dict(seq_lens=(9000, 40000), decode_query_len=4, identity=False), + ), + Cell( + why="selected blocks, decode without speculation at batch 8: the long-context route at T=8", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + identity=False, + ), + ), + ) + + note = """ + Certified only as the target calls it: one KV head of 16 query heads, + head_dim 128, FP8 q, packed NVFP4 cache, BF16 output, `num_topk_chunks` + left to the dispatcher. The function also takes BF16 / FP8 caches, other + head geometries and an explicit split-K factor (the in-tree tests drive + those); none is certified here. + + The MXFP8 hand-off each route publishes is certified on the long-context + kernel's own entry (`m3_sparse_decode_lc`); the Triton merge writes it with + the same recipe (`_store_mxfp8_head`). + """ + + def __call__( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + kv_scale=kv_scale, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + num_topk_chunks=num_topk_chunks, + identity_blocks=identity_blocks, + ) + return output + + def reference( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + """``output`` as the call is supposed to write it: float64 block-sparse attention over the + dequantized NVFP4 blocks (``sparse_attention_reference``), rounded once to ``output``'s dtype. + + Only the packed NVFP4 cache with one KV head is mirrored -- the target's + only configuration. ``kv_scale``, ``num_topk_chunks`` and + ``identity_blocks`` do not change what is computed. + """ + assert k_block_scale is not None and k_paged.shape[1] == 1, ( + "mirrors the one-KV-head NVFP4 cache only" + ) + attn = sparse_attention_reference( + q, + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.to(output.dtype) + + def is_valid( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) + + +m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py new file mode 100644 index 000000000000..76fc68fa5903 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py @@ -0,0 +1,436 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch long-context sparse decode over the selected NVFP4 blocks (16-CTA clusters). + +Also the plain-torch block-sparse attention and MXFP8 hand-off mirrors the +Triton sparse decode entry (``m3_sparse_decode``) reuses: both compute the same +thing over the same cache. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The TP4 geometry the kernel is compiled for. +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +#: NVFP4 elements per E4M3 block scale. +SF_VEC = 16 +#: MXFP8 elements per UE8M0 scale, and the o_proj width the hand-off is laid out for. +MX_BLOCK = 32 +WIDTH = NUM_HEADS * HEAD_DIM + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, block_scale: torch.Tensor, pages: torch.Tensor, quad_scales: bool +) -> torch.Tensor: + """Dequantize KV head 0 of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, 0].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, 0].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """Block-sparse decode attention over one KV head of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages(k, k_block_scale, pages, quad_scales=False).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True).reshape(-1, HEAD_DIM) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, + sf [num_rows_alloc * WIDTH // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) + num_sf_k_tiles = WIDTH // 128 + sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +class M3SparseDecodeLc(OpWrapper): + """Sparse decode attention, split-K merge and MXFP8 hand-off of one TP4 decode step in one launch. + + One cluster of 16 CTAs per query token: CTA ``r`` attends the token's + ``r``-th selected block (``topk[tok, r]``), and the cluster merges the 16 + partials over distributed shared memory into the BF16 row ``out[tok]``. With + ``q8`` / ``sf`` it also writes the row's MXFP8 copy and swizzled UE8M0 + scales for the o_proj GEMM. All operands are the byte views + ``m3_attn_lc.m3_lc_attn`` hands over (E4M3 ``q``, packed E2M1 pages, E4M3 + block scales). Returns ``out``. + + Per block: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16 + (exact), QK and PV on FP16 tensor cores with FP32 accumulation, the softmax + probabilities rounded to FP16 for PV, log2-domain softmax with ``ex2.approx``. + ``linear_softmax`` picks the denominator recurrence (the two agree to + rounding); the target picks it per batch shape. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation, one request: T=1, log-domain recurrence, MXFP8 hand-off", + spec=dict(seq_lens=(9000,), decode_query_len=1, linear_softmax=False, mxfp8=True), + ), + Cell( + why="Eagle3 verify at one request: T=4, the log-domain recurrence the target keeps for local batch 1", + spec=dict(seq_lens=(20000,), decode_query_len=4, linear_softmax=False, mxfp8=True), + ), + Cell( + why="Eagle3 verify at two requests: T=8 on the linear denominator", + spec=dict(seq_lens=(5000, 12345), decode_query_len=4, linear_softmax=True, mxfp8=True), + ), + Cell( + why="Eagle3 verify at batch 8 (T=32): short requests with -1 tails and zero-length padding rows", + spec=dict( + seq_lens=(16384, 2049, 130, 0, 7777, 4, 30000, 0), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + ), + ), + Cell( + why="decode without speculation at batch 8: T=8 rows of distinct requests", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + linear_softmax=False, + mxfp8=True, + ), + ), + Cell( + why="no hand-off buffers (the variant the call site takes past 128 rows), driven at T=8", + spec=dict(seq_lens=(5000, 12345), decode_query_len=4, linear_softmax=True, mxfp8=False), + ), + ) + + note = """ + Certified: the four instantiated variants (`linear_softmax` x hand-off on / + off), all with `early_trigger=True` -- the op refuses `False`. The + hand-off cells are certified bit-exactly against `mxfp8_handoff_reference` + applied to the op's own BF16 rows; the BF16 rows themselves within the band + the test derives. + + `topk` is read for the first `min(16, ceil(kv_len / 128))` entries of each + token; a -1 among them skips that CTA's block, entries past them are never + read. `use_pdl` only changes the launch attribute. The hand-off assumes + `num_sf_k_tiles = 16` (a 2048-wide o_proj input) and buffers with at least + as many rows as the step; `MERGE_MXFP8` sizes them at 128 rows. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + bt_prewait: bool = False, + ) -> torch.Tensor: + # merge_once: this target's copy of the op (minimaxM3Sm107Tp4/sparseDecodeLc4.cu), whose split-K merge forms the + # 16 weights once per CTA (same values, same order: bitwise the image op's output). + lc = ( + torch.ops.trtllm.minimax_m3_sm107_tp4_sparse_decode_lc + if merge_once + else torch.ops.trtllm.minimax_m3_sparse_decode_lc + ) + # prefetch: up to two read-only tensors the TP4 op pulls into L2 (a hint; the image op takes none). + # next_kv: (this layer's id buffer, the next sparse layer's, its K / V pools and block scales), TP4 op only. + extra = () + if merge_once and (prefetch or next_kv): + extra = tuple(prefetch or ()) + (None,) * (2 - len(prefetch or ())) + if next_kv: + extra += tuple(next_kv) + lc( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + decode_query_len, + sm_scale, + linear_softmax, + use_pdl, + early_trigger, + *extra, + **({"identity": True} if identity else {}), + **({"identity_step": True} if identity_step else {}), + **({"one_div": True} if one_div else {}), + # bt_prewait: the TP4 op stages the block-table row in shared memory before its grid wait (same output). + **({"bt_prewait": True} if bt_prewait and merge_once else {}), + ) + return out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + bt_prewait: bool = False, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 2048] the op is supposed to write: ``sparse_attention_reference`` in float64, + rounded once. ``linear_softmax``, ``use_pdl``, ``early_trigger``, ``one_div`` and ``bt_prewait`` do not change + what is computed; the hand-off buffers are checked separately with ``mxfp8_handoff_reference``. + """ + attn = sparse_attention_reference( + q.view(-1, NUM_HEADS, HEAD_DIM), + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.reshape(out.shape).to(out.dtype) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + early_trigger: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + bt_prewait: bool = False, + ) -> None: + """The inputs this op takes and answers wrongly (``m3_attn_lc.lc_applies`` is the target's form of the + layout half). + + The kernel addresses a page's token rows at a fixed 64-byte (data) and + 8-byte (scale) pitch and loads them 16 / 8 bytes at a time, so only the + page stride is read from the views. It reads the selected block ids, + the block table and the hand-off buffers unchecked. Reads ``topk`` and + ``seq_lens`` back to the host: a test-only cost. + """ + for name, t in (("k", k), ("v", v)): + assert t.dim() == 4 and t.shape[1] == 1 and t.shape[2:] == (PAGE, HEAD_DIM // 2), ( + f"{name} must be the [pages, 1, 128, 64] byte view of one KV head" + ) + assert t.stride(3) == 1 and t.stride(2) == HEAD_DIM // 2, ( + f"{name} rows must be dense 64-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0, ( + f"{name} pages must be 16-byte aligned" + ) + for name, t in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + assert t.dim() == 4 and t.stride(3) == 1 and t.stride(2) == HEAD_DIM // SF_VEC, ( + f"{name} rows must be dense 8-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 8 == 0 and t.stride(0) % 8 == 0, ( + f"{name} pages must be 8-byte aligned" + ) + assert q.data_ptr() % 16 == 0 and q.stride(0) % 16 == 0 and q.stride(1) % 16 == 0, ( + "q rows are loaded 16 bytes at a time and must be 16-byte aligned" + ) + assert k_global_scale.numel() == 1 and v_global_scale.numel() == 1, ( + "global scales are single values" + ) + total_q = q.shape[0] + assert ( + total_q % decode_query_len == 0 and seq_lens.numel() >= total_q // decode_query_len + ), ( + "every query token must belong to a request: seq_lens is read at tok // decode_query_len unchecked" + ) + seq = seq_lens.tolist() + num_cols = block_table.shape[1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk[tok, : min(TOPK, nblocks)] + assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( + f"token {tok}: selected block ids must be -1 or below the request's block count; the block " + "table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + f"num_sf_k_tiles must be {WIDTH // 128}: the hand-off rows are {WIDTH} wide and a smaller " + "value overlaps them silently" + ) + assert ( + q8.numel() >= total_q * WIDTH + and sf is not None + and sf.numel() >= ((total_q + 127) // 128 * 128 * (WIDTH // MX_BLOCK)) + ), "the hand-off buffers must hold every row of the step: they are written unchecked" + + +m3_sparse_decode_lc = M3SparseDecodeLc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..43166b472eca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The MSA pool's page: one sparse block is one 128-token KV page. +TOKENS_PER_BLOCK = 128 + + +class MsaKvUpdate(OpWrapper): + """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. + + ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its + own buffers at the head of every captured speculative decode step. One + single-CTA launch writes, for ``batch`` requests and ``total_q`` query + tokens (all int32): + + * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, + ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); + * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` + attends up to, through ``block_table``; + * ``n_valid_blocks[t]``, when given, the block count top-k selection is + bounded by. + + It returns None; the cells read the four buffers. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation at one request: qo = 1, every buffer written", + spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="decode without speculation at the captured batch 8", + spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", + spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), + ), + Cell( + why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", + spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), + ), + Cell( + why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", + spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), + ), + Cell( + why="positions crossing a page boundary: the slot moves to the previous page of the request", + spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), + ), + ) + + note = """ + The kernel covers any batch up to the block size's loop (it strides over + 256 threads), and any page size; only `tokens_per_block = 128`, the MSA + pool's page, is certified. + + The op does not bounds-check the request index a token reads + (`q_batch_row`) or the block-table column it derives: both come from the + metadata's own staging, which keeps them in range by construction, and an + out-of-range value reads whatever memory sits there. `is_valid` states both. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.minimax_m3_msa_kv_update( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + + def reference( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], + n_valid_blocks[:total_q])`` as the op is supposed to write them. + + The in-tree metadata's own torch re-derivation + (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook + call), statement for statement, returning new tensors rather than + writing the buffers. The buffer arguments only say which outputs exist: + a ``None`` buffer gives a ``None`` result. + """ + page = int(tokens_per_block) + kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) + cu_kv = None + if cu_kv_lens is not None: + cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) + cu_kv[1:] = torch.cumsum(kv_true, 0) + qbr = q_batch_row[:total_q].to(torch.long) + qo_dev = seq_lens[:batch] + pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(block_table.shape[1]) - 1 + ) + slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + n_valid = None + if n_valid_blocks is not None: + n_valid = torch.div( + (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" + ).to(torch.int32) + return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid + + def is_valid( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + """The inputs this op takes and answers wrongly. + + Every operand but the block table is read as a dense run from its first + element, so a strided view reads the wrong elements. The request index of + each token and the table row it selects are not bounds-checked. Reads + ``q_batch_row`` back to the host: a test-only cost. + """ + for name, t in ( + ("kv_lens", kv_lens), + ("kv_lens_staged", kv_lens_staged), + ("seq_lens", seq_lens), + ("q_batch_row", q_batch_row), + ("q_intra", q_intra), + ("msa_seq_lens", msa_seq_lens), + ("out_cache_loc", out_cache_loc), + ("cu_kv_lens", cu_kv_lens), + ("n_valid_blocks", n_valid_blocks), + ): + assert t is None or t.is_contiguous(), ( + f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" + ) + rows = q_batch_row[:total_q] + assert total_q == 0 or ( + int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) + ), ( + "q_batch_row must index a request below batch (and a block-table row): the op reads the " + "request's lengths and table row unchecked" + ) + + +msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/__init__.py new file mode 100644 index 000000000000..5da86fe24bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective-communication entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..d67b040e8bfd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams + + +class M3AllReduce(OpWrapper): + """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the + plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: + ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; + returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the + communicator's per call. + + The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps + beyond the multicast kernel's rows. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="tp4, a prefill chunk", spec=dict(world=4, m=2048, hidden=6144)), + Cell(why="tp4, a 64-row generation step", spec=dict(world=4, m=64, hidden=6144)), + Cell( + why="tp4, plain sum of eight drafter rows", + spec=dict(world=4, m=8, hidden=6144, plain=True), + ), + ) + + note = """ + The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then + adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum + to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. + """ + + def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): + return comm(x, all_reduce_params=params) + + def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): + """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" + total = torch.stack([x.float() for x in xs]).sum(0) + if params is None: + return total.to(xs[0].dtype) + updated = (params.residual.float() + total).to(xs[0].dtype) + h = updated.float() + normed = ( + h + * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) + * params.norm_weight.float() + ) + return normed.to(xs[0].dtype), updated + + +m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..0b999befbadc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,329 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode layer boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer + fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; + only the stored residual is rounded to BF16. + + The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are + summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic + and the normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + pre = total.to(torch.bfloat16).float() + residual.float() + updated = pre.to(torch.bfloat16) + u = pre.double() if drafter else updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +class M3McAllreduceNorm(OpWrapper): + """One TP group's boundary: ``updated = bf16(residual + sum_r bf16(x_r + x2_r))``, ``normed = RMSNorm(updated)``. + + Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls + its own view ``ws_uc`` (re-arming the words it read) and rotates the per-CTA call parities in ``flags``; the + workspace is the ``McAllReduceWorkspace`` bucket of ``maxt`` tokens. All ranks must make the same calls in the + same stream order. ``prepushed`` 2 means a producer already pushed ``x`` (and launched this grid after its + pushes); 3 that two producers pushed ``x`` into exchange set A and ``x2`` into set B. Returns + ``[normed, updated]`` or, with ``quant``, also the MXFP8 copy of ``normed`` (E4M3 rows, swizzled UE8M0). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The tp4 decode boundaries (hidden 6144, 4 ranks, eps 1e-6). Buckets 8 / 16 / 48 are separate workspaces; the + # 8-token bucket runs epilogue option 7 at 1, 2, 7, 8 tokens and 15 at 3..6, the others 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token at the attention boundary, x only, bucket 8 option 7", + spec=dict(m=1, maxt=8, opts=7, x2=False, quant=False, prepushed=0), + ), + Cell( + why="MoE boundary at 4 tokens: shared + routed summed in-kernel, MXFP8 copy for the next QKV, option 15", + spec=dict(m=4, maxt=8, opts=15, x2=True, quant=True, prepushed=0), + ), + Cell( + why="a full 8-token bucket with the second addend", + spec=dict(m=8, maxt=8, opts=7, x2=True, quant=False, prepushed=0), + ), + Cell( + why="attention boundary whose rows the o_proj GEMV already pushed (prepushed 2)", + spec=dict(m=8, maxt=8, opts=7, x2=False, quant=False, prepushed=2), + ), + Cell( + why="routed engine pushed set A, shared down projection set B (prepushed 3), with the MXFP8 copy", + spec=dict(m=6, maxt=8, opts=15, x2=True, quant=True, prepushed=3), + ), + Cell( + why="12 tokens pushed by the o_proj GEMV into the 16-token bucket", + spec=dict(m=12, maxt=16, opts=15, x2=False, quant=False, prepushed=2), + ), + Cell( + why="the 9..13-token MoE kernel's sum into the 16-token bucket, MXFP8 copy", + spec=dict(m=16, maxt=16, opts=15, x2=False, quant=True, prepushed=0), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=30, maxt=48, opts=15, x2=False, quant=False, prepushed=0), + ), + Cell( + why="the 48-token bucket full, second addend and MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, x2=True, quant=True, prepushed=0), + ), + # Every other row count of the decode boundaries (modeling.MC_ALLREDUCE): the bucket and epilogue options the + # target picks for it (ops/m3_mc_allreduce), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} tokens: a boundary row count between the ones above", + spec=dict( + m=m, + maxt=8 if m <= 8 else 16 if m <= 16 else 48, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1] if m <= 8 else 15, + x2=m % 2 == 0, + quant=m % 3 == 0, + prepushed=0, + ), + ) + for m in range(1, 49) + if m not in (1, 4, 8, 16, 30, 48) + ), + # Every other row count pushed by the o_proj (1..16) or the 1..8-token MoE engine (with the MXFP8 copy): + # modeling.ATTN_PUSH / modeling.MOE_PUSH. + *( + Cell( + why=f"{m} tokens prepushed: a pushed row count between the ones above", + spec=dict( + m=m, + maxt=8 if m <= 8 else 16, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1] if m <= 8 else 15, + x2=False, + quant=m <= 8, + prepushed=2, + ), + ) + for m in range(1, 17) + if m not in (8, 12) + ), + # The Eagle3 drafter boundaries (modeling.DRAFT_MC_AR): option 16 (the row kernel's drafter form) on 23 / 31 + # at every draft-pass row count 1..8 (pass 0: 4 x requests, later passes: requests), x only. + *( + Cell( + why=f"Eagle3 drafter boundary at {m} rows (opts {23 if m in (1, 2, 7, 8) else 31})", + spec=dict( + m=m, + maxt=8, + opts=23 if m in (1, 2, 7, 8) else 31, + x2=False, + quant=False, + prepushed=0, + drafter=True, + ), + ) + for m in range(1, 9) + ), + # The same boundaries on the 8-CTA cluster kernel (opts bit 128, modeling.DRAFT_MC_AR_CLUSTER, launched without + # the carveout: 23 | 128), alone and fed by the drafter projection's push (prepushed 2, DRAFT_AR_PUSH). + *( + Cell( + why=f"Eagle3 drafter boundary on the cluster kernel at {m} rows, prepushed {pre}", + spec=dict( + m=m, + maxt=8, + opts=23 | 128, + x2=False, + quant=False, + prepushed=pre, + drafter=True, + ), + ) + for pre in (0, 2) + for m in range(1, 9) + ), + ) + + note = """ + The epilogue option bits (1 early residual / weight load, 2 st.async cluster exchange, 4 early parity rotation, + 8 max carveout) are bit-identical by construction; only the four (bucket, option) pairs the target selects are + built. `prepushed` 1 (pushed, grid wait kept) is accepted by the op but no caller passes it. `pdl` changes only + launch overlap and is driven off in the cells. + + `normed` is gated to one BF16 rounding step: the kernel's RMS statistic is a warp / cluster tree of FP32 sums + and its `rsqrtf` is approximate, so a row whose scaled value lands next to a BF16 midpoint can round either + way. `updated` and the MXFP8 copy of the op's own `normed` are bit-exact. + """ + + def __call__( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_sm107_tp4_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + None, + None, + maxt, + quant, + prepushed, + opts, + cap_out, + cap_w, + cap_eps, + ) + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant, drafter) + + def is_valid( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + ) -> None: + h = residual.shape[-1] + # The op checks norm_w's element count, not its layout: a strided view is read as packed memory. + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + # The workspace is addressed as [set][parity][maxt][world][H / 2] words and never bounds-checked: a smaller + # buffer (another bucket's, or one built for fewer ranks) is polled and re-armed out of range. + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + # Flags are per (token, cluster CTA) of the bucket; the op checks the count, not that they belong to the + # workspace. Flags of another bucket hold another call history, and the parity they select is not the one the + # peers push into: the poll then waits on the other half forever or reads last call's rows. + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + + +m3_mc_allreduce_norm = M3McAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..0d457022c38f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize + multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode MoE boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_mc_allreduce_norm import boundary + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class M3McFinalizeAllreduceNorm(OpWrapper): + """The MoE boundary with the routed experts' finalize folded in. + + Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, + ``scales`` fp32 or bf16), then the group's all-reduce, residual add and RMSNorm exactly as + ``m3_mc_allreduce_norm``: same workspace, flags and call-order contract. Returns ``[normed, updated]`` or, with + ``quant``, also the MXFP8 copy of ``normed``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The tp4 target takes this path from 9 tokens (top-4 of 128 experts, fp32 routing weights, the shared expert + # always present): the 16- and 48-token buckets, epilogue option 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="the smallest finalize call (9 tokens, 16-token bucket) with the next QKV's MXFP8 copy", + spec=dict(m=9, maxt=16, opts=15, topk=4, quant=True, dropped=0), + ), + Cell( + why="a full 16-token bucket with dropped expert slots", + spec=dict(m=16, maxt=16, opts=15, topk=4, quant=False, dropped=5), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=24, maxt=48, opts=15, topk=4, quant=False, dropped=0), + ), + Cell( + why="the 48-token bucket full, MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, topk=4, quant=True, dropped=3), + ), + # Every other row count of the unfinalized MoE (modeling.MC_FINALIZE): the bucket and options the target picks. + *( + Cell( + why=f"{m} tokens: a finalize row count between the ones above", + spec=dict( + m=m, + maxt=16 if m <= 16 else 48, + opts=15, + topk=4, + quant=m % 2 == 0, + dropped=m % 4, + ), + ) + for m in range(9, 49) + if m not in (9, 16, 24, 48) + ), + ) + + note = """ + BF16 routing weights and the 8-token bucket are accepted by the op (option 7 there) but the target only enters + this path at 9 tokens and up, with the built-in MoE's fp32 weights; neither is a cell. The rest is as in + `m3_mc_allreduce_norm`'s note. + """ + + def __call__( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_sm107_tp4_mc_allreduce_norm( + fc2, + shared, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + 0, + opts, + ) + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + ) -> None: + h = residual.shape[-1] + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + # The kernel gathers fc2[perm] without a bound: an index past the FC2 rows reads whatever follows them. + lo, hi = int(perm.min()), int(perm.max()) + assert lo >= -1 and hi < fc2.shape[0], ( + f"perm spans [{lo}, {hi}]; entries must be -1 (dropped) or a row of fc2 (< {fc2.shape[0]})" + ) + + +m3_mc_finalize_allreduce_norm = M3McFinalizeAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/__init__.py new file mode 100644 index 000000000000..e9f4d8e20b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedding ops of this target: the generic token-embedding gather.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..3060c17acdff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="a 48-row decode verify, int64 ids", + spec=dict(rows=48, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ) + + note = """ + An exact gather: `reference` must match bit for bit. + """ + + def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) + + def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] + return rows * valid.unsqueeze(-1).to(table.dtype) + + +m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/__init__.py new file mode 100644 index 000000000000..4c6fc92cd49c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 decode embedding gather + first-layer RMSNorm as one CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py new file mode 100644 index 000000000000..78563b4db40d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. + +128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original +K3 wrapper this target does not promise FlashInfer's bit-identical reduction. +""" + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +NORM_THREADS = 128 +NORM_VEC = 8 +NORM_WARPS = 4 # capacity: fixed: the CTA's warps + + +@cute.kernel +def m3_embed_norm_kernel( + ids: cute.Tensor, # int32 or int64 [N] + table: cute.Tensor, # bf16 [V, H] + weight: cute.Tensor, # bf16 [H] + raw: cute.Tensor, # bf16 [N, H], out: the rows + out: cute.Tensor, # bf16 [N, H], out: the normed rows + vocab: cutlass.Int64, + eps: cutlass.Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + hidden: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 + ) + + tok = cutlass.Int64(ids[bidx]) + valid = (tok >= cutlass.Int64(0)) & (tok < vocab) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + + gX = cute.local_tile(table, tiler_mn, (row, 0)) + gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) + gY = cute.local_tile(out, tiler_mn, (bidx, 0)) + w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 + ) + thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + tXgR = thr_copy_O.partition_D(gR) + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. + tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) + if valid: + cute.copy(copy_atom_load, tXgX, tXrX) + cute.copy(copy_atom_load, tWgW, tWrW) + cute.copy(copy_atom_store, tXrX, tXgR) + + # flashinfer's RMSNormKernel arithmetic, in its order (see the header). + x = tXrX.load().to(cutlass.Float32) + x_sq = x * x + sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) + for i in cutlass.range_constexpr(5): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + if lane == 0: + reduction_buffer[0, warp] = sum_sq + cute.arch.barrier() + total = cutlass.Float32(0.0) + if lane < NORM_WARPS: + total = reduction_buffer[0, lane] + for i in cutlass.range_constexpr(5): + total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) + mean_sq = total / cutlass.Float32(hidden) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + w = tXrW.load().to(cutlass.Float32) + y = x * rstd * (w + cutlass.Float32(0.0)) + tXrO.store(y.to(out.element_type)) + cute.copy(copy_atom_store, tXrO, tXgO) + + +@cute.jit +def m3_embed_norm( + ids: cute.Tensor, + table: cute.Tensor, + weight: cute.Tensor, + raw: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int64, + eps: cutlass.Float32, + n_tokens: cutlass.Constexpr[int], + hidden: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. + blocks = hidden // (NORM_VEC * NORM_THREADS) + tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), + stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip + m3_embed_norm_kernel( + ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden + ).launch( + grid=[n_tokens, 1, 1], + block=[NORM_THREADS, 1, 1], + smem=NORM_WARPS * 4, + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py new file mode 100644 index 000000000000..276f8c32023b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host side of the MiniMax-M3 decode embedding + first-layer RMSNorm (m3_embedding_firstnorm_kernel.py): H 6144, +1..8 rows, into caller-owned buffers. + +Compiled on the first call per (rows, id dtype, table shape, PDL), which must happen outside CUDA-graph capture +(:class:`DecodeEmbedding` compiles at construction). +""" + +import threading +from typing import Dict, Tuple + +import cuda.bindings.driver as cuda_driver +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +MAX_ROWS = 8 # capacity: fixed: the kernel's built rows (decode steps) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def arg(tensor): + return from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _check(ids, table, weight, raw, out) -> None: + rows = ids.shape[0] if ids.dim() == 1 else 0 + if not ( + 1 <= rows <= MAX_ROWS + and ids.dtype in (torch.int32, torch.int64) + and ids.is_contiguous() + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN,) + and all( + t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() + for t in (raw, out) + ) + ): + raise ValueError( + "m3_embed_norm: expects ids int32/int64 [1..8], a contiguous bf16 table [V, 6144], bf16 weight [6144] " + "and contiguous bf16 raw / out [rows, 6144]" + ) + + +def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) + fn = _compiled[key] = cute.compile( + _m3_embed_norm_jit, + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + _check(ids, table, weight, raw, out) + fn = _compile(ids, table, weight, eps, raw, out, use_pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return out, raw + + +class DecodeEmbedding: + """Caller-owned buffers of one (rows, id dtype) decode embedding: ``raw`` (the gathered rows, the first residual) + and ``output`` (their RMSNorm), over ``table`` and the first layer's plain (non-Gemma) input norm ``weight`` / + ``eps``. Built (and compiled) before CUDA-graph capture.""" + + def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): + assert not torch.cuda.is_current_stream_capturing() + assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN + assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) + assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( + table, + weight, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(eps) + self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) + + def __call__(self, ids): + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return m3_embed_norm( + ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/__init__.py new file mode 100644 index 000000000000..9cc374bdf1cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMM entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..4fb5af737dc0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3Bf16Linear(OpWrapper): + """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. + + The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), + o_proj (this rank's input columns), gate_up and down projections run it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=2304, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=2048)), + Cell(why="drafter gate_up, one row", spec=dict(m=1, n=9216, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=4608)), + ) + + note = """ + cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x, w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return (x.float() @ w.float().t()).to(x.dtype) + + +m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py new file mode 100644 index 000000000000..1fb7a1d759ae --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense MLP's swigluoai activation as one C++ launch with PDL and an early release (DENSE_SWIGLU_CPP).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# The dense layers' swigluoai parameters, which the kernel hard-codes (modeling_minimaxm3._build_swiglu_oai_dense_mlp). +LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 + + +class M3DenseSwiglu(OpWrapper): + """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, + rounded once to bf16: trtllm::minimax_m3_sm107_tp4_swiglu_oai (minimaxM3Sm107Tp4/swigluOai.cu), bit for bit the + Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the + grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights + during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even + half width; ``__call__`` returns None otherwise or for other activation parameters (the caller runs the Triton + op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP4 dense MLP shard: gate | up 6144 wide (d = 3072), at the decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP4 dense MLP at 1 row", spec=dict(rows=1, d=3072)), + Cell(why="the TP4 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=3072)), + Cell(why="the TP4 dense MLP at 8 rows", spec=dict(rows=8, d=3072)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the + alpha-scaled ex2.approx sigmoid, div.full), checked over every bf16 gate and every bf16 up pattern. + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float] = LIMIT, + alpha: Optional[float] = ALPHA, + beta: Optional[float] = BETA, + use_pdl: bool = True, + ) -> Optional[torch.Tensor]: + if ( + (limit, alpha, beta) != (LIMIT, ALPHA, BETA) + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 4 != 0 + or x.data_ptr() % 4 != 0 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp4_swiglu_oai(x, out, bool(use_pdl)) + return out + + def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_dense_swiglu = M3DenseSwiglu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py new file mode 100644 index 000000000000..e4cd998d587d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 x BF16 -> BF16 low-M GEMM of the Eagle3 drafter's linears at fixed FlashInfer tactics (SM10x).""" + +from __future__ import annotations + +from typing import Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: (N, K) of the TP4 drafter shards -> ((max rows, kind, tactic), ...), rows ascending. Kind "d" is FlashInfer's direct +#: SIMT kernel with tactic (block, outputs per block, rows per block); kind "s" is its split-K tensor-core kernel with +#: tactic (mma_m, mma_n, split_k, ab_stages). A row count is served by the first entry whose max rows holds it; the +#: rows modeling.DRAFT_LOWM covers are only the ones measured (the fastest tactic at no larger error vs fp64 than +#: cuBLAS). Rubin SM107 TP4 sweep (mclk locked, >= 3x L2 of rotating weights, CUDA-graph timing), us cuBLAS -> +#: low-M at 1 / 4 rows: qkv 8.15 -> 5.90 / 8.32 -> 6.71, o 5.36 -> 3.40 / 5.44 -> 3.66, gate_up 12.93 -> 9.93 / +#: 13.03 -> 12.05, down 7.68 -> 6.16 / 7.64 -> 6.15. The fc is replicated (the TP2 shape): its TP2 tactic and row. +TACTICS: Mapping[Tuple[int, int], tuple] = { + # qkv + (2304, 12288): ((1, "d", (256, 2, 1)), (4, "s", (64, 8, 4, 10))), + # o_proj + (6144, 2048): ((1, "d", (128, 4, 1)), (4, "s", (64, 8, 1, 10))), + # gate_up + (9216, 6144): ((1, "d", (128, 2, 1)), (4, "s", (64, 8, 1, 6))), + # down + (6144, 4608): ((4, "s", (64, 8, 2, 10)),), + # the Eagle3 fc over the three captured hidden states + (6144, 18432): ((1, "d", (384, 2, 1)),), +} +#: Rows each (N, K) is certified at (a subset of what its table would serve). +ROWS: Mapping[Tuple[int, int], Tuple[int, ...]] = { + (2304, 12288): (1, 4), + (6144, 2048): (1, 4), + (9216, 6144): (1, 4), + (6144, 4608): (1, 4), + (6144, 18432): (1,), +} + + +def tactic_for(n: int, k: int, m: int) -> Optional[Tuple[str, tuple]]: + """(kind, tactic) of an [m, k] x [n, k]^T call, or None when the table does not serve it.""" + for rows, kind, tactic in TACTICS.get((n, k), ()): + if m <= rows: + return kind, tactic + return None + + +class M3DraftLowM(OpWrapper): + """``out`` [M, N] BF16 = ``x`` [M, K] BF16 @ ``weight`` [N, K]^T BF16, FP32 accumulation. + + ``kind`` / ``tactic`` pick FlashInfer's kernel: "d" runs ``run_direct_dense`` (SIMT, one dot product per output in + K order), "s" runs ``run_splitk_dense`` (tcgen05 MMA, K split over ``split_k`` CTAs and summed in FP32). Both + round once to BF16. The kernel is compiled on its first call per (M, N, K, tactic, PDL); the target compiles every + covered call before warmup (ops/m3_draft_lowm.prebuild). + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter qkv, 1 rows: d tactic", spec=dict(nk=(2304, 12288), m=1)), + Cell(why="drafter qkv, 4 rows: s tactic", spec=dict(nk=(2304, 12288), m=4)), + Cell(why="drafter o_proj, 1 rows: d tactic", spec=dict(nk=(6144, 2048), m=1)), + Cell(why="drafter o_proj, 4 rows: s tactic", spec=dict(nk=(6144, 2048), m=4)), + Cell(why="drafter gate_up, 1 rows: d tactic", spec=dict(nk=(9216, 6144), m=1)), + Cell(why="drafter gate_up, 4 rows: s tactic", spec=dict(nk=(9216, 6144), m=4)), + Cell(why="drafter down, 1 rows: s tactic", spec=dict(nk=(6144, 4608), m=1)), + Cell(why="drafter down, 4 rows: s tactic", spec=dict(nk=(6144, 4608), m=4)), + Cell(why="drafter fc, 1 rows: d tactic", spec=dict(nk=(6144, 18432), m=1)), + ) + + note = """ + The table is the TP4 drafter's: its shapes are the per-rank shards of the Eagle3 layer at 4-way tensor + parallelism (32 query and 2 KV heads per rank, intermediate 9216 per rank) and the replicated fc, at the b1 + decode rows (1 row in draft passes 1-2, 4 in pass 0). Certified at those shapes and rows only; FlashInfer's + kernels take other shapes, which are not measured here. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + if kind == "d": + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + DirectTactic, + run_direct_dense, + ) + + run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) + elif own_splitk: + # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # it also pushes the rows into a drafter boundary all-reduce (modeling.DRAFT_AR_PUSH). + from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push + + if ar is not None: + run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) + else: + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + else: + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + return out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """bf16(x @ weight^T) summed in fp64 and rounded once; ``out``, ``kind``, ``tactic`` and ``use_pdl`` only + choose where and how the op writes.""" + return (x.double() @ weight.double().t()).to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> None: + """The kernels index ``x`` and ``out`` densely: a strided view is read as if it were packed.""" + assert x.dim() == 2 and x.is_contiguous(), "x must be a contiguous [M, K] matrix" + assert weight.dim() == 2 and weight.is_contiguous(), ( + "weight must be a contiguous [N, K] matrix" + ) + assert out.is_contiguous() and tuple(out.shape) == (x.shape[0], weight.shape[0]), ( + "out must be a contiguous [M, N] matrix" + ) + assert ar is None or (own_splitk and kind == "s" and len(ar) == 5), ( + "ar (the boundary all-reduce push) runs only on this target's split-K copy" + ) + + +m3_draft_lowm = M3DraftLowM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..9702cc0778a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class M3Fp32Linear(OpWrapper): + """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the + fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies + (the engine leaves torch's default).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), + Cell( + why="a generation step above the tensor-core router's 48 rows", + spec=dict(m=64, e=128, k=6144), + ), + ) + + note = """ + Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against + the true fp32 reference. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x.to(torch.float32), w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return x.float() @ w.t() + + +m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..59bfa0780a00 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 x MXFP8 GEMV (tcgen05, block-scaled, CuTe DSL) on an activation its producer already quantized.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_dequant +from ..mx_gemv import op as _gemv + +__all__ = ["PROJECTIONS", "gemv_is_valid", "m3_mx_gemv", "served_cells"] + +# (N, K) of this target's MXFP8 projections at tp4, the GEMV's tuned shapes. +QKV_SPARSE = (2560, 6144) +QKV = (2304, 6144) +O_PROJ = (6144, 2048) +SHARED_GATE_UP = (1536, 6144) +SHARED_DOWN = (6144, 768) +DENSE_GATE_UP = (6144, 6144) +DENSE_DOWN = (6144, 3072) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "shared gate_up": SHARED_GATE_UP, + "shared down": SHARED_DOWN, + "dense gate_up": DENSE_GATE_UP, + "dense down": DENSE_DOWN, +} + + +def served_cells(named: tuple[Cell, ...], x_bf16: bool, max_m: int) -> tuple[Cell, ...]: + """A cell for every (projection, rows) the GEMV serves up to ``max_m`` rows that ``named`` does not list. + + Certification is per projection: a projection's split-K, ring and MMA width differ, so each (N, K, M) the + decode dispatch can send here is its own cell. + """ + listed = {(cell.spec["nk"], cell.spec["m"]) for cell in named} + return tuple( + Cell(why=f"{name}, {m} rows", spec=dict(nk=nk, m=m)) + for name, nk in PROJECTIONS.items() + for m in range(1, max_m + 1) + if (nk, m) not in listed and _gemv.serves(nk[0], nk[1], m, x_bf16) + ) + + +def gemv_reference( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """bf16(x @ dequant(weight)^T) for fp32 activation values x [M, K], summed in fp64 and rounded once.""" + return (x.double() @ mxfp8_dequant(weight, weight_scale).double().t()).to(torch.bfloat16) + + +def gemv_is_valid(weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + n, k = weight.shape + # The kernel reads both from data_ptr() as dense buffers at the layout it was compiled for. + assert weight.dtype == torch.float8_e4m3fn and weight.is_contiguous(), ( + "weight must be contiguous e4m3 [N, K]" + ) + assert weight_scale.numel() == n * k // 32, ( + "weight_scale must be the 128x4-swizzled UE8M0 buffer [N * K / 32]" + ) + + +class M3MxGemv(OpWrapper): + """bf16 [M, N] = xq [M, K] (e4m3, 128x4-swizzled UE8M0 scales ``xsf``) x weight [N, K]^T (e4m3, swizzled + ``weight_scale``), M <= 48; None when the shape is not served (the caller falls back). + + ``trigger_early``: release the PDL dependents before the epilogue. ``cfg``: a (split, ring[, contig, nacc, a_tmem, + st_async]) tuning override; default the table's. ``use_pdl``: default the environment's. ``ar`` = (multicast + workspace address, flags address, rank, world, bucket tokens, mode): also push the rows into the next boundary + all-reduce's exchange (then ``op.last_pushed`` is True). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The projections' (N, K) at the token counts each build differs in: the MMA width (8 / 16 / 32 / 48 columns) a + # call uses is the smallest that holds M, and the shared gate_up takes another split at 2 tokens and above 8. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 4)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell(why="sparse-layer QKV, 8 tokens", spec=dict(nk=QKV_SPARSE, m=8)), + Cell( + why="sparse-layer QKV, 48 verify rows: the widest MMA", spec=dict(nk=QKV_SPARSE, m=48) + ), + Cell(why="dense-attention QKV, one token", spec=dict(nk=QKV, m=1)), + Cell(why="dense-attention QKV, 13 rows: the 16-column MMA", spec=dict(nk=QKV, m=13)), + Cell(why="o_proj, one token (split 2)", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 24 rows: the 32-column MMA", spec=dict(nk=O_PROJ, m=24)), + Cell( + why="shared gate_up, one token (A from shared memory)", + spec=dict(nk=SHARED_GATE_UP, m=1), + ), + Cell(why="shared gate_up, 2 tokens: its split-4 config", spec=dict(nk=SHARED_GATE_UP, m=2)), + Cell( + why="shared gate_up, 12 rows: the wide split-4 config", + spec=dict(nk=SHARED_GATE_UP, m=12), + ), + Cell(why="shared down, one token (unsplit, K 768)", spec=dict(nk=SHARED_DOWN, m=1)), + Cell(why="shared down, 8 tokens", spec=dict(nk=SHARED_DOWN, m=8)), + Cell(why="dense gate_up, one token", spec=dict(nk=DENSE_GATE_UP, m=1)), + Cell(why="dense down, 8 tokens", spec=dict(nk=DENSE_DOWN, m=8)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling.MX_GEMV_ROWS). + CELLS = CELLS + served_cells(CELLS, x_bf16=False, max_m=48) + + note = """ + The kernel sums in fp32 in split-K order where `reference` sums in fp64: the bf16 outputs agree within the band the + test derives. Not covered: the all-reduce push (`ar`), which needs a multicast exchange; `cfg` overrides beyond the + table (the target passes none). Shapes outside the table, M > 48, and 17..32 rows at K < 1024 return None. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv(xq, xsf, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The dequantized product; the launch options are ignored.""" + return gemv_reference(mxfp8_dequant(xq, xsf), weight, weight_scale) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + m, k = xq.shape + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous(), ( + "xq must be contiguous e4m3 [M, K]" + ) + # A linear-layout [M, K / 32] scale buffer is read as the swizzled one: wrong scales, no error. + assert xsf.numel() >= (m + 127) // 128 * 128 * (k // 32), ( + "xsf must be the 128x4-swizzled scale buffer (rows padded to 128), not the linear one" + ) + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv = M3MxGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..615a6b2e6785 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mx_gemv_bf16.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 GEMV (CuTe DSL) on a bf16 activation, quantized to MXFP8 in the kernel's prologue.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_quant_dequant +from ..mx_gemv import op as _gemv +from .m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + gemv_is_valid, + gemv_reference, + served_cells, +) + +__all__ = ["m3_mx_gemv_bf16"] + + +class M3MxGemvBf16(OpWrapper): + """bf16 [M, N] = mxfp8_quantize(x) x weight [N, K]^T for a contiguous, 16-byte aligned bf16 ``x`` [M, K]; the + quantization is mxfp8_quantize's (UE8M0 scale per 32 = the smallest power of two >= amax / 448). None when the + shape is not served. The other arguments as ``m3_mx_gemv``'s. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The prologue quantizes one 32-element block per thread and round: the projections at the token counts the + # target sends here (no producer-quantized activation), up to the widest each serves. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 16 rows: the most its split-4 prologue serves", + spec=dict(nk=QKV_SPARSE, m=16), + ), + Cell(why="dense-attention QKV, 8 tokens", spec=dict(nk=QKV, m=8)), + Cell(why="o_proj, one token", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 32 rows: the 32-column MMA", spec=dict(nk=O_PROJ, m=32)), + Cell(why="shared gate_up, one token", spec=dict(nk=SHARED_GATE_UP, m=1)), + Cell(why="shared gate_up, 2 tokens: its split-4 config", spec=dict(nk=SHARED_GATE_UP, m=2)), + Cell(why="shared down, 8 tokens (K 768)", spec=dict(nk=SHARED_DOWN, m=8)), + Cell( + why="dense gate_up, 8 tokens: the most split 2 at K 6144 serves", + spec=dict(nk=DENSE_GATE_UP, m=8), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling.MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=32) + + note = """ + As `m3_mx_gemv`. The in-kernel quantization is mxfp8_quantize's bit for bit except where the kernel's + rcp.approx(448) and an exact 1/448 put a block's amax on different sides of a power of two (not observed on the + cells). Not covered: `ar`, `cfg` overrides. Returns None for rows beyond 8 rounds of the prologue per CTA. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv_bf16(x, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The MXFP8 round trip of x times the dequantized weight; the launch options are ignored.""" + return gemv_reference(mxfp8_quant_dequant(x), weight, weight_scale) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv_bf16 = M3MxGemvBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..0b422a7167e5 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned + +from .m3_mxfp8_linear import mxfp8_reference + + +class M3Mxfp8Gemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight + [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key + and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The decode dispatch hands it the boundary / merge MXFP8 activation on rows the GEMV does not serve: the + # unsplit shared down projection at 17..32 rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=768)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=768)), + ) + + note = """ + Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + """The dequantized product; ``tune`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..d7a70d725189 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned + +_BLOCK = 32 + + +def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" + pad_r = (rows + 127) // 128 * 128 + pad_c = (cols + 3) // 4 * 4 + s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) + return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] + + +def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, k = x.shape + s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) + factor = torch.pow(2.0, (s - 127).to(torch.float32)) + return ( + x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) + ).view(rows, k) + + +def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: + """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" + with true_fp32_matmul(): + return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=2560, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=2048, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=3072, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=1536, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=2048, path="tuned")), + ) + + note = """ + Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` + quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and + multiplies dequantized operands in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + if tuned_quantize_gemm: + return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + if flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, output_dtype + ) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are + ignored.""" + m, k = x.shape + blocks = x.float().view(m, k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(-1, keepdim=True) + exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) + scale = torch.pow(2.0, exp) + q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale + with true_fp32_matmul(): + out = q.view(m, k) @ _dequant(weight, weight_scale).t() + return (out * alpha.float()).to(output_dtype) + + +m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py new file mode 100644 index 000000000000..ee781b16ff42 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM at a fixed FlashInfer tactic (MX-TAC): the same choice on every rank and every start.""" + +from __future__ import annotations + +from typing import Dict, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + +from .m3_mxfp8_linear import mxfp8_reference + +# One runner per output dtype (a runner holds no per-call state). +_RUNNERS: Dict[Tuple[type, torch.dtype], object] = {} + + +def _runner(dtype: torch.dtype): + runner = _RUNNERS.get((FlashInferMXFP8GemmRunner, dtype)) + if runner is None: + runner = _RUNNERS[(FlashInferMXFP8GemmRunner, dtype)] = FlashInferMXFP8GemmRunner(dtype) + return runner + + +class M3Mxfp8PinnedGemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + swizzled UE8M0 scales) and an MXFP8 weight [N, K] + (swizzled scales) on FlashInfer's MXFP8 GEMM at ``tactic``, no AutoTuner; ``output_dtype`` out.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The decode graph's MXFP8 GEMMs on rows the GEMVs do not serve: the shared down projection on its boundary + # activation (17..32 rows) and the bf16 projections past the GEMV rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=768)), + Cell(why="packed QKV + index projection, 48 rows", spec=dict(m=48, n=4992, k=6144)), + ) + + note = """ + The MXFP8 product (fp32 accumulation, one rounding to the output dtype): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. The tactic only picks the kernel. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tactic: int, + ) -> torch.Tensor: + return _runner(output_dtype).forward([act, act_scale, weight, weight_scale], tactic=tactic) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tactic: int, + ) -> torch.Tensor: + """The dequantized product; ``tactic`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_pinned_gemm = M3Mxfp8PinnedGemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py new file mode 100644 index 000000000000..3f35af65048a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_quantize.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 activation quantize at a fixed tactic (MX-TAC), ahead of ``m3_mxfp8_pinned_gemm``.""" + +from __future__ import annotations + +from typing import Dict, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8QuantizeRunner + +_RUNNERS: Dict[torch.dtype, object] = {} + + +def _runner(dtype: torch.dtype): + runner = _RUNNERS.get(dtype) + if runner is None: + runner = _RUNNERS[dtype] = MXFP8QuantizeRunner(dtype) + return runner + + +class M3Mxfp8PinnedQuantize(OpWrapper): + """A bf16 activation [M, K] to MXFP8: (e4m3 [M, K], swizzled UE8M0 scales), on the runner's ``tactic``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="packed QKV + index projection input, 48 rows", spec=dict(m=48, k=6144)), + Cell(why="shared expert down projection input, 17 rows", spec=dict(m=17, k=768)), + ) + + note = """ + Bit-exact: `reference` is the stock MXFP8 quantize (``torch.ops.trtllm.mxfp8_quantize``, swizzled scales) of the + same input; every tactic computes the same block maxima and roundings. + """ + + def __call__(self, x: torch.Tensor, tactic: int) -> Tuple[torch.Tensor, torch.Tensor]: + return _runner(x.dtype).forward([x], tactic=tactic) + + def reference(self, x: torch.Tensor, tactic: int) -> Tuple[torch.Tensor, torch.Tensor]: + """The stock quantize; ``tactic`` is ignored.""" + return torch.ops.trtllm.mxfp8_quantize(x, True) + + +m3_mxfp8_pinned_quantize = M3Mxfp8PinnedQuantize() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py new file mode 100644 index 000000000000..f8deb858c843 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, +and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). + +The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values +``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 +marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a +per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: +per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 +code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved +it decodes to the BF16 weight bit for bit. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_target_head"] + +#: Verify rows per call (requests x (draft length + 1)). +MAX_ROWS = 8 +#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). +CTA_WARPS = 16 +#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. +TARGET_SHARD = 50016 +HIDDEN = 6144 +#: Codes per window: 0..14 decode, 15 marks an exception. +WINDOW = 15 +#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. +MAX_BASE = 254 - WINDOW +#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). +ROWS_PER_CHUNK = 4096 + + +def _pad_rows(weight: torch.Tensor) -> torch.Tensor: + n = weight.shape[0] + tiles = (n + 15) // 16 + if tiles * 16 == n: + return weight + return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) + + +def _fragment_view(w: torch.Tensor) -> torch.Tensor: + """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" + n, k = w.shape + # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) + return ( + w.view(n // 16, 2, 8, k // 16, 2, 4, 2) + .permute(0, 3, 2, 5, 4, 1, 6) + .reshape(n // 16, k // 16, 32, 8) + ) + + +def _exponents(weight: torch.Tensor) -> torch.Tensor: + return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF + + +def _window_base(hist: torch.Tensor) -> int: + """The base of the 15-exponent window holding the most weights (the lowest such base).""" + window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) + return int(window.argmax()) + + +def encode(weight: torch.Tensor, base: Optional[int] = None): + """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's + own window. + + Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile + << 16 | k), exc_w fp32 [E, 2] ({w, 0})). + """ + w = _pad_rows(weight) + n, k = w.shape + tiles, ks = n // 16, k // 16 + assert ks % 4 == 0, "K must be a multiple of 64" + bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF + expo = (bits >> 7) & 0xFF + if base is None: + base = _window_base(torch.bincount(expo.flatten(), minlength=256)) + code = expo - base + exc = (code < 0) | (code >= WINDOW) + code = torch.where(exc, torch.full_like(code, WINDOW), code) + sm = ((bits >> 8) & 0x80) | (bits & 0x7F) + sm = torch.where(exc, torch.zeros_like(sm), sm) + # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. + rows, cols = exc.nonzero(as_tuple=True) + order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) + rows, cols = rows[order], cols[order] + exc_w = torch.stack( + [ + w[rows, cols].float(), + torch.zeros((rows.numel(),), device=w.device), + ], + dim=1, + ) + exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) + exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) + exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( + torch.int32 + ) + # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. + smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) + cf = ( + _fragment_view(code.to(torch.uint8).view(n, k)) + .view(tiles, ks // 4, 4, 32, 8) + .to(torch.int32) + ) + # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. + cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) + for p in range(4): + cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) + cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 + # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. + seg01 = ( + smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) + ) # [.., 2, 32, 16] + seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) + packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] + return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() + + +def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): + """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" + n, _ = weight.shape + assert rows_per_chunk % 16 == 0 + hist = torch.zeros(256, dtype=torch.int64, device=weight.device) + for r in range(0, n, rows_per_chunk): + hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) + base = _window_base(hist) + packed, offs, poss, ws = [], [], [], [] + n_exc = 0 + for r in range(0, n, rows_per_chunk): + p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) + packed.append(p) + offs.append(off[:-1] + n_exc) + poss.append(pos) + ws.append(ew) + n_exc += pos.numel() + exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) + return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) + + +def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: + """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" + tiles, ks = (n + 15) // 16, k // 16 + p = packed.view(tiles, ks // 4, 3, 32, 16) + sm = ( + p[:, :, :2] + .permute(0, 1, 3, 2, 4) + .reshape(tiles, ks // 4, 32, 4, 8) + .permute(0, 1, 3, 2, 4) + .to(torch.int32) + ) + cw = ( + p[:, :, 2] + .contiguous() + .view(torch.int32) + .view(tiles, ks // 4, 32, 4) + .permute(0, 1, 3, 2) + .long() + & 0xFFFFFFFF + ) + codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( + torch.int32 + ) + bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) + bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) + frag = bits.reshape(tiles, ks, 32, 8) + # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). + w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) + return w.to(torch.int16).view(torch.bfloat16)[:n] + + +def decode( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + n: int, + k: int, + row0: int = 0, + row1: Optional[int] = None, +) -> torch.Tensor: + """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). + ``row0`` must be a multiple of 16.""" + row1 = n if row1 is None else row1 + assert row0 % 16 == 0 + t0, t1 = row0 // 16, (row1 + 15) // 16 + unit_bytes = (k // 64) * 3 * 32 * 16 + dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() + counts = (exc_off[1:] - exc_off[:-1]).long() + tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) + rows = tile_of * 16 + (exc_pos >> 16).long() + cols = (exc_pos & 0xFFFF).long() + sel = (rows >= row0) & (rows < row1) + dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) + return dec + + +def round_trip_ok( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + weight: torch.Tensor, + rows_per_chunk: int = ROWS_PER_CHUNK, +) -> bool: + """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent + window below the BF16 infinity exponent.""" + n, k = weight.shape + if not 0 <= base <= MAX_BASE or k % 64: + return False + if ( + packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 + or exc_off.numel() != (n + 15) // 16 + 1 + ): + return False + for r0 in range(0, n, rows_per_chunk): + r1 = min(n, r0 + rows_per_chunk) + dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) + if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): + return False + return True + + +class M3TargetHead(OpWrapper): + """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit + the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in + another summation order). + + ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` + (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are + the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and + # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with + # a shard whose rows are not a multiple of 16. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{m} verify row(s) of the TP4 target head shard", + spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), + ) + for m in range(1, MAX_ROWS + 1) + ), + Cell( + why="4 rows with outliers and exact zeros in many tiles (the exception lists)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), + ), + Cell( + why="2 rows over a shard of 4000 rows (250 tiles), K 1024", + spec=dict(m=2, n=4000, k=1024, exceptions=True), + ), + ) + + note = """ + The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception + adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs + from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); + integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp + order: the result is deterministic. + """ + + def __call__( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_target_head( + w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl + ) + return out + + def reference( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. + + ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" + weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + return logits.to(torch.bfloat16) + + def is_valid( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> None: + # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count + # stores a partial sum. + assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + assert out.data_ptr() != x.data_ptr(), "out must not alias x" + + +m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/long_select.py new file mode 100644 index 000000000000..193cc96e72bf --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/long_select.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. + +Same contract and bit-identical output as ``torch.ops.trtllm.minimax_m3_select_blocks`` on per-KV-head scores +[num_kv_heads, n_blocks, total_q]: init / local forcing, per-query valid-block count, ties to the smaller block index, +forced -inf picks emitted as -1, ascending block ids with -1 tail padding. + +One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups ``block % 16`` is a lower +bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared memory +and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every candidate +against the whole row (re-read). + +The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp4_select_long`` (this target's copy of ``minimax_m3_select_long`` +with a choice of where it releases the sparse decode), called through this target's ``m3_select_long`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long +from ..ops.attention.m3_select_long_dsl import m3_select_long_dsl + + +def long_select_blocks( + max_score_kv: torch.Tensor, + *, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + fast_tail: bool = True, + dsl: bool = False, +) -> Optional[torch.Tensor]: + """Top-16 select over per-KV-head scores [num_kv_heads, n_blocks, total_q] with 129..9216 blocks. + + Returns the head-major backed [total_q, num_kv_heads, 16] int32 table, or None outside this kernel's shapes + (callers keep the reference op). ``early_trigger`` releases the dependent sparse decode right after the select's + grid wait (``SELECT_EARLY_MAX_ROWS`` in modeling.py); otherwise ``trigger_phase`` 1 releases it once every score + of the row is loaded and 0 at the select's exit (``SELECT_LONG_TRIGGER_PHASE``). ``fast_tail`` ranks rows of at + most 64 candidates on warp 0 (``SELECT_LONG_SL3``; the same picks). ``dsl`` runs the 384-thread kernel + (``SELECT_LONG_DSL``; the same table; ``fast_tail`` does not apply). + """ + if ( + max_score_kv.dtype != torch.float32 + or not max_score_kv.is_cuda + or max_score_kv.dim() != 3 + or not 128 < max_score_kv.shape[1] <= MAX_BLOCKS + or n_valid_blocks.dtype != torch.int32 + or not n_valid_blocks.is_cuda + or n_valid_blocks.shape[0] < max_score_kv.shape[2] + ): + return None + num_kv_heads, _, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, 16), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + return (m3_select_long_dsl if dsl else m3_select_long)( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + fast_tail, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_attn_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_attn_lc.py new file mode 100644 index 000000000000..8e5027b71ffe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_attn_lc.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 long-context sparse decode attention in one launch. + +Replaces the generic chain's split-K producer (``_gqa_sparse_decode_kernel``, one block per chunk) and the LSE +merge (``_merge_topk_attn_out_kernel``, with its MXFP8 copy for the o_proj) once the top-k block table exists. + +Grid: one cluster of 16 CTAs per query token; CTA ``r`` attends the token's ``r``-th selected block (the Triton +producer's chunk ``r`` at 16 chunks), keeps the partial in shared memory, and after a cluster barrier merges head +``r`` from all 16 partials over DSMEM, writing the BF16 row and, optionally, its MXFP8 copy and swizzled UE8M0 +scales exactly as ``_store_mxfp8_head`` lays them out. + +Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and +PV on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +log-domain recurrence. Accumulation order inside the MMAs and the reductions differs, so results match to rounding, +not bit for bit. + +Schedule: the request length and the query rows (FP8 widened to FP16 in shared memory) are read before the PDL grid +wait, which overlaps them with the select ahead when it releases this grid early (modeling's +``SELECT_EARLY_MAX_ROWS``); +the split-K partials travel by ``st.async`` into each receiving CTA's own mbarrier instead of a cluster-wide barrier. +Neither changes a value or the merge order. + +Geometry: 16 query heads per KV head (TP4), one KV head, head_dim 128, page 128, top-k 16. The kernel is +``torch.ops.trtllm.minimax_m3_sparse_decode_lc``, called through this target's ``m3_sparse_decode_lc`` entry. +""" + +import torch + +from ..ops.attention.m3_sparse_decode_lc import m3_sparse_decode_lc + + +class PrefetchSlot: + """The read-only tensors the next long-context sparse decode prefetches into L2 (the attention layer's o_proj weight + and scales), set by the layer around its decode call and taken by the decode; empty means no prefetch.""" + + def __init__(self) -> None: + self.tensors = None + + def take(self): + tensors, self.tensors = self.tensors, None + return tensors + + +PREFETCH = PrefetchSlot() +# (this layer's id buffer, the next sparse layer's id buffer, its K / V pools, its K / V block scales) +NEXT_KV = PrefetchSlot() + + +def lc_applies(q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output) -> bool: + """Whether the step's operands fit the kernel: TP4 geometry, packed NVFP4 pages, FP8 q, BF16 [T, 16, 128] out.""" + if q.dim() != 3 or tuple(q.shape[1:]) != (16, 128) or q.dtype != torch.float8_e4m3fn: + return False + if q.stride(2) != 1 or q.stride(1) % 16 or q.stride(0) % 16 or q.data_ptr() % 16: + return False + if int(k_paged.shape[1]) != 1 or int(topk_idx.shape[0]) != 1 or int(topk_idx.shape[-1]) != 16: + return False + for t in (k_paged, v_paged): + if ( + t.stride(3) != 1 + or t.stride(2) != 64 + or int(t.shape[3]) != 64 + or t.data_ptr() % 16 + or t.stride(0) % 16 + ): + return False + for t in (k_block_scale, v_block_scale): + if t is None or t.stride(3) != 1 or t.stride(2) != 8 or t.data_ptr() % 8 or t.stride(0) % 8: + return False + return ( + output.dtype == torch.bfloat16 + and output.dim() == 3 + and output.is_contiguous() + and tuple(output.shape) == tuple(q.shape) + ) + + +def m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + linear_softmax=True, + mxfp8=None, + use_pdl=True, + early_trigger=True, + merge_once=False, + prefetch=None, + next_kv=None, + identity=False, + one_div=False, + identity_step=False, +): + """Drop-in for ``minimax_m3_sparse_attn_decode`` on the TP4 NVFP4 shape; callers check ``lc_applies`` first. + + ``topk_idx`` is the head-major [1, total_q, 16] view the Triton producer takes. ``mxfp8`` is ``(q8, sf)`` from + ``MERGE_MXFP8.buffers(2048, dev)`` to also emit the o_proj hand-off. The op is built with the early PDL + trigger only (``early_trigger=True``). ``merge_once`` runs this target's op with the merge weights formed once + per CTA (bitwise the same output); with it, ``prefetch`` (up to two contiguous tensors) is pulled into L2 by the + op's CTAs before their grid dependency. + """ + total_q = int(q.shape[0]) + q8, sf = ( + (mxfp8[0].view(torch.uint8), mxfp8[1].view(torch.uint8)) + if mxfp8 is not None + else (None, None) + ) + m3_sparse_decode_lc( + q.view(torch.uint8), + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + output.view(total_q, 2048), + q8, + sf, + 2048 // 128, + int(decode_query_len), + float(sm_scale), + bool(linear_softmax), + bool(use_pdl), + bool(early_trigger), + merge_once=bool(merge_once), + identity=bool(identity) and bool(merge_once), + identity_step=bool(identity_step) and bool(merge_once), + one_div=bool(one_div) and bool(merge_once), + prefetch=prefetch if merge_once else None, + next_kv=( + tuple(next_kv[:2]) + tuple(t.view(torch.uint8) for t in next_kv[2:]) + if merge_once and next_kv is not None + else None + ), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py new file mode 100644 index 000000000000..e89f2de8a8fb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M kernels (``DRAFT_LOWM``). + +The drafter's bias-free BF16 projections (qkv, o_proj, gate_up, down, the fc; not the head) call ``draft_linear`` +first: a covered row count runs ``ops/gemm/m3_draft_lowm`` at the tactic of its (N, K) table, every other call the +plain GEMM. ``prebuild`` compiles every covered call before the executor's warmup; a forward never compiles one, and a +covered call it finds uncompiled is a miss. +""" + +from __future__ import annotations + +from typing import FrozenSet, Iterable, Optional, Set, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm.logger import logger + +from ..ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for +from . import m3_report + +PATH = "m3_draft_lowm" + + +# ((N, K), M) of every call the target covers (its DRAFT_LOWM keys), set by prebuild. +_KEYS: FrozenSet[Tuple[Tuple[int, int], int]] = frozenset() +# (N, K, M, PDL) of every call prebuild compiled. +_BUILT: Set[Tuple[int, int, int, bool]] = set() +# (N, K, M) of every call that logged its engagement line. +_ENGAGED: Set[Tuple[int, int, int]] = set() +# DRAFT_AR_PUSH: the weight whose next low-M split-K call also pushes its rows into a boundary all-reduce +# (data pointer, the bucket's multicast words, its flags, rank, world, max tokens), set by arm_push; and whether the +# last armed call pushed: the data pointer of the output it pushed, 0 if none (read once by take_pushed). Host state +# set and read in one forward's issue order. +_PUSH_ARM: Optional[tuple] = None +_PUSHED = 0 + + +def arm_push(weight: torch.Tensor, workspace, rows: int) -> None: + """Arm ``weight``'s next covered split-K call to push into ``workspace``'s bucket for ``rows`` (see _PUSH_ARM).""" + global _PUSH_ARM, _PUSHED + from . import m3_mc_allreduce + + bucket = m3_mc_allreduce._bucket(workspace, rows) + flags = None if bucket is None else bucket.flags_for(8) + _PUSHED = 0 + _PUSH_ARM = None + if bucket is not None and flags is not None and bucket.max_tokens == 8: + _PUSH_ARM = ( + weight.data_ptr(), + bucket.mc.data_ptr(), + flags.data_ptr(), + bucket.rank, + bucket.world_size, + bucket.max_tokens, + ) + + +def take_pushed(x: torch.Tensor) -> bool: + """Whether the armed call pushed its rows; disarms. ``x`` is what the boundary all-reduce is about to reduce: it + must be the pushing call's own output (the prepushed all-reduce reads only the pushed words, never ``x``), so an + operation between that call and the all-reduce raises instead of silently reducing the wrong rows.""" + global _PUSH_ARM, _PUSHED + pushed, _PUSHED, _PUSH_ARM = _PUSHED, 0, None + if pushed and x.data_ptr() != pushed: + raise RuntimeError( + "the drafter's pushed rows are not the boundary all-reduce's input (an operation ran in between)" + ) + return bool(pushed) + + +def draft_linear( + weight: torch.Tensor, x: torch.Tensor, bias: Optional[torch.Tensor] +) -> Optional[torch.Tensor]: + """x [M, K] @ weight^T through the low-M kernel when the target covers it (``prebuild``'s keys); None for the + plain GEMM. The split-K ('s') tactics run this target's copy of FlashInfer's kernel (ops/m3_lowm_splitk, bit + for bit the same output at every certified cell), which also carries the drafter-boundary push epilogue.""" + if bias is not None or x.dim() != 2: + return None + n, k = weight.shape + m = int(x.shape[0]) + if ((n, k), m) not in _KEYS: + m3_report.uncovered(PATH, f"the drafter's {n}x{k} linear past its tactic table's rows") + return None + if not x.is_contiguous(): + # The first pass's fc may read a column window of the captured hidden states; the plain GEMM takes strides. + m3_report.uncovered(PATH, "a strided activation (the plain GEMM reads it in place)") + return None + use_pdl = get_env_enable_pdl() + if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + m3_report.missed(PATH, m, "the activation or weight is not BF16") + return None + if (n, k, m, use_pdl) not in _BUILT: + m3_report.missed(PATH, m, f"the {n}x{k} kernel at {m} rows was not compiled before warmup") + return None + kind, tactic = tactic_for(n, k, m) + if (n, k, m) not in _ENGAGED: + _ENGAGED.add((n, k, m)) + logger.info( + f"[M3 sm_107 tp4] fast path engaged: m3_draft_lowm {n}x{k} at {m} rows (device {x.device.index})" + ) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + global _PUSHED + ar = None + if _PUSH_ARM is not None and _PUSH_ARM[0] == weight.data_ptr() and kind == "s": + ar = _PUSH_ARM[1:] + _PUSHED = out.data_ptr() + return m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, True, ar) + + +def installed(weight: torch.Tensor) -> bool: + """Whether the target covers ``weight`` (bias-free BF16 [N, K]) at some row count (``prebuild``'s keys).""" + if weight.dtype != torch.bfloat16 or weight.dim() != 2: + return False + return any(nk == tuple(weight.shape) for nk, _ in _KEYS) + + +def prebuild(weights: Iterable[torch.Tensor], keys: Iterable[Tuple[Tuple[int, int], int]]) -> int: + """Take the target's covered ((N, K), M) ``keys``, then compile (and run once) every covered row count of every + covered drafter weight, outside capture, on the current stream. Returns the number of calls compiled.""" + global _KEYS + _KEYS = frozenset(keys) + use_pdl = get_env_enable_pdl() + built = 0 + shapes = {tuple(weight.shape): weight for weight in weights if installed(weight)} + for (n, k), weight in sorted(shapes.items()): + for nk, m in sorted(key for key in _KEYS if key[0] == (n, k)): + if (n, k, m, use_pdl) in _BUILT: + continue + kind, tactic = tactic_for(n, k, m) + x = torch.zeros(m, k, dtype=torch.bfloat16, device=weight.device) + out = torch.empty(m, n, dtype=torch.bfloat16, device=weight.device) + m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, True) + if kind == "s": + # The push variant a boundary all-reduce may arm (o_proj / down): compiled only, never run here + # (its pushes need the all-reduce's exchange and a matching call). + from .m3_lowm_splitk import compile_push + + compile_push(x, weight.t(), out, use_pdl, tactic) + _BUILT.add((n, k, m, use_pdl)) + built += 1 + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py new file mode 100644 index 000000000000..45a64dfe9210 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 draft pass for the MiniMax-M3 drafter's vocab-sharded BF16 head (TP decode, <= 8 requests). + +One kernel replaces a draft pass's head and the draft loop's per-pass bookkeeping, at the in-tree precision: + +* The drafter's BF16 lm_head shard is kept a second time, built by ``prepare_head`` before the executor's warmup: in + the lossless 12-bit coding of the target head (``ops.gemm.m3_target_head``; 75% of the bytes, kept only when it + decodes to the BF16 shard bit for bit), else in ``mma.m16n8k16`` A-fragment order (``bf16_fragments``, the same + values). The kernel streams it with BF16 MMAs and FP32 accumulation over the BF16 head input rows; a coded + exception decodes to +0 and its exact product x * w is added at its tile's flush. +* Each logit is rounded to BF16 (the in-tree logits are the BF16 head GEMM's output) and keyed + ``ordered(value) << 32 | ~global_index`` with -0 as +0, so the larger value wins and ties take the lowest index, as + the in-tree ``torch.max`` / ``torch.argmax`` over the gathered logits (``greedy_sample_draft_with_tp_gather``). +* The last CTA exchanges each row's key across the TP ranks through a multicast mapping (Lamport slots, see + ``m3_shard_accept``), then stores the draft token, the next pass's input id and position, and advances the draft KV + lengths as the in-tree draft loop does between passes. + +The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu) and called through the +catalog entry ``ops.spec.m3_draft_step``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok +from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments, max_tile_pieces +from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op +from . import m3_report +from .m3_shard_accept import AcceptWorkspace, create_workspace + +# The draft step launches this many CTAs instead of one per SM (3126 tiles: 200 CTAs x 15 + 9 x 14, at most 2 warp +# pieces per tile), so the CTAs own near-equal whole-tile counts and the stream does not end on the tiles' straggler +# tail (-1.04 / -1.18 us per pass at 1 / 4 rows vs 212 on Rubin TP4). A different static warp segmentation: +# exact-class. +GRID = 209 # capacity: fixed: 3126 head tiles of the TP4 shard, at most 2 warp pieces per tile on 209 CTAs +# The draft head kernel sums a tile that several warps stream in warp order (fixed by the launch shape) instead of by +# shared-memory atomics in arrival order: with fewer than 16 tiles per CTA a tile can take 3-4 pieces and the atomics' +# order changed the fp32 sum run to run; with 2 pieces the sum is the atomics' bit for bit. +FIXED_ORDER_PATH = "m3_draft_step_fixed_order" + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the 12-bit coded stream + with its exception lists (``coded``), or the fragment-ordered BF16 copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + coded: bool = False + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + exp_base: int = 0 + + +def head_supported(lm_head) -> bool: + """A bias-free BF16 head whose local shard the kernel can own: whole 16-row tiles, K a multiple of 16, at most + MAX_TILES_PER_CTA tiles per CTA of a one-CTA-per-SM grid.""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + if getattr(lm_head, "bias", None) is not None: + return False + n, k = weight.shape + if n % 16 or k % 16: + return False + return (n // 16 + _grid() - 1) // _grid() <= MAX_TILES_PER_CTA + + +def prepare_head(lm_head, tp_rank: int, coded: bool) -> DraftHead: + """The streamed copy of ``lm_head``'s local shard (global ids from tp_rank * rows): its 12-bit coding when + ``coded``, K is a multiple of 64 and the coding decodes to the shard bit for bit, else (fail + closed) the fragment-ordered BF16 copy. + Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_draft_head``).""" + weight = lm_head.weight + offset = tp_rank * weight.shape[0] + with torch.no_grad(): + w = weight.detach() + if coded and w.shape[1] % 64 == 0: + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return DraftHead( + weight=weight, + wf=packed, + vocab_offset=offset, + coded=True, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + exp_base=base, + ) + del packed, exc_off, exc_pos, exc_w + wf = bf16_fragments(w) + return DraftHead(weight=weight, wf=wf, vocab_offset=offset) + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _grid() -> int: + """The draft step's CTA count: ``GRID`` (at most one per SM).""" + return min(GRID, _num_sms()) + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace ([2][MAX_TOKENS][world] slots), created collectively on the first eager call; None + during a CUDA-graph capture that finds it missing.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = create_workspace(mapping, device, MAX_TOKENS) + return _WS + + +def _fixed_order( + rows: int, head: DraftHead, x: torch.Tensor, num_ctas: int, split: bool = False +) -> bool: + """The fixed warp-order tile sum where the launch lets a tile take more than 2 warp pieces (max_tile_pieces; a + split or equal-chunk launch always), with its engagement line; False keeps the atomics, deterministic on this + shape.""" + eq = getattr(head, "bnd", None) is not None + if not (split or eq): + chunks = x.shape[1] // (16 if head.exc_off is None else 64) + if ( + max_tile_pieces( + head.wf.shape[0] if head.exc_off is None else head.exc_off.numel() - 1, + num_ctas, + chunks, + ) + <= 2 + ): + return False + m3_report.taken( + FIXED_ORDER_PATH, + rows, + "the draft head sums shared tiles in warp order", + ) + return True + + +def draft_step( + head: DraftHead, + x: torch.Tensor, + rows: int, + x_rows: Optional[torch.Tensor], + mapping, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, +) -> bool: + """One fused greedy draft pass over ``rows`` requests (see the module docstring); False when the exchange + workspace is missing during a capture. Collective over the TP group. + + Args: + head: the prepared head of this rank (``prepare_head``), coded or BF16 fragments. + x: [*, K] BF16 head input rows (the drafter's final-norm output); row ``x_rows[b]`` (int64) or ``b`` is + request b's. + draft_out: int32 [rows, >= step + 1], unit column stride: column ``step`` receives the tokens. + ids_next / pos_next: int32 [rows]: the tokens, and ``pos_src[x_rows[b] or b] + 1`` (pos_src int32). + kv_lens: int32 draft KV lengths: ``-= draft_len - num_acc`` with ``seq_lens = 1`` on pass 0 when ``num_acc`` + is given, ``+= 1`` on later passes. + """ + ws = workspace(mapping, x.device) + if ws is None: + return False + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + _grid(), + True, + None, + head.exc_off, + head.exc_pos, + head.exc_w, + head.exp_base, + None, # no fold tail: the forward's tail runs the standalone one + fixed_order=_fixed_order(rows, head, x, _grid()), + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py new file mode 100644 index 000000000000..7936d2497731 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py @@ -0,0 +1,405 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's decode attention on Rubin: ``M3DrafterFullCtxFmha``. + +The drafter layer is a plain ``TrtllmAttention`` over its own FP8 paged cache (the full context, every draft pass). +On SM 10.7 its generation phase runs ``thop.attention``'s in-tree trtllm-gen runner, whose split is capped at 8 CTAs +per (request, KV head): at TP4 (1 KV head per rank) 45.4 us per pass at 165k tokens on 212 SMs. + +This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append +into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context +PrimTS split-KV decode of ``ops/m3_fullctx_decode`` at the same pinned recipes as the dense layers. + +With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead +(``minimax_m3_sm107tp4_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page +pool and block table taken from the KV cache pool pointers and mapping the preprocess reads. + +It serves generation-only steps of the covered (query rows per request, requests) shapes; the target installs it +first in the drafter layer's FMHA list (``install``), so every other step falls through to the stock libraries. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + FlashInferTrtllmGenFmha, +) +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode + +from . import m3_fullctx_decode as fullctx + +# With PDL, the fused RoPE launch takes the maximum shared-memory carveout so the drafter decode's large-smem CTAs +# co-reside with it (timing only). +ROPE_CARVEOUT = True + + +def fwd_has_no_mrope(fwd) -> bool: + return getattr(fwd, "mrope_position_deltas", None) is None + + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + +class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): + """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + + def __init__( + self, + attn: "TrtllmAttention", + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + pdl: bool = False, + on_rope_taken=None, + ) -> None: + super().__init__(attn) + self._recipes = recipes + self._max_kv_len = max_kv_len + self._own_reducer = own_reducer + self._fused_rope_mode = fused_rope_mode + self._pdl = pdl + self._on_rope_taken = on_rope_taken + self._on_taken = on_taken + self._on_uncovered = on_uncovered + self._on_missed = on_missed + + def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: + nreq = int(metadata.num_generations) + if nreq <= 0 or int(q.shape[0]) % nreq: + return None + return int(q.shape[0]) // nreq, nreq + + def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: + if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: + return False + if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: + return False + attn = self.attn + if ( + attn.num_heads != fullctx.HEADS + or attn.num_kv_heads != fullctx.KV_HEADS + or attn.head_dim != fullctx.HEAD_DIM + or not QuantMode(attn.quant_mode).has_fp8_kv_cache() + ): + self._on_uncovered("not the TP4 FP8 drafter geometry") + return False + # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's + # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. + if ( + forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant + or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig + ): + self._on_uncovered("per-step KV cache scales") + return False + if int(metadata.max_seq_len) > self._max_kv_len: + self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") + return False + shape = self._shape(q, metadata) + if shape is None or shape not in self._recipes: + self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") + return False + # Compiled before warmup (``install``); the forward never compiles one. + if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: + self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") + return False + return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) + + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + attn = params.attn + width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim + rotary_dim = int(attn.rope_params.dim) + return ( + self._fused_rope_mode is not None + and params.qkv_input is not None + and attention_input.dim() == 2 + and attention_input.dtype == torch.bfloat16 + and attention_input.is_contiguous() + and int(attention_input.shape[1]) == width + and not params.is_cross + and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) + and attn.rotary_cos_sin is not None + and attn.rotary_cos_sin.dtype == torch.float32 + and fwd_has_no_mrope(params.fwd) + and 0 < rotary_dim <= attn.head_dim + and rotary_dim % 16 == 0 + and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) + # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo + # runs for several requests only) and reads the workspace's, so it keeps that path. + and not (self._multi_token(params) and int(params.batch_size) == 1) + ) + + @staticmethod + def _multi_token(params: FmhaParams) -> bool: + return ( + params.spec_decoding_generation_lengths is not None + and int(params.attn.predicted_tokens_per_seq) > 1 + ) + + def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" + if not self._fused_rope_applies(params, attention_input): + return False + attn = params.attn + meta = params.meta + qlen, nreq = int(params.input_seq_length), int(params.batch_size) + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + if plan is None: + return False + # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary + # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the + # block table is the pool's [requests, 2, blocks] K / V sub-page indices. + pool_index, layer_in_pool = ( + int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] + ) + tpb = int(params.tokens_per_block) + sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim + base = ( + int(meta.host_kv_cache_pool_pointers[pool_index, 0]) + + layer_in_pool * int(params.kv_factor) * sub_page_bytes + ) + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) + ) + tables = meta.kv_cache_block_offsets[ + pool_index, params.seq_offset : params.seq_offset + nreq + ] + multi = self._multi_token(params) + if not torch.cuda.is_current_stream_capturing(): + # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess + # path checks; the op itself writes V at the table's V sub-page. + tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(tables.shape[-1], device=tables.device) + in_use = cols[None, :] < tpb_held[:, None] + offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + q = torch.empty( + int(params.num_tokens), + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=attention_input.device, + ) + + torch.ops.trtllm.minimax_m3_sm107tp4_drafter_rope_fp8_kv( + attention_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + int(attn.rope_params.dim), + tpb, + self._fused_rope_mode, + self._pdl, + ROPE_CARVEOUT, + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + tables[:, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + if self._on_rope_taken is not None: + self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") + self._on_taken(params.num_tokens) + return True + + def run_generation(self, params: FmhaParams) -> None: + attention_input = params.qkv_input if params.qkv_input is not None else params.query_input + if attention_input is None: + raise RuntimeError("the drafter decode requires QKV or query input.") + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + if self._run_fused_rope(params, attention_input): + return + # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 + # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). + fp8_context_fmha = False + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + _fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + _window_left, + is_multi_token_gen, + ) = thop.trtllm_gen_generation_preprocess( + attention_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + fp8_context_fmha, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + ) + qlen, nreq = params.input_seq_length, params.batch_size + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the + # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row + # pass is never multi-token. + if plan is None or (is_multi_token_gen and qlen == 1): + raise RuntimeError( + f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" + ) + rows = params.num_tokens * attn.num_heads * attn.head_dim + q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) + q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) + # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 + # of the pool (the two views have the same page count, as the kernel requires). + if not torch.cuda.is_current_stream_capturing(): + # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. + lengths = params.spec_decoding_generation_lengths + if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): + raise RuntimeError( + f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" + ) + # Only the blocks a request holds are laid out; the table's tail entries are unset. + tpb = int(params.tokens_per_block) + held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(block_tables.shape[-1], device=block_tables.device) + in_use = cols[None, :] < held[:, None] + offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. + # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + block_tables[:nreq, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + self._on_taken(params.num_tokens) + + +def install( + draft_model: torch.nn.Module, + recipes, + max_kv_len: int, + own_reducer: bool, + pdl: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + on_rope_taken=None, +) -> int: + """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention + backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it + was installed on.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + count = 0 + device = next(draft_model.parameters()).device + # The backend is not a module: each attention module of the drafter holds it as ``attn``. + for module in draft_model.modules(): + backend = getattr(module, "attn", None) + if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): + continue + if not bool((backend.kv_cache_scaling_factor == 1.0).all()): + on_uncovered("a scaled FP8 KV cache") + continue + fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) + libs = backend._fmha_manager.fmha_libs + if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): + libs.insert( + 0, + M3DrafterFullCtxFmha( + backend, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode=fused_rope_mode, + pdl=pdl, + on_rope_taken=on_rope_taken, + ), + ) + count += 1 + return count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py new file mode 100644 index 000000000000..79aff5cb80ce --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The speculative step's store and next-step staging at the tail of the forward (one Triton launch). + +For each of the first ``n`` rows (``n`` read from device memory, so a captured graph whose step was not armed writes +nothing), with ``s = slots[i]``: + +* the spec sampler's per-slot store, exactly as ``m3_spec_tail.slot_scatter`` writes it (zero-padded to each store + width): ``new[s][c]``, ``next_new[s][c]``, ``draft[s][c]`` and ``lens[s]``; +* the engine's next-step inputs, exactly as ``m3_spec_tail.step_gather`` writes them from that store for a step whose + rows are these rows in this order (no new tokens, first previous row 0): ``ids[i*T + t] = next_new[s][t]``, + ``draft_ids[i*D + d] = draft[s][d]``, ``pos_off[i*T + t] = lens[s]``, ``kv_off[i] = lens[s] - T``. + +Integer copies only: the values are the ones the sampler's and the engine's own copies write. + +With ``cpp`` the same writes run as this target's C++ op (trtllm::minimax_m3_sm107_tp4_fold_tail, +minimaxM3Sm107Tp4/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fold_tail_kernel( + ctrl, + slots, + src_new, + sn_s0, + sn_s1, + SN_W: tl.constexpr, + dst_new, + dn_s0, + dn_s1, + DN_W: tl.constexpr, + src_next, + sx_s0, + sx_s1, + SX_W: tl.constexpr, + dst_next, + dx_s0, + dx_s1, + DX_W: tl.constexpr, + src_draft, + sd_s0, + sd_s1, + SD_W: tl.constexpr, + dst_draft, + dd_s0, + dd_s1, + DD_W: tl.constexpr, + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + T: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(ctrl) + if i < n: + s = tl.load(slots + i).to(tl.int64) + c = tl.arange(0, BLOCK) + new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) + tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) + nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) + tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) + drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) + tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) + ln = tl.load(src_lens + i) + tl.store(dst_lens + s, ln) + # The next step's inputs read the store's padded rows: zero past each source width. + tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) + tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) + tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) + tl.store(kv_off + i, ln - T) + + +def fold_tail( + ctrl: torch.Tensor, + slots: torch.Tensor, + rows: int, + src_new: torch.Tensor, + dst_new: torch.Tensor, + src_next: torch.Tensor, + dst_next: torch.Tensor, + src_draft: torch.Tensor, + dst_draft: torch.Tensor, + src_lens: torch.Tensor, + dst_lens: torch.Tensor, + ids: torch.Tensor, + draft_ids: torch.Tensor, + pos_off: torch.Tensor, + kv_off: torch.Tensor, + tokens_per_step: int, + cpp: bool = False, +) -> bool: + """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs + are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width].""" + two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) + one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) + if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): + return False + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) + or slots.dtype != torch.long + ): + return False + t_step = tokens_per_step + d_step = t_step - 1 + if ( + rows <= 0 + or slots.numel() < rows + or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) + or ids.numel() < rows * t_step + or pos_off.numel() < rows * t_step + or draft_ids.numel() < rows * d_step + or kv_off.numel() < rows + ): + return False + widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) + widths += (src_draft.shape[1], dst_draft.shape[1], t_step) + if cpp and max(widths) <= 32: + torch.ops.trtllm.minimax_m3_sm107_tp4_fold_tail( + ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, + src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, + ) # fmt: skip + return True + block = triton.next_power_of_2(max(widths)) + _fold_tail_kernel[(rows,)]( + ctrl, + slots, + src_new, + src_new.stride(0), + src_new.stride(1), + src_new.shape[1], + dst_new, + dst_new.stride(0), + dst_new.stride(1), + dst_new.shape[1], + src_next, + src_next.stride(0), + src_next.stride(1), + src_next.shape[1], + dst_next, + dst_next.stride(0), + dst_next.stride(1), + dst_next.shape[1], + src_draft, + src_draft.stride(0), + src_draft.stride(1), + src_draft.shape[1], + dst_draft, + dst_draft.stride(0), + dst_draft.stride(1), + dst_draft.shape[1], + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + t_step, + d_step, + block, + num_warps=1, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py new file mode 100644 index 000000000000..be04f26f9aff --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's +``force_num_accepted_tokens``), in one launch. + +Triton: ``force_vec_kernel``: one CTA bumps the shared step counter and draws each request's count from the fixed-seed +pool with the in-tree (counter, slot) hash, exactly as ``_apply_force_accepted_tokens`` does. ``force_vec_req_kernel`` +(``FORCE_AL_PER_REQUEST``): each request draws from the same pool at its own key, its KV length at this verify. +Reached through the ``ops/spec`` entries ``m3_force_vec`` and ``m3_force_vec_req``. + +One compiled variant for every batch and accept count: specializing them (== 1, divisible by 16) would compile +mid-serving the first time a step reaches a new one, past the warm-up. +""" + +import triton +import triton.language as tl + + +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_kernel( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK: tl.constexpr, +): + c = tl.load(counter) + 1 + slots = tl.arange(0, BLOCK) + tl.store(counter, c) + idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask + live = slots < batch_size + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) + + +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_req_kernel( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK: tl.constexpr, +): + # One draw per request from its own key (its KV length at this verify): no shared counter, so a request's forced + # counts depend only on its own prompt and progress, not on what else the server ran before or beside it. + slots = tl.arange(0, BLOCK) + live = slots < batch_size + k = tl.load(keys + slots, mask=live, other=0).to(tl.int64) + idx = (k * key_stride) & pool_mask + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py new file mode 100644 index 000000000000..3dc0fc4a3888 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 full-context FP8 decode attention on Rubin: ``fullctx_decode``. + +The dense layers (0-2) attend every token of their FP8 E4M3 (unscaled) paged KV cache. On SM 10.7 the stock call +(flashinfer trtllm-gen, ``minimax_m3_trtllm_gen_dense_decode``) caps its split at 16 cluster CTAs per (request, KV +head): at TP4 (16 query heads and 1 KV head per rank) one request runs on 16 CTAs of 212 and takes 33.4-33.8 us +per call at 165k tokens. + +This runs the same arithmetic -- FP8 Q (the same E4M3 rounding of q), FP8 K/V, P quantized to E4M3 per split -- on +the vendored PrimTS decode kernel (``attention/backends/prims_ts``, SM 10.7 through the sm_100f family) at a pinned +split-KV recipe, which fans one request out over 96-108 CTAs. The split partials are merged by +``torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce`` (``cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/ +fullCtxReduce.cu``, head-count generic, reused by name): unnormalized fp32 accumulation over the active splits, one +final division, one BF16 rounding, the merge form of stock trtllm-gen. Only the split changes against stock +(``modeling.FULLCTX_DECODE_RECIPES``). + +A plan is compiled per (device, query rows per request, requests) by ``prebuild`` before the executor's warmup; the +forward only reads plans (``plan_for``). The plan is built from PrimTS's own config, compile and workspace helpers at +an explicit recipe; the PrimTS wrapper's planning policy is not used, so no PrimTS state is replaced. With +``own_reducer`` off the PrimTS reducer merges into FP16 and the result is cast to BF16. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Mapping, Optional, Tuple + +import torch + +# The TP4 per-rank geometry: 16 query heads over 1 KV head. +HEADS = 16 +KV_HEADS = 1 # capacity: fixed: the tp4 KV heads per rank +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim +PAGE = 128 +# The per-request K/V length PrimTS's config selection sizes its policy for (Q/KV tiles, launch mode); the pinned +# recipes were measured with it. The compiled capacity (``max_kv_len``) may be larger. +POLICY_KV_LEN = ( + 1 << 20 +) # capacity: fixed: the PrimTS policy's K/V length bound (the plans' compile shape) +HEADS_PER_KV = HEADS // KV_HEADS # capacity: fixed: the tp4 query heads per KV head + +_plans: Dict[Tuple[int, int, int, bool], "FullCtxPlan"] = {} +_lock = threading.Lock() + + +class FullCtxPlan: + """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" + + def __init__( + self, + device: torch.device, + qlen: int, + nreq: int, + max_kv_len: int, + recipe: Mapping, + own_reducer: bool, + pdl: bool, + ): + from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec + from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( + get_max_active_clusters_for_cluster_size, + make_decode_config, + ) + + self._tsdec = tsdec + self.qlen = qlen + self.nreq = nreq + self.pdl = pdl + self.recipe = dict(recipe) + splits = int(self.recipe["splits"]) + fp8 = tsdec._cutlass_dtype("float8_e4m3fn") + # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer + # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents + # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an + # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the + # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own + # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge + # (minimax_m3_sm107tp4_fullctx_reduce) waits as its first instruction. + cfg = make_decode_config( + headdim=HEAD_DIM, + args={"use_pdl": pdl, **self.recipe.get("config", {})}, + seq_len_q=qlen, + seq_len_kv=min(max_kv_len, POLICY_KV_LEN), + batch_size=nreq, + num_heads_q=HEADS, + num_heads_kv=KV_HEADS, + q_dtype=fp8, + k_dtype=fp8, + v_dtype=fp8, + o_dtype=tsdec._cutlass_dtype("float16"), + qkv_layout="pagedKv", + num_tokens_per_page=PAGE, + storage_tokens_per_page=PAGE, + split_kv_mode="gmem_reduction_with_separate_kernel", + splits_kv=splits, + max_splits_kv=splits, + min_loop_iters_per_split=1, + mask_type="causal" if qlen > 1 else "dense", + auto_tuner=True, + split_kv=True, + ) + _check_layout(cfg, qlen, splits) + with torch.cuda.device(device): + max_active = get_max_active_clusters_for_cluster_size(1) + spec = tsdec._decode_launch_spec_from_config( + cfg, + batch_size=nreq, + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + seq_len_q=qlen, + max_active_clusters=int(max_active), + ) + compile_spec = tsdec._make_decode_compile_spec( + spec, + device_index=_index(device), + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + page_size=PAGE, + max_kv_len=max_kv_len, + seq_len_q=qlen, + q_dtype_key="float8_e4m3fn", + k_dtype_key="float8_e4m3fn", + v_dtype_key="float8_e4m3fn", + output_dtype_key="float16", + use_packed_q=False, + kv_prefix_mode="dynamic", + kv_lengths_mode="dynamic", + ) + self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) + layout = tsdec._make_decode_workspace_layout( + spec.scratch_shapes, + torch.float16, + use_separate_reduction_kernel=True, + use_split_kv=True, + ) + self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) + self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) + self.workspace.split_kv_counter.zero_() + self.workspace.cu_seqlens_q.zero_() + self.workspace.attention_sinks.zero_() + rows = HEADS_PER_KV * qlen + if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): + raise RuntimeError( + f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " + f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" + ) + self.own_reducer = own_reducer + self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) + self.tile_kv = int(cfg.tile_size_kv) + self.insts_kv = int(cfg.num_insts_kv) + self.causal = qlen > 1 + shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) + self.staging = torch.empty(shape, dtype=torch.float16, device=device) + + def run( + self, + q: torch.Tensor, + k_pages: torch.Tensor, + v_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + mxfp8: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> bool: + """Decodes into output; with mxfp8 (E4M3 rows, UE8M0 scales) and a merge op that takes them, the merge + also writes the MXFP8 copy of output. Returns whether it did.""" + # The kernel takes a compact query; a column view of the fused QKV projection is copied once. + q = q.reshape(self.staging.shape).contiguous() + runtime = self._tsdec._prepare_decode_runtime_unchecked( + q, + (k_pages, v_pages), + None, + output_dtype=torch.float16, + bmm1_scale=sm_scale, + bmm2_scale=1.0, + out=self.staging, + ) + ws = self.workspace + self._main( + runtime.q, + runtime.k_cache, + runtime.v_cache, + runtime.k_sf_cache, + runtime.v_sf_cache, + runtime.out, + seq_lens, + ws.cu_seqlens_q, + block_table, + block_table, # Dense decode constexpr-elides the membership table. + ws.partial_o, + ws.partial_stats, + ws.split_kv_counter, + ws.attention_sinks, + runtime.num_physical_pages, + runtime.k_page_stride, + runtime.k_head_stride, + runtime.k_token_stride, + runtime.v_page_stride, + runtime.v_head_stride, + runtime.v_token_stride, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + if self.own_reducer: + args = ( + ws.partial_o, + ws.partial_stats, + seq_lens, + output, + self.qlen, + self.q_tokens_per_cta, + self.tile_kv, + self.insts_kv, + self.causal, + self.pdl, + ) + if mxfp8 is not None and emits_mxfp8(): + torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce(*args, mxfp8[0], mxfp8[1]) + return True + torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce(*args) + return False + self._primts_reducer( + runtime.out, + seq_lens, + ws.cu_seqlens_q, + ws.partial_o, + ws.partial_stats, + ws.attention_sinks, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + output.view(self.staging.shape).copy_(self.staging) + return False + + +def _check_layout(cfg, qlen: int, splits: int) -> None: + """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. + + Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is + ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding + window, attention sinks or packed Q. + """ + problems = [] + if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): + problems.append("not split-KV with the separate GMEM reducer") + if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: + problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") + if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: + problems.append("cluster-SMEM reduction or a persistent scheduler") + if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: + problems.append("packed Q, a sliding window or attention sinks") + if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): + problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") + if qlen > 1: + if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): + problems.append("multi-token Q without the grouped [token, head] tile layout") + elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): + problems.append( + f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" + ) + if problems: + raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + + +def emits_mxfp8() -> bool: + """Whether the loaded merge op can also write the o_proj MXFP8 hand-off (its schema takes q8 / sf).""" + op = torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce + return "q8" in str(op.default._schema) + + +def has_own_reducer() -> bool: + """Require the BF16 split-KV merge op supplied by this target build.""" + torch.ops.trtllm.minimax_m3_sm107tp4_fullctx_reduce + return True + + +def plan_for( + device: torch.device, qlen: int, nreq: int, own_reducer: bool = True +) -> Optional[FullCtxPlan]: + """The prebuilt plan for this shape, or None.""" + return _plans.get((_index(device), qlen, nreq, own_reducer)) + + +def prebuild( + device: torch.device, + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + pdl: bool, +) -> int: + """Compile the plan of every (query rows per request, requests) in ``recipes``. Call before warmup.""" + built = 0 + for (qlen, nreq), recipe in recipes.items(): + key = (_index(device), qlen, nreq, own_reducer) + with _lock: + if key not in _plans: + _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) + built += 1 + return built + + +def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: + """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" + return pool[0::subpages_per_slot], pool[1::subpages_per_slot] + + +def _index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py new file mode 100644 index 000000000000..fad75d44694a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py @@ -0,0 +1,638 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with two schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer``) through ``IndexScoreRunner`` +and called through this target's ``m3_index_score`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, griddepcontrol_wait + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # No explicit launch_dependents: the only consumer is the block selector, + # which reduces over every block this grid writes, so triggering it early + # would leave it resident and spinning in its own wait. The implicit + # trigger as the CTAs exit still gives it the launch overlap. + early_q = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + hint_blocks = cute.ceil_div(seq_lens[batch_id], BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, BLOCK_K * head_dim * (dtype.width // 8) + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # seq_lens must be read after the wait. It is a graph-stable buffer that + # on_update_kv_lens patches every step, and num_blocks bounds an unmasked + # block_table load and an unmasked score store, so a stale length is an + # out-of-bounds write rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score_t2sc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score_t2sc.py new file mode 100644 index 000000000000..52a79fbffca9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score_t2sc.py @@ -0,0 +1,682 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule: T2SC, this target's copy of +the sm_107 TP2 / TP8 targets' scorer (``modeling.INDEX_SCORE_T2SC``). The sm_103-derived ``ops/m3_index_score`` +stays the kill path. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with three schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. +* ``prestage_old`` (with ``early_issue``, ``modeling.INDEX_SCORE_T2SC_PRESTAGE_OLD``): when the CTA's first block holds + only tokens older than this step (block < (seq_lens hint - decode query length) / 128), its K tile is TMA-loaded into + its stage before the wait; Q (and a first block that is not all old) is still loaded after it. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer_t2sc``) through +``IndexScoreRunner`` and called through this target's ``m3_index_score_t2sc`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import ( + TRTLLM_ENABLE_PDL, + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + prestage_old: bool = False, + early_trigger: bool = False, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K + # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's + # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. + self.prestage_old = prestage_old and early_issue + # early_trigger (on for this target): release the dependent selector right after this CTA's + # setup, so its launch and residency overlap the scoring; it still waits for this whole grid. + self.early_trigger = early_trigger + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # Without early_trigger the selector launches at this grid's implicit trigger (the CTAs' exit). With it, the + # selector is already resident in its own grid-dependency wait, which still returns only once every block + # this grid writes is complete. + early_q = False + early_k = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + seqlen_hint = seq_lens[batch_id] + hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + if cutlass.const_expr(self.prestage_old): + # Stage the first K tile now when every token of the block predates this step. + if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: + early_k = True + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + if not early_k: + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # The length that bounds this grid's work is read after the wait (the + # pre-wait reads above are prefetch / early-issue hints only, clamped to + # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens + # patches every step, and num_blocks bounds an unmasked block_table load + # and an unmasked score store, so a stale length is an out-of-bounds write + # rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + prestage_old=early, + early_trigger=True, + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py new file mode 100644 index 000000000000..11ef3a496132 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py @@ -0,0 +1,1249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +This target's copy of FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the Eagle3 +drafter's low-M linears, ``ops/m3_draft_lowm``): the same kernel and the same split order at every tactic, +bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's final +cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, the +SwiGLU prologue that follow. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 # capacity: fixed: a kernel tile constant + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 # capacity: fixed: an alignment + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 # capacity: fixed: pipeline stages +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 # capacity: fixed: pipeline stages + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= _MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES + base = SplitKTactic(mma_m, mma_n, split_k, min_stages) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(min_stages, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 # capacity: fixed: a split index, not a size + + +# The drafter-boundary push (``ar_push``, modeling.DRAFT_AR_PUSH): the exchange words of +# ops/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.s32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + llvm.inline_asm( + None, + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + ar_push: bool = False, + ) -> None: + self.ar_push = ar_push + self.acc_dtype = cutlass.Float32 + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64 = cutlass.Int64(0), + ar_flags: cutlass.Int64 = cutlass.Int64(0), + ar_rank: Int32 = Int32(0), + ar_world: Int32 = Int32(1), + ar_maxt: Int32 = Int32(8), + ar_tokens: Int32 = Int32(0), + ar_n: Int32 = Int32(0), + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel( + a, b, c, bias, ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_tokens, ar_n + ).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + ar_ws: cutlass.Int64, # ar_push: the boundary all-reduce's multicast exchange words (int32) + ar_flags: cutlass.Int64, # ar_push: its per-(token, cluster CTA) call flags (parity = flag & 1) + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, # the bucket's tokens + ar_tokens: Int32, # rows pushed (public M) + ar_n: Int32, # public N (the hidden size): the exchange row width + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + # ar_push: the owner's BF16 tile (kernel-M = output column, kernel-N = token), M-major. + sD = cute_ext.allocate( + mC.element_type, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=16, + ) + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sD, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl and not self.ar_push): + # ar_push: the dependent is a prepushed all-reduce, which skips its grid wait: it is released only after + # this grid's own grid wait and its pushes, in the owner's _store_tile (the prepushed >= 2 contract). + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + 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) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sD: cute.Tensor, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + self._store_tile(thr_t2r, rAcc, rD, gD_epi, c_dtype) + if cutlass.const_expr(self.ar_push): + self._push_tile( + thr_t2r, + rD, + sD, + epi_tid, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + @cute.experimental.jit + def _store_tile(self, thr_t2r, rAcc, rD, gD_epi, c_dtype: cutlass.Constexpr): + """The owner rank's reduced FP32 tile: one cast to ``c_dtype`` and the store to global memory.""" + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + @cute.experimental.jit + def _push_tile( + self, + thr_t2r, + rD, + sD, + epi_tid: cutlass.Int32, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output + would store them (ops/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" + cute_ext.partition_and_copy(thr_t2r, rD, sD) + cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + pairs = self.cta_m // 2 + col0 = bidx * self.cta_m + half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) + for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): + word = epi_tid + w0 + t = word // pairs + p = word % pairs + if t < ar_tokens: + n = col0 + 2 * p + par = _ld_volatile_s32( + ar_flags + cutlass.Int64(4) * cutlass.Int64(t * 8 + n // (ar_n // 8)) + ) & Int32(1) + off = ( + cutlass.Int64(par) * half + + cutlass.Int64(t * ar_world + ar_rank) * cutlass.Int64(ar_n // 2) + + cutlass.Int64(n // 2) + ) + _st_global_b32( + ar_ws + off * cutlass.Int64(4), + _bf16x2_word( + sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + ), + ) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +import torch as _torch # noqa: E402 (as in FlashInfer) + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ) + + +@cute.experimental.jit +def _bmm_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ar_tokens: Int32, + ar_n: Int32, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + stream, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D 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, b, out, and bias must share BF16 or FP16 dtype") + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a, + b, + bias, + out, + pdl: bool, + tactic: SplitKTactic, +): + """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +@functools.cache +def _get_compiled_push( + dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] +): + kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, ar_push=True) + compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + z64, z32 = cutlass.Int64(0), Int32(0) + return cute_ext.compile( + _bmm_push, + kernel, + *compile_tensors[:3], + stream, + z64, + z64, + z32, + Int32(1), + Int32(8), + z32, + Int32(6144), + ) + + +def compile_push(a, b, out, pdl: bool, tactic) -> None: + """Compile (only) the push variant of ``run_splitk_dense_push`` for these operands' shapes and layouts.""" + cute_tensors = _to_cute_swap(a, b, out, None) + _get_compiled_push(a.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) + + +def run_splitk_dense_push(a, b, out, pdl: bool, tactic: SplitKTactic, ar: tuple): + """``run_splitk_dense`` without bias that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange; + ``ar`` = (multicast words address, flags address, rank, world, bucket tokens) of the 8-token bucket.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_push(a.dtype, tactic, pdl, cute_tensors[4]) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + ws, flags, rank, world, maxt = ar + compiled( + *cute_tensors[:3], + stream, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + Int32(out.shape[0]), + Int32(out.shape[1]), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py new file mode 100644 index 000000000000..142bcb9a9133 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py @@ -0,0 +1,422 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-latency TP all-reduce + residual + RMSNorm for decode (M <= 8 tokens) over a multicast mapping. + + partial = bf16(x + x2) (optional second input, e.g. shared + routed experts) + updated = bf16(residual + sum_r partial_r) (ranks summed in FP32, rank order) + normed = bf16(updated * rsqrt(mean(updated^2) + eps) * w) + +One cluster of CLUSTER CTAs per token; CTA c owns columns [c * H / CLUSTER, ...), one 16-byte chunk (8 BF16) +per thread. Each CTA pushes its slice of this rank's partial through the multicast mapping (one store lands in +every rank's buffer), releases dependents, polls its own buffer until all ranks' slices arrived, sums them, adds +the residual and reduces the RMS statistic across the cluster through DSMEM. + +Exchange protocol as in the K3 sandwich: int32 words, 0x80000000 = empty (pushed -0.0 halves become +0.0), each +CTA keeps a call counter whose parity selects the buffer half, readers re-arm the words they read. All ranks of +the TP group must issue the same calls in the same stream order. + +The workspace holds two exchange sets. With prepushed 3 two producers pushed this rank's x (set A) and x2 (set B), +e.g. the routed experts' finalize and the shared expert's down projection, and the kernel forms bf16(x + x2) per +rank from the two slots, exactly as the unpushed path does. + +Producers that push (the o_proj / shared-expert GEMV ``ar`` tuple, the routed-expert engine's ``arp`` descriptor) are +pointed at a bucket's exchange by its ``mc`` address, ``flags_for(8)`` counters, ``rank``, ``world_size`` and +``max_tokens``; ``get_workspace`` and ``_bucket`` select it. + +The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp4_mc_allreduce_norm``, reached through the ``ops/comm`` entries +``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``; this module owns the workspaces, the bucket and +epilogue-option choice and the MXFP8 publication. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op +from ..ops.comm.m3_mc_finalize_allreduce_norm import ( + m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, +) + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +BUCKETS = (8, 16, 48) # token capacity of each workspace (grid = bucket * cluster CTAs) +EMPTY_WORD = -2147483648 + +# Attention / MoE boundary epilogue options (each bit-identical): 1 residual + norm weight loaded at entry (before +# the poll); 2 cluster RMS exchange by st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated +# right after the poll (no trailing __syncthreads); 8 max-smem carveout so the next kernel's CTAs co-reside early. +# The 8-token bucket, by token count 1..8 (GB300 TP4 tuning). +AR_OPTS_BY_M = (7, 7, 15, 15, 15, 15, 7, 7) +# The 16- and 48-token buckets. +AR_OPTS_BY_BUCKET = {16: 15, 48: 15} +# The MoE-finalize all-reduce, per bucket. +AR_FIN_OPTS_BY_BUCKET = {8: 7, 16: 15, 48: 15} + + +# Option 32: the 8-CTA cluster kernel in the stock TP4 oneshot reduction order instead of the row kernel (not finalize). +_OPT_ONESHOT_CLUSTER = 32 + +# Option 64: the row kernel's max-smem-carveout instance (as option 8 is for the cluster kernel). +_OPT_ROW_CARVEOUT = 64 + +# Option 128: the drafter boundary (option 16) on the 8-CTA cluster kernel instead of the world-4 row kernel (routing +# only; mcAllreduceNormTp4.cu invoke). Exact-class to the row kernel's FlashInfer order: the RMS statistic is the +# cluster's sum (modeling.DRAFT_MC_AR_CLUSTER). +_OPT_DRAFTER_CLUSTER = 128 + +# Option 2 is also the safe ordering of the cluster RMS exchange: its stores into peer CTAs' shared memory follow a +# cluster barrier arrive / wait pair. Without it they run before any cluster barrier, when a peer need not have +# started, so every call keeps it. +_OPT_CLUSTER_BARRIER = 2 + + +def _ar_opts(m: int, bucket: int, finalize: bool = False, oneshot_cluster: bool = False) -> int: + """Epilogue options for an m-row call on the given bucket; every option is bit-identical. A bucket missing from + the tables gets 0, and every result keeps _OPT_CLUSTER_BARRIER. ``oneshot_cluster`` adds _OPT_ONESHOT_CLUSTER.""" + if oneshot_cluster and not finalize: + return _ar_opts(m, bucket) | _OPT_ONESHOT_CLUSTER + if finalize: + opts = AR_FIN_OPTS_BY_BUCKET.get(bucket, 0) + elif bucket == 8: + opts = AR_OPTS_BY_M[m - 1] + else: + opts = AR_OPTS_BY_BUCKET.get(bucket, 0) + return opts | _OPT_CLUSTER_BARRIER + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags + + +class McBuckets: + """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" + + def __init__(self, workspaces): + self.workspaces = workspaces + self.hidden = workspaces[0].hidden + + def for_tokens(self, m: int): + for ws in self.workspaces: + if m <= ws.max_tokens: + return ws + return None + + +def _bucket(workspace, m): + return workspace.for_tokens(m) if isinstance(workspace, McBuckets) else workspace + + +def decline_reason(x, residual, workspace, cluster=8, x2=None) -> str: + """Why ``m3_mc_allreduce_norm`` returned None for these operands (for the dispatch's miss message).""" + if workspace is None: + return "no multicast workspace: it is created collectively on the first eager decode step" + bucket = _bucket(workspace, x.shape[0]) + if bucket is None: + return f"no workspace bucket holds {x.shape[0]} tokens (buckets {BUCKETS})" + if bucket.flags.get(cluster) is None: + return f"no {cluster}-CTA call counters for the {bucket.max_tokens}-token bucket" + return ( + f"operands outside the kernel: x {tuple(x.shape)} {x.dtype} contiguous={x.is_contiguous()}, " + f"residual contiguous={residual.is_contiguous()}, hidden {bucket.hidden}" + + ("" if x2 is None else f", x2 {tuple(x2.shape)} contiguous={x2.is_contiguous()}") + ) + + +def m3_mc_allreduce_norm( + x, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + x2=None, + quant=False, + prepushed=0, + oneshot_cluster=False, + drafter_norm=False, + carveout=False, + capture=None, + drafter_cluster=False, +): + """Return (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture`` = (out [>= M, H] view, fc_norm weight + [H], eps): the boundary also writes the Eagle3 capture of the updated residual, bit for bit + m3_spec_tail.capture_norm (CAPTURE_NORM_IN_BOUNDARY; not with ``drafter_norm``). ``drafter_cluster``: with + ``drafter_norm``, the 8-CTA cluster kernel's drafter form (option 128; its carveout is option 8) instead of the row + kernel's. + + A None makes the caller take another collective, so it must be the same on every TP rank. It is: it depends only + on the shapes, dtypes and layouts of the operands (the same code builds them on every rank) and on whether the + workspace's flags exist, which ``get_workspace`` creates collectively, before any capture, for cluster 8.""" + workspace = _bucket(workspace, x.shape[0]) + if workspace is None: + return None + if ( + x.dim() != 2 + or x.shape[0] > MAX_TOKENS + or x.dtype != torch.bfloat16 + or x.shape[1] != workspace.hidden + or not x.is_contiguous() + or not residual.is_contiguous() + or (x2 is not None and (not x2.is_contiguous() or x2.shape != x.shape)) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + # The Eagle3 drafter boundary (opts bit 16, the row kernel's FlashInfer fused_add_rmsnorm order): 8-token bucket, + # x only, no MXFP8 copy, never the cluster kernel. + if drafter_norm and ( + workspace.max_tokens != 8 or x2 is not None or quant or prepushed not in (0, 2) + ): + return None + if drafter_norm and drafter_cluster: + opts = ( + (16 | _ar_opts(x.shape[0], workspace.max_tokens)) & ~8 + | (8 if carveout else 0) + | _OPT_DRAFTER_CLUSTER + ) + elif drafter_norm: + opts = ( + 16 | _ar_opts(x.shape[0], workspace.max_tokens) | (_OPT_ROW_CARVEOUT if carveout else 0) + ) + else: + opts = _ar_opts(x.shape[0], workspace.max_tokens, oneshot_cluster=oneshot_cluster) + if carveout: + # the max-smem carveout on whichever kernel runs: option 8 on the cluster kernel, 64 on the row kernel + opts |= 8 if oneshot_cluster else _OPT_ROW_CARVEOUT + outs = _allreduce_norm_op( + x, + x2, + residual, + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and x.shape[1] % 128 == 0), + int(prepushed), + opts, + *_capture_args(capture, drafter_norm), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def _capture_args(capture, drafter_norm: bool) -> tuple: + """(cap_out, cap_w, cap_eps) for the op, from ``capture`` = (out, fc_norm weight, eps) or None.""" + if capture is None: + return (None, None, 0.0) + assert not drafter_norm, "the capture epilogue is not built for the drafter boundary" + out, weight, eps = capture + return (out, weight, float(eps)) + + +def m3_mc_finalize_allreduce_norm( + fc2, + perm, + scales, + shared, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + quant=False, +): + """MoE finalize (sum_k scale * fc2[perm] + shared) + all-reduce + residual + RMSNorm in one kernel. + Returns (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. (It runs the shared GB300 op, which has no capture + epilogue: a capture layer finalized here keeps the capture_norm kernel.)""" + m = residual.shape[0] + workspace = _bucket(workspace, m) + if workspace is None: + return None + if ( + residual.dim() != 2 + or m > MAX_TOKENS + or residual.shape[1] != workspace.hidden + or fc2.dtype != torch.bfloat16 + or not fc2.is_contiguous() + or fc2.dim() != 2 + or perm.dtype != torch.int32 + or perm.dim() != 2 + or perm.shape[0] != m + or not perm.is_contiguous() + or scales.dtype not in (torch.float32, torch.bfloat16) + or not scales.is_contiguous() + or (shared is not None and (shared.shape != residual.shape or not shared.is_contiguous())) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _finalize_allreduce_norm_op( + fc2, + perm, + scales, + shared, + residual.contiguous(), + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and residual.shape[1] % 128 == 0), + _ar_opts(m, workspace.max_tokens, finalize=True), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +_WORKSPACE: dict = {} +_AGREE_COMMS: dict = {} + + +def _tp_comm(mapping): + """A communicator over mapping's TP group, created once (collective over that group only).""" + key = tuple(mapping.tp_group) + comm = _AGREE_COMMS.get(key) + if comm is None: + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + comm = mapping.tp_group_pg + else: + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + _AGREE_COMMS[key] = comm + return comm + + +def tp_agree(mapping, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of mapping's TP group (one blocking host collective). + + For a fast path whose availability is decided from rank-local state (a kernel build that can fail on one rank): + every rank must take the same path, so that they issue the same collectives with the same exchange protocol and + the same kernels. Call it once per decision, at build time, at the same point on every rank, outside CUDA-graph + capture; never per step. With one rank, or attention DP (whose ranks run different batches), it returns ``ok``. + """ + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return ok + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("tp_agree must run outside CUDA-graph capture") + from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded + + return _mnnvl_workspace_all_succeeded(_tp_comm(mapping), bool(ok)) + + +def tp_agree_built(mapping, built, what: str): + """``built`` when it is not None on every TP rank, else None on every rank (see ``tp_agree``).""" + if tp_agree(mapping, built is not None) or built is None: + return built + from tensorrt_llm.logger import logger + + logger.warning(f"{what}: built here but not on another TP rank; every rank falls back") + return None + + +def get_workspace(mapping, hidden: int): + """Process-wide workspace (created collectively on first eager use); None while capturing before creation. + + The first eager call for a hidden size (the warmup forward, before any CUDA-graph capture) creates one + McAllReduceWorkspace per bucket in BUCKETS on every TP rank (MPI collective over the TP group) plus each + bucket's 8-CTA-cluster call counters; later calls, captured or not, return the same McBuckets. The workspaces + live until process exit. + """ + ws = _WORKSPACE.get(hidden) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + return None + wss = [McAllReduceWorkspace.create(mapping, hidden, b) for b in BUCKETS] + for w in wss: + w.flags_for(8) + ws = McBuckets(wss) + _WORKSPACE[hidden] = ws + return ws diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/__init__.py new file mode 100644 index 000000000000..f818a0e6c9f1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/__init__.py @@ -0,0 +1,16 @@ +# 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. +"""MiniMax-M3 MoE layer of the sm_107 TP4 target for low-batch decode as a CuTe DSL kernel: ``m3_moe`` for 1..8 tokens +(op.py) at the TP4 routed and shared intermediate of 768 per rank.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe16_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe16_kernel.py new file mode 100644 index 000000000000..ff609e78ef29 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe16_kernel.py @@ -0,0 +1,3871 @@ +# 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. +"""MiniMax-M3 routed + shared experts for 9..16 decode tokens on one TP rank: ``m3_moe16`` (v2). + +Same math as the 1..8-token engine (m3_moe_kernel.py), operation for operation for the routed part (bitwise equal to +TRTLLM-Gen's fused MoE on the fp32 router-GEMV logits); a different schedule for 9..16 rows: + +- Front: the hidden states land in smem by bulk copies in K chunks (the idle ring); all 16 warps run the router + GEMV (m3_router_gemv's per-thread order) and the NVFP4 quantization of B1 from smem; the logits cross CTAs as + tagged words; top-k one warp per token. +- Ring: FC1 / FC2 weight stages by TMA; their scale words go by 4-byte cp.async straight into the instruction atom + layout (SFA_CP; no raw scale box, no splicer warp), completing on the stage's full barrier. +- FC1 -> FC2 hand-off (B2IMG): the FC1 epilogue writes each unit's FP4 values and scale byte into a per-group global + image laid out as the FC2 B tile + its scale atoms; one counter per group (release add per finished unit, every + counter gains exactly 192 per call); warp 5 loads a completed image with two bulk copies. +- Shared expert on NSH more CTAs (MXFP8 tcgen05), as the 1..8 engine. +Configuration via M3_CONFIG (m_max, num_ctas, stages, nb2_cap, lag, shared, sfa_cp, b2buf, ...). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +H = 6144 # capacity: fixed: the checkpoint's hidden size +I_TP = 768 +TWO_I = 2 * I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 12) +assert 9 <= M_MAX <= 16, M_MAX +NCTA = _cfg("num_ctas", 128) +FUSED_ROUTER = True +NROUTE = E_LOCAL +NGEMV = (NROUTE + NCTA - 1) // NCTA +assert NGEMV == 1, "m3_moe16: one router expert per CTA (num_ctas >= 128)" +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +LAG = _cfg("lag", 3) +G_CAP = TOP_K * M_MAX +NG_CAP = M_MAX +THREADS = 512 # capacity: fixed: the CTA size +NWARPS = THREADS // 32 +GEMV_THREADS = 256 +SHARED = bool(_cfg("shared", 0)) +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +SFA_CP = bool( + _cfg("sfa_cp", 1) +) # A scales by 4-byte cp.async into the atom layout (else raw box + splicer warp) +B2IMG = True +PERM = bool(_cfg("perm", 1)) +PERM_MIN_NG = max( + 3, _cfg("perm_min_ng", 5) +) # (measured: helps at random routing G >= 23, hurts at spec G 13-14) +# a partial last slot group becomes group 1 (the plan must match: op16) +DBG_DELAY = bool( + _cfg("dbg_delay", 0) +) # (stress builds) pseudo-random producer / epilogue / release stalls +DBG_NOSTORE = bool(_cfg("dbg_nostore", 0)) # (timing experiment only: wrong results) +EPI_FENCE = bool( + _cfg("epi_fence", 0) +) # (rel_warp 0) writer lanes' fence.acq_rel.gpu before the release add +EPI_PFENCE = bool( + _cfg("epi_pfence", 1) +) # producer-side fence.proxy.async.global before the release add +REL_WARP = bool( + _cfg("rel_warp", 1) +) # warp 8 does the per-group release adds (off the epilogue's path) +NSH = 24 if SHARED else 0 +SH_I = 768 +SH_KB = H // 32 # capacity: fixed: a kernel tile constant +SH_DKB = SH_I // 32 # capacity: fixed: a kernel tile constant +SH_SPLIT = 2 +SH_GT = SH_KB // 4 // SH_SPLIT # capacity: fixed: a kernel tile constant +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * (SH_I // 128) # capacity: fixed: pipeline stages +SHX_BAR = 6 +SH_PRETILE = False +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = _cfg("n1", 48 if _cfg("m_max", 12) <= 12 else 64) # B rows QS q + t +N1_LD = 1 << (N1 - 1).bit_length() # tcgen05.ld .num (a power of two) covering N1 columns +QS = N1 // 4 +assert QS >= M_MAX +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 +K1_TILES = H // 256 +S1 = K1_TILES // C1 # capacity: fixed: pipeline stages +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 +KT2 = I_TP // 256 +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = _cfg("n2", (4 * M_MAX + 15) // 16 * 16) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +IMG_B2 = KT2 * NUM_BYTES_B2 # one group's B2 tile image (global and smem) +IMG_SF = KT2 * NUM_KB * 512 # its scale atoms +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) # every group counter gains 192 per call +B1IMG_OFF = None # set below (after N1): the activation image after the NG_CAP group images +TK_WARPS = [0, 1, 2, 3, 6, 8, 9, 10, 11, 12, 13, 14, 5, 7, 15, 4][ + : _cfg("m_max", 12) +] # top-k warp of token t (4 last: early0) +B1_WARP = 15 if _cfg("m_max", 12) <= 14 else -1 # (-1: warp 0 after the routing) +NB2_CAP = _cfg("nb2_cap", 2) +B1IMG_OFF = M_MAX * IMG_BYTES +B1IMG_BYTES = (S1 * N1 * 128 + M_MAX * (H // 16) + 15) // 16 * 16 +IMG_TOTAL = B1IMG_OFF + B1IMG_BYTES +UNITS_ROW = TWO_I // R1 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +TOPK_BAR = 4 +ALL_BAR = 5 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 16) +LG_REPL = _cfg("lg_repl", 64) +LG_STRIDE = M_MAX * E_LOCAL + 32 +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # B2 image loader; slots + plan # capacity: fixed: a warp role index +S_WARP = 6 # (SFA_CP = 0) scale splicer # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +EARLY0 = bool(_cfg("early0", 1)) +EVICT_FIRST = 0x12F0000000000000 +EVICT_LAST = 0x14F0000000000000 +_LOG2E = 1.4426950408889634 +NB2BUF = _cfg("b2buf", 2) +SPB = 1 +# x staging: K chunks of XCK elements, NXB smem buffers (rotating if fewer than the chunks) +XCK = _cfg("xck", 2048) +NXC = H // XCK +XCB = M_MAX * XCK * 2 # bytes of one chunk (all tokens) +STAGE_A = NUM_BYTES_A +STAGE_SF = NUM_BYTES_SFA if SFA_CP else NUM_BYTES_SFA_RAW + NUM_BYTES_SFA +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + NB2BUF * IMG_B2 + + IMG_SF + + NB2_CAP * M_MAX * TOP_K * R2 * 2 + + SPB * M_MAX * C1 * R1 * 4 + + (G_CAP + 1) * (NG_CAP + 1) * 4 + + 6 * 1024 +) +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // (STAGE_A + STAGE_SF) +ACC1_COL = 0 +ACC2_COL = 2 * N1 +ACC2_STRIDE = 1 << (N2 - 1).bit_length() +SFB1_COL = ACC2_COL + 2 * ACC2_STRIDE +SFB1_WSTEP = 1 +SFB2_COL = ( + SFB1_COL + NUM_SFB1_ATOMS * 4 +) # one TMEM slot (loaded after the previous group's FC2 MMAs) +SFA1_COL = SFB2_COL + KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# smem: [ring A | ring SF | B2 bufs | B2 SF staging | B1]; x chunks alias [0, _OFF_B1) during the front +_OFF_SFA = STAGES * NUM_BYTES_A +_OFF_B2 = (_OFF_SFA + STAGES * STAGE_SF + 1023) // 1024 * 1024 +_OFF_SFB2 = _OFF_B2 + NB2BUF * IMG_B2 +_OFF_B1 = (_OFF_SFB2 + IMG_SF + 1023) // 1024 * 1024 +BIG = _OFF_B1 + S1 * NUM_BYTES_B1 +NXB = min(NXC, _OFF_B1 // XCB) +assert NXB >= 2, (NXB, _OFF_B1, XCB) +# shared CTA layout in BIG: one pass over all M tokens; B rows SHN q + t (gate_up, 2 K slices), t (down) +SHN = 16 # capacity: fixed: a kernel tile constant +SH_NPASS = 1 +SH_HINT = EVICT_FIRST +SH_PF_POST = bool( + _cfg("sh_pf_post", 1) +) # L2 prefetch of the shared weights after the grid dependency (else before) +_SB_A = 0 +_SH_FIX = ( + SH_GT * 2 * SHN * 128 + + 6 * SHN * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + 64 * SHN * 4 + + 2 * 32 * SHN * 4 + + 2048 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * 2 * SHN * 128 +_SB_XC = _SB_BH + 6 * SHN * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + 64 * SHN * 4 +_SB_END = _SB_GU + 2 * 32 * SHN * 4 +SH_ACC_GU = 0 +SH_ACC_DN = 2 * SHN +SH_SFA = SH_ACC_DN + 2 * SHN +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT # capacity: fixed: a kernel tile constant +SH_TMEM = 256 # capacity: fixed: TMEM columns +if SHARED: + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) +TD_STRIDE = 64 +PFD = _cfg("pfd", 0) # L2 prefetch distance in ring stages (0: off) +PF_WARP = 12 # capacity: fixed: a warp role index +TICKET = bool(_cfg("ticket", 0)) # CTA roles by entry order +# priority tokens: token 0 and PT1 (M / 2: the second request's first row at spec c2) get their logits published first +NPT = _cfg("npt", 0) +PT1 = M_MAX // 2 +PF_PRIO = bool( + _cfg("pf_prio", 0) +) # L2 prefetch of the priority tokens' experts (this CTA's 1/NCTA of each) +PF_W1 = TWO_I * (H // 2) // NCTA +PF_S1 = TWO_I * (H // 16) // NCTA +PF_W2 = H * (I_TP // 2) // NCTA +PF_S2 = H * (I_TP // 16) // NCTA +assert PF_W1 % 16 == 0 and PF_S1 % 16 == 0 and PF_W2 % 16 == 0 and PF_S2 % 16 == 0 + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + # B rows n >= 32 (N1 = 64) go to the atom's next 32-bit column (row n: lane n % 32, column n // 32) + for w in cutlass.range_constexpr((N1 + 31) // 32): + n = lane + cutlass.Int32(32 * w) + q = n // cutlass.Int32(QS) + tl = n - q * cutlass.Int32(QS) + sf_ok = (n < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr( + base_row + warp * cutlass.Int32(32), + base_col + SFB1_COL + a * 4 + w * SFB1_WSTEP, + ), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + gpos, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = gpos % cutlass.Int32(SH_NST) # ring position (stage g of this pass) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue( + lane, + sid, + cutlass.Int32(g), + s_big, + sh_full, + tma_gu, + tma_dn, + gusf, + dnsf, + cutlass.Int32(g), + ) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + elif cutlass.const_expr(not SH_PF_POST): + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(SH_PF_POST): + if warp == 0: # after the grid dependency: DRAM is idle during the routed front + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + for ps in cutlass.range_constexpr(SH_NPASS): + TT0 = 0 # one pass: B rows SHN q + t (gate_up), t (down) + MS = M_MAX + GB = ps * SH_STAGES # ring positions before this pass + PAR = ps & 1 # parity of the per-pass barriers (B ready, accumulators) + if warp == 0: + for g in cutlass.range(SH_NST if ps == 0 else 0, SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) # ring position over the passes + _wait( + sh_empty.subview(gg % cutlass.Int32(SH_NST)), + ((gg // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf, gg) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=2 * SHN, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SHN, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) + st = gg % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (gg // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32((2 * SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32((SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(SHN), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + if cutlass.const_expr(ps > 0): + _wait( + acc_dn, PAR ^ 1 + ) # the previous pass's MMAs are done: its B tiles and TMEM scales are free + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((MS * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(MS * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=(t + cutlass.Int32(TT0)) * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(SHN) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + + it * cutlass.Int32(2 * SHN * 128) + + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of + # k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(SHN - 1) + ok = (n < cutlass.Int32(2 * SHN)) & (t < cutlass.Int32(MS)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n // cutlass.Int32(SHN)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t + # at column 8 q + t + _wait(acc_gu, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=2 * SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(MS): + bigf.store( + cutlass.Float32(av[SHN + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(MS): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((32 * MS + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(32 * MS): + p = pq // cutlass.Int32(MS) + tt = pq - p * cutlass.Int32(MS) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(SHN) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(32) + p) * cutlass.Int32(SHN) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview( + (tt + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + + sid * cutlass.Int32(32) + + p + ) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((MS * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(MS * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = (t + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + kb * cutlass.Int32( + 32 + ) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + + kt * cutlass.Int32(SHN * 128) + + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (t & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + ( + ((c0 + cutlass.Int32(1)) ^ (t & cutlass.Int32(7))) + << cutlass.Int32(2) + ), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(MS) + for kt in cutlass.range_constexpr(6): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * SHN, + cutlass.Float32, + ), + num=SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(MS): + _st_relaxed_u32( + dnx.subview(cutlass.Int32((tt + TT0) * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +# ============================================================================= +# v2 helpers: global release / acquire, proxy fences. +# ============================================================================= +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _atom_add_u64(addr, val, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr), cutlass.Int64(val)], + "atom.relaxed.gpu.global.add.u64 $0, [$1], $2;", + "=l,l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.proxy.async.global;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_acq_rel_gpu(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.acq_rel.gpu;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _dbg_delay(salt): + """(DBG_DELAY stress builds) a pseudo-random ~0..4 us stall: protocol robustness under skew.""" + if cutlass.const_expr(DBG_DELAY): + r = cutlass.Int32(_gtime() >> cutlass.Int64(6)) * cutlass.Int32(1103515245) + salt + if ((r >> cutlass.Int32(12)) & cutlass.Int32(7)) == cutlass.Int32(0): + for _ in cutlass.range((r >> cutlass.Int32(16)) & cutlass.Int32(31), unroll=1): + _nanosleep(128) + + +@cute.jit +def _waitt(bar: cutlass.Array, parity, acc): + """_wait that (TIMING) adds the waited ns to acc.""" + t0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + if cutlass.const_expr(TIMING): + acc = acc + (_gtime() - t0) + return acc + + +@cute.jit +def _tstore(tdbg: cutlass.Array, bx, i: cutlass.Constexpr, v): + if cutlass.const_expr(TIMING): + tdbg.store(cutlass.Int64(v), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@cute.jit +def _issue_fc1( + lane, + stage, + e, + c, + j, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a1, + tma_s1, + s1w: cutlass.Array, +): + """FC1 stage j of unit c of expert e into ring stage `stage` (the caller waited its empty barrier): 4 TMA boxes + (rows 32 c .., k-tiles 4 j + q) and the stage's scale atoms (SFA_CP: word [r][q] of atom kb <- raw atom + (k-tile 4 j + q, kb) word [r][c % 4], by this warp's 32 lanes).""" + row = c * cutlass.Int32(R1) + kx = j * cutlass.Int32(C1 * KT_BYTES) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview(stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(q * NUM_BYTES_BOX1)), + tma_a1, + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32(SF1_KA) + + j * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview(stage * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1, + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + atom0 = (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32( + SF1_KA + ) + j * cutlass.Int32(16) + src0 = atom0 * cutlass.Int32(128) + lane * cutlass.Int32(4) + (c & cutlass.Int32(3)) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _issue_fc2( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a2, + tma_s2, + s2w: cutlass.Array, +): + """FC2 k-tile t2 of output block blk for the group's n_here slots: box jj = rows 32 blk .. of slot jj's w2 at + k-tile t2; scale word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) word [r][blk % 4].""" + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32(NUM_BYTES_BOX1 if SFA_CP else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2, + (cutlass.Int32(t2 * KT_BYTES), blk * cutlass.Int32(R2), e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2, + ( + cutlass.Int32(0), + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _quant_words(w, sc, rsc, rsix): + """NVFP4 of one 16-value block given as 8 bf16-pair words: (lo word, hi word, scale e4m3 bits).""" + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.jit +def _quant_block(xs32: cutlass.Array, widx, sc, rsc, rsix): + """NVFP4 of one 16-value block (8 bf16 words at widx): (lo word, hi word, scale e4m3 bits), fp4_quantize's ops.""" + w = [] + for h in cutlass.range_constexpr(2): + v4 = xs32.load(idx=widx + cutlass.Int32(4 * h), vector_size=4, alignment=16) + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_gu_desc: cutlass.GridConstant[cuda.TensorMap], + tma_dn_desc: cutlass.GridConstant[cuda.TensorMap], + bias: cutlass.Array, # f32 [128] + escal: cutlass.Array, # f32 [6, 128] + plan: cutlass.Array, # int32 [4 M + 1, NCTA, M + 1] + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + img32: cutlass.Array, # int32 [NG_CAP, IMG_BYTES / 4]: per-group FC2 B image (tile + scale atoms) + cnt: cutlass.Array, # int32 [NG_CAP * 8]: per-group unit counters (index 8 g) + s1w: cutlass.Array, # int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # int32 words of w2_weight_scale + w1b: cutlass.Array, # int8 bytes of w3_w1_weight (L2 prefetch) + w2b: cutlass.Array, # int8 bytes of w2_weight (L2 prefetch) + epochs: cutlass.Array, # int32 [NCTA + NSH] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8] + tdbg: cutlass.Array, # int64 [NCTA + NSH, 64] + xb32: cutlass.Array, # int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # f32 router weight [128, 6144] + sfsc: cutlass.Array, # f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # int64 [LG_REPL, LG_STRIDE] + guw: cutlass.Array, + gusf: cutlass.Array, + dnw: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if cutlass.const_expr(TICKET): + # role by entry order: the CTAs that become resident last (behind the predecessor's CTAs) take the shared + # expert roles (most slack), the first NCTA the routed ones (the router GEMV / logits exchange) + s_tk = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + if tidx == 0: + old = _atom_add_u64( + cnt.subview(cutlass.Int32(NG_CAP * 8 + 8)).data_ptr().toint(), cutlass.Int64(1) + ) + s_tk.store(cutlass.Int32(old % cutlass.Int64(NCTA + NSH)), idx=0) + prims.barrier_cta_sync(0) + bx = s_tk.load(idx=0) + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + # ring scale words: SFA_CP: the instruction atoms [STAGES][4 x 512 B]; else raw boxes [STAGES][8 KB] + spliced + sSF = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(STAGES * STAGE_SF // 4,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA + (0 if SFA_CP else STAGES * NUM_BYTES_SFA_RAW)), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=512, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(NB2BUF * IMG_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), shape=(IMG_SF // 4,), dtype=cutlass.Int32, alignment=1024 + ) + ys = cutlass.Array( + cutlass.BFloat16, + NB2_CAP * M_MAX * TOP_K * R2, + space=cutlass.AddressSpace.smem, + alignment=16, + ) # bf16 values + s_part = cutlass.Array( + cutlass.Float32, SPB * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array( + cutlass.Float32, 9 * M_MAX + 16, space=cutlass.AddressSpace.smem, alignment=16 + ) + # the activation's scale codes [M, 384] until they are in TMEM: aliases ys (first written by the FC2 epilogue) + s_xsf = cutlass.Array( + ys.data_ptr(0), shape=(M_MAX * (H // 64),), dtype=cutlass.Int32, alignment=16 + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + img_bar = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + sfb2_free = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_full = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + x_empty = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + route_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + ys_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + grp_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_load = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + topk_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tok0_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + px_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + s_prog = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array(cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_gn = cutlass.Array( + cutlass.Int32, NG_CAP, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots per group + s_emap = cutlass.Array(cutlass.Int32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array(cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32) + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + if tidx < NB2BUF: + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), 1) + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < NXB: + prims.mbarrier_init(x_full.subview(tidx), 1) + prims.mbarrier_init(x_empty.subview(tidx), NWARPS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + prims.mbarrier_init(sfb2_free.subview(0), 1) + prims.mbarrier_init(route_bar.subview(0), 32) + prims.mbarrier_init(ys_done.subview(0), EPI_THREADS) + prims.mbarrier_init(b1_load.subview(0), 1) + prims.mbarrier_init(topk_bar.subview(0), M_MAX * TOP_K) + prims.mbarrier_init(tok0_bar.subview(0), TOP_K) + prims.mbarrier_init(px_full.subview(0), 1) + if tidx < 4: + prims.mbarrier_init(grp_done.subview(tidx), EPI_THREADS) + if tidx == 0: + s_prog.store(cutlass.Int32(0), idx=0) + if warp == X_WARP: # slot tables: empty (the slot warp fills them after the top-k) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + for z in cutlass.range_constexpr(E_LOCAL // 32): + s_emap.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + if cutlass.const_expr(not SFA_CP): + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + # router weights (expert bx) in registers: thread tg = tid % 256 owns elements 4 tg + 1024 j + tg = tidx % cutlass.Int32(GEMV_THREADS) + half = tidx // cutlass.Int32(GEMV_THREADS) + rwv = [ + rw.load( + idx=bx * cutlass.Int32(H) + tg * cutlass.Int32(4) + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + # ---- x chunks into the idle ring (warp 6): chunk c = K [XCK c, XCK (c + 1)) of every token + if warp == S_WARP: + if cutlass.const_expr(NPT > 0): + # the priority tokens' rows first (their logits are published ahead of the others) + if lane == 0: + prims.mbarrier_arrive_expect_tx(px_full.subview(0), NPT * H * 2) + cute.arch.sync_warp() + if lane < NPT: + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(lane * cutlass.Int32(H * 2)), + xb32.subview( + cutlass.select_( + lane == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1) + ) + * cutlass.Int32(H // 2) + ), + px_full.subview(0), + H * 2, + ) + for c in cutlass.range_constexpr(NXB): + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(c), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(c * XCB) + lane * cutlass.Int32(XCK * 2)), + xb32.subview(lane * cutlass.Int32(H // 2) + cutlass.Int32(c * XCK // 2)), + x_full.subview(c), + XCK * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + img8f = cutlass.Array( + img32.data_ptr(0), shape=(IMG_TOTAL,), dtype=cutlass.Int8, alignment=16 + ) + if cutlass.const_expr(NPT > 0): + # ---- priority tokens (0 and PT1): their logits first (m3_router_gemv's per-thread / warp order) + _wait(px_full.subview(0), 0) + ph = cutlass.select_(half < cutlass.Int32(NPT), half, cutlass.Int32(0)) + pacc = cutlass.Float32(0.0) + px32 = cutlass.Array( + sB1.data_ptr(0), shape=(NPT * H // 2,), dtype=cutlass.Int32, alignment=16 + ) + for j in cutlass.range_constexpr(H // 1024): + pxw = px32.load( + idx=ph * cutlass.Int32(H // 2) + cutlass.Int32(512 * j) + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + px0 = cutlass.Int32(pxw[0]) + px1 = cutlass.Int32(pxw[1]) + pacc = _fma_rn( + (px0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][0]), + pacc, + ) + pacc = _fma_rn( + (px0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][1]), + pacc, + ) + pacc = _fma_rn( + (px1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][2]), + pacc, + ) + pacc = _fma_rn( + (px1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][3]), + pacc, + ) + for off in (16, 8, 4, 2, 1): + pacc = _add_rn(pacc, cute.arch.shuffle_sync_bfly(pacc, off)) + if lane == 0: + s_red.store(pacc, idx=cutlass.Int32(9 * M_MAX) + warp) # [half][8 warps] + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if tidx < cutlass.Int32(NPT * LG_REPL): + hh = tidx // cutlass.Int32(LG_REPL) + rr = tidx - hh * cutlass.Int32(LG_REPL) + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn( + sacc, s_red.load(idx=cutlass.Int32(9 * M_MAX + w) + hh * cutlass.Int32(8)) + ) + tp = cutlass.select_(hh == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1)) + _st_relaxed_u64( + lgbuf.subview(rr * cutlass.Int32(LG_STRIDE) + tp * cutlass.Int32(E_LOCAL) + bx) + .data_ptr() + .toint(), + ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)), + ) + if tidx == 0: + _stamp(tdbg, bx, 43) + # ---- this CTA's share of the NVFP4 activation (QB blocks, from global) for the global B1 image: loads first + QB = (M_MAX * (H // 16) + NCTA - 1) // NCTA + qb = bx * cutlass.Int32(QB) + tidx + q_ok = (tidx < cutlass.Int32(QB)) & (qb < cutlass.Int32(M_MAX * (H // 16))) + qb = cutlass.select_(q_ok, qb, cutlass.Int32(0)) + qtok = qb // cutlass.Int32(H // 16) + qkb = qb - qtok * cutlass.Int32(H // 16) + xq_v = [ + xb32.load( + idx=qtok * cutlass.Int32(H // 2) + qkb * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + for h in range(2) + ] + # ---- router GEMV from the staged chunks, all 512 threads (rolled: the front runs once per call) + MH = (M_MAX + 1) // 2 # tokens per GEMV half + acc_r = cute.make_rmem_tensor((MH,), cutlass.Float32) + for i in cutlass.range_constexpr(MH): + acc_r[i] = cutlass.Float32(0.0) + big32 = cutlass.Array( + s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16 + ) + JPC = XCK // 1024 # router j-steps per chunk + for xc in cutlass.range(NXC, unroll=1): + buf = xc % cutlass.Int32(NXB) + _wait(x_full.subview(buf), (xc // cutlass.Int32(NXB)) & cutlass.Int32(1)) + if (tidx == 0) & (xc == cutlass.Int32(0)): + _stamp(tdbg, bx, 41) + xw0 = buf * cutlass.Int32(XCB // 4) # chunk word base in s_big + for jj in cutlass.range_constexpr(JPC): + jg = xc * cutlass.Int32(JPC) + cutlass.Int32(jj) + wsel = [cutlass.Float32(rwv[0][z]) for z in range(4)] + for jx in cutlass.range_constexpr(1, H // 1024): + for z in cutlass.range_constexpr(4): + wsel[z] = cutlass.select_( + jg == cutlass.Int32(jx), cutlass.Float32(rwv[jx][z]), wsel[z] + ) + for i in cutlass.range_constexpr(MH): + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + xw = big32.load( + idx=xw0 + + t * cutlass.Int32(XCK // 2) + + cutlass.Int32(512 * jj) + + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + acc_g = acc_r[i] + acc_g = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[0], acc_g + ) + acc_g = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[1], acc_g + ) + acc_g = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[2], acc_g + ) + acc_g = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[3], acc_g + ) + acc_r[i] = acc_g + if cutlass.const_expr(NXB < NXC): + if xc + cutlass.Int32(NXB) < cutlass.Int32(NXC): + # every thread's reads of this chunk are done (CTA barrier, then a generic -> async proxy fence by + # the refilling warp: the bulk copy overwrites what the generic proxy read); racecheck-clean + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == S_WARP: + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(buf), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + buf * cutlass.Int32(XCB) + lane * cutlass.Int32(XCK * 2) + ), + xb32.subview( + lane * cutlass.Int32(H // 2) + + (xc + cutlass.Int32(NXB)) * cutlass.Int32(XCK // 2) + ), + x_full.subview(buf), + XCK * 2, + ) + if tidx == 0: + _stamp(tdbg, bx, 42) + for i in cutlass.range_constexpr(MH): + v = acc_r[i] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + if (lane == 0) & (t < cutlass.Int32(M_MAX)): + s_red.store(v, idx=(warp % cutlass.Int32(8)) * cutlass.Int32(M_MAX) + t) + # this CTA's quantization blocks into the global B1 image (B1 tile bytes, then the [M, 384] scale codes) + if q_ok: + qw = [] + for h in cutlass.range_constexpr(2): + for z in cutlass.range_constexpr(4): + qw.append(cutlass.Int32(xq_v[h][z])) + qlo, qhi, qsf = _quant_words(qw, sc, rsc, rsix) + img32.store( + (qlo, qhi), idx=cutlass.Int32(B1IMG_OFF // 4) + _b1_addr(qtok, qkb), alignment=8 + ) + img8f.store(cutlass.Int8(qsf), idx=cutlass.Int32(B1IMG_OFF + S1 * NUM_BYTES_B1) + qb) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == cutlass.Int32(0): + _dbg_delay(bx * cutlass.Int32(3)) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX + t))) + lw = ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)) + _st_relaxed_u64( + lgbuf.subview(tidx * cutlass.Int32(LG_STRIDE) + cutlass.Int32(t * E_LOCAL) + bx) + .data_ptr() + .toint(), + lw, + ) + if tidx == 0: + _stamp(tdbg, bx, 21) + if tidx == cutlass.Int32( + THREADS - 32 + ): # B1 share released after the logits (its gpu fence stalls warp 15 only) + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() # once, after the join (ALL_BAR) of the writer threads + _red_release_add( + cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint(), cutlass.Int32(1) + ) + # ---- the B1 image (every CTA's share) into smem (warp 13; the top-k warps are 0 .. M - 1) + if cutlass.const_expr(B1_WARP >= 0): + if warp == cutlass.Int32(B1_WARP): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + + # ---- top-k: token t on warp t (keys sigmoid + bias, ties to the lower id), TRTLLM-Gen MiniMax2 ops + tk_t = cutlass.Int32(-1) # this warp's top-k token (TK_WARPS[t] = warp) + for i_tk in cutlass.range_constexpr(M_MAX): + tk_t = cutlass.select_(warp == cutlass.Int32(TK_WARPS[i_tk]), cutlass.Int32(i_tk), tk_t) + if tk_t >= cutlass.Int32(0): + t = tk_t + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 22) + lgv = [cutlass.Int32(res[z]).bitcast(cutlass.Float32) for z in range(4)] + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hk = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hk[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + if lane < TOP_K: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 20) + if cutlass.const_expr(PF_PRIO): + if (t == cutlass.Int32(0)) | (t == cutlass.Int32(PT1)): + if lane < TOP_K: + # this CTA's 1/NCTA of each of the token's experts into L2 (DRAM is idle until the routing is + # complete; every byte is streamed later anyway) + _l2_prefetch( + w1b.subview( + e_l * cutlass.Int32(TWO_I * H // 2) + bx * cutlass.Int32(PF_W1) + ) + .data_ptr() + .toint(), + PF_W1, + ) + _l2_prefetch( + s1w.subview( + (e_l * cutlass.Int32(TWO_I * H // 16) + bx * cutlass.Int32(PF_S1)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S1, + ) + _l2_prefetch( + w2b.subview( + e_l * cutlass.Int32(H * I_TP // 2) + bx * cutlass.Int32(PF_W2) + ) + .data_ptr() + .toint(), + PF_W2, + ) + _l2_prefetch( + s2w.subview( + (e_l * cutlass.Int32(H * I_TP // 16) + bx * cutlass.Int32(PF_S2)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S2, + ) + cute.arch.sync_warp() + if lane < TOP_K: # every writer lane arrives (its own release of s_ids / s_w) + prims.mbarrier_arrive(topk_bar.subview(0)) # this token's ids / weights are in smem + if t == cutlass.Int32(0): + prims.mbarrier_arrive( + tok0_bar.subview(0) + ) # token 0's experts: the early group-0 stages + cute.arch.sync_warp() + + if cutlass.const_expr(EARLY0): + if warp == W_WARP: + # slots 0..3 = token 0's four distinct experts; group 0's stage ranges are the same for every G + _wait(tok0_bar.subview(0), 0) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(STAGES): + _issue_fc1( + lane, + n0, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + n0 = n0 + cutlass.Int32(1) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if warp == X_WARP: + # slots: entry i < 4 M is (token i / 4, top-k i % 4); distinct experts in entry order + NCH = (TOP_K * M_MAX + 31) // 32 + _wait(topk_bar.subview(0), 0) + gcnt = cutlass.Int32(0) + rec = [] # per chunk: (first, rank, slot rank of the entry, expert, entry valid) + for ch in cutlass.range_constexpr(NCH): + ei = lane + cutlass.Int32(32 * ch) + valid = ei < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, ei, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + known = cutlass.Boolean(False) + prev = cutlass.Int32(-1) + key = my_e + if cutlass.const_expr(ch > 0): + prev = cutlass.select_( + valid, + s_emap.load(idx=cutlass.select_(valid, my_e, cutlass.Int32(0))), + cutlass.Int32(-1), + ) + known = prev >= cutlass.Int32(0) + key = cutlass.select_(known, cutlass.Int32(-1000) - lane, my_e) + cute.arch.sync_warp() # every lane's s_emap read before any lane's write + same = _match_any(key) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (known == cutlass.Boolean(False)) & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = gcnt + cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + slot_l = cutlass.select_(known, prev, slot_l) + if first: + s_emap.store(rank, idx=my_e) + if valid: + if bx == 0: + route_dbg.store( + my_e, + idx=(ei // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + ei % cutlass.Int32(TOP_K), + ) + rec.append((first, rank, slot_l, my_e, valid)) + cute.arch.sync_warp() + gcnt = gcnt + cute.arch.popc(bal) + # slot of rank r: with a partial last group (tail < 4) and >= 3 groups, that group becomes group 1 (its + # FC2 then is not the last hand-off behind a short FC1 phase); the full groups 1.. shift up by one + ngr = (gcnt + cutlass.Int32(3)) >> cutlass.Int32(2) + tail = gcnt - (ngr - cutlass.Int32(1)) * cutlass.Int32(4) + do_perm = ( + cutlass.Boolean(PERM) + & (tail < cutlass.Int32(4)) + & (ngr >= cutlass.Int32(PERM_MIN_NG)) + ) + for z in cutlass.range_constexpr((G_CAP + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP): + s_el.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + for ch in cutlass.range_constexpr(NCH): + first, rank, slot_r, my_e, valid = rec[ch] + ei = lane + cutlass.Int32(32 * ch) + for which in cutlass.range_constexpr(2): + r = rank if which == 0 else slot_r + sl_m = cutlass.select_( + r < cutlass.Int32(4), + r, + cutlass.select_( + r >= (ngr - cutlass.Int32(1)) * cutlass.Int32(4), + r - (ngr - cutlass.Int32(2)) * cutlass.Int32(4), + r + cutlass.Int32(4), + ), + ) + sl_m = cutlass.select_(do_perm, sl_m, r) + if cutlass.const_expr(which == 0): + if first: + s_el.store(my_e, idx=sl_m) + else: + if valid: + s_kpos.store( + ei % cutlass.Int32(TOP_K), + idx=sl_m * cutlass.Int32(M_MAX) + ei // cutlass.Int32(TOP_K), + ) + # slots per group + if lane < cutlass.Int32(NG_CAP): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + nh = cutlass.select_(do_perm & (lane == cutlass.Int32(1)), tail, nh) + nh = cutlass.select_( + do_perm & (lane == ngr - cutlass.Int32(1)), cutlass.Int32(4), nh + ) + s_gn.store(nh, idx=lane) + cute.arch.sync_warp() + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + prims.mbarrier_arrive( + route_bar.subview(0) + ) # every lane: its slots / kpos / plan-row writes are in smem + cute.arch.sync_warp() + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = s_gn.load(idx=lane) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add(cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad) + else: + _wait(route_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(B1_WARP < 0): + if warp == cutlass.Int32(0): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + if warp < 4: # B1 scales to TMEM (the MMA needs them only with the first ring stage) + _wait(b1_load.subview(0), 0) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + n_iss = cutlass.Int32( + 0 + ) # stages issued (early ones included): the L2-prefetch warp's throttle + w_emp = cutlass.Int64(0) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _dbg_delay(bx) + _issue_fc1( + lane, + stage, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _issue_fc2( + lane, + stage, + g, + blk, + t2, + n_here, + s_el, + sA, + sSF, + ab_full, + tma_a2_desc.get_ptr(), + tma_s2_desc.get_ptr(), + s2w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if lane == 0: + _stamp(tdbg, bx, 3) + _tstore(tdbg, bx, 30, w_emp) + + # ---- (SFA_CP = 0) scale splicer (warp 6): atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if cutlass.const_expr(not SFA_CP): + if warp == S_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt_ = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt_ = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt_ = cutlass.Int32(KT2) + for _j in cutlass.range(cnt_, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32( + 4 + ) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSF.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSF.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSF.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSF.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + # ---- FC2 B loader (warp 5): group g's image once its counter is complete + if warp == X_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for g in range(ng): + buf = g % cutlass.Int32(NB2BUF) + if g >= cutlass.Int32(NB2BUF): + _waitd( + b2_free.subview(buf), + ((g // cutlass.Int32(NB2BUF)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 29) + caddr = cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2_POLL_NS) + cv = _ld_acquire_u32(caddr) + cute.arch.sync_warp() + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 51) + _fence_proxy_async_global() + if g >= cutlass.Int32(1): + _waitd( + sfb2_free.subview(0), (g - cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 107 + ) + if lane == 0: + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(0), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), tdbg, 108 + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 52) + if g >= cutlass.Int32(1): + # the single TMEM scale slot: the previous group's FC2 MMAs are done + gp = g - cutlass.Int32(1) + _waitd( + b2_free.subview(gp % cutlass.Int32(NB2BUF)), + (gp // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), + tdbg, + 109, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 53) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for aa in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFB2_COL + aa * 4), + sd0 + cutlass.Int64(aa * (512 >> 4)), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + prims.tcgen05_commit(sfb2_free.subview(0)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + + # ---- (PFD) L2 prefetch warp (12): the same stage sequence as the producer, PFD stages ahead of it + if cutlass.const_expr(PFD > 0): + if warp == cutlass.Int32(PF_WARP): + pf = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + kx = j * cutlass.Int32(C1 * KT_BYTES) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_prefetch( + tma_a1_desc.get_ptr(), + [ + kx + cutlass.Int32(q * KT_BYTES), + c * cutlass.Int32(R1), + e, + ], + [], + ) + atom0 = ( + e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2)) + ) * cutlass.Int32(SF1_KA) + j * cutlass.Int32(16) + _l2_prefetch( + s1w.subview(atom0 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 8192, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load( + idx=g * cutlass.Int32(4) + cutlass.Int32(jj) + ) + prims.cp_async_bulk_tensor_prefetch( + tma_a2_desc.get_ptr(), + [ + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ], + [], + ) + a2 = ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) * cutlass.Int32(SF2_KA) + cutlass.Int32(4 * t2) + _l2_prefetch( + s2w.subview(a2 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 2048, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + + # ---- (REL_WARP) group releases (warp 8): this CTA's finished units of group p, one release add per phase + if cutlass.const_expr(REL_WARP): + if warp == cutlass.Int32(8): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + grp_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + _dbg_delay(bx * cutlass.Int32(13) + p) + if (lane == 0) & (n_done > cutlass.Int32(0)): + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + _red_release_add( + cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + if (lane == 0) & (p == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 54) + cute.arch.sync_warp() + + # ---- MMA (warp 7) + if warp == MMA_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + m_b2 = cutlass.Int64(0) + m_full = cutlass.Int64(0) + m_acc = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + m_acc = _waitt( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + cute.arch.sync_warp() + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g % cutlass.Int32(NB2BUF) + m_b2_0 = m_b2 + m_b2 = _waitt( + b2_ready.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), m_b2 + ) + if cutlass.const_expr(TIMING): + if (lane == 0) & (g < cutlass.Int32(4)): + tdbg.store( + m_b2 - m_b2_0, + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(44) + g, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + m_acc = _waitt( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB2_COL + cutlass.Int32(t2 * NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + cute.arch.sync_warp() + k2 = k2 + cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit( + b2_free.subview(buf) + ) # this group's B2 tile and TMEM scales are free + cute.arch.sync_warp() + if lane == 0: + _tstore(tdbg, bx, 26, m_b2) + _tstore(tdbg, bx, 27, m_full) + _tstore(tdbg, bx, 28, m_acc) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + img8 = cutlass.Array( + img32.data_ptr(0), shape=(NG_CAP * IMG_BYTES,), dtype=cutlass.Int8, alignment=16 + ) + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if qi < cutlass.Int32(6 * G_CAP): + e_sl = s_el.load(idx=sl_e) + if e_sl >= cutlass.Int32(0): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32(escal.load(idx=fld * cutlass.Int32(E_LOCAL) + e_sl)), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + e_wait = cutlass.Int64(0) + e_busy = cutlass.Int64(0) + e_a = cutlass.Int64(0) + e_pw = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + unit_done = (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)) + e_wait = _waitt( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + tb0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1_LD, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=(t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + tc0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_a = e_a + (tc0 - tb0) + g_u = slot >> cutlass.Int32(2) + jrow = slot & cutlass.Int32(3) + # this warp's tokens t = warp + 4 tr, all in flight together (latency-bound otherwise) + NT = (M_MAX + 3) // 4 + xr = cute.make_rmem_tensor((NT,), cutlass.Float32) + for tr in cutlass.range_constexpr(NT): + tq = cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + xr[tr] = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=(tq * C1 + 0) * R1 + lane), + s_part.load(idx=(tq * C1 + 1) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 2) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 3) * R1 + lane), + ) + if (j0 == cutlass.Int32(0)) & (j1 != cutlass.Int32(S1)): + # head of a unit split with the next CTA: its partial sums + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + _st_relaxed_u32( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint(), + xr[tr].bitcast(cutlass.Int32), + ) + else: + if j0 != cutlass.Int32(0): + # tail: the head's partial sums (head first, as TRTLLM-Gen's K order) + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + pbits = _ld_relaxed_u32(paddr) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xr[tr] = _add_rn(pbits.bitcast(cutlass.Float32), xr[tr]) + cute.arch.sync_warp() + _dbg_delay(bx * cutlass.Int32(7) + warp) + es = slot * cutlass.Int32(8) + ea = s_esc.load(idx=es) + ec = s_esc.load(idx=es + cutlass.Int32(1)) + eal = s_esc.load(idx=es + cutlass.Int32(2)) + ebe = s_esc.load(idx=es + cutlass.Int32(3)) + eli = s_esc.load(idx=es + cutlass.Int32(4)) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + n0 = jrow * cutlass.Int32(M_MAX) + t2 = cblk >> cutlass.Int32(4) + cc = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = g_u * cutlass.Int32(IMG_BYTES) + sfoff = ( + gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2 * cutlass.Int32(NUM_KB) + + ((cblk & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(512) + ) + + (cblk & cutlass.Int32(3)) + ) + for tr in cutlass.range_constexpr(NT): + # (skip) only tokens routed to this slot (warp-uniform); other B2 rows are never used + if s_kpos.load( + idx=slot * cutlass.Int32(M_MAX) + + cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + ) >= cutlass.Int32(0): + xv = xr[tr] + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + qv, sfb = _swiglu_quant( + xv, gate_raw, is_up, ea, ec, eal, ebe, eli + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + t = warp + cutlass.Int32(4 * tr) + if ( + (lane == 0) + & (t < cutlass.Int32(M_MAX)) + & cutlass.Boolean(not DBG_NOSTORE) + ): + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), row n + # at byte 16 (n % 32) + 4 (n / 32), byte c % 4 + n = n0 + t + img32.store( + (lo, hi), + idx=( + gbase + + t2 * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=sfoff + + (n & cutlass.Int32(31)) * cutlass.Int32(16) + + (n >> cutlass.Int32(5)) * cutlass.Int32(4), + ) + tc1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_pw = e_pw + (tc1 - tc0) + if cutlass.const_expr(not REL_WARP): + if unit_done: + if lane == 0: + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + if cutlass.const_expr(EPI_FENCE): + _fence_acq_rel_gpu() + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=EPI_BAR, number_of_threads=EPI_THREADS + ) # s_part reads done + if cutlass.const_expr(not REL_WARP): + if unit_done & (tidx == 0): + _red_release_add( + cnt.subview(g_u * cutlass.Int32(8)).data_ptr().toint(), + cutlass.Int32(1), + ) + if cutlass.const_expr(TIMING): + e_busy = e_busy + (_gtime() - tb0) + k1 = k1 + cutlass.Int32(1) + if cutlass.const_expr(REL_WARP): + prims.mbarrier_arrive( + grp_done.subview(p & cutlass.Int32(3)) + ) # this CTA's group-p units are out + if tidx == 0: + _stamp(tdbg, bx, 10) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + e_wait = _waitt( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + if cutlass.const_expr(TIMING): + if ( + (tidx == 0) + & (g == ng - cutlass.Int32(1)) + & (bb == cutlass.Int32(0)) + ): + _stamp(tdbg, bx, 8) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = [] + ya = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), + cutlass.Float32, + ), + num=ACC2_STRIDE, + ) + for t in cutlass.range_constexpr(M_MAX): + v = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + v = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(ya[w * M_MAX + t]), v + ) + yv.append(v) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if warp < s_gn.load(idx=g): + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + cutlass.BFloat16( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)) + ), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + if tidx == 0: + _stamp(tdbg, bx, 6) + _tstore(tdbg, bx, 31, e_wait) + _tstore(tdbg, bx, 34, e_busy) + _tstore(tdbg, bx, 36, e_a) + _tstore(tdbg, bx, 38, e_pw) + prims.mbarrier_arrive(ys_done.subview(0)) # this thread's ys values are in smem + + # ---- finalize on all 16 warps: (block bb, token t) pairs, the fma chain over t's top-4 in order, + y_sh + NPAIR = NB2_CAP * M_MAX + NPW = (NPAIR + NWARPS - 1) // NWARPS + wyl = [] + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + ok = (pair < cutlass.Int32(NPAIR)) & (bb < nb) + ch = cutlass.select_( + ok, + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wyl.append( + _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + if cutlass.const_expr(SHARED) + else cutlass.Int32(0) + ) + _wait(ys_done.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 19) + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if (pair < cutlass.Int32(NPAIR)) & (bb < nb): + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ), + acc_o, + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SHARED): + wy = wyl[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + bias: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + img32: cute.Tensor, + cnt: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + w1b: cute.Tensor, + w2b: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + bias, + escal, + plan, + partials, + img32, + cnt, + s1w, + s2w, + w1b, + w2b, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..81122a06c18c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3914 @@ +# 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. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP4 rank (routed and shared intermediate 768): +``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 # capacity: fixed: the checkpoint's hidden size +# routed intermediate per rank: 1536 at TP2 (this target: FC2 K = 6 full 256-K tiles), 768 at TP4, 384 at TP8 (FC2 K +# = 1.5 256-K tiles: the last tile's A box is half out of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 +# MMAs / scale atoms are issued; no padding, no extra bytes) +I_TP = _cfg("i_tp", 768) +TWO_I = 2 * I_TP +assert I_TP % 128 == 0 and 128 <= I_TP <= 1536, I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 # capacity: fixed: the CTA size +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is the low-M MXFP8 GEMM's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# TAIL_TRIGGER (with LATE_TRIGGER, no all-reduce push): the CTAs that own more FC2 blocks than the rest (the first +# H / 32 mod NCTA, which end last) release the dependents right after the grid dependency, so the dependents launch +# when the other CTAs tear down, onto the SMs they leave, instead of after the last of the heavy CTAs. +TAIL_TRIGGER = bool(_cfg("tail_trigger", 1)) +NHEAVY = (H // 32) % NCTA if NCTA < H // 32 else 0 +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +# the shared intermediate per rank: 1536 (TP2), 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate +# positions: its gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP2: 1 x (64 + +# 64), full K in one accumulator, the order of the stock tcgen05 MXFP8 GEMM; TP4: 2 x (32 + 32), the order of +# m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. +SH_I = _cfg("sh_i", 768) +SH_ROWS = SH_I // 24 +# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) +SH_SPLIT = 128 // (2 * SH_ROWS) +SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N +SH_KB = H // 32 # capacity: fixed: a kernel tile constant (gate_up K blocks) (gate_up K blocks) +SH_DKB = SH_I // 32 # down K blocks +SH_DKT = SH_I // 128 # down k-tiles +SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the +# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction +# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) +) # (in-chain: a win at M <= 2, not 4..8) +if SHARED: + assert SH_I in (384, 768, 1536), SH_I + assert SH_PRETILE or SH_SPLIT == 2, ( + "the un-pretiled shared stream is the 32-row (TP4) layout only" + ) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # 6 stages per unit # capacity: fixed: pipeline stages +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) +# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) +KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 + + +def _nkb2(t2: int) -> int: + """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" + return NUM_KB if t2 < KT2 - 1 else KB2_LAST + + +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy +# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few +# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. +HB_REPL = _cfg("hb_repl", 1) +assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) +HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +# Hang guard: every cross-CTA spin (the kernel needs all its CTAs resident at once) traps after this long, so a lost +# co-residency or a missing producer surfaces as a device exception instead of a silent hang. The timer is read only +# inside a spin iteration; the arithmetic is unchanged. +HANG_NS = _cfg("hang_ms", 2000) * 1000000 +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # capacity: fixed: a warp role index +S_WARP = 6 # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +# E0_M4: at M = 4 the MMA warp routes +# token 3 (NRW0 = 4), so the early group-0 stages go to FC2 B builder warp 8 instead (idle until the routing is done; +# RW_EXTRA = 0 at M = 4). Issue order only: exact. +E0_M4 = ( + bool(_cfg("e0_m4", 1)) and NRW0 == 4 and RW_EXTRA == 0 and THREADS // 32 > 8 +) # M == 4 only: at M > 4 warp 8 routes (RW_EXTRA) +E0_WARP = ( + 8 if E0_M4 else MMA_WARP +) # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and (NRW0 <= 3 or E0_M4) and bool(_cfg("e0_off", 1)) +# B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing +# warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the +# two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. +if B1_AFTER_TOPK and XSTAGE and EARLY0 and not E0_OFFLOAD: + raise ValueError( + "m3_moe: b1_after_topk=1 with xstage, early0 and no early-stage offload deadlocks (warp 4 waits on x_bar[1] " + "before TOPK2_BAR, the B1 build waits on TOPK2_BAR before x_bar[1]); set b1_after_topk=0 or early0=0" + ) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# TP2 at 8 tokens: the 98 KB of hidden states fit no 5-stage ring; the router GEMV then reads them from global (the +# kernel's non-XSTAGE path, same math) +if XSTAGE and M_MAX * H * 2 > STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA): + XSTAGE = False +assert not XSTAGE or M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 +_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 +_SH_FIX = ( + SH_GT * SH_NB * 128 + + SH_DKT * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + _SH_XCH_BYTES + + 2 * SH_ROWS * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * SH_NB * 128 +_SB_XC = _SB_BH + SH_DKT * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + _SH_XCH_BYTES +_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = SH_NB +SH_SFA = SH_ACC_DN + 16 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +# pre-tiled stages beyond the ring that warp 0 prefetches into L2, per lane (TP2: 72 stages) +SH_PF = max(2, (SH_STAGES - SH_NST + 31) // 32) if SHARED else 2 +# TP2 (SH_GT 48 gate_up k-tiles, SH_DKT 12) needs more than 256 columns: allocate the next power of two +SH_TMEM = 256 if SH_SFB_DN + 4 * SH_DKT <= 256 else 512 +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 32 * SH_PF, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 4 * SH_DKT <= SH_TMEM, ( + SH_NST, + _SB_END, + BIG, + ) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _trap_after(elapsed, *, loc=None, ip=None): + """trap when elapsed (ns, int64) exceeds HANG_NS.""" + _llvm.inline_asm( + None, + [cutlass.Int64(elapsed).ir_value(loc=loc, ip=ip)], + f"{{ .reg .pred p; setp.gt.s64 p, $0, {HANG_NS}; @p trap; }}", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _hang_tick(t0): + """One spin iteration of a cross-CTA wait: the wait's start (latched on the first call; t0 = 0 before it), after + trapping once the wait has lasted longer than HANG_NS.""" + now = _gtime() + t1 = cutlass.select_(t0 == cutlass.Int64(0), now, t0) + _trap_after(now - t1) + return t1 + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + # (a partial last k-tile has no atoms past I_TP) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(SH_PF): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(SH_PF) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SH_NB, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if cutlass.const_expr(SH_SPLIT == 1): + # TP2: one K slice; quadrant qd = rows 32 qd .. (gate positions 0..63, then up 0..63), no partials: + # the owner sum is (0 + acc), then bf16, at _SB_GU word (u SH_ROWS + p) 8 + t = (32 qd + l) 8 + t + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + elif cutlass.const_expr(SH_SPLIT == 2): + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + else: + # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are + # accumulator columns 8 qd + t + av = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_GU + qd * cutlass.Int32(8), + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(1): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd == cutlass.Int32(0): + # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + for s_ in cutlass.range_constexpr(1, SH_SPLIT): + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + lane * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(SH_ROWS * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + t0h = cutlass.Int64(0) + while pend != cutlass.Int32(0): + t0h = _hang_tick(t0h) + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(SH_DKT): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(LATE_TRIGGER and TAIL_TRIGGER and NHEAVY > 0): + if ar_off & (bx < cutlass.Int32(NHEAVY)): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + t0h = cutlass.Int64(0) + while done == cutlass.Int32(0): + t0h = _hang_tick(t0h) + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + t0h = cutlass.Int64(0) + while (cv - cnt_target) < cutlass.Int32(0): + t0h = _hang_tick(t0h) + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + for rr in cutlass.range_constexpr(HB_REPL): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4) + + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + hb_off = cutlass.Int32(0) + if cutlass.const_expr(HB_REPL > 1): + # this CTA's copy of the records + hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + t0h = cutlass.Int64(0) + while done == cutlass.Int32(0): + t0h = _hang_tick(t0h) + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview( + (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) + ), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + t0h = cutlass.Int64(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + t0h = _hang_tick(t0h) + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if cutlass.const_expr(HB_REPL == 1): + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + else: + if lane < HB_REPL: # copy `lane` + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4) + + lane * cutlass.Int32(HB_STRIDE), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + t0h = cutlass.Int64(0) + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + t0h = _hang_tick(t0h) + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py new file mode 100644 index 000000000000..ab60e7d96370 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py @@ -0,0 +1,779 @@ +# 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. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP4 rank (routed and shared intermediate 768 per rank), +over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm.logger import logger + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +INTER = 768 # capacity: fixed: the routed experts' width per tp4 rank +NUM_EXPERTS = 128 +SHARED_INTER = 768 # capacity: fixed: the shared expert's width per tp4 rank +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size (m3_moe's shared-expert CTAs (beside num_ctas routed CTAs)) +TOP_K = 4 +U1 = 2 * INTER // 32 # FC1 units per expert +# The router GEMV takes at most 2 experts per CTA (router weights held in registers). +MIN_CTAS = 64 # capacity: fixed: an SM split, not a size + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # FC1 stages (items) per unit # capacity: fixed: pipeline stages +_ITEM_BYTES = _UNIT_BYTES // _S1 # capacity: fixed: a ring item's bytes +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters. Off at TP4 (at 4 / 8 tokens +# 0.5 / 0.8 us/call slower than the tagged records in the TP4 sweep). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 0) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit). 5 at 8 tokens (with a 50 ns FC2 record poll below) +# measured 0.8 us/call faster than 6 in the TP4 sweep (same MMAs, bit-identical outputs); at 4 tokens the default is +# fastest with heavy_fc2_only. +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 5) +# The FC2 B builders' record poll back-off in ns (0 = the kernel's default, 200); 50 at 8 tokens. +B2_POLL_NS_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 50) +# ``early0``: warp 4 streams group 0 right after its own top-k. Off at 8 tokens (0.6 us/layer faster in the chain: +# there it is on every routing warp's critical path and the ring holds only 5 stages). +EARLY0_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 0) +# The plan's ``heavy_fc2_only``: 192 FC2 output blocks over 188 routed CTAs leave 4 CTAs with two blocks, whose +# trailing FC2 after the last hand-off otherwise ends the kernel 3-7 us after the rest; they stream no FC1 instead +# (1.7 / 0.9 us/call faster at 4 / 8 tokens in the TP4 sweep; bit-identical outputs). +HEAVY_FC2_ONLY_BY_M: Tuple[int, ...] = (0, 0, 0, 1, 0, 0, 0, 1) +# ``hb_repl``: copies of the tagged FC1 -> FC2 records (applied when > 1); CTA bx polls copy bx % hb_repl, which +# spreads the polls of the few hot record lines over L2. 7 and 8 tokens hand off through the B images instead. +HB_REPL_BY_M: Tuple[int, ...] = ( + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, +) # off at TP4 (0.2 us/call slower at 4 tokens) +# ``lg_repl``: replicas of the router logit exchange at up to LG_REPL_MAX_TOKENS tokens (fewer stores before the +# exchange; bit-identical outputs). Above that the kernel's default applies. +LG_REPL = 32 +LG_REPL_MAX_TOKENS = 0 # capacity: fixed: measured crossover (off at TP4: flat in the sweep) +# The TP4 FLAT stack: the TP2 FLAT values (no B2 images, 16 router-logit replicas, AR push parities loaded in the +# finalize, PDL release right after the grid wait) over the per-m tables above. Bitwise-identical output (m 1..8 x 4 +# routings, eager + graph); standalone m=4 -0.49 us/call (late_trigger). The values are part of the engine's module +# key. +FLAT_STACK = {"b2img": 0, "lg_repl": 16, "ar_pre": 0, "late_trigger": 0} +# FLAT only up to 4 tokens: at 8 tokens the per-m tables are faster (TP2: 2.96 us/call; TP4 in-model at 2 requests +# +175 us/step with FLAT). Bit-identical either way. +FLAT_MAX_TOKENS = 4 # capacity: fixed: measured crossover + + +def flat_applies(num_tokens: int) -> bool: + """Whether an engine state for ``num_tokens`` tokens runs the FLAT stack.""" + return num_tokens <= FLAT_MAX_TOKENS + + +# Byte-balanced FC1 ranges below one unit per CTA (:func:`_fine_ranges`) for groups too small to give every CTA a +# whole unit; False gives such groups whole units on the CTAs with the least FC2 work. +FINE_PLAN = True +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 + = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row + R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then + down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per + layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + P = I_ // n_cta + split = 128 // (2 * P) + gt = H_ // 128 // split + dkt = I_ // 128 + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + octant = r // P + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(2 * dkt): + jt, kt = gd // dkt, gd % dkt + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, gt + gd, :16384] = img.reshape(-1) + out[s, gt + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def _fine_ranges(weights, total: int): + """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item + counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary + strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit + runs at least to that unit's end. Returns the item counts.""" + n = len(weights) + items = [0] * n + s = 0 + acc = 0.0 + for c in range(n): + acc += weights[c] + e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) + if s % _S1 and e < total: # started inside a unit: run to (at least) its end + e = max(e, min(total, (s // _S1 + 1) * _S1)) + items[c] = e - s + s = e + assert s == total + return items + + +def make_phase_plan( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = FINE_PLAN, + heavy_fc2_only: bool = False, +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + # (heavy_fc2_only) CTAs with more FC2 blocks than the minimum stream no FC1: they run each group's FC2 as soon as + # its hand-off lands, so their trailing FC2 after the last hand-off matches the others' two lagged groups + min_nb = min(nbs) + no_fc1 = [heavy_fc2_only and nbs[c] > min_nb for c in range(num_ctas)] + BIG_W = 1 << 40 + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + for c in range(num_ctas): + if no_fc1[c]: + fc2[c] += BIG_W + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [ + 0 if no_fc1[c] else max(_S1, (limit - f) // _ITEM_BYTES) + for c, f in enumerate(fc2) + ] + + lo, hi = 0, max(f for c, f in enumerate(fc2) if not no_fc1[c]) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 + lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] + sw = sum(wts) + items = _fine_ranges([w * total / sw for w in wts], total) + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + inter: int = INTER, + ): + """``inter``: the routed intermediate per rank; this engine is built for INTER (768) only.""" + if inter != INTER: + raise ValueError( + f"m3_moe (TP4): routed intermediate {inter} per rank is not supported (only {INTER})" + ) + self.inter = inter + logger.info(f"[M3 MoE] Rubin m3_moe engine (tp4 tables): state for {num_tokens} tokens") + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + if self.num_ctas > HIDDEN // 32: + # a routed CTA that owns no 32-row FC2 output block deadlocks the FC2 hand-off on its second call + raise ValueError( + f"m3_moe runs at most {HIDDEN // 32} routed CTAs (one FC2 output block each), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, + num_tokens, + M3_MOE_LAG, + 0.0 if num_tokens == 1 else 1.0, + True, + heavy_fc2_only=bool(HEAVY_FC2_ONLY_BY_M[num_tokens - 1]), + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "i_tp": INTER, + } + if HB_REPL_BY_M[num_tokens - 1] > 1: + self.config["hb_repl"] = HB_REPL_BY_M[num_tokens - 1] + if num_tokens <= LG_REPL_MAX_TOKENS: + self.config["lg_repl"] = LG_REPL + if shared: + self.config["sh_i"] = SHARED_INTER + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + if B2_POLL_NS_BY_M[num_tokens - 1] > 0: + self.config["b2_poll_ns"] = B2_POLL_NS_BY_M[num_tokens - 1] + if not EARLY0_BY_M[num_tokens - 1]: + self.config["early0"] = 0 + if flat_applies(num_tokens): + self.config.update(FLAT_STACK) + # engagement line (every rank; once per engine state = once per token count) + from tensorrt_llm.logger import logger as _flog + + _flog.info( + "MiniMax-M3 MoE engine FLAT (TP4): " + + " ".join(f"{k}={self.config[k]}" for k in FLAT_STACK) + + f" tokens={num_tokens}" + ) + else: + from tensorrt_llm.logger import logger as _flog + + _flog.info( + f"MiniMax-M3 MoE engine FLAT (TP4) off above {FLAT_MAX_TOKENS} tokens (per-m tables): " + f"tokens={num_tokens}" + ) + # E0@m4: at 4 tokens the early group-0 FC1 stages are issued by FC2 B builder warp 8 instead of the routing + # warp 4 (issue order only, bitwise; TP4 -0.19 / TP8 -0.21 us/call). + if num_tokens == 4: + from tensorrt_llm.logger import logger as _e0log + + _e0log.info( + "MiniMax-M3 MoE engine E0M4 (%s): early group-0 stages on warp 8, tokens=4" % "TP4" + ) + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + # Every CTA waits on others (cross-CTA spins on partials, hand-off records, the shared expert's h and y): the + # grid must be resident at once, one CTA per SM. The module's shared-memory / TMEM budget assertion makes one + # CTA fit an SM; the grid must not exceed the SMs. (The kernel also traps after HANG_NS in any such spin.) + if self.num_ctas + mod.NSH > sms: + raise ValueError( + f"m3_moe: a grid of {self.num_ctas} + {mod.NSH} CTAs exceeds the {sms} SMs; its CTAs wait on each " + "other and must all be resident at once" + ) + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full( + (mod.G_CAP * mod.U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw + ) + # the FC1 -> FC2 records; with HB_REPL > 1 copies, HB_STRIDE int32 words apart + hb_words = mod.G_CAP * num_tokens * mod.U1 * 4 + if mod.HB_REPL > 1: + hb_words = mod.HB_REPL * mod.HB_STRIDE + self.hbuf32 = torch.zeros(hb_words, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [3072, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 768], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + # keyed on the config the kernel module was built with (not self.config), so a module built with any other + # trace-time config can never reuse another build's compiled kernel + key = (_ENGINE_PATH,) + tuple(sorted(self.mod.M3_CONFIG.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py new file mode 100644 index 000000000000..446776419315 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py @@ -0,0 +1,288 @@ +# 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. +"""Host side of ``m3_moe16`` (m3_moe16_kernel.py): the MiniMax-M3 MoE layer (fused router + routed experts + shared +expert) for 9..16 decode tokens on one TP4 rank. The 1..8-token engine (op.M3MoeState) is untouched; select this +state only for num_tokens > 8. + +Output: bf16 [M, 6144] = bf16(shared + routed), the routed part bit for bit TRTLLM-Gen's fused MoE on the fp32 router +GEMV logits (m3_router_gemv's math), the shared part as the in-model m3_mx_gemv shared expert. + + state = M3Moe16State(device, num_tokens=12) # before CUDA-graph capture; one per token count + layer = state.layer(w3_w1_weight, w3_w1_weight_scale, w2_weight, w2_weight_scale, fc31_alpha, fc31_scale_c, + act_alpha, act_beta, act_clamp, fc2_alpha, e_score_correction_bias, + router_weight=router_weight, fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf)) + out = layer(hidden) # first call compiles: call once before capture + +Protocol state (per state object, shared by the layers that run on it in stream order): per-group FC2-B images and +unit counters, the B1 image + counter, per-CTA epochs. Every call adds exactly 192 to each group counter and NCTA to +the B1 counter, so nothing is reset between calls (graph replay safe). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch + +from .op import ( + HIDDEN, + SHARED_CTAS, + SHARED_INTER, + U1, + M3MoeLayer, + _compiled, + _kernel_module, + _shared_views, + _view, + is_supported, +) + + +def _group_slots(g: int, p: int, perm: bool, min_ng: int = 5) -> int: + """Slots of group p when the call has g distinct experts (the kernel's slot order).""" + ng = (g + 3) // 4 + tail = g - 4 * (ng - 1) + if p >= ng: + return 0 + if perm and tail < 4 and ng >= max(3, min_ng): + return tail if p == 1 else 4 + return min(4, g - 4 * p) + + +def make_phase_plan16(num_ctas: int, m_max: int, lag: int = 2, perm: bool = True, min_ng: int = 5): + """op.make_phase_plan (early0, tail_w 1) with the kernel's group sizes (PERM: a partial last group is group 1).""" + S1 = 6 + UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) + ITEM = UNIT_BYTES // S1 + blk_bytes = 32 * (768 // 2 + 768 // 16) + gcap = 4 * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + for g in range(4, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = _group_slots(g, p, perm, min_ng) + total = n_p * U1 * S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * _group_slots(g, p - lag, perm, min_ng) * blk_bytes + if p == ng - 1 and not (p == 0 and m_max >= 2): + for q in range(max(0, ng - lag), ng): + fc2[c] += nbs[c] * _group_slots(g, q, perm, min_ng) * blk_bytes + items = [0] * num_ctas + if total >= S1 * num_ctas: + + def counts_for(limit): + return [max(S1, (limit - f) // ITEM) for f in fc2] + + lo, hi = 0, max(fc2) + total * ITEM + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > S1), + key=lambda c: (-(fc2[c] + items[c] * ITEM), c), + ) + items[c] -= 1 + extra -= 1 + else: + for c in sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // S1]): + items[c] = S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +_ENGINE16_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe16_kernel.py") +M16_MIN_TOKENS, M16_MAX_TOKENS = 9, 16 +# Default routed CTAs (beside the SHARED_CTAS shared-expert CTAs). +M16_NUM_CTAS = 128 +# Phase lag: phase p streams FC1 of group p and FC2 of group p - M16_LAG (the kernel's default lag). +M16_LAG = 3 + + +class M3Moe16State: + """Workspace + compiled kernel of ``m3_moe16`` for ``num_tokens`` (9..16) tokens per call.""" + + fused_router = True + shared = True + + def __init__(self, device: torch.device, num_tokens: int, num_ctas: Optional[int] = None): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3Moe16State allocates and compiles: build it before CUDA-graph capture" + ) + if not M16_MIN_TOKENS <= num_tokens <= M16_MAX_TOKENS: + raise ValueError( + f"m3_moe16 takes {M16_MIN_TOKENS}..{M16_MAX_TOKENS} tokens per call, not {num_tokens}" + ) + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe16: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or M16_NUM_CTAS + if not 128 <= self.num_ctas <= sms - SHARED_CTAS: + raise ValueError( + f"m3_moe16 runs 128 .. {sms - SHARED_CTAS} routed CTAs + {SHARED_CTAS} shared (all resident)" + ) + self.m = num_tokens + plan, nb_cap = make_phase_plan16(self.num_ctas, num_tokens, M16_LAG) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "shared": 1, + } + try: + self.mod = mod = _kernel_module(self.config, _ENGINE16_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe16: num_tokens={num_tokens} does not fit the smem / TMEM budget: {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + # split units' head sums: every slot empty (0xFFFFFFFF) between calls + self.partials = torch.full((mod.G_CAP * U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw) + # FC2-B images per group + the B1 image; counters: group g at 8 g, B1 at 8 NG_CAP, CTA tickets at 8 NG_CAP + 8 + self.img = torch.zeros(mod.IMG_TOTAL // 4, dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 16, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros( + mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw + ) # logits: tag << 32 | f32 + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> M3MoeLayer: + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + layer.shared_weights = _shared_views(shared_weights, "m3_moe16") + layer.shared_mxfp8 = tuple(shared_weights) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError(f"m3_moe16: expects contiguous bf16 hidden states [{m}, {HIDDEN}]") + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + layer.bias, + layer.escal, + self.plan.view(-1), + self.partials, + self.img, + self.cnt, + layer.w3_w1_weight_scale.reshape(-1).view(torch.int32), + layer.w2_weight_scale.reshape(-1).view(torch.int32), + layer.w3_w1_weight.reshape(-1).view(torch.int8), + layer.w2_weight.reshape(-1).view(torch.int8), + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + *layer.shared_weights, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_moe16 compiles on its first call: call it once before capture" + ) + import cutlass.cute as cute + + sig = [_view(t, 16 if t.element_size() * t.numel() % 16 == 0 else 4) for t in args] + key = (_ENGINE16_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py new file mode 100644 index 000000000000..48a7827f59b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hand-off of an MXFP8-quantized activation from the kernel that produced it to the next MXFP8 GEMM. + +A producer that already holds a BF16 row in registers (the multicast all-reduce + RMSNorm, the split-K attention +merge) also writes its MXFP8 form (E4M3 rows + UE8M0 scales in the 128x4-swizzled layout of +``trtllm::mxfp8_quantize(x, True)``) and publishes it here keyed by the BF16 tensor. The MXFP8 linear whose input +is that tensor takes it and skips its own quantize kernel. Every take clears the slot, so a hand-off is only ever +consumed by the first MXFP8 linear that runs after the producer. + +Module state shared across files (one instance each per process): + +* ``SLOT`` (PrequantSlot): published by ``m3_mc_allreduce`` (``quant=True``), the sparse-decode merge and the fused + decode attention (MXFP8 output); cleared by the sparse decode before it runs; taken by the MXFP8 linear dispatch + and by the one-launch QKV + attention + o_proj path (which also peeks ``pending``). +* ``AR_PUSH`` (ArPushSlot): armed by the fast forward around the attention call; the o_proj GEMV takes the target + and marks ``pushed``; the fast forward checks ``take_pushed`` before the attention boundary all-reduce. +* ``MOE_PUSH`` (MoePushSlot): armed by the fast forward before the MoE; the routed-expert engine takes ``engine`` + and the shared expert's down projection takes ``shared``, each marks its pushed output; the fast forward + checks ``take_pushed`` before the MoE boundary all-reduce. +""" + + +class PrequantSlot: + def __init__(self): + self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) + + def publish(self, x, xq, xsf): + self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) + + def clear(self): + self.pending = None + + def take(self, x): + """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" + pending, self.pending = self.pending, None + if ( + pending is None + or x.dim() != 2 + or x.data_ptr() != pending[0] + or x.shape[0] != pending[1] + or x.shape[1] != pending[2] + ): + return None + return pending[3], pending[4] + + +SLOT = PrequantSlot() + + +class ArPushSlot: + """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. + + The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output + it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" + + def __init__(self): + self.target = None + self.pushed = None # data_ptr of the pushed output + + def arm(self, target): + self.target = target + self.pushed = None + + def take_target(self): + target, self.target = self.target, None + return target + + def take_pushed(self, x): + pushed, self.pushed = self.pushed, None + return pushed is not None and x.data_ptr() == pushed + + +AR_PUSH = ArPushSlot() + + +class MoePushSlot: + """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, + ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv + ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" + + def __init__(self): + self.engine = None + self.shared = None + self.engine_pushed = None # data_ptr of the pushed routed output + self.shared_pushed = None # data_ptr of the pushed shared output + + def arm(self, engine, shared): + self.engine, self.shared = engine, shared + self.engine_pushed = self.shared_pushed = None + + def disarm(self): + self.engine = self.shared = None + + def take_engine(self): + engine, self.engine = self.engine, None + return engine + + def take_shared(self): + shared, self.shared = self.shared, None + return shared + + def take_pushed(self, shared, routed): + """(routed pushed, shared pushed) for these outputs; clears the marks.""" + e = routed is not None and self.engine_pushed == routed.data_ptr() + s = shared is not None and self.shared_pushed == shared.data_ptr() + self.engine_pushed = self.shared_pushed = None + return e, s + + +MOE_PUSH = MoePushSlot() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_report.py new file mode 100644 index 000000000000..7f864bd240c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_report.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast-path reports a kernel module makes from inside a call: ``taken`` (an engagement line), ``missed`` (a covered +shape that cannot run: raises under ``modeling_v2='require'``), ``uncovered`` and ``capture_missed``. + +The policy (logging once, raising under require) is the target's: ``modeling.py`` installs its functions here at import, +so a kernel reports without importing the target module. Until then every report is a no-op. +""" + +from __future__ import annotations + +from typing import Callable, Optional + + +def _noop(*args, **kwargs) -> None: + return None + + +taken: Callable[..., None] = _noop +missed: Callable[..., None] = _noop +uncovered: Callable[..., None] = _noop +capture_missed: Callable[..., None] = _noop + + +def install( + *, + on_taken: Callable[[str, Optional[int], str], None], + on_missed: Callable[[str, Optional[int], str], None], + on_uncovered: Callable[[str, str], None], + on_capture_missed: Callable[[str, Optional[int], str], None], +) -> None: + """Route this module's reports to the target's policy.""" + global taken, missed, uncovered, capture_missed + taken, missed, uncovered, capture_missed = on_taken, on_missed, on_uncovered, on_capture_missed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_sampled_accept.py new file mode 100644 index 000000000000..dbeb5c8850a6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_sampled_accept.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). + +Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +active must stay on the generic path; the caller decides. + +Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +(the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens heavier +than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the nucleus +with probability >= top_p, so most rows decide on their second exchange. + +Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based Philox +keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one 4-word +record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and reads +the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the same +order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing and +two call halves keep a rank one call ahead from overwriting a slot still being read. +""" + +from typing import Optional + +import torch + +from ..ops.spec.m3_sampled_accept import m3_sampled_accept as _sampled_accept_op +from ..ops.spec.m3_sampled_accept_ws_words import m3_sampled_accept_ws_words as _ws_words_op + +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu: a row's CTAs +# launch as one thread-block cluster, bounded record polls) and reached through the ops/spec entry m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) + + +class SampledAcceptWorkspace: + """Exchange slots behind one multicast mapping (two call halves), the (finished rows, call) counters, the per-row + token scratch, an extra-attempt counter and the SA-1R round-0 decided counter. Created collectively, outside + CUDA-graph capture.""" + + def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): + self.uc, self.mc_ptr, self.rank, self.world = uc, mc_ptr, rank, world + self.handle, self.comm = handle, comm + dev = uc.device + self.counters = torch.zeros(2, dtype=torch.int32, device=dev) + self.tokens = torch.zeros(MAX_ROWS, dtype=torch.int32, device=dev) + self.extra_attempts = torch.zeros(1, dtype=torch.int32, device=dev) + self.decided = torch.zeros(1, dtype=torch.int32, device=dev) + + @classmethod + def create(cls, mapping, device, one_round_k: int) -> "SampledAcceptWorkspace": + # SA-1R's records carry the round-0 list words: the size depends on the variant (``one_round_k``). + words = int(_ws_words_op(max(mapping.tp_size, 1), one_round_k)) + if mapping.tp_size == 1: + uc = torch.zeros(words, dtype=torch.int64, device=device) + return cls(uc, uc.data_ptr(), 0, 1) + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 8, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("SampledAcceptWorkspace.create failed on another rank") + return cls(uc, mc.data_ptr(), mapping.tp_rank, mapping.tp_size, handle, comm) + + +_WS: Optional[SampledAcceptWorkspace] = None + + +def workspace(mapping, device, one_round_k: int) -> Optional[SampledAcceptWorkspace]: + """The process-wide workspace for the SA-1R variant ``one_round_k`` (the same on every call of a process).""" + global _WS + if _WS is None: + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = SampledAcceptWorkspace.create(mapping, device, one_round_k) + return _WS + + +def sampled_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: Optional[torch.Tensor], + seeds: torch.Tensor, + offsets: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, + ws: Optional[SampledAcceptWorkspace] = None, + use_pdl: bool = True, + *, + one_round_k: int, + prewait_params: bool, + early_trigger: bool, + ilp_walk: bool, +): + """Sampled verification (temperature + top-p) from this rank's bf16 logits shard [nreq * (draft_len + 1), V / tp]: + outputs as shard_accept (accepted [nreq, D+1], num_accepted [nreq], gather_ids [nreq]); None + when unavailable. Per-row sampling parameters and Philox (seed, offset) are the spec metadata's, one per row. + top_ks / min_ps are not applied: rows with top-k or min-p active must be routed to the generic path by the + caller. ``one_round_k`` (SA-1R: 0, 9 or 82), ``prewait_params``, ``early_trigger`` and ``ilp_walk`` select the + kernel variant (the target's constants).""" + del top_ks, min_ps + nreq, draft_len = draft_tokens.shape + rows = nreq * (draft_len + 1) + if logits.dim() != 2 or logits.shape[0] != rows or rows > MAX_ROWS or logits.shape[1] > 51200: + return None + if logits.dtype != torch.bfloat16 or logits.stride(1) != 1: + return None + ws = ws if ws is not None else workspace(mapping, logits.device, one_round_k) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + _sampled_accept_op( + logits, + int(vocab_offset), + temperatures[:rows].float().contiguous(), + top_ps[:rows].float().contiguous(), + seeds[:rows].long().contiguous(), + offsets[:rows].long().contiguous(), + ws.uc, + ws.mc_ptr, + ws.counters, + ws.tokens, + ws.rank, + ws.world, + draft_tokens.int().contiguous(), + batch_idx.int().contiguous(), + accepted, + num_accepted, + gather_ids, + int(force_total), + next_new, + force_vec, + ws.extra_attempts, + bool(use_pdl), + prewait_params=bool(prewait_params), + early_trigger=bool(early_trigger), + one_round_k=int(one_round_k), + decided=ws.decided, + ilp_walk=bool(ilp_walk), + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py new file mode 100644 index 000000000000..49bcef462239 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). + +Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp4_select_long``, +one 1024-thread CTA per row), from a CTA of 384 threads: + +* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and + mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's + valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned + key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped + block and masks the result, so a thread's loads issue back to back. +* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of + the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. +* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by + that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate + against the whole row instead. +* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. + +The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait +(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are +compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and +loads them before the executor's warmup, so a decode step never compiles one. +""" + +from typing import Dict, Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.utils import ( + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +SELECT_THREADS = 384 +SELECT_WARPS = SELECT_THREADS // 32 +# 1M tokens + draft tokens at 128 tokens per block. +MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row +MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging +TOPK = 16 # capacity: fixed: MSA's top-k +NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) +# The key of -inf (the reference's 0xff800000): never emitted. +NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 +# (keys per thread, rows up to this many valid blocks) +ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) + +# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare +# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group +# minima [16], per-warp candidate counts [SELECT_WARPS]. +_CAND = 0 +_PICK_KEY = _CAND + 2 * MAX_CAND +_PICK_BLOCK = _PICK_KEY + TOPK +_WARP_MIN = _PICK_BLOCK + TOPK +_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 +_WARP_CNT = _GROUP_MIN + 16 +SMEM_WORDS = _WARP_CNT + SELECT_WARPS +SMEM_BYTES = 4 * SMEM_WORDS + + +def _key(score): + """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" + bits = cutlass.Float32(score).bitcast(cutlass.Int32) + return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) + + +def _imin(a, b): + return cutlass.select_(b < a, b, a) + + +def _imax(a, b): + return cutlass.select_(b > a, b, a) + + +def _forced_key(scores, head, block, query, init_blocks, local_start): + raw = scores[head, block, query] + s = cutlass.select_( + block >= local_start, + cutlass.Float32(LOCAL_SCORE), + cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), + ) + return _key(s) + + +@cute.jit +def _select_row( + row, + scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides + n_valid: cute.Tensor, # int32 [>= total_q] + out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] + init_blocks, + local_blocks, + num_kv_heads, + smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory + bar_id: cutlass.Constexpr, + loaded_trigger: cutlass.Constexpr, + raw_prewait, + n_valid_prewait: cutlass.Constexpr, +): + """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier + ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the + row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by + the caller before its grid wait; otherwise it is read here.""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % cutlass.Int32(32) + warp = tid // cutlass.Int32(32) + query = row // num_kv_heads + head = row % num_kv_heads + if cutlass.const_expr(n_valid_prewait): + raw = raw_prewait + else: + raw = n_valid[query] + valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) + local_start = _imax(raw - local_blocks, cutlass.Int32(0)) + cand = cute.make_tensor( + cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) + ) + + if tid < TOPK: + smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) + smem[_PICK_BLOCK + tid] = cutlass.Int32(0) + + if valid <= TOPK: + # Every valid block is picked (thread t holds block t). + if tid < valid: + smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) + smem[_PICK_BLOCK + tid] = tid + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): + items = ITEM_CLASSES[ci][0] + lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] + hi = ITEM_CLASSES[ci][1] + if (valid > lo) & (valid <= hi): + key = cute.make_rmem_tensor(items, cutlass.Int32) + local_min = cutlass.Int32(NO_KEY) + for s in cutlass.range_constexpr(items): + # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. + block = tid + cutlass.Int32(s * SELECT_THREADS) + k = _forced_key( + scores, head, _imin(block, valid - 1), query, init_blocks, local_start + ) + key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) + local_min = _imin(local_min, key[s]) + local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) + if lane < 16: + smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + if tid < 16: + g = cutlass.Int32(NO_KEY) + for w in cutlass.range_constexpr(SELECT_WARPS): + g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) + smem[_GROUP_MIN + tid] = g + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + bound = cutlass.Int32( + -2147483648 + ) # the smallest key: signed keys, unlike the reference's unsigned ones + for g in cutlass.range_constexpr(16): + bound = _imax(bound, smem[_GROUP_MIN + g]) + + # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(items): + cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) + incl = cnt + for d in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync( + incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) + ) + incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) + if lane == 31: + smem[_WARP_CNT + warp] = incl + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + before = cutlass.Int32(0) + total = cutlass.Int32(0) + for w in cutlass.range_constexpr(SELECT_WARPS): + c = smem[_WARP_CNT + w] + before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) + total += c + slot = before + incl - cnt + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + if slot < MAX_CAND: + cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( + tid + cutlass.Int32(s * SELECT_THREADS) + ) + slot += 1 + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if total <= MAX_CAND: + for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): + mine = cand[i] + rank = cutlass.Int32(0) + for j in cutlass.range(total, unroll=8): + rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) + smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) + else: + # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + block = tid + cutlass.Int32(s * SELECT_THREADS) + rank = cutlass.Int32(0) + for j in cutlass.range(valid, unroll=1): + kj = _forced_key(scores, head, j, query, init_blocks, local_start) + less = (kj < key[s]) | ((kj == key[s]) & (j < block)) + rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = key[s] + smem[_PICK_BLOCK + rank] = block + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if warp == 0: + # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. + v = cutlass.Int32(NO_KEY) + if lane < TOPK: + pk = smem[_PICK_KEY + lane] + if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): + v = smem[_PICK_BLOCK + lane] + for size_log in cutlass.range_constexpr(1, 6): + for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): + other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) + ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) + lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) + v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) + if lane < TOPK: + out[query, head, lane] = cutlass.select_( + v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v + ) + + +class _SelectLongDslKernel: + """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count + before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because + ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's + first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it + (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" + + def __init__( + self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool + ): + self.use_pdl = use_pdl + self.early_trigger = early_trigger + self.loaded_trigger = (not early_trigger) and trigger_phase == 1 + self.n_valid_prewait = n_valid_prewait + + @cute.kernel + def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): + row, _, _ = cute.arch.block_idx() + smem = SmemAllocator() + words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) + raw_prewait = cutlass.Int32(0) + if cutlass.const_expr(self.n_valid_prewait): + raw_prewait = n_valid[row // num_kv_heads] + griddepcontrol_wait() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + _select_row( + row, + scores, + n_valid, + out, + init_blocks, + local_blocks, + num_kv_heads, + words, + bar_id=1, + loaded_trigger=self.loaded_trigger, + raw_prewait=raw_prewait, + n_valid_prewait=self.n_valid_prewait, + ) + + @cute.jit + def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): + num_kv_heads = scores.shape[0] + self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( + grid=(scores.shape[2] * num_kv_heads, 1, 1), + block=(SELECT_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + +# The select reads its row's valid-block count before its grid wait instead of after it, which takes one dependent +# global load off the post-wait path; the selection is unchanged bit for bit. Exact because the count +# (MsaMetadata.msa_n_valid_blocks) is written only before the step's first layer, behind a plain launch +# (attention/test_modeling_v2_m3_n_valid_contract.py). +N_VALID_PREWAIT = True +#: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). +kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} + + +def kernel_key( + use_pdl: bool, early_trigger: bool, trigger_phase: int +) -> Tuple[bool, bool, int, bool]: + """The compile key of a launch (the trigger phase does not matter with the early trigger). The kernel reads its + row's valid-block count before its grid wait (``N_VALID_PREWAIT``).""" + return ( + bool(use_pdl), + bool(early_trigger), + 0 if early_trigger else int(trigger_phase), + N_VALID_PREWAIT, + ) + + +def _compile(key: Tuple[bool, bool, int, bool]): + if key not in kernel_cache: + heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() + scores = make_fake_tensor( + cutlass.Float32, + (heads, blocks, rows), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) + out = make_fake_tensor( + cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) + ) + kernel_cache[key] = cute.compile( + _SelectLongDslKernel(*key), + scores, + n_valid, + out, + cutlass.Int32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + return kernel_cache[key] + + +def select_long_dsl( + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool, + early_trigger: bool, + trigger_phase: int, +) -> None: + """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" + _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( + scores, n_valid_blocks, out, init_blocks, local_blocks + ) + + +def prebuild(keys, device) -> list: + """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" + scores = torch.zeros( + 1, 2048, 4, dtype=torch.float32, device=device + ) # capacity: fixed: a dummy prebuild row, not a size + n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) + out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) + built = [] + for key in keys: + key = kernel_key(*key) + _compile(key)(scores, n_valid, out, 0, 1) + built.append(key) + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py new file mode 100644 index 000000000000..2b7d05ad920f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 verification for the MiniMax-M3 target head (vocab-sharded TP head, verify rows <= 48). + +Replaces the all-gather of the target logits, the argmax and the strict-acceptance ops of a greedy verify step with +one kernel over this rank's BF16 logits shard (the target lm_head's own local GEMM output): + +* Every logit becomes a 64-bit key ``ordered(value) << 32 | ~global_index`` (ties pick the lowest index, as + ``torch.argmax``), max-reduced per row: warp shuffles, then ``atomicMax``. +* The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping + (Lamport: 0 = empty, the key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the + slots and runs the strict acceptance. Call n uses slot half n % 2. + +The argmax is taken over the BF16 logits without arithmetic, so the tokens are those of the in-tree path's argmax +over the gathered logits. The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/ +shardAccept.cu) and reached through the ``ops/spec`` entry ``m3_shard_accept``. +""" + +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from ..ops.spec.m3_shard_accept import m3_shard_accept as _shard_accept_op + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 # capacity: fixed: DECODE_MAX_TOKENS verify rows + + +@dataclass(eq=False) +class AcceptWorkspace: + """Per-rank exchange slots ([2 halves][SHARD_ROWS rows][world] uint64) behind one multicast mapping, the local + key accumulators and the (finished CTAs, call) counters.""" + + uc: torch.Tensor + mc_ptr: int + keys: torch.Tensor + counters: torch.Tensor + rank: int + world_size: int + handle: Any = None + comm: Any = None + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + if not _mnnvl_workspace_all_succeeded(comm, True): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = mc.data_ptr(), mapping.tp_rank + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None + return AcceptWorkspace( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) + + +def _workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace, created collectively outside CUDA-graph capture on first use.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + # A capture that finds it missing is the caller's miss (it is built on the eager steps). + return None + _WS = create_workspace(mapping, device, SHARD_ROWS) + return _WS + + +def shard_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, +): + """Greedy verification from this rank's BF16 logits shard. + + Args: + logits: [nreq * (draft_len + 1) <= SHARD_ROWS, V / tp] BF16, unit column stride: this rank's columns + (global ids ``vocab_offset ..``) of the verify rows, request-major. + draft_tokens: [nreq, draft_len] int32, the drafts being verified. + batch_idx: [nreq] int32, each request's row in the verify layout. + force_total / force_vec: forced acceptance (see ``m3_shard_accept``). + next_new: optional int32 [nreq, >= 1]; column 0 receives the last accepted token. + + Returns: + (accepted_tokens [nreq, draft_len + 1] int32, num_accepted [nreq] int32, first-pass gather ids [nreq] + int64), or None when the shape is not served or the workspace is missing during a capture. Collective over + the TP group. + """ + nreq, draft_len = draft_tokens.shape + if ( + logits.dim() != 2 + or logits.shape[0] != nreq * (draft_len + 1) + or logits.shape[0] > SHARD_ROWS + or logits.dtype != torch.bfloat16 + or logits.stride(1) != 1 + ): + return None + ws = _workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + _shard_accept_op( + logits, + vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_tokens.contiguous(), + batch_idx, + accepted, + num_accepted, + gather_ids, + force_total, + next_new, + _num_sms(), + force_vec, + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py new file mode 100644 index 000000000000..0d4dc9cd1b64 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small fused kernels of the MiniMax-M3 Eagle3 step (GB300, decode). + +``chunk_norm`` / ``capture_norm`` apply the drafter's per-layer fc_norm RMSNorms to the captured target hidden states +(the in-tree ``apply_eagle3_fc`` norms, same fp32 arithmetic); ``draft_input`` is a fused draft pass's input (the +embedding lookup, the drafter layer's input_layernorm / hidden_norm and their concat, same arithmetic); +``step_gather`` / ``slot_scatter`` move the step's integer outputs between the spec sampler's store and the engine's +input buffers. +""" + +from typing import Optional, Tuple + +import torch + +# The kernels are compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu) and reached through the +# ops/spec entries. +from ..ops.spec.m3_capture_norm import m3_capture_norm as _capture_norm_op +from ..ops.spec.m3_chunk_norm import m3_chunk_norm as _chunk_norm_op +from ..ops.spec.m3_draft_input import m3_draft_input as _draft_input_op +from ..ops.spec.m3_slot_scatter import m3_slot_scatter as _slot_scatter_op +from ..ops.spec.m3_step_gather import m3_step_gather as _step_gather_op + + +def draft_input( + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + eps_embed: float, + w_hidden: torch.Tensor, + eps_hidden: float, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(cat([rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden], -1) [M, 2H], hidden[rows] [M, H]) for + int32 ids [M] and int64 rows [M] (rows None: hidden[:M]); ids outside the table embed to zero.""" + m, h = ids.shape[0], table.shape[1] + out = torch.empty((m, 2 * h), dtype=table.dtype, device=table.device) + residual = torch.empty((m, h), dtype=table.dtype, device=table.device) + return _draft_input_op( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, use_pdl + ) + + +def chunk_norm( + src: torch.Tensor, weights: torch.Tensor, eps: float, use_pdl: bool = True +) -> torch.Tensor: + """Per-chunk RMSNorm of src [M, >= C*H] (row stride may exceed C*H) with weights [C, H] -> [M, C*H].""" + c, h = weights.shape + out = torch.empty((src.shape[0], c * h), dtype=src.dtype, device=src.device) + return _chunk_norm_op(src, weights, eps, out, use_pdl) + + +# The overlap scheduler's per-step gathers and the spec sampler's slot stores as one kernel each (step_gather / +# slot_scatter) instead of ~10 and 4 small torch ops; called by the target's SpecStepFold. + + +def step_gather( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 +) -> bool: + """See step_gather_kernel; False when the inputs are outside what the kernel covers (caller falls back).""" + ts = (new_tok, next_draft, lens, slots, ppi, ids, draft, pos_off, kv_off) + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or next_draft.dim() != 2 + or new_tok.dim() < 2 + ): + return False + if (new_tok.dim() == 3 and new_tok.shape[2] != 1) or not ( + lens.is_contiguous() + and ids.is_contiguous() + and draft.is_contiguous() + and pos_off.is_contiguous() + and kv_off.is_contiguous() + ): + return False + _step_gather_op( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 + ) + return True + + +def slot_scatter( + slots, src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft +) -> bool: + """See slot_scatter_kernel (2-D int32 views, zero padding); False to fall back.""" + ts = (src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft) + if any(t.dtype != torch.int32 or not t.is_cuda for t in ts) or slots.dtype != torch.long: + return False + two = (src_new, dst_new, src_next_new, dst_next_new, src_draft, dst_draft) + if any(t.dim() != 2 for t in two) or src_lens.dim() != 1 or dst_lens.dim() != 1: + return False + _slot_scatter_op( + slots, + src_new, + dst_new, + src_next_new, + dst_next_new, + src_lens, + dst_lens, + src_draft, + dst_draft, + ) + return True + + +def capture_norm( + h: torch.Tensor, + residual: Optional[torch.Tensor], + weight: torch.Tensor, + eps: float, + dst: torch.Tensor, +) -> None: + """dst [M, H] (row stride may exceed H) = rmsnorm(h + residual) * weight with chunk_norm's rounding.""" + _capture_norm_op(h, residual, weight, eps, dst) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py new file mode 100644 index 000000000000..f5a1b229c777 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard of a verify step of at most 8 rows, streamed from a lossless 12-bit coding of its BF16 +weight (75% of the bytes; ``ops.gemm.m3_target_head`` has the format). + +``prepare_head`` codes this rank's shard once, before the executor's warmup, and keeps it only when it decodes to the +BF16 weight bit for bit (``round_trip_ok``); a shard that does not keeps the in-tree BF16 GEMM. ``logits`` runs the +kernel: BF16 x BF16 MMAs with fp32 accumulation over the full K and one BF16 rounding per logit, the in-tree head's +arithmetic in another summation order. The kernel is compiled into TensorRT-LLM +(cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu) and called through the catalog entry ``ops.gemm.m3_target_head``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from ..ops.gemm.m3_target_head import CTA_WARPS, MAX_ROWS, encode_chunked, round_trip_ok +from ..ops.gemm.m3_target_head import m3_target_head as _target_head_op + + +@dataclass(eq=False) +class TargetHead: + """One rank's target lm_head shard in the 12-bit coding, and the kernel's split-tile scratch. + + ``coded`` is False when the coding did not reproduce the weight bit for bit (the in-tree GEMM then runs).""" + + weight: torch.Tensor + coded: bool + w: Optional[torch.Tensor] = None + base: int = 0 + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + part: Optional[torch.Tensor] = None + counter: Optional[torch.Tensor] = None + grid: int = 0 + + +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def head_supported(lm_head) -> bool: + """A bias-free 2-D BF16 head of whole 16-row tiles whose K the coding covers (a multiple of 64).""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + return ( + getattr(lm_head, "bias", None) is None + and weight.shape[0] % 16 == 0 + and weight.shape[1] % 64 == 0 + ) + + +def prepare_head(lm_head, grid: Optional[int] = None) -> TargetHead: + """Code ``lm_head``'s local shard and check it decodes to the BF16 weight bit for bit (fails closed: a shard that + does not is returned with ``coded=False`` and nothing kept). Built once, eagerly, before the executor's warmup + (``M3Eagle3Worker.prebuild_target_head``).""" + weight = lm_head.weight + with torch.no_grad(): + w = weight.detach() + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if not round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return TargetHead(weight=weight, coded=False) + if grid is None: + grid = _num_sms() + n = weight.shape[0] + return TargetHead( + weight=weight, + coded=True, + w=packed, + base=base, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + part=torch.zeros(grid * CTA_WARPS * 2 * 4 * 32, dtype=torch.float32, device=weight.device), + counter=torch.zeros((n + 15) // 16, dtype=torch.int32, device=weight.device), + grid=grid, + ) + + +def logits(head: TargetHead, x: torch.Tensor, use_pdl: bool = True) -> torch.Tensor: + """This rank's BF16 logits shard [rows, n] of the rows x [rows <= MAX_ROWS, K] (BF16) from the coded shard.""" + assert head.coded and 1 <= x.shape[0] <= MAX_ROWS, (head.coded, x.shape) + if x.stride(1) != 1 or x.stride(0) % 8 or x.data_ptr() % 16: + x = x.contiguous() + n = head.weight.shape[0] + out = torch.empty((x.shape[0], n), dtype=torch.bfloat16, device=x.device) + return _target_head_op( + head.w, + x, + out, + head.part, + head.counter, + head.exc_off, + head.exc_pos, + head.exc_w, + head.base, + n, + head.grid, + use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py new file mode 100644 index 000000000000..db0077117fe3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The token map of ``thop.trtllm_gen_generation_preprocess`` for one request with several tokens. + +For multi-token generation the op's QKV preprocessing places every row by the workspace token map (``tokensInfo``, +(request, token) per row), but builds that map only when ``isBuildDecoderInfoKernelNeeded()`` holds, which at one +request with the default RoPE it does not. The op then applies whatever map the last multi-request call left in the +workspace: a row can land on another request's length and pages, and its RoPE position and K/V slot are wrong. This +module wraps the op so that, for one request and several tokens, the map the build would write, (0, t) for each row t, +is written into the op's own map slice of the workspace (the offset from the op's layout function) just before the +call. Every other call passes through unchanged. The rows' positions still come from the request's generation length, +exactly as with a freshly built map. + +The write is a device copy from a map prebuilt at install (no allocation, no host sync), so it is captured into the +decode graphs with the call it precedes. +""" + +from __future__ import annotations + +from typing import Callable, Dict, Optional, Tuple + +import torch + +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.logger import logger + +#: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). +_ARG_NAMES = ( + "qkv_input", + "workspace", + "sequence_lengths", + "spec_decoding_generation_lengths", + "spec_decoding_position_offsets", + "kv_cache_block_offsets", + "host_kv_cache_pool_pointers", + "host_kv_cache_pool_mapping", + "kv_scale_orig_quant", + "kv_scale_quant_orig", + "attention_output_orig_quant", + "rotary_inv_freq", + "rotary_cos_sin", + "mrope_position_deltas", + "layer_idx", + "seq_offset", + "num_heads", + "num_kv_heads", + "head_size", + "tokens_per_block", + "kv_cache_quant_mode", + "max_attention_window_size", + "cyclic_attention_window_size", + "num_tokens", + "batch_beam", + "input_seq_length", + "max_past_kv_length", + "rotary_embedding_dim", + "rotary_embedding_base", + "rotary_embedding_scale_type", + "rotary_embedding_scale", + "rotary_embedding_max_positions", + "position_embedding_type", + "bmm1_scale", + "bmm2_scale", + "fp8_context_fmha", + "predicted_tokens_per_seq", + "attention_chunk_size", + "multi_processor_count", + "total_num_blocks", + "kv_factor", + "need_build_kv_cache_metadata", + "cross_attention", + "skip_fmha_workspace", +) +_INT2_BYTES = 8 # capacity: fixed: an element size + +_original: Optional[Callable] = None +_max_rows = 0 +_maps: Dict[torch.device, torch.Tensor] = {} +_offsets: Dict[Tuple, int] = {} + + +def _no_taken(rows: int) -> None: + pass + + +def _no_missed(rows: int, reason: str) -> None: + pass + + +_on_taken: Callable[[int], None] = _no_taken +_on_missed: Callable[[int, str], None] = _no_missed + + +def _indexed(device: torch.device) -> torch.device: + """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: + """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" + device = _indexed(device) + m = _maps.get(device) + if m is None or m.shape[0] < rows: + reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" + _on_missed(rows, reason) + raise RuntimeError(f"m3_tokmap_guard: {reason}") + return m[:rows] + + +def _map_offset(args: dict, rows: int) -> int: + x = args["qkv_input"] + skip = bool(args.get("skip_fmha_workspace", False)) + key = ( + x.dtype, + rows, + int(args["num_heads"]), + int(args["head_size"]), + int(args["rotary_embedding_dim"]), + int(args["num_kv_heads"]), + skip, + ) + off = _offsets.get(key) + if off is None: + layout = thop.get_trtllm_gen_generation_workspace_layout( + x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip + ) + off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) + if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: + raise RuntimeError( + f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " + f"(offset {off}, {size} bytes): the layout changed under this guard" + ) + _offsets[key] = off + return off + + +def _guarded(*args, **kwargs): + a = dict(zip(_ARG_NAMES, args)) + a.update(kwargs) + rows = int(a["num_tokens"]) + if ( + a.get("spec_decoding_generation_lengths") is not None + and int(a["predicted_tokens_per_seq"]) > 1 + and int(a["batch_beam"]) == 1 + and rows > 1 + and not bool(a.get("cross_attention", False)) + ): + ws = a["workspace"] + src = _one_request_map(ws.device, rows) + off = _map_offset(a, rows) + # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. + dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] + dst.view(torch.int32).view(rows, 2).copy_(src) + _on_taken(rows) + return _original(*args, **kwargs) + + +def install( + device: torch.device, + max_rows: int, + on_taken: Callable[[int], None], + on_missed: Callable[[int, str], None] = _no_missed, +) -> bool: + """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to + ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this + call installed the wrapper.""" + global _original, _max_rows, _on_taken, _on_missed + device = _indexed(device) + if device not in _maps or _maps[device].shape[0] < max_rows: + t = torch.arange(max_rows, dtype=torch.int32, device=device) + _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() + _max_rows = max(_max_rows, max_rows) + _on_taken = on_taken + _on_missed = on_missed + if _original is not None: + return False + _original = thop.trtllm_gen_generation_preprocess + thop.trtllm_gen_generation_preprocess = _guarded + logger.info( + "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " + f"(up to {max_rows} rows)" + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/__init__.py new file mode 100644 index 000000000000..77d7b84d9bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py new file mode 100644 index 000000000000..258d45007a9b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe16_engine.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 9..13 decode rows in one CuTe DSL launch (``m3_moe16``): the router, the routed NVFP4 +experts and the shared MXFP8 expert, on caller-owned state (an ``M3Moe16State`` and one ``M3MoeLayer`` per MoE +layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..m3_moe.op import M3MoeLayer +from ..m3_moe.op16 import M3Moe16State +from .m3_moe_engine import moe_reference + +__all__ = ["M3Moe16State", "m3_moe16_engine"] + + +class M3Moe16Engine(OpWrapper): + """One MiniMax-M3 MoE layer for M (9..13 here; the state takes 9..16) decode rows: + ``layer.state.run(layer, hidden, out, arp)``, bf16(shared + routed) [M, 6144]. + + The same arguments and the same math as ``m3_moe_engine`` (the routed part operation for operation, the shared + expert always folded in) on an ``M3Moe16State``: a different schedule for more rows. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # One build per row count (m_max is a trace-time constant); the target routes 9..13 rows here (spec-decode verify + # of 2 x 6 rows lands at 12). + CELLS: tuple[Cell, ...] = ( + Cell(why="9 rows: the smallest build, a partial last slot group", spec=dict(tokens=9)), + Cell(why="10 rows: its own build (m_max is a trace-time constant)", spec=dict(tokens=10)), + Cell(why="11 rows: its own build, an odd row count", spec=dict(tokens=11)), + Cell(why="12 rows: two requests verifying 5-token drafts", spec=dict(tokens=12)), + Cell(why="13 rows: the largest the target routes here", spec=dict(tokens=13)), + ) + + note = """ + Bands as `m3_moe_engine` (same approximations, same quantization boundaries). Not covered: 14..16 rows, which the + state takes and this target never routes here; the all-reduce push (`arp`). No `is_valid`, as `m3_moe_engine`. + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return layer.state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``m3_moe_engine``'s reference with the shared expert; `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=True) + + +m3_moe16_engine = M3Moe16Engine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..7fe1bcb50976 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py @@ -0,0 +1,326 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 1..8 decode tokens in one CuTe DSL launch (``m3_moe``): the router, the routed NVFP4 +experts and the shared MXFP8 expert, on caller-owned state (an ``M3MoeState`` and one ``M3MoeLayer`` per MoE layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ..m3_moe.op import ( + HIDDEN, + INTER, + SHARED_INTER, + TOP_K, + M3MoeLayer, + M3MoeState, + ar_push_desc, + is_supported, +) + +__all__ = [ + "M3MoeLayer", + "M3MoeState", + "ar_push_desc", + "is_supported", + "m3_moe_engine", + "moe_reference", + "moe_route", + "mxfp8_dequant", + "mxfp8_quant_dequant", + "unswizzle_sf", +] + +# The model's constants the kernels are built with: routed_scaling_factor, the shared expert's SwiGLU-OAI +# (swiglu_alpha, swiglu_limit, the (up + 1) offset). +ROUTED_SCALING = 2.0 +SHARED_ALPHA = 1.702 +SHARED_LIMIT = 7.0 + +_LOG2E = 1.4426950216293334961 +_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) +_F32_MAX = 3.40282346638528859812e38 + + +def _f32(x: float) -> torch.Tensor: + return torch.tensor(x, dtype=torch.float32) + + +def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: + """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" + lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) + lo = lut[(packed & 0xF).long()] + hi = lut[(packed >> 4).long()] + return torch.stack((lo, hi), dim=-1).flatten(-2) + + +def _e2m1_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" + a = x.abs() + q = torch.full_like(a, 6.0) + for hi, value, inclusive in ( + (5.0, 4.0, True), + (3.5, 3.0, False), + (2.5, 2.0, True), + (1.75, 1.5, False), + (1.25, 1.0, True), + (0.75, 0.5, False), + (0.25, 0.0, True), + ): + q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) + return torch.where(x < 0, -q, q) + + +def _e4m3_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" + return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + + +def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: + """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" + bits = x.contiguous().view(torch.int32) + code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + return code.clamp(0, 254) + + +def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: + """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" + flat = sf.reshape(-1).view(torch.uint8) + r = torch.arange(rows, device=flat.device).view(-1, 1) + b = torch.arange(blocks, device=flat.device).view(1, -1) + padded_blocks = (blocks + 3) // 4 * 4 + idx = ( + (r // 128) * (128 * padded_blocks) + + (b // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (b % 4) + ) + return flat[idx] + + +def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" + rows, k = q.shape + codes = unswizzle_sf(sf, rows, k // 32).float() + return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( + 32, dim=1 + ) + + +def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: + """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to + fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" + rows, k = x.shape + blocks = x.float().view(rows, k // 32, 32) + code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) + scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) + return (_e4m3_round(blocks / scale) * scale).view(rows, k) + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16).float() + + +def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" + return (a.double() @ b.double().t()).float() + + +def _row_perm(rows: int, gated: bool) -> torch.Tensor: + """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle + of epilogue tile 128).""" + probe = torch.empty(rows, 1) + perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) + if gated: + perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] + return perm + + +def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: + """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row + order (the global scale is not applied: it is the per-expert alpha).""" + rows, half = w.shape + codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() + vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) + return vals[perm_inv] + + +def moe_route(x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor): + """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" + logits = _gemm(x, router_weight) + sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() + keys = sig + bias + ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] + sk = keys.gather(1, ids) - bias[ids] + total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) + w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) + return ids, _bf16(w) + + +def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: + """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" + rows, k = x.shape + blocks = x.view(rows, k // 16, 16) + vmax = blocks.abs().amax(dim=-1, keepdim=True) + sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) + # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element + # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). + osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) + return (_e2m1_round(blocks * osc) * sff).view(rows, k) + + +def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: + """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 + input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with + act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" + a, c, alpha, beta, lim = (esc[i] for i in range(5)) + gc = torch.where(gate > lim, lim, gate) + gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) + uc = torch.minimum(torch.maximum(up, -lim), lim) + c1 = a * (alpha * _f32(_LOG2E).to(up.device)) + sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) + h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) + blocks = h.view(-1, 16) + amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) + sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) + osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) + return (_e2m1_round((blocks * osc) * c) * sff).view(-1) + + +def _shared_expert(x: torch.Tensor, shared_mxfp8) -> torch.Tensor: + """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the + model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" + gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 + gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) + gate, up = gu[:, :SHARED_INTER], gu[:, SHARED_INTER:] + a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) + ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 + sg = 1.0 / ( + torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 + ) + h = _bf16(ub * (sg * a)) + return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) + + +def moe_reference(hidden: torch.Tensor, layer: M3MoeLayer, shared: bool) -> torch.Tensor: + """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" + x = hidden.float() + m = x.shape[0] + ids, weights = moe_route(x, layer.router_weight, layer.bias) + xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) + inv31 = torch.argsort(_row_perm(2 * INTER, gated=True)).to(x.device) + inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) + y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) + for e in torch.unique(ids).tolist(): + w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) + w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) + esc = layer.escal[:, e] + for t, k in (ids == e).nonzero().tolist(): + acc1 = _gemm(xr[t : t + 1], w31)[0] + # The buffer holds [w3 (up); w1 (gate)]. + h = _fc1_epilogue(acc1[:INTER], acc1[INTER:], esc) + y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) + acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) + for k in range(TOP_K): + acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() + routed = acc.float() + if not shared: + return routed.to(torch.bfloat16) + return (_shared_expert(x, layer.shared_mxfp8) + _bf16(routed)).to(torch.bfloat16) + + +class M3MoeEngine(OpWrapper): + """One MiniMax-M3 MoE layer for M (1..8) decode tokens: ``layer.state.run(layer, hidden, out, arp)``. + + ``hidden`` bf16 [M, 6144] contiguous, M = the state's token count. ``layer`` (from ``M3MoeState.layer``) carries the + TRTLLM-Gen NVFP4 expert buffers read in place, the per-expert scalars, the router weight / bias, the activation's + NVFP4 global scale and, on a ``shared`` state, the shared expert's MXFP8 weights; its state carries the compiled + kernel and the workspace every layer on it shares. Returns bf16 [M, 6144]: bf16(shared + routed) on a ``shared`` + state, else the routed partial. ``out``: write there instead of a new tensor. ``arp`` (``ar_push_desc``): also + push the rows into a multicast all-reduce exchange. + + The routed part is TRTLLM-Gen's NVFP4 MoE on the fp32 router logits: sigmoid (tanh form) + correction bias top-4 + of 128 (ties to the lower id), weights sigmoid / sum * 2.0 rounded to bf16; the activation quantized to NVFP4 with + the layer's global scale; per expert the raw FC1 sums, SwiGLU-OAI (alpha / beta / clamp), the NVFP4 requant with + fc31_scale_c, FC2, bf16(fc2_alpha x sum); the fma chain over the top-4 in order. The shared expert is the target's + MXFP8 SwiGLU-OAI MLP. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # One cell per token count: m_max is a trace-time constant, so each is its own build, and the target builds all + # eight (decode batches of 1..8 requests, spec-decode verify rows up to 8). The whys name what the build differs in. + CELLS: tuple[Cell, ...] = ( + Cell( + why="1 token: pre-tiled shared expert, last phase balanced without trailing FC2", + spec=dict(tokens=1), + ), + Cell( + why="2 tokens: pre-tiled shared expert, the one build without AR_PRE", + spec=dict(tokens=2), + ), + Cell( + why="3 tokens: 4 FC2-B builder warps, the shared expert from swizzled scales", + spec=dict(tokens=3), + ), + Cell( + why="4 tokens: one token per routing warp, logits replicated 128x", spec=dict(tokens=4) + ), + Cell( + why="5 tokens: 5 builder warps, a routing warp takes two tokens, logits 32x", + spec=dict(tokens=5), + ), + Cell(why="6 tokens: one request verifying a 5-token draft", spec=dict(tokens=6)), + Cell(why="7 tokens: the B2IMG global FC1 -> FC2 hand-off", spec=dict(tokens=7)), + Cell( + why="8 tokens: B2IMG with a 6-stage ring, 32 distinct experts at most", + spec=dict(tokens=8), + ), + ) + + note = """ + The kernel's elementwise math uses the PTX approximations (rcp.approx, ex2.approx, tanh) where `reference` uses + torch's correctly rounded ones, and its GEMMs sum in tensor-core order where `reference` sums in fp64: a + quantization boundary (NVFP4 activation / intermediate code, routing near-tie) can fall on the other side, so the + two agree within the band the test derives, not bit for bit. The routed part is bit for bit the stock TRTLLM-Gen + NVFP4 MoE (fp4_quantize + fp4_block_scale_moe_runner) on the same routing; the test checks that too. + + Not covered: the all-reduce push (`arp`), which needs a multicast exchange (the decode forward's boundary all-reduce + exercises it); the non-shared state (`shared=False`), which this target never builds. Every layer on one state + must run in one stream order: the workspace is shared and tagged per call, never reset. + + No `is_valid`: the compiled call checks every buffer's dtype, shape and alignment against its signature and + raises on a mismatch (a misaligned `hidden`, an `out` that is not bf16 [M, 6144], a malformed `arp`). + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return layer.state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Plain torch over the layer's buffers (see the class docstring); `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=layer.state.shared) + + +m3_moe_engine = M3MoeEngine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..6e0a44f38b30 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TF32 tensor-core router of the unfused decode MoE (9..16 generation rows), optionally with the NVFP4 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) + + +def nvfp4_quantize( + x: torch.Tensor, sf_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], + the pre-rounding scaled values [M, K]). + + Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. + """ + gs = sf_scale.float().reshape(()) + m, k = x.shape + blocks = x.float().view(m, k // 16, 16) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +def tf32_rna(w: torch.Tensor) -> torch.Tensor: + """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" + bits = w.contiguous().view(torch.int32) + rounded = (bits + 0x1000) & ~0x1FFF + finite = (bits & 0x7F800000) != 0x7F800000 + return torch.where(finite, rounded, bits).view(torch.float32) + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation (``fp4_quantize``'s bytes): returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target routes the unfused MoE's pure generation steps of 9..16 rows here (ROUTER_TC_GEN), logits only. + CELLS: tuple[Cell, ...] = ( + Cell( + why="9 rows: a 3-request verify past the MoE engine", spec=dict(m=9, k=6144, fp4=False) + ), + Cell(why="10 rows", spec=dict(m=10, k=6144, fp4=False)), + Cell(why="11 rows", spec=dict(m=11, k=6144, fp4=False)), + Cell(why="12 rows: 3 requests x 4 tokens", spec=dict(m=12, k=6144, fp4=False)), + Cell(why="13 rows: a partial second token tile", spec=dict(m=13, k=6144, fp4=False)), + Cell(why="14 rows", spec=dict(m=14, k=6144, fp4=False)), + Cell(why="15 rows", spec=dict(m=15, k=6144, fp4=False)), + Cell( + why="16 rows: 4 requests x 4 tokens, two full token tiles", + spec=dict(m=16, k=6144, fp4=False), + ), + ) + + note = """ + The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both + sum the same 16 slice partials in the same order. Only the one the test device picks is driven. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> tuple[torch.Tensor, ...]: + """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" + with true_fp32_matmul(): + logits = x.float() @ tf32_rna(w).t() + if sf_scale is None: + return (logits,) + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return logits, packed, sf + + def is_valid( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + if sf_scale is not None: + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py new file mode 100644 index 000000000000..8c2a67356d3f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One kernel for MSA's device-side KV-length correction (MsaAttentionMetadata.on_update_kv_lens). + +The torch formulation is ~23 one-CTA kernels (min, cumsum, gathers, div / clamp / remainder, the block-table lookup, +the valid-block counts) and runs at the head of every captured decode step, twice. This kernel computes the same +int32 results in one launch: + + kv_true[b] = min(kv_lens[b], staged[b]) -> seq_lens_out[b], cu_kv[b + 1] = prefix sum (cu_kv[0] = 0) + pos[t] = kv_true[r] - qo[r] + intra[t] (r = q_batch_row[t]) + out_loc[t] = table[r, min(max(pos, 0) // page, ncols - 1)] * page + max(pos, 0) % page + n_valid[t] = (max(pos + 1, 0) + page - 1) // page + +``msa_kv_update`` is the target's ``MsaKvUpdate`` (``tensorrt_llm._torch.pyexecutor.runtime_hooks``): the MSA +metadata calls it with its own buffers and falls back to the torch formulation when it returns False. The kernel is +``torch.ops.trtllm.minimax_m3_msa_kv_update``, called through this target's ``msa_kv_update`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.msa_kv_update import msa_kv_update as _kv_update_op + + +def msa_kv_update( + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], +) -> bool: + """The fused correction (see module doc); False, with nothing written, when an operand is outside what the + kernel covers. Arguments in ``MsaKvUpdate`` order; ``seq_lens`` are the per-request query lengths (``qo``).""" + ts = [ + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + msa_seq_lens, + out_cache_loc, + ] + ts += [t for t in (cu_kv_lens, n_valid_blocks) if t is not None] + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or block_table.dim() != 2 + or block_table.stride(1) != 1 + ): + return False + if any(not t.is_contiguous() for t in ts if t is not block_table): + return False + if ( + kv_lens.numel() < batch + or kv_lens_staged.numel() < batch + or seq_lens.numel() < batch + or msa_seq_lens.numel() < batch + or q_batch_row.numel() < total_q + or q_intra.numel() < total_q + or out_cache_loc.numel() < total_q + or (cu_kv_lens is not None and cu_kv_lens.numel() < batch + 1) + or (n_valid_blocks is not None and n_valid_blocks.numel() < total_q) + ): + return False + _kv_update_op( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/__init__.py new file mode 100644 index 000000000000..ef3a4c5847ea --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) as a CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..932971278c49 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1299 @@ +# 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. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows per CTA = the MMA's M # capacity: fixed: a kernel tile constant +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row # capacity: fixed: a kernel tile constant +MMA_INST_K = 32 +K_BLOCKS = ( + CTA_K // MMA_INST_K +) # 4 MMAs per k-tile, one per scale byte of the atom # capacity: fixed: a kernel tile constant +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +# The swigluoai activation of MiniMax-M3's dense MLP, gate * sigmoid(1.702 * gate) * (up + 1) with gate clamped to <= 7 +# and up to [-7, 7], in exactly the instruction sequence the Triton silu_and_mul kernel compiles to (modules/swiglu.py +# on SM107): the gate clamp on the packed bf16 pair, the up clamp as min.xorsign.abs, two rounded multiplies into +# ex2.approx, div.full for the reciprocal, two rounded multiplies, round to bf16. No multiply feeds an add, so the +# explicit .rn forms compile to what the Triton kernel's contractable forms do. +SWIGLU_OAI_ALPHA = 1.702 +SWIGLU_OAI_LIMIT = 7.0 + + +@dsl_user_op +def _swiglu_oai_bf16x2(gate_word, up_word, *, loc=None, ip=None): + """The activation of one packed bf16 pair of gate and up values: its two bf16 results as one int32 word (the low + half from the low halves).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [cutlass.Int32(w).ir_value(loc=loc, ip=ip) for w in (gate_word, up_word)], + "{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, $1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, $2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 $0, r1, r0;\n\t}", "=r,r,r", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[ + bool + ], # x_bf16: x_words is [M, 2K] gate | up; the activation is swigluoai(gate, up) +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in if x_swiglu else k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + if cutlass.const_expr(x_swiglu): + # The up half of the same columns, K values after the gate in the row. + qz_wvecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(k_in // 2 + qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + if cutlass.const_expr(x_swiglu): + qz_word = _swiglu_oai_bf16x2( + qz_word, cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + ) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[bool] = False, +) -> None: + # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. + assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( + "x_swiglu needs x_bf16 and excludes din" + ) + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, x_swiglu, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py new file mode 100644 index 000000000000..3e6b07326658 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py @@ -0,0 +1,521 @@ +# 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. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) for already-quantized activations of at most 8 tokens. + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring): GB300 exposed time behind a 4.25 us PDL predecessor at M 1 / 8, vs the fused low-M GEMM: +# QKV 2.78 / 3.16 vs 4.58 / 6.21 us, o_proj 2.23 / 2.41 vs 4.40 / 5.09, shared gate_up 2.43 / 2.74 vs 3.56 / 4.30, +# shared down 1.74 / 1.66 vs 2.92 / 3.38. A config is (split, ring[, contig, nacc, a_tmem, st_async]); the optional +# fields default to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + (2560, 6144): (4, 12), + (6144, 2048): (2, 8), + # Shared-expert gate_up: split 2, ring 12, single accumulator, A operand from shared memory. + (1536, 6144): (2, 12, 1, 1, 0, 1), + # Shared-expert down: unsplit, ring 3, single accumulator, A operand from shared memory. + (6144, 768): (1, 3, 1, 1, 0, 1), + (2304, 6144): (4, 12), + (6144, 3072): (2, 12), + # Dense-layer gate_up: 6.4 / 6.3 us at M 1 / 8 vs FlashInfer 8.8. + (6144, 6144): (2, 12), +} +# Per-token-count configs (N, K, M): the grid that suits the concurrent routed-expert engine depends on how long +# that engine runs. At 2 tokens the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = { + (1536, 6144, 2): (4, 12, 1, 1, 0, 1), +} +# Above 8 tokens (spec-decode verify rows): the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = { + (1536, 6144): (4, 12, 1, 1, 0, 1), +} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # GB300 exposed time vs the FlashInfer prequantized GEMM (behind a PDL predecessor), at 24 / 48 tokens: QKV 3.76 / + # 4.63 vs 7.46 / 11.44 us, o_proj 2.74 / 3.33 vs 4.46 / 5.84, shared gate_up 4.57 / 5.71 vs 7.16 / 11.25. Behind + # only for K = 768 (unsplit shared down) at 17-32 tokens: 3.1 vs 2.9. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], + x_swiglu: bool = False, +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + if nacc != 1: + # More than one TMEM accumulator is broken on SM107 (nondeterministic or garbage + # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + x_swiglu, + ) + return fn + + +def _views( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight, + weight_scale, + y, + x_bf16: bool, + x_swiglu: bool = False, +) -> tuple: + # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. + k = x.shape[1] // 2 if x_swiglu else x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, + x_swiglu: bool = False, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + if x_swiglu: + k //= 2 + n = weight.shape[0] + if weight.shape[1] != k: + return None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu) + if resolved is None: + return None + key, runtime_ar = resolved + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, + x_swiglu: bool = False, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros( + m, + 2 * k if x_swiglu else k, + dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, + device=dev, + ) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_swiglu( + gate_up: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``m3_mx_gemv_bf16(swiglu(gate_up), ...)`` in one launch: ``gate_up`` bf16 [M <= 8, 2K] (contiguous, gate then + up) is activated in the kernel's quantize prologue with the swigluoai arithmetic of the Triton silu_and_mul kernel + (alpha 1.702, limit 7, the (up + 1) offset; bit for bit), then quantized as the bf16 path does.""" + if ( + gate_up.dtype != torch.bfloat16 + or gate_up.dim() != 2 + or not gate_up.is_contiguous() + or gate_up.data_ptr() % 16 != 0 + or gate_up.shape[1] % 2 != 0 + ): + return None + return _run(gate_up, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar, True) + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/__init__.py new file mode 100644 index 000000000000..acec8fdbb26f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalization entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..31c499cab2a6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step embedding gather and the first layer's input RMSNorm in one CuTe DSL launch, into caller buffers.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm + +__all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] + +VOCAB = 200064 # MiniMax-M3's vocabulary; the decode table is the full (replicated) embedding + + +class M3EmbeddingFirstnorm(OpWrapper): + """Write ``table[ids]`` into ``raw`` (bf16 [M, 6144], a zero row for an id outside ``[0, V)``) and + ``raw * rsqrt(mean(raw^2) + eps) * weight`` into ``out``; returns ``(out, raw)`` -- the first layer's normed input + and its residual. ``ids`` int32 or int64 [M], M = 1..8; ``table`` bf16 [V, 6144] contiguous; ``weight`` bf16 + [6144]. The first call per (M, id dtype, table shape, PDL) compiles (``DecodeEmbedding`` does it at build).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target builds every (rows 1..8, id dtype) owner; the kernel is one CTA per row, so the rows differ only in + # grid size: the ends of the range and both id dtypes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, int32 ids", spec=dict(rows=1, id_dtype=torch.int32)), + Cell(why="eight requests, int32 ids", spec=dict(rows=8, id_dtype=torch.int32)), + Cell(why="one request, int64 ids", spec=dict(rows=1, id_dtype=torch.int64)), + Cell( + why="six rows (one request verifying a draft), int64 ids", + spec=dict(rows=6, id_dtype=torch.int64), + ), + # Every other row count the decode dispatch takes (modeling.EMBED_FIRSTNORM), both id dtypes. + *( + Cell( + why=f"{rows} rows: a decode row count between the ones above", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in (2, 3, 4, 5, 7) + ), + ) + + note = """ + The gathered rows are exact. The norm is flashinfer's RMSNorm arithmetic in its reduction order with a fast + rsqrt, where `reference` sums in fp64 and takes torch's rsqrt: within the band the test derives, not bit for bit. + An id outside [0, V) gives a zero row and a zero norm (torch's embedding would raise). No accepted input is known + to compute wrongly: the host side rejects any other dtype, shape or a strided buffer. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return m3_embed_norm(ids, table, weight, eps, raw, out, use_pdl) + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x * rstd * weight.double()).to(torch.bfloat16), rows + + +m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/__init__.py new file mode 100644 index 000000000000..d25c41dc9557 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries this target alone calls: the Eagle3 verify ops, the capture / fc_norm RMSNorms and +the step-input staging.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..be5aa46fe774 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_chunk_norm import rmsnorm_rows + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = ( + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ) + + note = """ + fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional + one-ulp move against torch's statistic. + """ + + def __call__( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) + return dst + + def reference( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" + x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h + return rmsnorm_rows(x, w, eps) + + def is_valid( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> None: + # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read + # past its elements without an error. + assert w.is_contiguous(), "w must be dense [H]" + + +m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..6699a55fc3e6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states +(``m3_spec_tail.chunk_norm``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ) + + note = """ + fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then + against torch's statistic; the cells' band states it. + """ + + def __call__( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) + return dst + + def reference( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" + c, h = w.shape + x = src[:, : c * h].reshape(src.shape[0], c, h) + return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) + + def is_valid( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ) -> None: + c, h = w.shape + assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + + +m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..8f0f69ee812b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat +(``m3_spec_tail.draft_input``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3DraftInput(OpWrapper): + """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = + hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of + a vocab-local embedding). Returns ``(out, residual)``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling.DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the + # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests + # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", + spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", + spec=dict(m=n, h=6144, gather=True, invalid_id=False), + ) + for n in range(1, 9) + ), + Cell( + why="third pass of 3 requests: one row each, hidden read in place", + spec=dict(m=3, h=6144, gather=False, invalid_id=False), + ), + Cell( + why="one row with an out-of-table id: the embedding half is zero", + spec=dict(m=1, h=6144, gather=False, invalid_id=True), + ), + ) + + note = """ + The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's + summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, + also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the + drafter's embedding is the whole table, so its ids are always in range. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl + ) + return out, residual + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / + ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" + valid = (ids >= 0) & (ids < table.shape[0]) + emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) + emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) + hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] + ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) + return ref, (hid.clone() if residual is not None else None) + + def is_valid( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> None: + h = table.shape[1] + # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. + assert table.is_contiguous() and hidden.is_contiguous(), ( + "table / hidden must be dense [*, H]" + ) + if rows is not None: + assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" + else: + assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" + # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + assert residual is None or ( + residual.is_contiguous() and residual.shape == (ids.shape[0], h) + ), "residual must be contiguous [M, H]" + + +m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py new file mode 100644 index 000000000000..75626e570244 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter MLP's plain SwiGLU as one C++ launch (``modeling.DRAFT_SILU_CPP``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3DraftSilu(OpWrapper): + """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: + trtllm::minimax_m3_sm107_tp4_silu_and_mul (minimaxM3Sm107Tp4/siluAndMul.cu), bit for bit the Triton + ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP + runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs + the Triton op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP4 drafter MLP shard: intermediate 18432 / 4 = 4608, gate | up 9216 wide, at the drafter's decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP4 drafter MLP at 1 row", spec=dict(rows=1, d=4608)), + Cell(why="the TP4 drafter MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=4608)), + Cell(why="the TP4 drafter MLP at 8 rows", spec=dict(rows=8, d=4608)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (the same PTX instruction sequence), checked over every + bf16 gate and every bf16 up pattern. + """ + + def __call__(self, x: torch.Tensor) -> Optional[torch.Tensor]: + if ( + x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 16 != 0 + or x.data_ptr() % 16 != 0 + or x.shape[0] > 65535 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp4_silu_and_mul(x, out) + return out + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul(x) + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_draft_silu = M3DraftSilu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..5dbb324a7d50 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One pass of the fused greedy Eagle3 draft loop: the drafter's vocab-sharded BF16 head (streamed from its lossless +12-bit coding, ``ops.gemm.m3_target_head``, or as BF16 fragments), the cross-rank argmax and the loop's per-pass +bookkeeping in one kernel (``m3_draft_step.draft_step``).""" + +from __future__ import annotations + +import functools +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ..gemm.m3_target_head import decode as decode_coded +from .m3_shard_accept import exchange_tokens, local_argmax, make_keys + +#: Requests (head rows) per pass, and exchange rows per call half of the workspace. +MAX_TOKENS = 8 +#: 16-row weight tiles one CTA can own (the kernel's shared-memory accumulator). +MAX_TILES_PER_CTA = 32 +KERNEL_WARPS = 16 # warps per CTA (draftStep.h kWarps) + + +@functools.lru_cache(maxsize=64) +def max_tile_pieces(n_tiles: int, num_ctas: int, chunks_per_tile: int) -> int: + """The most warps that stream one weight tile under the kernel's whole-tile split: CTA c owns tiles + [c * n / grid, (c + 1) * n / grid), and its (tile, chunk) range is cut into KERNEL_WARPS equal contiguous + segments. The shared-memory atomics add a tile's pieces in arrival order: deterministic for at most 2 pieces (two + addends from 0 commute), not for more (fixed_order then sums them in warp order).""" + worst = 1 + for c in range(num_ctas): + lo, hi = c * n_tiles // num_ctas, (c + 1) * n_tiles // num_ctas + span = (hi - lo) * chunks_per_tile + bounds = [w * span // KERNEL_WARPS for w in range(KERNEL_WARPS + 1)] + for t in range(hi - lo): + a, b = t * chunks_per_tile, (t + 1) * chunks_per_tile + worst = max( + worst, sum(1 for w in range(KERNEL_WARPS) if bounds[w] < b and bounds[w + 1] > a) + ) + return worst + + +#: Draft length of the fused loop, and the drafter's TP4 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 50016 +HIDDEN = 6144 + + +def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the + 12-bit coded stream with its exception lists decoded.""" + if exc_off is None: + return from_fragments(wf) + n = (exc_off.numel() - 1) * 16 + return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + + +def shard_logits( + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() + + +class M3DraftStep(OpWrapper): + """One fused greedy draft pass over ``rows <= 8`` requests (one head row each). + + Head input ``x[x_rows[b]]`` (pass 0: the last accepted verify row of each request) or ``x[b]``. The head is this + rank's BF16 lm_head shard in fragment order (``m3_draft_step.bf16_fragments``, global ids ``vocab_offset ..``), + or with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base`` its lossless 12-bit coding (``wf`` = the uint8 stream + of ``m3_target_head.encode_chunked``, which must round-trip to the BF16 shard); each logit is the fp32 dot product + rounded to BF16 and the row's token is the global argmax over all ranks' + shards (ties, -0 / +0 included, to the lowest global id), exchanged through the multicast slots + ``[2][8][world]``. The last CTA then writes ``draft_out[b, step]`` and ``ids_next[b]`` (the token), ``pos_next[b] + = pos_src[x_rows[b] or b] + 1``, and advances the draft KV lengths as the in-tree draft loop does between passes + (pass 0 with ``num_acc``: ``kv_lens -= draft_len - num_acc`` and ``seq_lens = 1``; later passes: ``kv_lens += + 1``). ``out_keys`` (optional) receives each row's winning key. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling.DRAFT_STEP): draft_len 3 at 1..8 requests, the drafter's + # full-vocabulary BF16 head (50016 rows per rank of 4, hidden 6144). Pass 0 reads the verify rows' hidden (4 per + # request) through x_rows and rolls the KV back; the later passes read one row per request. + _PASSES: tuple[Cell, ...] = ( + *( + Cell( + why=f"pass 0 at {n} request(s): head rows gathered from {4 * n} verify rows, KV rollback, seq_lens = 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=0, + rank=n % 4, + world=4, + vs=DRAFT_SHARD, + k=HIDDEN, + ), + ) + for n in range(1, MAX_TOKENS + 1) + ), + Cell( + why="pass 1 at two requests: one row each, kv_lens + 1, rank 3's ids", + spec=dict( + nreq=2, draft_len=DRAFT_LEN, step=1, rank=3, world=4, vs=DRAFT_SHARD, k=HIDDEN + ), + ), + Cell( + why="last pass (step 2) at 8 requests, rank 0", + spec=dict( + nreq=8, draft_len=DRAFT_LEN, step=2, rank=0, world=4, vs=DRAFT_SHARD, k=HIDDEN + ), + ), + Cell( + why="pass 2 at 5 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=5, + draft_len=DRAFT_LEN, + step=2, + rank=1, + world=4, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + ), + ), + ) + # The coded head (the fused draft loop's default) at every pass shape, and the BF16 fragments (the coding's + # fail-closed path) at the first and last pass shapes and with ties. + CELLS: tuple[Cell, ...] = ( + *(Cell(why=f"coded: {c.why}", spec=dict(c.spec, coded=True)) for c in _PASSES), + *( + Cell(why=f"bf16: {c.why}", spec=dict(c.spec, coded=False)) + for c in (_PASSES[0], _PASSES[MAX_TOKENS - 1], _PASSES[-1]) + ), + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit; + the coded head's weight is the BF16 shard bit for bit; its exceptions decode to +0 in the MMAs and their exact + products are added at the tile's flush) + and the argmax is the in-tree one over BF16 logits, so a token differs from the in-tree path's only where the fp32 + sum, accumulated in another order, rounds to the other side of a BF16 boundary and that decides the row's maximum. + The cells' rows are decidable and the peers' keys sit 0.25 away (or tie exactly); the edge tests (forced BF16 + ties, -0 / +0, top-2 one fp32 step apart that round to one BF16 value, adjacent BF16 values) use integer data + whose fp32 sums are exact in any order and hold the kernel to the in-tree path bit for bit. The cross-rank + exchange is certified on one device with the peers' keys written into the slots (see m3_shard_accept). + """ + + def __call__( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> None: + # fold: (the 14 fold-tail tensors in buildParams' order, rows, tokens per step) -- the step's forward-tail store + # and next-step staging, run by the pass's last CTA after its tokens (unused by this target, which passes None). + tensors, fold_rows, fold_t = ( + (list(fold[0]), int(fold[1]), int(fold[2])) if fold is not None else (None, 0, 0) + ) + torch.ops.trtllm.minimax_m3_sm107tp4_draft_step( + x, + x_rows, + rows, + wf, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + num_ctas, + pdl, + out_keys, + exc_off, + exc_pos, + exc_w, + exp_base, + tensors, + fold_rows, + fold_t, + fixed_order, + ) + + def reference( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> dict[str, torch.Tensor]: + """Every buffer the pass writes, computed on copies: "logits" (this rank's BF16 logits as fp32), "local_keys" + (this rank's key per row), "keys" (the winning key), "tokens", "draft_out", "ids_next", "pos_next", "kv_lens" + and "seq_lens". + + Reads ``ws_uc``, ``counters`` and ``kv_lens`` before the call. ``keys``, ``ws_mc_ptr``, ``num_ctas``, ``pdl``, + ``out_keys`` and ``fixed_order`` do not change the result; ``fold`` adds the fold tail's writes (m3_fold_tail), + not these. + """ + h = x[x_rows.long()] if x_rows is not None else x[:rows] + logits = shard_logits(h, wf, exc_off, exc_pos, exc_w, exp_base) + best, idx = local_argmax(logits, vocab_offset) + local_keys = make_keys(best, idx) + tokens = exchange_tokens(local_keys, ws_uc, counters, rank, world, MAX_TOKENS) + half = int(counters[1].item()) & 1 + slots = ws_uc.view(torch.int64)[ + half * MAX_TOKENS * world : (half * MAX_TOKENS + rows) * world + ] + slots = slots.view(rows, world).clone() + slots[:, rank] = local_keys + sign = torch.iinfo(torch.int64).min + win = (slots ^ sign).max(dim=1).values ^ sign + d_out = draft_out.clone() + d_out[:rows, step] = tokens + src = x_rows.long() if x_rows is not None else torch.arange(rows, device=x.device) + lens = kv_lens.clone() + sl = seq_lens.clone() if seq_lens is not None else None + if step == 0 and num_acc is not None: + lens[:rows] -= draft_len - num_acc[:rows] + sl[:rows] = 1 + elif step > 0: + lens[:rows] += 1 + out = { + "logits": logits, + "local_keys": local_keys, + "keys": win, + "tokens": tokens, + "draft_out": d_out, + "ids_next": tokens.clone(), + "pos_next": pos_src[src] + 1, + "kv_lens": lens, + } + if sl is not None: + out["seq_lens"] = sl + return out + + def is_valid( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> None: + assert not bool(keys[:rows].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + # The kernel reads the head row x + x_rows[b] * K and the position pos_src[x_rows[b]] with no bound. + if x_rows is not None: + assert bool(((x_rows[:rows] >= 0) & (x_rows[:rows] < x.shape[0])).all()), ( + "x_rows must index x" + ) + assert bool((x_rows[:rows] < pos_src.shape[0]).all()), ( + "pos_src must cover the head rows" + ) + else: + assert pos_src.shape[0] >= rows, "pos_src must cover the head rows" + # Each CTA accumulates at most MAX_TILES_PER_CTA tiles in shared memory. + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + assert -(-tiles // num_ctas) <= MAX_TILES_PER_CTA, "too many weight tiles per CTA" + # The atomics' arrival-order sum is deterministic only while no tile takes more than 2 warp pieces. + assert ( + fixed_order + or max_tile_pieces(tiles, num_ctas, x.shape[1] // (16 if exc_off is None else 64)) <= 2 + ), ( + "a tile takes more than 2 warp pieces on this grid: the atomics' sum is order-dependent (fixed_order)" + ) + if exc_off is not None: + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), ( + "exc_off must end at the exception count" + ) + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + # A NaN logit maps to a key above every finite one and wins the row silently. + h = x[x_rows.long()] if x_rows is not None else x[:rows] + assert not bool(torch.isnan(h).any()), "NaN head inputs win the argmax" + if fold is not None: + assert len(fold) == 3 and len(fold[0]) == 14, ( + "fold is (the 14 fold-tail tensors, rows, tokens per step)" + ) + if num_acc is not None: + na = num_acc[:rows] + assert bool(((na >= 1) & (na <= draft_len + 1)).all()), ( + "num_acc counts must lie in [1, draft_len + 1]" + ) + + +m3_draft_step = M3DraftStep() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..b980aeefc8d8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..m3_force_vec import force_vec_kernel + + +class M3ForceVec(OpWrapper): + """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` + (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 + ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + force_vec_kernel[(1,)]( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + """The counts (a new tensor; ``counter`` and ``out`` are not written).""" + c = int(counter.item()) + 1 + slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) + idx = (c * counter_stride + slots * slot_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec = M3ForceVec() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py new file mode 100644 index 000000000000..416425f9a01d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, each request drawn at its own key.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..m3_force_vec import force_vec_req_kernel + + +class M3ForceVecReq(OpWrapper): + """For each of the ``batch_size`` requests ``b`` write ``min(base_total + (pool[(keys[b] * key_stride) & + pool_mask] < frac), max_total)`` into ``out`` (int32): the forced-acceptance draw from the fixed-seed fp32 ``pool`` + at the request's own key (int32 ``keys``, its KV length at this verify), so a request's draws depend on neither the + server's history nor the other requests. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or + 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it, with ``FORCE_AL_PER_REQUEST``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + force_vec_req_kernel[(1,)]( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + """The counts (a new tensor; ``out`` is not written).""" + idx = (keys[:batch_size].to(torch.int64) * key_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec_req = M3ForceVecReq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..7623f5fe582b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampled (temperature + top-p) Eagle3 verification over the TP-sharded vocabulary in one kernel +(``m3_sampled_accept.sampled_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_shard_accept import strict_accept + +#: Verify rows per call. +MAX_ROWS = 48 +#: The fused sampling kernel's 1 / T guard. +TEMP_EPSILON = 1e-6 + + +def nucleus_probs( + logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +class M3SampledAccept(OpWrapper): + """Sampled verification of ``rows = nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard. + + Each row draws one token from its top-p nucleus of softmax(logits / T) over the full (all-rank) vocabulary by + distributed rejection sampling (Philox keyed by the row's seed / offset, the row, the attempt and the CTA), then + the strict acceptance runs against ``draft`` as in ``m3_shard_accept``: ``acc_tokens``, ``num_acc``, + ``gather_ids``, optional ``next_new[:, 0]``. ``tokens`` receives the sampled token per row; ``extra_attempts`` + accumulates exchange rounds beyond the second. The exchange goes through the multicast workspace (two call + halves of ``minimax_m3_sm107tp4_sampled_accept_ws_words(world)`` words) with the (finished rows, call) ``counters``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3 (T = 4 rows per request) at the TP4 caps: 1 / 2 / 8 requests (4 / 8 / 32 rows, 8 CTAs per row) + # and the 48-row maximum (4 CTAs per row); the 50016-column target shard. Single rank (see note). + CELLS: tuple[Cell, ...] = ( + Cell( + why="cap 1 (4 rows), top_p 0.8 over a 12-token head: the draws follow the nucleus distribution", + spec=dict( + nreq=1, vs=50016, vocab_offset=0, check="distribution", top_p=0.8, temperature=0.7 + ), + ), + Cell( + why="cap 8 (32 rows, 256 co-resident CTAs): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=8, + vs=50016, + vocab_offset=2 * 50016, + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ), + Cell( + why="cap 2 (8 rows) with top_p = 1 (no cut: the first draw stands), next_new and forced counts", + spec=dict( + nreq=2, + vs=50016, + vocab_offset=50016, + check="distribution", + top_p=1.0, + temperature=1.0, + force="vec", + ), + ), + Cell( + why="48 rows (4 CTAs per row): the second launch geometry, nucleus distribution at top_p 0.5", + spec=dict( + nreq=12, + vs=50016, + vocab_offset=3 * 50016, + check="distribution", + top_p=0.5, + temperature=1.3, + force="total", + ), + ), + ) + + note = """ + What is sampled is gated statistically, the rest exactly. Every drawn token must lie in the reference nucleus; + the empirical distribution over many calls (distinct Philox offsets) must sit within a stated total-variation + band of the reference's; the acceptance outputs must equal the strict acceptance of the drawn tokens exactly. + The Philox stream itself is not reproduced in torch: the kernel's draw uses the hardware exp / log + approximations (__expf, __logf) and a CTA-order inverse CDF, which torch does not reproduce bit for bit. + + Certified at world 1 only. The decision is replicated from every rank's per-round records, which another rank + computes from its own shard on the fly, so peers cannot be stood in for on one device; the cross-rank path + needs the multi-device job and is not covered here. Rows with top-k or min-p must not reach this op (it applies + neither); is_valid states it. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + prewait_params: bool = True, + early_trigger: bool = False, + one_round_k: int = 0, + decided: Optional[torch.Tensor] = None, + ilp_walk: bool = False, + ) -> None: + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept( + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + prewait_params, + early_trigger, + one_round_k, + decided, + ilp_walk, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + prewait_params: bool = True, + early_trigger: bool = False, + one_round_k: int = 0, + decided: Optional[torch.Tensor] = None, + ilp_walk: bool = False, + ) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + prewait_params: bool = True, + early_trigger: bool = False, + one_round_k: int = 0, + decided: Optional[torch.Tensor] = None, + ilp_walk: bool = False, + ) -> None: + rows = logits.shape[0] + # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the + # wrong rows. + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + # Per-row parameters are read at the row index, unchecked. + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), ( + "per-row parameters must cover every row" + ) + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + # The finished-row counter elects the CTA that runs the acceptance. + assert int(counters[0].item()) == 0, ( + "counters[0] (finished rows) must be zero between calls" + ) + + +m3_sampled_accept = M3SampledAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..9193b3dce61d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept_ws_words.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Size of the sampled-verify exchange workspace (``m3_sampled_accept.SampledAcceptWorkspace.create``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Exchange rounds per row, verify rows per call, CTAs per row (at most), uint64 words per record. +MAX_ATTEMPTS = 16 +MAX_ROWS = 48 +MAX_SPLITS = 8 +RECORD_WORDS = 4 +#: SA-1R (one_round_k 9 / 82): 9 listed bf16 values ride each record, 3 per extra word. +LIST_WORDS = {0: 0, 9: 3, 82: 3} + + +class M3SampledAcceptWsWords(OpWrapper): + """uint64 words of the multicast workspace ``minimax_m3_sm107tp4_sampled_accept`` exchanges through for + ``world`` ranks: two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="the tp4 target's workspace", spec=dict(world=4)), + Cell(why="a single rank (local workspace, no multicast)", spec=dict(world=1)), + Cell( + why="the tp4 target's workspace with SA-1R pooled (82)", + spec=dict(world=4, one_round_k=82), + ), + Cell( + why="the tp4 target's workspace with SA-1R top-9 (9)", spec=dict(world=4, one_round_k=9) + ), + ) + + note = """ + A host computation (no device work); it is an op so the workspace size and the kernel's indexing come from one + definition. + """ + + def __call__(self, world: int, one_round_k: int = 0) -> int: + return torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept_ws_words(world, one_round_k) + + def reference(self, world: int, one_round_k: int = 0) -> int: + """Two halves of attempts x rows x (world x splits) records x record words.""" + return ( + 2 + * MAX_ATTEMPTS + * MAX_ROWS + * world + * MAX_SPLITS + * (RECORD_WORDS + LIST_WORDS[one_round_k]) + ) + + +m3_sampled_accept_ws_words = M3SampledAcceptWsWords() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..8beba4373b02 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Greedy Eagle3 verification from this rank's logits shard: per-row argmax, the cross-rank (value, index) exchange +and the strict acceptance in one kernel (``m3_shard_accept.shard_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Verify rows per call, and exchange rows per call half of its workspace. +SHARD_ROWS = 48 + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out + + +class M3ShardAccept(OpWrapper): + """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` + (global ids ``vocab_offset ..``). + + Each row's token is its global argmax over all ranks' shards: keys ``ordered(value) << 32 | ~index`` (ties take + the lowest index, as ``torch.argmax``) are max-reduced over the shard, pushed into every rank's exchange slot + through the multicast slots ``[2][48][world]`` (Lamport: 0 = empty) and max-reduced over the ranks by the last + CTA, which then runs the strict acceptance: ``acc_tokens [nreq, D + 1]``, ``num_acc [nreq]`` (``force_vec`` / + ``force_total`` override the counts), ``gather_ids [nreq]`` (the first draft pass's row per request) and, when + given, ``next_new[:, 0]`` (the last accepted token). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every greedy verify step of the target (modeling.SHARD_ACCEPT): draft_len 3 at 1..12 requests (4..48 rows), the + # target lm_head's BF16 shard logits (50016 columns per rank at TP4). The logits come from the head GEMM, so the + # op only reduces. + CELLS: tuple[Cell, ...] = ( + Cell( + why="cap 1: 4 verify rows (1 request x T = 4), rank 0 of 4, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=50016, world=4, rank=0, force="none", strided=False), + ), + Cell( + why="cap 2: 8 verify rows, rank 2, global ids offset by 2 x 50016", + spec=dict(nreq=2, draft_len=3, vs=50016, world=4, rank=2, force="none", strided=False), + ), + Cell( + why="cap 4: 16 verify rows (4 requests x T = 4), rank 1 of 4, next_new staged", + spec=dict(nreq=4, draft_len=3, vs=50016, world=4, rank=1, force="none", strided=False), + ), + Cell( + why="cap 8: 32 verify rows, rank 3, global ids offset by 3 x 50016", + spec=dict(nreq=8, draft_len=3, vs=50016, world=4, rank=3, force="none", strided=False), + ), + Cell( + why="12 rows (3 requests) read from a row-strided view of a wider logits buffer, forced counts", + spec=dict(nreq=3, draft_len=3, vs=50016, world=4, rank=0, force="vec", strided=True), + ), + Cell( + why="the 48-row maximum (12 requests at draft 3), force_total, rows with exact cross-rank ties", + spec=dict( + nreq=12, + draft_len=3, + vs=50016, + world=4, + rank=2, + force="total", + strided=False, + ties=True, + ), + ), + # Every other request count of the verify range, on the 4-rank 50016-column shards. + *( + Cell( + why=f"{n} requests ({4 * n} verify rows): a request count between the ones above", + spec=dict( + nreq=n, draft_len=3, vs=50016, world=4, rank=n % 4, force="none", strided=False + ), + ) + for n in (5, 6, 7, 9, 10, 11) + ), + ) + + note = """ + The op reduces BF16 logits exactly, so the reference is bit-exact, including cross-rank ties (lowest global + index wins). The cross-rank exchange is certified on one device: the other ranks' keys are written into this + rank's slots before the call, which is what their multicast pushes would have put there; the multicast mapping + itself (ws_mc_ptr = this rank's slots, as at world 1) is not exercised across devices here. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + num_ctas, + force_vec, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Lowest-index maximum of each row's BF16 logits, the largest key among it and the peers' keys in + ``ws_uc``, then the strict acceptance. Reads ``ws_uc`` / ``counters`` before the call; the scratch and + output buffers and ``num_ctas`` do not change the result.""" + best, idx = local_argmax(logits.float(), vocab_offset) + tokens = exchange_tokens(make_keys(best, idx), ws_uc, counters, rank, world, SHARD_ROWS) + out = strict_accept(tokens, draft, batch_idx, force_total, next_new, force_vec) + out["tokens"] = tokens + return out + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[: logits.shape[0]].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + assert ws_uc.numel() >= 2 * SHARD_ROWS * world, "ws_uc must hold [2][48][world] slots" + # A NaN logit maps to a key above every finite one and wins the row silently. + assert not bool(torch.isnan(logits).any()), "NaN logits win the argmax" + nreq = draft.shape[0] + assert bool(((batch_idx[:nreq] >= 0) & (batch_idx[:nreq] < nreq)).all()), ( + "batch_idx must index the nreq requests" + ) + if force_vec is not None: + fv = force_vec[:nreq] + assert bool(((fv >= 1) & (fv <= draft.shape[1] + 1)).all()), ( + "force_vec counts must lie in [1, D + 1]" + ) + + +m3_shard_accept = M3ShardAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..ff07c54de1f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> +``m3_spec_tail.slot_scatter``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: + """src rows zero-padded (or cut) to ``width`` columns.""" + out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) + w = min(width, src.shape[1]) + out[:, :w] = src[:, :w] + return out + + +class M3SlotScatter(OpWrapper): + """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, + ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The + destinations are the store's transposed [slots, width] views.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of + # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) + # is zero-padded. + CELLS: tuple[Cell, ...] = ( + Cell( + why="8 rows into shuffled slots, widths equal to the store's", + spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), + ), + Cell( + why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", + spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), + ), + ) + + note = """ + Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. + """ + + def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): + torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) + + def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): + """The four destinations after the call, on copies.""" + s = slots.long() + out = {} + for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): + r = dst.clone() + r[s] = padded_rows(src[: s.numel()], dst.shape[1]) + out[name] = r + lr = dl.clone() + lr[s] = sl[: s.numel()] + out["dl"] = lr + return out + + def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): + assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( + "slots must index the store" + ) + # Two rows into one slot race (which one lands is unordered). + assert slots.unique().numel() == slots.numel(), "slots must be distinct" + assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + + +m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py new file mode 100644 index 000000000000..36e421412f30 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's speculative step correction as one launch (``RuntimeHooks.spec_offsets_op``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger + +# The engagement line is logged once per process (module state: __call__ reads no instance state). +_ENGAGED: list = [] + + +class M3SpecOffsets(OpWrapper): + """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp4_spec_offsets (minimaxM3Sm107Tp4/specOffsets.cu), the + same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other + layouts (the caller runs the ATen adds).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="1 request, 4 tokens (draft length 3)", spec=dict(tokens=4, seqs=1)), + Cell(why="8 requests, 32 tokens", spec=dict(tokens=32, seqs=8)), + Cell(why="wrap-around values", spec=dict(tokens=8, seqs=2, wrap=True)), + ) + + note = "Integer adds: the reference (the two in-place ATen adds) is exact." + + def __call__(self, pos, pos_off, kv, kv_off) -> bool: + ts = (pos, pos_off, kv, kv_off) + if ( + any( + t.dtype != torch.int32 or t.dim() != 1 or t.stride(0) != 1 or not t.is_cuda + for t in ts + ) + or pos.numel() != pos_off.numel() + or kv.numel() != kv_off.numel() + ): + return False + torch.ops.trtllm.minimax_m3_sm107_tp4_spec_offsets(pos, pos_off, kv, kv_off) + if not _ENGAGED: + _ENGAGED.append(True) + logger.info("[M3 host] spec step correction: one launch taken (spec_offsets_op)") + return True + + def reference(self, pos, pos_off, kv, kv_off) -> None: + pos += pos_off + kv += kv_off + + def is_valid(self, pos, pos_off, kv, kv_off) -> None: + assert pos.numel() == pos_off.numel() and kv.numel() == kv_off.numel(), ( + "equal lengths per pair" + ) + + +m3_spec_offsets = M3SpecOffsets() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..16c4000199da --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch +(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS: tuple[Cell, ...] = ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + ) + + note = """ + Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid + states them. + """ + + def __call__( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ) + + def reference( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" + nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) + s = slots[:P].long() + ids_r, draft_r = ids.clone(), draft.clone() + ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) + draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) + pos_r = torch.zeros_like(pos_off) + pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] + kv_r = torch.zeros_like(kv_off) + kv_r[b0 : b0 + P] = lens[s] - T + return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} + + def is_valid( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or + # writes another buffer. + s = slots[:P] + assert bool( + ( + (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) + ).all() + ), "slots must index the store" + assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( + "ppi must index lens" + ) + assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" + assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" + assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( + "offset buffers too short" + ) + + +m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py new file mode 100644 index 000000000000..122d43fb175a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode graph's step-input prologue in one launch (this target's ``RuntimeHooks.pack_prologue_op``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepPrologue(OpWrapper): + """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch + (trtllm::minimax_m3_sm107_tp4_step_prologue, minimaxM3Sm107Tp4/stepPrologue.cu). + + ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one + ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the + window, a device copy's a device address). The launch writes up to ``max_desc`` copies of each kind, every host + copy before any device copy, then zeroes the header; with ``epoch`` it increments ``epoch[0]`` once. The default + prologue does the same with two parallel-scatter launches (``step_inputs._graph_scatter_kernel``) and a header + fill.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP4 decode step's packs: a few dozen host copies of 4 to a few hundred bytes (positions, KV lengths, slots, + # the spec sampler's uploads), a handful of device copies (the graph runner's static inputs), some reading what a + # host copy of the same window wrote. + CELLS: tuple[Cell, ...] = ( + Cell(why="an empty window (a replay with no new upload)", spec=dict(h2d=0, d2d=0)), + Cell(why="host copies only, odd sizes and alignments", spec=dict(h2d=40, d2d=0)), + Cell( + why="host copies feeding device copies (a device copy reads a host copy's destination)", + spec=dict(h2d=24, d2d=8, chained=True), + ), + Cell( + why="every descriptor slot (128 + 128 of the pack's 256)", spec=dict(h2d=128, d2d=128) + ), + ) + + note = """ + Byte copies: the reference (the default prologue's Triton launches and fill on the same window) is exact. + Destinations are raw addresses; is_valid cannot check them. + """ + + def __call__(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + torch.ops.trtllm.minimax_m3_sm107_tp4_step_prologue(window, epoch, max_desc) + + def reference(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + """The default prologue on the same arguments: it writes the same destinations.""" + from tensorrt_llm._torch.pyexecutor.step_inputs import _graph_scatter_kernel + + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + bump = epoch is not None and not device_src + _graph_scatter_kernel[(max_desc,)]( + window, + stage_i64, + epoch if bump else stage_i64, + DEVICE_SRC=device_src, + BLOCK=256, + BUMP_EPOCH=bump, + num_warps=1, + ) + window[:16].zero_() + + def is_valid(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + assert window.dtype == torch.uint8 and window.is_contiguous(), "window: contiguous uint8" + assert window.numel() >= 16 + 24 * max_desc, "window: header and max_desc descriptors" + counts = window[:16].view(torch.int64).tolist() + assert counts[0] + counts[1] <= max_desc, "at most max_desc copies in all" + + +m3_step_prologue = M3StepPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py new file mode 100644 index 000000000000..d146b9d05d2e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py @@ -0,0 +1,1730 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +GB300/SM103 tuning targets GQA group 16, top-k 16 and four-token Eagle3 +verification. Split-K, query pairing and launch tables require four KV heads; +the linear-softmax policy also covers one and two KV heads. Other hardware +and shapes retain the default policies, including supported SM100 NVFP4. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* A packed NVFP4 cache is read in place, dequantized in-register from its E2M1 + nibbles and E4M3 block scales. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* A single top-k chunk writes the final output directly, avoiding scratch and + the otherwise redundant LSE merge. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. + +This is the MiniMax-M3 NVFP4 / sm_107 / tp4 target's copy. On top of the +in-tree decode it routes long-context TP4 steps to the one-launch +``m3_attn_lc`` kernel, short-sequence verify steps of at least +``MQ_MIN_TOKENS`` rows to the multi-query producer, releases PDL dependents +early, and has the merge emit the MXFP8 o_proj hand-off (``m3_prequant.SLOT``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + check_decode_span_shape, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +from . import m3_attn_lc, m3_prequant, m3_report + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +_TL_DTYPES = { + torch.bfloat16: tl.bfloat16, + torch.float16: tl.float16, + torch.float32: tl.float32, +} + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size + +# NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale +# byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. +NVFP4_SF_VEC_SIZE = 16 + + +# Fewest query rows for which the merge writes the MXFP8 o_proj hand-off: every decode step, so the MXFP8 o_proj +# never runs its own quantize kernel after attention. +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and swizzled UE8M0 + scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips its own + quantization. One 128-row buffer per width serves every layer: each o_proj consumes it before the next layer's + merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + +MERGE_MXFP8 = MergeMxfp8() + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + if IDENTITY: + # Slot `it` of the identity selection, bounded by this row's own kv extent above. + blk = it + else: + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, + # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + if IDENTITY: + blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) + blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) + else: + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for cand in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + if IDENTITY: + blk_0 = tl.where(valid_0, idx_0, sentinel) + blk_1 = tl.where(valid_1, idx_1, sentinel) + else: + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the + # pool, reads nothing. + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) +# Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) + + +# Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query +# tokens; below it the per-token producer is as fast. +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover +# Warps per multi-query producer program. +MQ_NUM_WARPS = 8 +# The long-context decode (``_longattn_lc_decode``) in its engagement line (no row range: every step the LC +# applies to takes it). +LC_PATH = "m3_sparse_decode_lc" + + +def _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + identity=False, + identity_step=False, +) -> None: + """minimax_m3_sparse_attn_decode through m3_attn_lc: same table, output and MXFP8 hand-off contract. ``identity`` + takes the short-sequence identity selection from the device lengths (the table is not read); ``identity_step`` + marks a step with no select (the kernel reads q after its grid wait), also when the staged table is read.""" + + use_pdl = get_env_enable_pdl() + total_q, num_heads, head_dim = q.shape + linear = _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=1, + gqa_group_size=num_heads, + max_topk=int(topk_idx.shape[-1]), + decode_query_len=decode_query_len, + ) + mxfp8 = None + m3_prequant.SLOT.clear() + if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(num_heads * head_dim, output.device) + # The early trigger lets the o_proj launch and stream its weights while attention runs. + m3_attn_lc.m3_lc_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + linear_softmax=linear, + mxfp8=mxfp8, + use_pdl=use_pdl, + early_trigger=True, + # This target's op (sparseDecodeLc4): the merge weights formed once per CTA, and the split-K merge dividing + # once per output (fixed order, no atomics; inside the same band of the float64 reference). + merge_once=True, + prefetch=m3_attn_lc.PREFETCH.take(), + next_kv=m3_attn_lc.NEXT_KV.take(), + identity=identity, + identity_step=identity_step, + one_div=True, + ) + m3_report.taken(ONE_DIV_MERGE_PATH, total_q) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), mxfp8[0][:total_q], mxfp8[1] + ) + + +# Engagement path of the identity-selection producers reading their blocks from the kv lengths (IDENTITY). +DEVICE_IDENTITY_PATH = "sparse_decode_device_identity" +# Engagement path of the LC's one-division split-K merge. +ONE_DIV_MERGE_PATH = "sparse_decode_lc_one_div_merge" + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + q may be FP8, as a fused QK-norm/RoPE producer emits it. The kernel widens + it to output.dtype in-register rather than making the caller materialize a + widened copy, which would cost a standalone kernel on every sparse layer. + The attention math is unchanged either way, since E4M3 widens exactly. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. + + Passing the four NVFP4 scale tensors selects the packed-NVFP4 cache instead: + k_paged/v_paged then hold two E2M1 elements per byte, the block scales are + paged E4M3 bytes in the layouts msa_scatter writes, and the two per-tensor + dequant scales are single-element FP32 device tensors. + + num_topk_chunks overrides the split-K factor and exists for tests: the + merged result must not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) + nvfp4_args = (k_block_scale, v_block_scale, k_global_scale, v_global_scale) + kv_nvfp4 = any(arg is not None for arg in nvfp4_args) + scale_cols = head_dim // NVFP4_SF_VEC_SIZE + if kv_nvfp4: + capability = torch.cuda.get_device_capability(q.device) + if capability not in _SM100F_CAPABILITIES: + raise NotImplementedError( + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; " + f"got SM{capability[0] * 10 + capability[1]}." + ) + _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) + # Two E2M1 elements per byte; Triton must see them unsigned so that the + # high-nibble shift does not sign-extend. + k_paged, v_paged = k_paged.view(torch.uint8), v_paged.view(torch.uint8) + k_block_scale, v_block_scale = ( + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused arguments. + scale_arg = kv_scale if use_scale else output + k_sc_arg, v_sc_arg = (k_block_scale, v_block_scale) if kv_nvfp4 else (output, output) + k_gs_arg, v_gs_arg = (k_global_scale, v_global_scale) if kv_nvfp4 else (output, output) + + widen_q = q.dtype in _FP8_DTYPES + if widen_q and output.dtype not in _TL_DTYPES: + raise ValueError( + f"MiniMax-M3 sparse decode cannot widen FP8 q into {output.dtype}; " + f"supported compute dtypes are {sorted(d.__str__() for d in _TL_DTYPES)}." + ) + # Packed NVFP4 dequantizes to FP16. The production producer supplies FP8 + # q, which widens exactly to FP16, so keep both dot operands there instead + # of converting every K/V tile to BF16. Accumulation remains FP32. + if widen_q: + compute_dtype = tl.float16 if kv_nvfp4 else _TL_DTYPES[output.dtype] + else: + compute_dtype = tl.float32 + + # Long-context TP4 steps: attention, split-K merge and MXFP8 hand-off in one launch. + + if ( + kv_nvfp4 + and widen_q + and num_topk_chunks is None + and m3_attn_lc.lc_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ) + ): + _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + # The identity selection from the device lengths (the staged table is not read). + identity=identity_blocks, + # Every identity step (no select ran): q after the grid wait. + identity_step=identity_blocks, + ) + m3_report.taken(LC_PATH, total_q) + return + + query_group_size = ( + _sm100f_nvfp4_query_group_size( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else 1 + ) + # The stock in-tree decode's split-K chunking at every row count (bit for bit its rows), as the SM107 TP2 target. + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = ( + _sm100f_nvfp4_num_topk_chunks( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else None + ) + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one + # program per (request, block) covers all of them and dequantizes each block once. + use_mq = ( + identity_blocks + and total_q >= MQ_MIN_TOKENS + and kv_nvfp4 + and num_kv_heads == 1 + and decode_query_len > 1 + and head_dim == 128 + and gqa_group_size * decode_query_len <= 128 + ) + if identity_blocks and kv_nvfp4 and widen_q: + # The identity-selection route is modeling.SPARSE_DECODE_PRODUCER's. + m3_report.taken( + "sparse_decode_producer", + total_q, + "multi_query" if use_mq else "per_row", + ) + if use_mq: + num_topk_chunks = max_topk + direct_output = num_topk_chunks == 1 + if direct_output: + # With one chunk, acc_o is already normalized and final. Point the + # producer at the caller's output and compile away the LSE store. + # Keeping a real dummy pointer preserves Triton's pointer contract. + o_arg = output + lse_arg = output + o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) + lse_strides = (0, 0, 0) + else: + # Persistent arena rather than torch.empty, so the partials keep one + # address across CUDA graph replays. FP32 keeps the split-K factor from + # perturbing the merged result. + reserve = torch.cuda.is_current_stream_capturing() + o_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + o_strides = tuple(o_arg.stride()) + lse_strides = tuple(lse_arg.stride()) + + use_pdl = get_env_enable_pdl() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = ( + _sm100f_nvfp4_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + if total_q % query_group_size or decode_query_len % query_group_size: + raise ValueError( + f"query group {query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." + ) + + producer_kernel = ( + _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel + ) + device_identity = identity_blocks + if device_identity and not use_mq: + m3_report.taken(DEVICE_IDENTITY_PATH, total_q) + single_kernel_args = ( + { + # No split-K arrival counters: the separate merge kernel combines the chunks. + "merge": (output, None), + "tail": (output.stride(0), output.stride(1), output.stride(2), num_kv_heads), + # An early release pays off when the dependent is the consumer of the + # final output (it can prefetch), or a merge grid that itself releases + # its dependent at entry: the o_proj then launches during this kernel and + # streams its weights while attention runs. Above 8 query rows the + # concurrent weight stream slows the larger attention grid more than the + # prefetch saves (GB300 TP4: -1.1% at 48 rows). + "options": {"FUSED_MERGE": False, "EARLY_TRIGGER": direct_output or total_q <= 8}, + } + if query_group_size == 1 + else {"merge": (), "tail": (), "options": {}} + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + else {} + ) + if use_mq: + from .triton_sparse_decode_mq import launch_mq_producer + + launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + sm_scale=sm_scale, + decode_query_len=decode_query_len, + use_pdl=use_pdl, + early_trigger=False, + num_warps=MQ_NUM_WARPS, + ) + else: + producer_kernel[ + (total_q // query_group_size) * num_topk_chunks, + num_kv_heads, + ]( + q, + k_paged, + v_paged, + scale_arg, + k_sc_arg, + v_sc_arg, + k_gs_arg, + v_gs_arg, + topk_idx, + o_arg, + lse_arg, + *single_kernel_args["merge"], + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + k_sc_arg.stride(0), + k_sc_arg.stride(1), + v_sc_arg.stride(0), + v_sc_arg.stride(1), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + *o_strides, + *lse_strides, + block_table.stride(0), + block_table.shape[1], + k_paged.shape[0], + *single_kernel_args["tail"], + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + WIDEN_Q=widen_q, + COMPUTE_DTYPE=compute_dtype, + USE_PDL=use_pdl, + KV_NVFP4=kv_nvfp4, + USE_LINEAR_SOFTMAX=use_linear_softmax, + SCALE_COLS=scale_cols, + DIRECT_OUTPUT=direct_output, + # Identity selection from the lengths this kernel reads, not the staged table (dispatch kill switch). + IDENTITY=device_identity, + **single_kernel_args["options"], + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + m3_prequant.SLOT.clear() + if not direct_output: + width = num_heads * head_dim + mxfp8 = ( + MERGE_MXFP8.buffers(width, output.device) + if MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS + and head_dim == 128 + and output.dim() == 3 + and output.is_contiguous() + and output.dtype == torch.bfloat16 + else None + ) + mxfp8_args = ( + { + "q8_ptr": mxfp8[0], + "sf_ptr": mxfp8[1], + "num_sf_k_tiles": width // 128, + "MXFP8_OUT": True, + } + if mxfp8 is not None + else {} + ) + merge_launch_options = ( + _sm100f_nvfp4_merge_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_arg, + lse_arg, + output, + head_dim, + *o_strides, + *lse_strides, + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + # Release the o_proj at the merge's entry so it streams its weights during the merge. + EARLY_TRIGGER=True, + **mxfp8_args, + **pdl_launch, + **merge_launch_options, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output.view(total_q, width), mxfp8[0][:total_q], mxfp8[1]) + + +__all__ = [ + "MERGE_MXFP8", + "MERGE_MXFP8_MIN_TOKENS", + "MergeMxfp8", + "NVFP4_SF_VEC_SIZE", + "SPARSE_BLOCK_SIZE", + "_sm100f_nvfp4_merge_launch_options", + "_sm100f_nvfp4_num_topk_chunks", + "_sm100f_nvfp4_query_group_size", + "_sm100f_nvfp4_use_linear_softmax", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py new file mode 100644 index 000000000000..048d31b82d68 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multi-query block-sparse decode producer for MiniMax-M3 short-sequence verify steps. + +In short-sequence mode every request attends at most ``topk`` KV blocks and block selection is the identity, so +all ``decode_query_len`` query tokens of a request read the same blocks. The per-token producer +(``_gqa_sparse_decode_kernel``) then dequantizes every block once per query token. This producer runs one program +per (request, block) for all of the request's query tokens: rows are (token, head) pairs, each block is +dequantized once, and the causal limit is a per-row mask. + +It writes the same per-block partials (output scaled by 1/l and the V global scale, log2 LSE) as the per-token +producer with one block per chunk, so ``_merge_topk_attn_out_kernel`` merges them unchanged. +""" + +import triton +import triton.language as tl + +from .triton_sparse_decode import _dequant_nvfp4_rows + + +@triton.jit( + do_not_specialize=["num_reqs"], do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"] +) +def _gqa_sparse_decode_mq_kernel( + q_ptr, # [num_reqs * QLEN, num_heads, head_dim], FP8 or FP16/BF16 + k_ptr, # packed NVFP4 K: [num_pages, 1, page_size, head_dim / 2] uint8 + v_ptr, # packed NVFP4 V: same layout + k_sc_ptr, # K block scales [num_pages, 1, page_size, head_dim / 16], token-linear + v_sc_ptr, # V block scales, 4x4 token-quad order + k_gs_ptr, + v_gs_ptr, + t_ptr, # topk table [1, total_q, topk]; the request's last token's row is read + o_ptr, # partials [topk, total_q, num_heads, head_dim] fp32 + lse_ptr, # partials [topk, total_q, num_heads] fp32 (log2) + block_table_ptr, + seq_lens, + num_reqs, + sm_scale, + stride_qn, + stride_qh, + stride_k_blk, + stride_k_pos, + stride_v_blk, + stride_v_pos, + stride_ksc_blk, + stride_vsc_blk, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + QLEN: tl.constexpr, + GROUP: tl.constexpr, # query heads per KV head (16) + BLOCK_R: tl.constexpr, # next power of two of QLEN * GROUP + HEAD_DIM: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SCALE_COLS: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr, +): + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) # top-k position == block index under identity selection + sm_scale_log2e = sm_scale * 1.4426950409 * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + + off_r = tl.arange(0, BLOCK_R) + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, HEAD_DIM) + off_dp = tl.arange(0, HEAD_DIM // 2) + off_s = tl.arange(0, SCALE_COLS) + row_t = off_r // GROUP + row_h = off_r % GROUP + r_mask = off_r < QLEN * GROUP + tok = pid_r * QLEN + row_t + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + pid_r) + # Token t of the request sits at seq_len - QLEN + t and attends positions [0, kv_len_t). + kv_len = tl.maximum(seq_len - QLEN + row_t + 1, 0) + blk = tl.load(t_ptr + (pid_r * QLEN + QLEN - 1) * stride_tn + pid_c * stride_tk) + blk_ok = blk >= 0 + blk = tl.where(blk_ok, blk, 0) + page = tl.load(block_table_ptr + pid_r * stride_bt_b + blk, mask=blk_ok, other=0).to(tl.int64) + # Rows past the last token's position are never read: their cache bytes may be stale (0 * NaN in the PV dot). + n_valid = tl.where( + blk_ok, tl.minimum(tl.maximum(seq_len - blk * BLOCK_SIZE_K, 0), BLOCK_SIZE_K), 0 + ) + n_mask = off_n < n_valid + + q = tl.load( + q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :], + mask=r_mask[:, None], + other=0.0, + ).to(tl.float16) + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + page * stride_k_blk + off_n[:, None] * stride_k_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + page * stride_ksc_blk + off_n[:, None] * SCALE_COLS + off_s[None, :], + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ) + valid = ( + (blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None]) & n_mask[None, :] & r_mask[:, None] + ) + qk = tl.dot(q, tl.trans(k_deq).to(tl.float16)) * sm_scale_log2e + qk = tl.where(valid, qk, float("-inf")) + m = tl.max(qk, axis=1) + m_safe = tl.where(m == float("-inf"), 0.0, m) + p = tl.exp2(qk - m_safe[:, None]) + p_sum = tl.sum(p, axis=1) + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + page * stride_v_blk + off_n[:, None] * stride_v_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ).to(tl.float16) + acc = tl.dot(p.to(tl.float16), v) + # An empty partial stores zero with LSE -inf, so the merge gives it weight 0. + scale = tl.where(p_sum > 0.0, 1.0 / p_sum, 0.0) * v_global_scale + lse = tl.where(p_sum > 0.0, m + tl.log2(p_sum), float("-inf")) + acc = acc * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + tok[:, None] * stride_o_b + + row_h[:, None] * stride_o_h + + off_d[None, :], + acc, + mask=r_mask[:, None], + ) + tl.store(lse_ptr + pid_c * stride_l_c + tok * stride_l_b + row_h * stride_l_h, lse, mask=r_mask) + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + *, + sm_scale, + decode_query_len, + use_pdl, + early_trigger, + num_warps=8, + num_stages=1, +): + """Launch the multi-query producer; arguments as minimax_m3_sparse_attn_decode after its NVFP4 views.""" + total_q, num_heads, head_dim = q.shape + num_reqs = total_q // decode_query_len + max_topk = int(topk_idx.shape[-1]) + rows = decode_query_len * num_heads + pdl_launch = {"launch_pdl": True} if use_pdl else {} + _gqa_sparse_decode_mq_kernel[(num_reqs, max_topk)]( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + o_arg, + lse_arg, + block_table, + seq_lens, + num_reqs, + sm_scale, + q.stride(0), + q.stride(1), + k_paged.stride(0), + k_paged.stride(2), + v_paged.stride(0), + v_paged.stride(2), + k_block_scale.stride(0), + v_block_scale.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + o_arg.stride(0), + o_arg.stride(1), + o_arg.stride(2), + lse_arg.stride(0), + lse_arg.stride(1), + lse_arg.stride(2), + block_table.stride(0), + QLEN=decode_query_len, + GROUP=num_heads, + BLOCK_R=triton.next_power_of_2(rows), + HEAD_DIM=head_dim, + BLOCK_SIZE_K=int(k_paged.shape[2]), + SCALE_COLS=head_dim // 16, + USE_PDL=use_pdl, + EARLY_TRIGGER=early_trigger, + num_warps=num_warps, + num_stages=num_stages, + **pdl_launch, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/weights.py new file mode 100644 index 000000000000..07fbe7f06468 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/weights.py @@ -0,0 +1,865 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_107 / tp4. + +Two tables, one per checkpoint the target reads: + +* ``MODEL_WEIGHTS`` (``MiniMaxM3Tp4Weights``): the vision-language checkpoint's language model, + ``language_model.*``. Every rank is handed the whole checkpoint, so every tensor-parallel split is this + table's: the Q heads, the K/V/index-Q heads (one each per rank), the dense MLPs and the shared expert's + intermediate, the o_proj / down-projection input columns, every routed expert's intermediate (all 128 + experts live on every rank), and the vocabulary of the shell's ``lm_head``. The index-K head is replicated. +* ``DRAFT_WEIGHTS`` (``MiniMaxM3Tp4DraftWeights``): the one-layer BF16 Eagle3 drafter's own checkpoint + (``speculative_model``), loaded through the shell's ``load_draft_weights`` when speculation is on. + +Storage is HF ``[out, in]`` row-major. Five families need a transform, each a layout a kernel demands: + +* **MXFP8 block scales** (every attention projection, the dense MLPs and the shared expert): the checkpoint + stores UE8M0 bytes ``[N, K / 32]``; the GEMMs read the 128x4-swizzled order + (``trtllm::block_scale_interleave``). A fused projection concatenates its shards' scales first and swizzles + once. +* **The five-way sparse-layer projection** ``[Q | K | V | index-Q | index-K]`` and the three-way dense one + ``[Q | K | V]``: the per-rank shards concatenated along N. +* **Gated MLPs** (dense and shared expert): ``[gate ; up]`` rows. +* **NVFP4 routed experts** (layers 3..59): ``[up (w3) ; gate (w1)]`` rows, then the trtllm-gen row shuffle of + the weight and of its E4M3 scales (gated-activation reorder + 128-row epilogue tile shuffle, + ``trtllmgen_maybe_get_cached_*_permute_indices``) and the 128x4 swizzle of the scales. +* **Gemma boundary norms**: ``input_layernorm``, ``post_attention_layernorm`` and ``model.norm`` are stored as + Gemma norms, ``(1 + w) * x``; the fused all-reduce + RMSNorm kernels apply ``w * x``, so ``1`` is added + (in BF16) at load. The per-head Q/K/index norms stay Gemma: their kernels apply the offset themselves. + +The routed experts' per-tensor scalars (input scales and second-level weight scales, every expert) load raw; +``modeling`` derives the kernels' global scales, alphas and activation constants from them once the weights +are real. + +Loading contract: ``load(model, weights)`` consumes the engine-provided checkpoint dict, fills every declared +parameter exactly once and asserts coverage in both directions: every declared parameter written, and every +checkpoint key either consumed or in ``expected_unconsumed`` (the vision tower and its projector, which this +text-only target does not build, and the language model's multi-token-prediction head). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) + +#: ``layers=ALL_LAYERS`` is the common case: a weight every decoder layer carries. +ALL_LAYERS = object() + +_LM = "language_model." +#: Checkpoint key families this target does not build: the vision tower and its two projectors (text only), and the +#: language model's multi-token-prediction head (speculation, when on, is the separate Eagle3 drafter). +_UNBUILT_PREFIXES = ( + "vision_tower.", + "multi_modal_projector.", + "patch_merge_mlp.", + _LM + "model.mtp.", +) + +MX_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 elements along K +NV_BLOCK = 16 # NVFP4: one E4M3 scale per 16 elements along K +EPILOGUE_TILE_M = ( + 128 # the trtllm-gen MoE runner's row tile: fixes its weight and scale row shuffles +) + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + ``shape`` takes the dimension bundle and returns the declared shape. ``src`` is a callable + ``(d, i) -> [(key or key tuple, index into the param | None), ...]`` with keys relative to the checkpoint + root; a key tuple hands the transform those tensors together. ``layers`` is ``ALL_LAYERS``, ``None`` (not + per-layer) or ``d -> layer indices``. ``transform(src, d)`` runs on the destination's device before the copy. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + +# ------------------------------------------------------------------------------------------------ shards + + +def _col( + t: torch.Tensor, d: SimpleNamespace, tp: int | None = None, rank: int | None = None +) -> torch.Tensor: + """This rank's rows (column-parallel output split), even split, sliced on the host and moved to the + destination device (only this rank's bytes are read).""" + tp = d.tp if tp is None else tp + rank = d.rank if rank is None else rank + n = t.shape[0] // tp + return _dev(t[rank * n : (rank + 1) * n], d) + + +def _row(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """This rank's input columns (row-parallel input split), even split, on the destination device.""" + n = t.shape[1] // d.tp + return _dev(t[:, d.rank * n : (d.rank + 1) * n], d) + + +def _dev(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """``t`` on the destination device (``d.dev``, set per fill).""" + return t.to(d.dev, non_blocking=True) + + +def _swizzle(scale: torch.Tensor) -> torch.Tensor: + """128x4 swizzle of a row-major ``[N, K / 32]`` UE8M0 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(scale.contiguous()) + + +# --------------------------------------------------------------------------------------------- transforms + + +def _gemma_fold(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A Gemma boundary norm as the plain norm the fused all-reduce epilogue applies: ``w + 1`` in BF16.""" + return t + 1.0 + + +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) + + +def _qkv_sparse_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[Q | K | V | index-Q | index-K]``: Q by attention heads, K/V/index-Q by KV heads, index-K replicated.""" + q, k, v, iq, ik = t + return torch.cat( + [_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0 + ).contiguous() + + +def _qkv_sparse_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v, iq, ik = t + return _swizzle( + torch.cat([_col(q, d), _col(k, d), _col(v, d), _col(iq, d), _dev(ik, d)], dim=0) + ) + + +def _qkv_dense_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +def _qkv_dense_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return _swizzle(torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0)) + + +def _index_qk_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[index-Q | index-K]`` in BF16: this rank's index-Q heads and the replicated index-K, the E4M3 values cast + (the scale tensors, the last two sources, are consumed only).""" + iq, ik = t[0], t[1] + return torch.cat([_col(iq, d), _dev(ik, d)], dim=0).to(d.dtype).contiguous() + + +def _gate_up_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return torch.cat([_col(gate, d), _col(up, d)], dim=0).contiguous() + + +def _gate_up_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return _swizzle(torch.cat([_col(gate, d), _col(up, d)], dim=0)) + + +def _row_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _row(t, d).contiguous() + + +def _row_s(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _swizzle(_row(t, d)) + + +def _col_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _col(t, d).contiguous() + + +#: Permute indices of the trtllm-gen expert relayouts, keyed by (shape, family, scale group): one expert's are +#: every expert's, so each is computed once per load. +_PERMUTE_CACHE: dict = {} + + +def _expert_w31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 operand: ``[up ; gate]`` packed-FP4 rows of this rank, then the gated-activation row + reorder and the epilogue-tile shuffle.""" + up, gate = t + w = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + w, _PERMUTE_CACHE, EPILOGUE_TILE_M, is_gated_act_gemm=True + ) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 E4M3 scales: ``[up ; gate]`` rows of this rank, the same row reorder at scale + granularity, then the 128x4 swizzle.""" + up, gate = t + s = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK, is_gated_act_gemm=True + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled.view(torch.uint8).reshape(shape)) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + + +def _expert_w2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 operand: this rank's input columns, then the epilogue-tile row shuffle.""" + w = _row(t, d).contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(w, _PERMUTE_CACHE, EPILOGUE_TILE_M) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 E4M3 scales: this rank's columns, the row shuffle at scale granularity, the swizzle.""" + s = _row(t, d).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w2_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled).view(torch.float8_e4m3fn).reshape(shape) + ) + + +# ------------------------------------------------------------------------------------------------ sources + + +def _one(template: str): + return lambda d, i, t=template: [(_LM + t.format(i=i), None)] + + +def _many(*templates: str): + return lambda d, i, ts=templates: [(tuple(_LM + t.format(i=i) for t in ts), None)] + + +def _per_expert(*suffixes: str): + """One source per routed expert, written at its expert index (all experts are local at tp4).""" + + def build(d, i, sx=suffixes): + out = [] + for e in range(d.num_experts): + keys = tuple(_LM + f"model.layers.{i}.block_sparse_moe.experts.{e}.{s}" for s in sx) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _dense_layers(d): + return range(d.first_sparse) + + +def _sparse_layers(d): + return range(d.first_sparse, d.num_layers) + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense layers, and every layer without the five-way packed + projection (``fuse_qkv_index_projection=False``).""" + return range(d.first_sparse) if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return range(d.first_sparse, d.num_layers) if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else range(d.first_sparse, d.num_layers) + + +_ATTN = "model.layers.{i}.self_attn" +_MOE = "model.layers.{i}.block_sparse_moe" + + +class MiniMaxM3Tp4Weights: + """MiniMax-M3 NVFP4 / sm_107 / tp4: the language model's table and the passes over it. + + ``declare`` allocates every weight it names, ``manifest`` maps checkpoint keys onto them, ``load`` fills + them; ``_targets`` is the one place that decides which parameters exist. + """ + + #: Free the per-expert relayout scratch after each expert stack. + RELEASE_AFTER = "_e_s2" + + WEIGHTS: tuple[W, ...] = ( + # --- the two boundary norms of every layer (Gemma, folded) --- + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + transform=_gemma_fold, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + transform=_gemma_fold, + ), + # --- attention: per-head Q/K norms (Gemma, unfolded), the packed QKV and the output projection --- + W("q_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".q_norm.weight")), + W("k_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".k_norm.weight")), + W( + "qkv_w", + shape=lambda d: (d.qkv_dense_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in "qkv")), + transform=_qkv_dense_w, + layers=_three_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_dense_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many(*(f"{_ATTN}.{p}_proj.weight_scale_inv" for p in "qkv")), + transform=_qkv_dense_s, + layers=_three_way_qkv_layers, + ), + W( + "qkv_w", + shape=lambda d: (d.qkv_sparse_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in ("q", "k", "v", "index_q", "index_k"))), + transform=_qkv_sparse_w, + layers=_five_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_sparse_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + *( + f"{_ATTN}.{p}_proj.weight_scale_inv" + for p in ("q", "k", "v", "index_q", "index_k") + ) + ), + transform=_qkv_sparse_s, + layers=_five_way_qkv_layers, + ), + # The separate BF16 index projection [index-Q | index-K] of the unfused sparse layers. The built-in model + # builds it as an unquantized Linear and fills it through the fused gate/up loader, which casts the + # checkpoint's E4M3 values to BF16 and does not apply their block scales: so does this table (the scale + # tensors are consumed, not applied). + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + src=_many( + _ATTN + ".index_q_proj.weight", + _ATTN + ".index_k_proj.weight", + _ATTN + ".index_q_proj.weight_scale_inv", + _ATTN + ".index_k_proj.weight_scale_inv", + ), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_layers, + ), + W( + "o_w", + shape=lambda d: (d.hidden, d.q_width), + dtype=torch.float8_e4m3fn, + src=_one(_ATTN + ".o_proj.weight"), + transform=_row_w, + ), + W( + "o_s", + shape=lambda d: (d.hidden * (d.q_width // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_ATTN + ".o_proj.weight_scale_inv"), + transform=_row_s, + ), + # --- dense SwiGLU-OAI MLP (layers 0..2) --- + W( + "mlp_gu_w", + shape=lambda d: (2 * d.dense_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight", "model.layers.{i}.mlp.up_proj.weight" + ), + transform=_gate_up_w, + layers=_dense_layers, + ), + W( + "mlp_gu_s", + shape=lambda d: (2 * d.dense_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight_scale_inv", + "model.layers.{i}.mlp.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_dense_layers, + ), + W( + "mlp_dn_w", + shape=lambda d: (d.hidden, d.dense_inter), + dtype=torch.float8_e4m3fn, + src=_one("model.layers.{i}.mlp.down_proj.weight"), + transform=_row_w, + layers=_dense_layers, + ), + W( + "mlp_dn_s", + shape=lambda d: (d.hidden * (d.dense_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one("model.layers.{i}.mlp.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_dense_layers, + ), + # --- MoE (layers 3..59): the fp32 router and its bias, the shared expert --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + dtype=torch.float32, + src=_one(_MOE + ".gate.weight"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "sh_gu_w", + shape=lambda d: (2 * d.shared_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + _MOE + ".shared_experts.gate_proj.weight", _MOE + ".shared_experts.up_proj.weight" + ), + transform=_gate_up_w, + layers=_sparse_layers, + ), + W( + "sh_gu_s", + shape=lambda d: (2 * d.shared_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _MOE + ".shared_experts.gate_proj.weight_scale_inv", + _MOE + ".shared_experts.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_sparse_layers, + ), + W( + "sh_dn_w", + shape=lambda d: (d.hidden, d.shared_inter), + dtype=torch.float8_e4m3fn, + src=_one(_MOE + ".shared_experts.down_proj.weight"), + transform=_row_w, + layers=_sparse_layers, + ), + W( + "sh_dn_s", + shape=lambda d: (d.hidden * (d.shared_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_MOE + ".shared_experts.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_sparse_layers, + ), + # --- the routed experts, NVFP4 (trtllm-gen layouts), all 128 on every rank --- + W( + "e_w31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_per_expert("w3.weight", "w1.weight"), + transform=_expert_w31, + layers=_sparse_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight"), + transform=_expert_w2, + layers=_sparse_layers, + ), + W( + "e_s31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w3.weight_scale", "w1.weight_scale"), + transform=_expert_s31, + layers=_sparse_layers, + ), + W( + "e_s2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w2.weight_scale"), + transform=_expert_s2, + layers=_sparse_layers, + ), + # The per-tensor scalars of every expert, raw; modeling derives the kernels' scalars from them. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + layers=_sparse_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_scale"), + ("isc3", "w3.input_scale"), + ("isc2", "w2.input_scale"), + ("ws2_1", "w1.weight_scale_2"), + ("ws2_3", "w3.weight_scale_2"), + ("ws2_2", "w2.weight_scale_2"), + ) + ), + # --- not per-layer --- + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=_one("model.embed_tokens.weight"), + layers=None, + ), + W( + "final_norm", + shape=lambda d: (d.hidden,), + src=_one("model.norm.weight"), + transform=_gemma_fold, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + sparse = cfg.sparse_attention_config + freq = [int(f) for f in sparse["sparse_attention_freq"]] + tp = mapping.tp_size + head_dim = int(cfg.head_dim) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + index_heads = max(int(sparse["sparse_num_index_heads"]) // tp, 1) + index_dim = int(sparse["sparse_index_dim"]) + return SimpleNamespace( + fused=bool( + getattr( + core.model_config.sparse_attention_config, "fuse_qkv_index_projection", False + ) + ), + num_layers=cfg.num_hidden_layers, + first_sparse=freq.index(1), + hidden=cfg.hidden_size, + head_dim=head_dim, + heads=heads, + kv_heads=kv_heads, + index_heads=index_heads, + index_dim=index_dim, + q_width=heads * head_dim, + qkv_dense_rows=(heads + 2 * kv_heads) * head_dim, + qkv_sparse_rows=(heads + 2 * kv_heads) * head_dim + index_heads * index_dim + index_dim, + dense_inter=int(cfg.dense_intermediate_size) // tp, + shared_inter=int(cfg.shared_intermediate_size) * int(cfg.n_shared_experts) // tp, + moe_inter=int(cfg.intermediate_size) // tp, + num_experts=int(cfg.num_local_experts), + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None)] + return [(f"l{i}_{entry.name}", i) for i in layers] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares (meta under the engine's meta init).""" + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key or key tuple, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, core, weights) -> set: + """The vision tower and its projectors, and the multi-token-prediction head: in the checkpoint, never built + by this target.""" + return {k for k in weights.keys() if k.startswith(_UNBUILT_PREFIXES)} + + def load(self, model, weights) -> None: + """Fill every declared weight and the shell's ``lm_head`` shard; assert coverage both ways.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + assert set(manifest) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest) ^ set(core.w.keys()), + ) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(core, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(core.w[param_key], ckpt, index, transform, weights, d, consumed) + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + # The shell owns the vocab-parallel head: this rank's rows of the full table. + _fill(model.lm_head.weight, _LM + "lm_head.weight", None, _col_w, weights, d, consumed) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + } + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the unbuilt families (vision, MTP head); " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _materialize(entry) -> torch.Tensor: + """``[:]`` on a lazy safetensors slice; a 0-dim tensor (the per-tensor scalars) rejects that index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +def _fill(param, ckpt, index, transform, weights, d, consumed) -> None: + keys = ckpt if isinstance(ckpt, tuple) else (ckpt,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + d.dev = dst.device + # Host tensors: a transform slices this rank's part on the host and moves only that to the device. + src_all = tuple(_materialize(weights[key]) for key in keys) + if transform is not None: + src = transform(src_all if isinstance(ckpt, tuple) else src_all[0], d) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + if src.dim() == 0 and dst.dim() == 0: + pass + elif src.numel() == 1 and dst.numel() == 1: + src = src.reshape(dst.shape) + src = src.to(dst.device, non_blocking=True) + assert tuple(dst.shape) == tuple(src.shape), (ckpt, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + _release(weights, keys) + + +def _release(weights, keys) -> None: + """Drop consumed checkpoint tensors from the engine's dict as the load goes (bounded host memory).""" + if hasattr(weights, "mark_consumed_keys"): + weights.mark_consumed_keys(keys) + else: + for key in keys: + weights.pop(key, None) + + +MODEL_WEIGHTS = MiniMaxM3Tp4Weights() + + +# -------------------------------------------------------------------------------------------- the drafter + + +def _draft_one(template: str): + return lambda d, i, t=template: [(t, None)] + + +def _draft_many(*templates: str): + return lambda d, i, ts=templates: [(tuple(ts), None)] + + +def _draft_fc_norms(d): + return range(d.capture_layers) + + +def _draft_qkv(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + return torch.cat([_col(q, d), _col(k, d), _col(v, d)], dim=0).contiguous() + + +class MiniMaxM3Tp4DraftWeights: + """The one-layer BF16 Eagle3 drafter (``speculative_model``), tp4. + + The fc over the three captured target layers and its per-layer fc_norms, the layer's three norms, its QKV + (input ``[embeds | hidden]``, so twice the hidden width; Q by heads, K/V by KV heads) and o_proj, its gated + MLP, the final norm and the full-vocabulary head (vocab-parallel). The checkpoint's ``embed_tokens`` is the + target's table (the drafter reads the target's); it is a predicted non-load. + """ + + WEIGHTS: tuple[W, ...] = ( + W( + "fc", + shape=lambda d: (d.hidden, d.capture_layers * d.hidden), + src=_draft_one("fc.weight"), + layers=None, + ), + W( + "fc_norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"fc_norm.{i}.weight", None)], + layers=_draft_fc_norms, + ), + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.input_layernorm.weight"), + layers=None, + ), + W( + "hidden_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.hidden_norm.weight"), + layers=None, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.post_attention_layernorm.weight"), + layers=None, + ), + W( + "qkv", + shape=lambda d: (d.qkv_rows, 2 * d.hidden), + src=_draft_many(*(f"layers.0.self_attn.{p}_proj.weight" for p in "qkv")), + transform=_draft_qkv, + layers=None, + ), + W( + "o", + shape=lambda d: (d.hidden, d.q_width), + src=_draft_one("layers.0.self_attn.o_proj.weight"), + transform=_row_w, + layers=None, + ), + W( + "gu", + shape=lambda d: (2 * d.inter, d.hidden), + src=_draft_many("layers.0.mlp.gate_proj.weight", "layers.0.mlp.up_proj.weight"), + transform=_gate_up_w, + layers=None, + ), + W( + "dn", + shape=lambda d: (d.hidden, d.inter), + src=_draft_one("layers.0.mlp.down_proj.weight"), + transform=_row_w, + layers=None, + ), + W("final_norm", shape=lambda d: (d.hidden,), src=_draft_one("norm.weight"), layers=None), + ) + + def dims(self, draft) -> SimpleNamespace: + cfg = draft.draft_config.pretrained_config + mapping = draft.draft_config.mapping + tp = mapping.tp_size + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + return SimpleNamespace( + hidden=cfg.hidden_size, + capture_layers=draft.capture_layers, + head_dim=head_dim, + q_width=heads * head_dim, + qkv_rows=(heads + 2 * kv_heads) * head_dim, + inter=cfg.intermediate_size // tp, + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + if entry.layers is None: + return [(entry.name, None)] + return [(f"{entry.name}{i}", i) for i in entry.layers(d)] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, draft) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + d = self.dims(draft) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, draft, weights) -> set: + """The drafter checkpoint's own embedding: the drafter reads the target's table.""" + return {"embed_tokens.weight"} & set(weights.keys()) + + def load(self, draft, weights) -> None: + d = self.dims(draft) + manifest = self.manifest(draft) + assert set(manifest) == set(draft.w.keys()), set(manifest) ^ set(draft.w.keys()) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(draft, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(draft.w[param_key], ckpt, index, transform, weights, d, consumed) + _fill(draft.lm_head.weight, "lm_head.weight", None, _col_w, weights, d, consumed) + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + f"drafter checkpoint coverage: unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = MiniMaxM3Tp4DraftWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py new file mode 100644 index 000000000000..856847d0b0fb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/__init__.py @@ -0,0 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 NVFP4 / sm_107 / tp8.""" + +from ..kernels import fake_ops # noqa: F401 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py new file mode 100644 index 000000000000..69ea6f145b80 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py @@ -0,0 +1,6140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ModelingV2 target: MiniMax-M3 NVFP4 / sm_107 / tp8 -- self-contained modeling code. + +MiniMax-M3's language model: 60 layers, hidden 6144, 64 query heads and 4 KV heads of dim 128. Layers 0-2 are dense +attention with a dense SwiGLU-OAI MLP; layers 3-59 are MSA block-sparse attention (an index branch scores 128-token +blocks and the top 16 are attended) with a MoE of 128 routed experts (top 4, sigmoid scores with a correction bias, +renormalized and scaled by 2) plus one shared expert. The checkpoint is a modelopt mixed-precision export: NVFP4 routed +experts, MXFP8 projections, fp32 router, bf16 embedding and head; the KV cache is NVFP4 for the sparse layers and FP8 +for the dense ones. + +``tp8`` is eight Rubin GPUs over two trays (multi-node NVLink): head-split attention (8 query heads, one KV head and +one index head per rank, each KV head shared by two ranks) and tensor-parallel experts (every expert on every rank, +intermediate 384 per rank), no attention DP. +Speculative decoding, when configured, is one-model Eagle3 (draft length 3): the BF16 drafter (``M3Eagle3Draft``) and +the BF16 target head, both at checkpoint precision; decode steps of up to 2 requests verify and draft through fused +kernels at that precision (``M3Eagle3Worker``). + +**The core** (``MiniMaxM3Tp8ModelingV2``) holds the weights (``weights.MODEL_WEIGHTS``) and the runtime resources the +steps use -- one MSA attention backend per layer, the TP all-reduce communicator, the RoPE table, the MXFP8 GEMM +dispatch state the engine tunes -- and dispatches each step, on the host, from its context count alone +(``_select_target``): + +* ``DecodeTarget`` takes every pure generation step. Up to ``DECODE_MAX_TOKENS`` tokens (up to 2 requests, each one + token or a draft verify of 1 + 3 rows) with 1-D input ids it runs the decode-step kernels and picks each layer + boundary's collective by the step's token bucket (``DECODE_PLANS``); it runs its own generic layer for the first + layer when the fused embedding does not take the step, and for every layer of a wider step. +* ``PrefillTarget`` takes every step with context rows: prefill chunks and mixed steps, every layer in its generic + form over the MSA backend. + +The two targets share no body: each owns its layer, attention and MoE compositions and its ``step_args``, the one +projection of the step metadata its forward reads. What both call are the module-level op adapters, each picking the +kernel for one call from shapes and layouts, which read no step metadata themselves. + +Every fast path the decode step can take, the rows each covers (``FAST_PATHS``, each paired with the entry whose +``CELLS`` certify it) and what a step that cannot take one does are below: **uncovered** (outside every covered +range: log once, run the generic path) or **missed** (inside a covered range but a precondition is missing: raise +``FastPathUnavailable`` under ``modeling_v2='require'``, else log once and fall back). A row range is empty unless a +landed cell has certified it at the TP8 geometry on Rubin. Every fast path a GEN rank takes prints a rank-tagged +engagement line once per (path, rows) (``taken``). Kernel modules report through ``ops/m3_report``, which this +module wires to the same policy at import. + +**Hand-offs between kernels.** A producer that already holds an activation in registers writes its MXFP8 form for the +next MXFP8 projection: the decode boundary all-reduce for the next layer's QKV, the sparse-decode merge and the +long-context decode for the o_proj. It publishes it in ``ops/m3_prequant.SLOT`` keyed by the BF16 tensor; the next +MXFP8 projection on a decode graph takes it, and every take clears the slot. The decode boundaries' push targets (the +o_proj / dense-down / shared-expert GEMVs and the MoE engine writing straight into the boundary all-reduce's exchange) +are passed explicitly from the layer to the producer and back. + +Compiled kernels (the MoE engine, the decode GEMVs with their push variants, the index scorer, the long-row select, +the full-context decode plans) are built before the executor's warmup; a forward never compiles one, and one it finds +missing is a miss. Other state a fast path needs (multicast workspaces, push descriptors) is built on the eager decode +steps; a CUDA-graph capture only reads it, so a capture that finds it missing is a miss (``capture_missed``). + +**What this target asserts rather than adapts**: SM 10.7, the topology and checkpoint (``routing.py``); the +MiniMax-M3 MSA sparse attention runtime with the fused five-way QKV/index projection, an FP8 indexer and an FP8 or +NVFP4 KV cache, no torch.compile and no piecewise prefill graphs (``_validate_runtime``); text-only input (the vision +tower is not built and a step carrying image input raises). +""" + +from __future__ import annotations + +import atexit +import copy +import dataclasses +import os +import threading +import time +from dataclasses import dataclass +from types import MappingProxyType +from typing import ( + Any, + Callable, + Dict, + FrozenSet, + Hashable, + Iterable, + List, + Mapping, + Optional, + Tuple, +) + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + ModelingV2Mode, + resolved_mode, + step_contract_enabled, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.moe.fp4_block_scale_moe_runner import ( + fp4_block_scale_moe_runner, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( + flashinfer_fused_add_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_rmsnorm import ( + flashinfer_rmsnorm, +) +from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( + fp4_quantize, +) +from tensorrt_llm._torch.attention.backends import AttentionMetadata +from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3MsaSparseAttention +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + MSA_REQUIRED_TOPK, + msa_paged_kv, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, +) +from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_speculative import ( + _ensure_draft_vocab_size, + _get_requested_draft_moe_backend, + _set_draft_kv_cache_quant_algo, + _slice_spec_position_ids, +) +from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + TensorParallelMode, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) +from tensorrt_llm._torch.modules.logits_processor import LogitsProcessor +from tensorrt_llm._torch.modules.multi_stream_utils import maybe_execute_in_parallel +from tensorrt_llm._torch.moe.fused_moe.routing import RoutingMethodType +from tensorrt_llm._torch.pyexecutor.breakable_cuda_graph import is_in_breakable_cuda_graph +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy, RuntimeHooks +from tensorrt_llm._torch.speculative import SpecMetadata +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.utils import ActType_TrtllmGen +from tensorrt_llm._utils import mpi_rank, prefer_pinned +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +from . import weights as _weights +from .ops import m3_draft_step as m3_draft_step_kernel +from .ops import ( + m3_drafter_attn, + m3_fc_wait, + m3_fold_tail, + m3_mc_allreduce, + m3_prequant, + m3_report, + m3_spec_tail, +) +from .ops import m3_fullctx_decode as m3_fullctx_decode_kernel +from .ops import m3_index_score as _index_score_kernel +from .ops import m3_sampled_accept as m3_sampled_accept_kernel +from .ops import m3_shard_accept as m3_shard_accept_kernel +from .ops import msa_kv_update as _msa_kv_update_kernel +from .ops import triton_sparse_decode as _sparse_decode_kernel +from .ops.activation.m3_silu_and_mul import m3_silu_and_mul +from .ops.attention.m3_draft_attention import m3_draft_attention +from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope +from .ops.attention.m3_fullctx_decode import m3_fullctx_decode +from .ops.attention.m3_index_score import m3_index_score +from .ops.attention.m3_msa_attention import m3_msa_attention +from .ops.attention.m3_producer import m3_producer +from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 +from .ops.attention.m3_qkv_insert import m3_qkv_insert +from .ops.attention.m3_select_blocks import m3_select_blocks +from .ops.attention.m3_select_long import m3_select_long +from .ops.attention.m3_select_long_dsl import m3_select_long_dsl +from .ops.attention.m3_sparse_decode import m3_sparse_decode +from .ops.attention.m3_sparse_decode_lc8 import m3_sparse_decode_lc8 +from .ops.comm.m3_allreduce import m3_allreduce +from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm +from .ops.embedding.m3_embedding import m3_embedding +from .ops.gemm.m3_bf16_linear import m3_bf16_linear +from .ops.gemm.m3_dense_swiglu import ALPHA as DENSE_ALPHA +from .ops.gemm.m3_dense_swiglu import BETA as DENSE_BETA +from .ops.gemm.m3_dense_swiglu import LIMIT as DENSE_LIMIT +from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu +from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS +from .ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for +from .ops.gemm.m3_fp32_linear import m3_fp32_linear +from .ops.gemm.m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + m3_mx_gemv, +) +from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 +from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm +from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear +from .ops.long_select import long_select_blocks +from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine +from .ops.moe.m3_router_tc import m3_router_tc +from .ops.mx_gemv import op as mx_gemv +from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm +from .ops.spec.m3_capture_norm import m3_capture_norm +from .ops.spec.m3_chunk_norm import m3_chunk_norm +from .ops.spec.m3_draft_input import m3_draft_input +from .ops.spec.m3_draft_silu import m3_draft_silu +from .ops.spec.m3_draft_step import m3_draft_step +from .ops.spec.m3_force_vec import m3_force_vec, m3_force_vec_req +from .ops.spec.m3_sampled_accept import m3_sampled_accept +from .ops.spec.m3_shard_accept import m3_shard_accept +from .ops.spec.m3_spec_offsets import m3_spec_offsets +from .ops.spec.m3_step_prologue import m3_step_prologue + +_SM = (10, 7) +# -------------------------------------------------------------------------------------------------------------------- +# Row ranges +# -------------------------------------------------------------------------------------------------------------------- + +# The largest generation step the decode target takes: 2 requests verifying 1 + 3 drafts (the GEN worker's batch). +DECODE_MAX_TOKENS = 8 # capacity: fixed: the decode kernels' step bound + + +@dataclass(frozen=True) +class DecodePlan: + """The layer-boundary collectives of one decode token bucket (GB300 tp4 measurements in the comments).""" + + name: str + max_tokens: int + # The attention boundary's multicast all-reduce + residual + RMSNorm (5.05 vs 5.27 us native at 1 token). + attn_mc_ar: bool + # The same kernel at the MoE boundary with the shared + routed add folded in: wins up to 48 tokens. + moe_mc_ar: bool + # The MoE-boundary all-reduce also emits the next QKV activation in MXFP8 for the decode GEMV / GEMM. + qkv_prequant: bool + + +# tiny: every step up to 8 tokens (c1..c8 without speculation, the c1 verify, every draft pass); small: the c2 +# verify; verify: the c4 / c8 verifies. Each is one multicast all-reduce bucket (ops/m3_mc_allreduce.BUCKETS). +# The multicast all-reduce + residual + RMSNorm at both boundaries of every TP2 decode step (one 8-token bucket), the +# MoE boundary's also publishing the next QKV activation in MXFP8. SM107 TP2 hop: 3.5..4.4 us at 1..8 rows against +# 8.7..9.5 us for the stock AUTO (NCCL) all-reduce + fused residual RMSNorm, and 13.3..14.0 us at the MoE boundary +# (stock add + all-reduce + mxfp8_quantize); bitwise equal to stock (ops/comm/m3_mc_allreduce_norm). +# False puts both boundaries back on the built-in all-reduce. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +MC_AR_TP2 = True +DECODE_PLANS = ( + DecodePlan( + "tiny", DECODE_MAX_TOKENS, attn_mc_ar=MC_AR_TP2, moe_mc_ar=MC_AR_TP2, qkv_prequant=MC_AR_TP2 + ), +) + +# The dense layers' decode attention (trtllm-gen, msa_decode.dense_decode_max_seq_len): the max_seq_len it is handed +# per request count. It only sizes the multi-CTA KV split (every CTA still covers its share of the actual length; the +# result changes only by the split-merge order), and the model maximum over-splits AgentX-length caches at one +# request: GB300 TP4, 16 heads, FP8 KV, 64K context: 15.1 us at 16K vs 28.6 us at 1M. Two requests are neutral (the +# grid is unchanged) and keep the step's bound. + + +# TP2 skeleton: no cap (the built-in split) until a TP2 cap is measured and certified. +DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH: Mapping[int, int] = MappingProxyType({}) + +# The dense layers' decode attention on PrimTS split-KV decode (ops/m3_fullctx_decode) instead of the trtllm-gen +# call, per (query rows per request, requests): the Q tile (``config``) and the split-KV fanout (``splits``) with the +# BF16 merge. At TP8 (8 query heads, 1 KV head per rank) the stock split is capped at 16 cluster CTAs per (request, KV +# head), i.e. 16 of 212 CTAs at one request. Rubin TP8 geometry, FP8 KV, chained us per call at 165k / 323k, every +# recipe passing the fp32 oracle (microbenches tp8s 777734 and tp4k 777845, locked clock, cold L2): (4, 1) 12.65 / +# 15.88 at tile 16 / 96 splits / 12 K/V stages; (4, 2) 15.85 / 21.32 at tile 32 / 96 splits. At one request, 108 or +# more splits at tile 16 fall off a cliff at 165k, and fewer K/V stages are slower at both lengths. The drafter's +# one-row passes (the 8-row Q tile has no separate-reduction profile with FP16 partials, so tile 16): (1, 1) 10.75 us +# at 108 splits / 12 K/V stages (tp8b 777783); (1, 2) 13.35 at 64 splits / 12 K/V stages (tp8c 779066). +FULLCTX_DECODE_RECIPES: Mapping[Tuple[int, int], Mapping[str, Any]] = MappingProxyType( + { + (4, 1): MappingProxyType( + { + "config": {"tile_size_q": 16, "groups_tokens_heads_q": True, "kv_stages": 12}, + "splits": 96, + } + ), + (4, 2): MappingProxyType( + {"config": {"tile_size_q": 32, "groups_tokens_heads_q": True}, "splits": 96} + ), + (1, 1): MappingProxyType({"config": {"tile_size_q": 16, "kv_stages": 12}, "splits": 108}), + (1, 2): MappingProxyType({"config": {"tile_size_q": 16, "kv_stages": 12}, "splits": 64}), + } +) +# False puts the dense layers back on the trtllm-gen call at every shape. +FULLCTX_DECODE_ENABLED = True +# The dense layers' FP8 q/k/v (decode steps of the hybrid cache): the TP2 op rounds Q/K to BF16 and then to E4M3 in +# its store, the bits of the BF16 fused_qk_norm_rope output cast to FP8. False runs that BF16 kernel and a separate cast +# (the same bits, one more kernel). +DENSE_QKV_FP8_VIA_BF16 = True +# On a full-context decode step that op also stores the new K/V rows into the layer's paged pool (the bytes +# write_msa_phase_kv's indexed copies store, one launch instead of about ten). False writes them with +# write_msa_phase_kv. +DENSE_KV_INSERT_FUSED = True +# The engagement line of that fused store (dispatch.taken, one per row count and rank). +DENSE_KV_INSERT_PATH = "m3_dense_kv_insert" +# The drafter MLP's C++ SwiGLU (DRAFT_SILU_CPP, spec_worker._m3_draft_mlp): its own taken / not-taken lines. +DRAFT_SILU_PATH = "m3_draft_silu" +# The drafter boundary all-reduce fed by its projection's push (DRAFT_AR_PUSH): its own taken / not-taken lines. +DRAFT_AR_PUSH_PATH = "m3_draft_ar_push" +# The split-KV merge on the BF16 op (minimax_m3_sm107_tp2_fullctx_reduce) instead of PrimTS's FP16 reducer and a cast: +# Rubin, 165k, graph-us per call (1, 1) 14.5 vs 18.5, (4, 1) 17.3 vs 21.3, (1, 2) 18.5 vs 23.1, (4, 2) 23.3 vs 27.5. +FULLCTX_DECODE_OWN_REDUCER = True +# The TP8 merge also writes the MXFP8 copy of its BF16 rows (E4M3 + 128x4-swizzled UE8M0, bit for bit +# trtllm::mxfp8_quantize) into the MERGE_MXFP8 buffers and publishes it to m3_prequant.SLOT, so the dense layer's +# MXFP8 o_proj skips its own quantization. Kill: False leaves the o_proj to quantize its BF16 input. +FULLCTX_DECODE_MXFP8 = True +FULLCTX_DECODE_MXFP8_PATH = "m3_fullctx_merge_mxfp8" +# The K/V length bound the full-context plans are compiled for; a step whose cache length bound is longer is a miss. It +# covers the TP2 AgentX config's max_seq_len (1 << 20) plus the speculative rows the cache reserves past it. +FULLCTX_DECODE_MAX_KV_LEN = 1 << 21 # capacity: fixed: the full-context plans' compiled K/V bound +# The Eagle3 drafter layer's generation phase on the same full-context decode (ops/m3_drafter_attn: the in-tree +# RoPE + FP8 K/V append, then the pinned PrimTS recipe) instead of thop.attention's in-tree runner, which caps its +# split at 8 CTAs per (request, KV head) on SM 10.7: 45-46 us per draft pass at 165k. +DRAFTER_FULLCTX_ENABLED = True +# On those passes the drafter's RoPE, FP8 query and FP8 K/V append run as one launch (the op +# minimax_m3_sm107tp8_drafter_rope_fp8_kv: the in-tree preprocess's GPT-NeoX rotation and unit-scale E4M3 store, and +# the torch E4M3 cast of its BF16 query) instead of thop.trtllm_gen_generation_preprocess and the cast. Kill constant. +DRAFTER_ROPE_FUSED = True +# The engagement line of that launch (dispatch.taken, one per row count and rank). +DRAFTER_ROPE_PATH = "m3_drafter_rope" +# thop.trtllm_gen_generation_preprocess builds its workspace token map only for several requests, yet a one-request +# multi-token generation (the drafter's first pass at one request, any stock trtllm-gen fallback) places its rows by +# that map: without a rebuild it applies the last multi-request call's map (rows on another request's length and pages, +# token-0 drafts). ops/m3_tokmap_guard writes the one-request map (0, t) into the op's map slice right before such +# calls: what the build writes (padded rows included), so bitwise the op with a fresh map. False:. +PREPROCESS_TOKMAP_PATH = "m3_preprocess_tokmap" +# How the rotation's products are summed, as the in-tree preprocess build sums them (0: two products then the add, +# 1: fma(cos, x, sin * pair), 2: fma(sin, pair, cos * x)); the op's bitwise test pins it. +DRAFTER_ROPE_MODE = 1 + +# The sparse layers' NVFP4 QKV + indexer producer (ops/attention/m3_producer, every step its geometry fits) triggers its +# dependent kernel on entry, so the index scorer launches and streams its pre-wait index-K prefetch while the producer +# runs (its outputs are read only after the scorer's own grid wait). Rubin TP2, 57-layer PDL chain behind a 21 MB +# pre-wait L2 prefetch (the 165K scorer's shape): 3.59 (native) -> 1.93 us per layer at 4 rows, 3.64 -> 2.00 at 8. +PRODUCER_EARLY_TRIGGER = True +# The producer also reads its per-step inputs (position ids, cache slots, RoPE rows) before its grid wait, leaving only +# the packed QKV row after it. Their writers are host-to-device copies before the replay and plain-launched kernels +# earlier in the step (the input preprocessing, msa_kv_update or its torch fallback), so they are final before the +# producer can start; a PDL-launched writer of any of them would have to turn this off. The call site's +# position_ids.reshape(-1).contiguous().to(torch.int32) and packed.contiguous() launch a plain torch copy right before +# the producer when the dtype or layout differs: also a plain writer, so the contract holds. Bit-identical outputs. +# TP8: True runs this target's copy of the producer (trtllm::minimax_m3_sm107_tp8_producer_pw, +# minimaxM3Sm107Tp8/producerPw8.cu, the TP2 producerPw2 at 8 query heads) with the pre-wait reads; False (the kill +# value) runs the shared minimax_m3_producer_pw at 8 query heads, which reads them after its wait. +# Step-input writers added since: the step-input prologue's C++ op (STEP_PACK_CPP_PROLOGUE) and the C++ spec fold tail +# (STEP_FOLD_TAIL_CPP) are plain launches like the Triton kernels they replace (their tests assert it on the captured +# graph); STEP_KV_LENS_UPDATE_ONCE only drops an earlier writer whose outputs the surviving msa_kv_update rewrites. +# The spec step correction (STEP_SPEC_OFFSETS_ONE_LAUNCH: position_ids, kv_lens) is a plain launch too (its test +# asserts it). +PRODUCER_PREWAIT_STEP_INPUTS = True +# #117: on a decode step of at most 8 rows the sparse layer's QKV + index GEMV runs the producer in its owner epilogue +# (ops/mx_gemv m3_mx_gemv_producer, m3_layer.phase_producer.producer_head): one launch and one PDL edge fewer per +# sparse layer, bit for bit the GEMV + m3_producer chain (q, index-q and every cache byte). The QKV linear defers its +# GEMV (``_linear``'s ``defer``: a ``_DeferredGemv``) and the producer call site runs the fused kernel, or the plain +# GEMV when it declines. +# Edges bench 785463 (TP8 c1, 57 layers, cold L2, both orders): -0.61 .. -0.98 us/layer at 165K / 323K. Kill: False. +QKV_PRODUCER_FOLD = True +QKV_PRODUCER_FOLD_PATH = "m3_qkv_producer_fold" + +# The fused embedding + first RMSNorm kernel: one owner per row count and index dtype. SM107: 1.55..1.65 us against +# 3.4..6.0 us for the stock embedding + RMSNorm at 1..8 rows (the gather bitwise, the norm one ulp off on ~5e-6 of the +# elements: summation order); it also hands the first layer its residual, +# so layer 0 takes the decode layer instead of the generic one. False runs layer 0 on the generic layer. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +EMBED_FIRSTNORM_TP2 = True +EMBED_FIRSTNORM_MAX_TOKENS = DECODE_MAX_TOKENS if EMBED_FIRSTNORM_TP2 else 0 +# The o_proj pushes its rows into the attention boundary's multicast all-reduce (prepushed mode 2) up to this many +# tokens; wider steps run the GEMV and a plain multicast all-reduce. SM107 TP2: -0.5..-1.2 us per layer, bitwise equal +# to the unpushed boundary. False keeps the GEMV and the boundary separate. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +ATTN_PUSH_TP2 = True +ATTN_PUSH_MAX_TOKENS = DECODE_MAX_TOKENS if (MC_AR_TP2 and ATTN_PUSH_TP2) else 0 +# The dense layers' MLP down projection pushes its rows into the post-MLP boundary's multicast all-reduce (prepushed +# mode 2), as the o_proj does into the attention boundary: the all-reduce no longer waits for the whole down GEMV grid +# before loading and pushing the rows itself. SM107 TP2 (stack-ON nsys 773230): the two non-capture dense layers' +# post-MLP hops end 5.5-5.8 us after the down GEMV vs 2.4-2.5 us for a pushed hop. Bitwise equal to the unpushed +# boundary (the same bf16 words). False keeps the down GEMV and the boundary separate. +DENSE_PUSH_TP2 = True +DENSE_PUSH_MAX_TOKENS = DECODE_MAX_TOKENS if (MC_AR_TP2 and DENSE_PUSH_TP2) else 0 +# The generic layer's boundaries (prefill, mixed steps, Eagle3 capture layers): the most rows that take the multicast +# all-reduce + residual + RMSNorm, alone and with the shared + routed add folded in: the decode boundaries' certified +# rows, under the same constant (MC_AR_TP2; off = the built-in all-reduce). +GENERIC_MC_AR_MAX_TOKENS = DECODE_MAX_TOKENS if MC_AR_TP2 else 0 +GENERIC_MC_AR_FOLDED_MAX_TOKENS = DECODE_MAX_TOKENS if MC_AR_TP2 else 0 +AR_PUSH_MODE = 2 +# The 1..8-token MoE engine (router, routed experts and the shared expert in one launch, ops/moe/m3_moe_engine at +# the TP2 width 1536). Rubin TP2 ablation chain (57 layers behind a 58.4 MB PDL producer, row = (P + MoE) - P, both +# orders): 31.3 / 50.7 us/layer at 4 / 8 tokens against the built-in MoE's 42.4 / 55.4 (FlashInfer MXFP8 shared +# expert); bit for bit the built-in block's shared.add_(routed). KILL_MOE_ENGINE: the built-in MoE at every row count. +KILL_MOE_ENGINE = False +MOE_ENGINE_MAX_TOKENS = 0 if KILL_MOE_ENGINE else 8 +# The engine's push of its sum into the MoE boundary's multicast all-reduce (prepushed mode 2) on the rows the +# engine serves, with the multicast boundary (DECODE_PLANS' moe_mc_ar): the boundary reads the pushed words instead +# of loading and pushing the rows itself. The forward arms the push only on rows MOE_PUSH covers. +MOE_PUSH_MAX_TOKENS = MOE_ENGINE_MAX_TOKENS if MC_AR_TP2 else 0 +# When the engine pushes its sum into the boundary all-reduce, every CTA releases the PDL dependents right after its +# grid dependency wait instead of after the push, so the all-reduce (which polls the pushed tagged words, not the grid) +# becomes resident on the SMs the engine frees. Scheduling only: the same bits (world-8 push matrix). Off at TP8: Rubin +# TP8 greedy early column, same-job A,B + B,A, flat (period +20..+26 at 155K..323K); True turns it on. +MOE_EARLY_AR_RELEASE = False +# A greedy Eagle3 verify step of up to this many requests verifies with the fused shard argmax + acceptance kernel +# over the target lm_head's BF16 shard logits (ops/m3_shard_accept), at the configured draft length. +SHARD_ACCEPT_DRAFT_LEN = 3 # capacity: fixed: the built Eagle3 draft length +# 0: (the gathered logits + argmax run). Rubin TP2 greedy verify (microbench): stock 120.4 us at b1; fused +# with the coded head 117.0 us; the BF16-head + shard-accept arm is not measured yet. +SHARD_ACCEPT_MAX_REQUESTS = 2 # capacity: fixed: the TP2 GEN batch +# A sampled (temperature + top-p, advanced_sampling_mode no_topk) Eagle3 verify step of up to this many requests +# verifies with the fused sampled acceptance (ops/m3_sampled_accept, the TP2-shard kernel) over the target +# lm_head's BF16 shard logits, at the configured draft length (the AgentX contract). 0: (the gathered logits +# and the in-tree sampler run). Rubin TP2 b1/b2 (p5 microbench): 180.1 / 181.8 -> 115.1 / 113.9 us per verify. +SAMPLED_VERIFY_MAX_REQUESTS = 2 # capacity: fixed: the TP8 GEN batch +# A step verified by the in-tree path takes the fused first drafter inputs (the capture norms and the first pass's +# gather) up to this many requests. +FIRST_DRAFTER_INPUTS_MAX_REQUESTS = 0 # capacity: fixed: off at TP2 (0) +# A step verified by the fused kernels (dispatch.SHARD_ACCEPT / SAMPLED_ACCEPT) of up to this many requests takes the +# fused first drafter inputs (spec_worker._m3_first_drafter_inputs: the accepted tokens as they lie, the first pass's +# gather rows, the fc_norm chunk norm or the armed capture norm, then the in-tree fc). 0: (the in-tree first +# drafter inputs run, the captures stay in-tree and the fused draft loop steps aside). +FUSED_FIRST_INPUTS_MAX_REQUESTS = 2 # capacity: fixed: the TP2 GEN batch +# On a fused-verify step of up to this many verify rows the Eagle3 captures write rmsnorm(bf16(h + residual)) * +# fc_norm[i] in one kernel (ops/spec/m3_capture_norm) once the fc_norm weights are stacked; the fused first drafter +# inputs then skip their chunk norm. 0: (the in-tree capture, then the chunk norm). +CAPTURE_NORM_MAX_ROWS = 8 # capacity: fixed: the TP2 verify rows (2 requests x 4) +# The first draft pass's per-layer fc_norm over the captured columns in one kernel (ops/spec/m3_chunk_norm) at up +# to this many verify rows; wider first passes take the in-tree first drafter inputs. 0:. +CHUNK_NORM_MAX_ROWS = 8 # capacity: fixed: the TP2 verify rows (2 requests x 4) +# The fused draft loop's two all-reduce boundaries per pass (o_proj -> post-attention norm, down -> final norm) of up +# to this many rows run the multicast all-reduce + residual + RMSNorm (ops/m3_mc_allreduce, opts bit 16: the +# drafter's stock rounding, normalizing the unrounded FP32 sum as FlashInfer fused_add_rmsnorm does) instead of the +# built-in all-reduce + RMSNorm. Kill: 0. Rubin TP2 (p16 microbench): 10.0..11.0 -> 4.6..5.3 us per boundary. +DRAFT_MC_AR_MAX_ROWS = 8 # capacity: fixed: the TP8 verify rows +# The drafter's down -> final-norm boundary (the all-reduce right before each pass's draft step) runs without the +# max-smem carveout (opts bit 8, which AR_OPTS_BY_M sets for 3..6 rows): the carveout made the 212-CTA draft step wait +# for the reduce's CTAs (TP8 census r3: pass 0 at 4 rows ran its reduce at 233,472 B smem, passes 1-2 at 16,384, and +# pass 0's step +10 us; TP4's draft-step penalty, same mechanism). Bit-identical. Kill: True. +DRAFT_MC_AR_CARVEOUT = False +# The drafter's o_proj and down projection, on the copied split-K kernel (DRAFT_LOWM_OWN_SPLITK), also push their rows +# into the drafter boundary all-reduce's exchange (the canonical words its own push would store, at the call parity), +# which then runs prepushed 2 instead of pushing x itself; the copy releases its dependents only after its own grid +# wait and the pushes (the prepushed >= 2 contract). Bitwise. Rubin TP8 world 8, residuals chained boundary to +# boundary (push-3 782332): 8.56 -> 8.08 us per boundary (1 row), 9.21 -> 8.10 (4 rows), ~ -4 us/step. Kill: False. +DRAFT_AR_PUSH = True +# On a fused-verify step of up to this many verify rows the drafter's first fc is forked onto the aux stream right +# after the last captured target layer, overlapping the remaining target layers and the head (spec_worker). +FC_EARLY_MAX_ROWS = 0 # capacity: fixed: measured fork crossover +# A sampled fused verify of up to this many verify rows on the in-tree head GEMM forks the drafter's first fc after the +# head instead of beside it (spec_worker): the fc's split-K grid otherwise holds the SMs first and the head GEMM starts +# late. Kill: 0. Rubin TP2 c1 (locked nsys 768559): the head started 17.4 us after its input was ready. +FC_AFTER_HEAD_MAX_ROWS = 8 # capacity: fixed: measured fork crossover +# The forked fc's completion reaches the first draft input through a counter instead of a stream join: the fork bumps a +# device counter after the fc, a one-thread wait kernel queued after the acceptance spins on it, and the fork is joined +# after the first draft input. That input then has one producer on its own stream (a programmatic, PDL edge) instead of +# the acceptance on another graph stream: in-model the acceptance -> draft input edge cost 2.0-2.4 us (TP2 r4.4) and +# 3.2-3.8 us (TP8 tp8-full) per step. Ordering only: the same values. Kill: False (join before the draft input). +DI_FC_FLAG = True +DI_FC_FLAG_PATH = "m3_di_fc_flag" +# A decode-only Eagle3 step of up to this many requests drafts with the fused draft loop (spec_worker): per pass the +# draft input in one launch (ops/m3_spec_tail.draft_input), the drafter layer's in-tree attention, MLP and +# all-reduces, and one kernel for the draft head (its 12-bit coded copy, or the BF16 fragments when the coding does +# not decode bit for bit), the cross-rank argmax and the pass's bookkeeping (ops/m3_draft_step), at the +# configured draft length. The head kernel takes one row per request, at most 8. +# The fused draft loop streams the 12-bit coded copy of the draft head (when it decodes bit for bit); False keeps the +# BF16 fragments. +DRAFT_STEP_DRAFT_LEN = SHARD_ACCEPT_DRAFT_LEN # capacity: fixed: the built draft length +# 0: (the in-tree draft loop runs). Rubin TP2 per draft pass (head + argmax + exchange): 122.6 / 126.0 -> +# 102.8 / 102.8 us at b1 / b2 (BF16 fragments; the coded head is slower on sm_107: DRAFT_HEAD_CODED False). +DRAFT_STEP_MAX_REQUESTS = 2 # capacity: fixed: the TP8 GEN batch +# H-R2-3: on the fused draft loop's last pass, the step's forward-tail store and next-step staging (the C++ fold tail, +# STEP_FOLD_TAIL_CPP) run in that draft_step's last CTA after its tokens (foldTail.h foldTailWarp, the standalone tail's +# code bit for bit), removing the draft_step -> fold_tail launch edge; M3SpecStepFold.fused_tail_operands then makes +# store_in_forward skip its launch. Off until its in-model pair. Kill: False. +DRAFT_STEP_FOLD_TAIL = False +DRAFT_STEP_FOLD_TAIL_PATH = "m3_draft_step_fold_tail" +# The Eagle3 drafter's BF16 linears (qkv, o_proj, gate_up, down, fc) run FlashInfer's fixed-tactic low-M kernels +# (ops/gemm/m3_draft_lowm) at the rows its TP8 table is certified for, up to this many rows; other calls keep the +# in-tree cuBLAS GEMM. Kill: 0. Rubin TP8 shards (sweep 777890, mclk locked) at 1 row: qkv 6.15 -> 3.94, +# o 4.97 -> 2.42, gate_up 8.09 -> 5.78, down 5.62 -> 4.00 us; the replicated fc 21.50 -> 17.77 (TP2). +DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the certified low-M rows +# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (ops/m3_lowm_splitk): bit for bit the +# same output at every certified cell (spec/test_modeling_v2_m3_lowm_splitk.py); the seam for the drafter-boundary push +# epilogue and the down projection's SwiGLU prologue. Kill: False (FlashInfer's kernel). +DRAFT_LOWM_OWN_SPLITK = True +# On top of DRAFT_SILU_CPP: where the drafter MLP's down projection runs on a split-K tactic of the target's copy +# (DRAFT_LOWM_OWN_SPLITK), its activation is built in the GEMM's B load from gate | up instead of a separate launch +# (ops/m3_draft_lowm.draft_down_swiglu): bit for bit the activation + GEMM +# (spec/test_modeling_v2_m3_lowm_splitk.py), one launch fewer per draft pass. Kill: False. OFF in t8-3b: bitwise but +# slower (SpecTail push-6b 782695, world 8, chained: 10.14 / 10.81 -> 17.38 / 66.85 us per boundary at 1 / 4 rows, +# ~ +70 us/step); +# the down projection keeps its boundary push through draft_linear and the C++ SwiGLU (DRAFT_SILU_CPP) runs. +DRAFT_SILU_FOLD = False +DRAFT_LOWM_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = { + nk: frozenset(m for m in rows if m <= DRAFT_LOWM_MAX_ROWS) + for nk, rows in DRAFT_LOWM_TACTIC_ROWS.items() +} +# The identity-selection sparse decode moves from the per-row producer to the multi-query producer at this many rows +# (ops/triton_sparse_decode.MQ_MIN_TOKENS); the multi-query producer needs one KV head, which the TP8 heads have +# (8 query heads on 1 KV head per rank), so identity steps of at least 17 rows take it. +SPARSE_DECODE_MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query crossover +SPARSE_DECODE_KV_HEADS = 1 # capacity: fixed: the tp8 KV heads per rank +# The long-row top-16 selector takes engines of 129..LONG_SELECT_MAX_BLOCKS KV blocks of 128 tokens; it holds 2, 4 or +# 9 keys per thread by the row's valid-block count. At most 128 blocks, or past the maximum, the built-in selector +# runs; 0 turns the path off. Rubin TP2 (2 KV heads), 165,521 tokens, behind the scorer, ablation per layer: +# -0.9 us at 1 request and -1.2 us at 2 requests against the built-in selector. +LONG_SELECT_MIN_BLOCKS = 129 # capacity: fixed: past the built-in selector's 128 +LONG_SELECT_MAX_BLOCKS = 9216 # capacity: fixed: the select kernels' built blocks per row +LONG_SELECT_KEY_WIDTHS = (2, 4, 9) +# The long-row select releases the sparse decode after it right after its own grid wait (the decode stages q and the +# request length before its wait) up to this many query rows; past it the trigger stays at the select's exit. On Rubin +# TP2 the early-launched long-context sparse decode (SPARSE_DECODE_LC) crowds the select's CTAs off the SMs already at +# one request (4 rows): releasing at the exit is 336 us per decode step faster at 1 request and 350 at 2 (in-model, +# graph time, 165K-context mixed draws, locked memory clock), so the early release is off at every row count. +SELECT_EARLY_MAX_ROWS = 0 # capacity: fixed: off at every row count +# With the long-context decode's single merge (m3_attn_lc8.MERGE_ONCE), that op's CTAs also pull the layer's o_proj +# weight and scales into L2 before their grid dependency (cp.async.bulk.prefetch.L2, a hint: no value changes), so the +# o_proj GEMV after it finds them L2-warm (the SM107 TP4 LC's form, sparseDecodeLc4). Kill: False. +SPARSE_DECODE_LC8_PREFETCH_OPROJ = True +# With the long-context decode's single merge (m3_attn_lc8.MERGE_ONCE), that op also records each step's selected +# block ids per layer and pulls into L2 the K/V pages + scales of the blocks the NEXT sparse layer selected at the +# previous step (65-68% of a layer's distinct selected blocks recur one step later: MSA probe 778226 at TP4), so that +# layer's sparse decode reads them L2-warm. Hints only; up to N1_MAX_ROWS query rows. Kill: False. +SPARSE_DECODE_LC8_PREFETCH_NEXT_KV = True +N1_MAX_ROWS = 128 # capacity: fixed: the next-KV prefetch's row bound +# Without the early release, the long-row select releases the sparse decode once every score of its row is loaded +# (1, before the exact rank of the candidates) instead of at its exit (0). The release still follows +# the select's own grid wait, which the sparse decode's pre-wait q / length / block-table reads rely on. The decode's +# 16-CTA cluster launch and prologue then overlap the rank: Rubin TP2 in-model, locked memory clock, AgentX draws, +# 2 instances vs the r4.3 stored A: replay period -162 us fixed, -94 at a 323K-token request. +SELECT_LONG_TRIGGER_PHASE = 1 +# The long-row select runs as ops/m3_select_long_dsl (ops/attention/m3_select_long_dsl): one 384-thread CTA per row +# instead of one 1024-thread CTA, the same selection bit for bit, with 4, 11 or 24 keys per thread +# (LONG_SELECT_DSL_KEY_WIDTHS) and the same trigger choices. On the Rubin TP8 decode chain (57 layers behind the +# producer stream, 1 / 2 requests) the select is -0.15 / -0.12 us per layer at 165K tokens, 0.00 / -0.01 at 323K and +# -0.04 / -0.04 at 500K against the 1024-thread kernel. Kill: False (the 1024-thread kernel). +SELECT_LONG_DSL = True +LONG_SELECT_DSL_KEY_WIDTHS = (4, 11, 24) +# The MSA index scorer (ops/attention/m3_index_score) of a sparse decode step, up to INDEX_SCORE_MAX_REQUESTS requests +# (0 turns the path off; wider steps run the in-tree CuTe DSL scorer): its CTAs per request by request count, and up +# to INDEX_SCORE_EARLY_MAX_REQUESTS requests the early issue (Q and the first K tile right after the grid wait, one +# older block per CTA prefetched into L2 before it). The split is captured with the graph, so it is one value per +# request count over every context. Rubin TP2 (2 index heads), 165,521 tokens, behind a 58 MB PDL producer, ablation +# per layer: 1 request 3.27 us at 1536 CTAs (512: 3.76, in-tree 256-CTA scorer 6.07); 2 requests 4.97 us at 768 +# (in-tree 6.71). Within 0.2 us of each context's best split from 16K to 400K tokens except 2 requests at 400K (+0.6). +# The early issue reads the request's first page id from the MSA page table before the grid wait: exact because the +# page table is written only on the host side before the replay (attention/test_modeling_v2_m3_block_table_contract.py); +# the request length read there is a hint, every load and store is bounded by the one read after the wait. +INDEX_SCORE_MAX_REQUESTS = 2 # capacity: fixed: the TP8 GEN batch +# At 1 request the split is one wave of CTAs (1272 = 212 SMs x 6 resident at this CTA's shared memory): at the +# 1-index-head geometry (57-layer chain behind the producer, scores and the selection bitwise) -0.03 / -0.27 / +0.06 us +# per layer at 165K / 323K / 500K tokens against 1536. Kill: 1536. +INDEX_SCORE_SPLIT_K: Mapping[int, int] = MappingProxyType({1: 1272, 2: 768}) +INDEX_SCORE_EARLY_MAX_REQUESTS = 2 # capacity: fixed: measured early-issue crossover + +# Synth only (TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS, fractional): the fused verify's forced accepted counts are +# drawn per request, keyed by the request's KV length at the verify (its prompt plus what it has accepted), instead of +# from the server-wide counter, so a request's forced counts do not depend on what the server ran before (A/B draws +# become history-independent). Production sampling never forces counts. False: (the shared counter, the in-tree +# sequence). +FORCE_AL_PER_REQUEST = True + +# The decode MXFP8 GEMVs (ops/mx_gemv, SM107 TP2 table), per projection (N, K): the rows each serves, on an +# activation its producer quantized (the MoE boundary for the QKV, the attention merge for the o_proj) and on a bf16 +# one (quantized in the GEMV's prologue, bit-identical to mxfp8_quantize). SM107 graph us per projection at 1..8 rows: +# 2.2..8.6 against 4.7..10.5 for the stock MXFP8 quantize + FlashInfer CuTe DSL GEMM (ops/mx_gemv/op.py). +# False puts every projection back on the stock decode-graph GEMM. +# On by default since r4.2: the intermittent illegal memory access first seen with the GEMM/AR stack on was the +# built-in Triton sparse decode's (fixed in r4); with the stack on, GSM8K 4-in-flight ran crash-free on r4.1. +MX_GEMV_TP2 = True +# MX-TAC: the decode-graph MXFP8 projections the GEMVs do not serve take one fixed per-bucket decision instead of the +# AutoTuner's per-start choice (whose L2-hot profiling left c2 bimodal at +-6%): FlashInfer's CuTe DSL quantize and +# CuTe DSL GEMM. Stock's candidates per bucket are quantize {TRT-LLM, CuTe DSL} x GEMM {CUTLASS, CuTe DSL}; on SM107 +# every TP2 projection at 1..16 rows runs fastest on this pair, cold and L2-hot (graph us, cold: 4.7..10.3 against +# 5.0..11.5 with the TRT-LLM quantize and 9.7..58 with CUTLASS), and all four candidates give the same bits (quantize +# bytes and scales, GEMM outputs). A constant is rank-agreed and start-independent by construction. +# False returns these projections to the AutoTuner. +MXFP8_DECODE_QUANT_TACTIC = 0 # MXFP8QuantizeRunner.CUTE_DSL +MXFP8_DECODE_GEMM_TACTIC = 0 # FlashInferMXFP8GemmRunner.CUTE_DSL +_ALL_ROWS = frozenset(range(1, DECODE_MAX_TOKENS + 1)) +_GEMV_PROJECTIONS = ( + QKV_SPARSE, + QKV, + O_PROJ, + SHARED_GATE_UP, + SHARED_DOWN, + DENSE_GATE_UP, + DENSE_DOWN, +) +MX_GEMV_ROWS: Dict[Tuple[int, int], FrozenSet[int]] = ( + {nk: _ALL_ROWS for nk in _GEMV_PROJECTIONS} if MX_GEMV_TP2 else {} +) +MX_GEMV_BF16_MAX_TOKENS: Dict[Tuple[int, int], int] = ( + {nk: DECODE_MAX_TOKENS for nk in _GEMV_PROJECTIONS} if MX_GEMV_TP2 else {} +) + +Key = Hashable + + +@dataclass(frozen=True) +class RowRange: + """The shapes one fast path takes, and the catalog entry whose cells certify each of them. + + ``keys`` are the covered shapes: row counts for most paths, (query rows per request, requests) for the attention + kernels, a variant name for thresholds. ``certifies`` maps a cell's spec of ``entry`` to the keys that cell + drives (empty for a cell of another variant). + """ + + path: str + keys: FrozenSet[Key] + entry: OpWrapper + certifies: Callable[[Mapping[str, Any]], Iterable[Key]] + + def covers(self, key: Key) -> bool: + return key in self.keys + + def certified(self) -> FrozenSet[Key]: + return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + + +def _rows(lo: int, hi: int) -> FrozenSet[int]: + return frozenset(range(lo, hi + 1)) + + +def _spec_rows(field: str, **match: Any) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + """A cell's ``spec[field]`` when the cell's other fields match ``match``.""" + + def certifies(spec: Mapping[str, Any]) -> Iterable[Key]: + if any(spec.get(name) != value for name, value in match.items()): + return () + return (spec[field],) + + return certifies + + +def _const(*keys: Key) -> Callable[[Mapping[str, Any]], Iterable[Key]]: + return lambda spec: keys + + +def _sparse_decode_producer(spec: Mapping[str, Any]) -> Iterable[Key]: + if not spec["identity"]: + return () + rows = spec["decode_query_len"] * len(spec["seq_lens"]) + mq = rows >= SPARSE_DECODE_MQ_MIN_TOKENS and SPARSE_DECODE_KV_HEADS == 1 + return ("multi_query" if mq else "per_row",) + + +def index_score_schedule(requests: int) -> Tuple[int, bool]: + """The index scorer's (CTAs per request, early issue) at this many requests (a covered request count).""" + return INDEX_SCORE_SPLIT_K[requests], requests <= INDEX_SCORE_EARLY_MAX_REQUESTS + + +def _index_score_requests(spec: Mapping[str, Any]) -> Iterable[Key]: + """A cell certifies its request count when it drives the TP8 index head (1 per rank) with that count's schedule.""" + requests = len(spec["seq_lens"]) + if spec.get("heads") != 1 or requests not in INDEX_SCORE_SPLIT_K: + return () + if (spec["split_k"], spec["early"]) != index_score_schedule(requests): + return () + return (requests,) + + +def _long_select_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The key widths a cell's rows reach: every width up to its blocks' when rows of mixed lengths are drawn.""" + blocks = spec["num_blocks"] + widest = 2 if blocks <= 2048 else 4 if blocks <= 4096 else 9 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_KEY_WIDTHS if w <= widest) + return (widest,) + + +def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: + """The 384-thread select's key widths a cell's rows reach (every width up to its blocks' for mixed rows).""" + blocks = spec["num_blocks"] + widest = 4 if blocks <= 4 * 384 else 11 if blocks <= 11 * 384 else 24 + if spec["valid"] == "mixed": + return tuple(w for w in LONG_SELECT_DSL_KEY_WIDTHS if w <= widest) + return (widest,) + + +EMBED_FIRSTNORM = RowRange( + "embed_firstnorm", + _rows(1, EMBED_FIRSTNORM_MAX_TOKENS), + m3_embedding_firstnorm, + _spec_rows("rows"), +) +# The identity-selection Triton producer + merge (every request within 16 blocks): the per-row producer, and from +# SPARSE_DECODE_MQ_MIN_TOKENS rows the multi-query producer (one KV head), certified at the TP8 heads on SM107. +SPARSE_DECODE_PRODUCER = RowRange( + _sparse_decode_kernel.PRODUCER_PATH, + _sparse_decode_kernel.PRODUCER_KEYS, + m3_sparse_decode, + _sparse_decode_producer, +) +# The TP8 long-context sparse decode at 8 query heads over one KV head (ops/m3_attn_lc8): attention, split-K merge +# and the MXFP8 o_proj hand-off in one launch, for every decode step with a selected table; "prewait" up to +# m3_attn_lc8.PREWAIT_MAX_ROWS query rows, "plain" above. Exact-class (in the float64 reference's band; run-to-run +# deterministic). An empty range routes every step back to the Triton producer + merge. +SPARSE_DECODE_LC8 = RowRange( + _sparse_decode_kernel.LC8_PATH, + _sparse_decode_kernel.LC8_KEYS, + m3_sparse_decode_lc8, + lambda spec: ("prewait" if spec["prewait"] else "plain",), +) +# Every token count of a sparse layer on an NVFP4 KV cache at the TP8 heads (prefill and decode): the producer takes the +# geometry, not a row range. Kill switch: frozenset() routes every sparse layer back to the native producer. +PRODUCER = RowRange("m3_producer", frozenset({"nvfp4_tp8"}), m3_producer, _const("nvfp4_tp8")) +INDEX_SCORE = RowRange( + "m3_index_score", + _rows(1, INDEX_SCORE_MAX_REQUESTS), + m3_index_score, + _index_score_requests, +) +LONG_SELECT = RowRange( + "long_select", + frozenset(LONG_SELECT_KEY_WIDTHS) if LONG_SELECT_MAX_BLOCKS else frozenset(), + m3_select_long, + _long_select_widths, +) +LONG_SELECT_DSL = RowRange( + "long_select_dsl", + frozenset(LONG_SELECT_DSL_KEY_WIDTHS) + if LONG_SELECT_MAX_BLOCKS and SELECT_LONG_DSL + else frozenset(), + m3_select_long_dsl, + _long_select_dsl_widths, +) +MC_ALLREDUCE = RowRange( + "mc_allreduce_norm", + _rows(1, DECODE_MAX_TOKENS) if MC_AR_TP2 else frozenset(), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=0), +) +ATTN_PUSH = RowRange( + "attn_push", + _rows(1, ATTN_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +DENSE_PUSH = RowRange( + "dense_push", + _rows(1, DENSE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +# The router of a pure generation step of 9..16 rows on the generic layer (above the MoE engine's 8; the TP2 GEN batch 4 +# with 3 drafts): the TF32 tensor-core router (minimax_m3_router_tc, 16-CTA clusters, one reduction order), the +# arithmetic the built-in gate takes on SM107 at these rows (2.9-3.5 us/call against ~9.8 us for the cuBLAS TF32 +# split-K router; closer to the fp32 product: max |dlogit| 1.7e-6 vs 6.7e-5). Devices that cannot co-schedule 16-CTA +# clusters keep the fp32 router. Context steps never take it. +ROUTER_TC_GEN_MIN_ROWS = 9 # capacity: fixed: past the MoE engine's 8 rows +ROUTER_TC_GEN_MAX_ROWS = 16 # capacity: fixed: the shared gate's router_tc range +ROUTER_TC_GEN = RowRange( + "m3_router_tc", + _rows(ROUTER_TC_GEN_MIN_ROWS, ROUTER_TC_GEN_MAX_ROWS), + m3_router_tc, + _spec_rows("m"), # every cell checks the logits, with or without the NVFP4 activation +) +MOE_PUSH = RowRange( + "moe_push", + _rows(1, MOE_PUSH_MAX_TOKENS), + m3_mc_allreduce_norm, + _spec_rows("m", prepushed=AR_PUSH_MODE), +) +MOE_ENGINE = RowRange( + "m3_moe_engine", _rows(1, MOE_ENGINE_MAX_TOKENS), m3_moe_engine, _spec_rows("tokens") +) +# The GEMM paths are keyed (projection (N, K), rows): each projection runs its own split-K / ring / MMA +# configuration, so a row count is certified per projection. +MX_GEMV = RowRange( + "m3_mx_gemv", + frozenset((nk, m) for nk, rows in MX_GEMV_ROWS.items() for m in rows), + m3_mx_gemv, + lambda spec: ((spec["nk"], spec["m"]),), +) +MX_GEMV_BF16 = RowRange( + "m3_mx_gemv_bf16", + frozenset((nk, m) for nk, hi in MX_GEMV_BF16_MAX_TOKENS.items() for m in _rows(1, hi)), + m3_mx_gemv_bf16, + lambda spec: ((spec["nk"], spec["m"]),), +) +# (draft length, requests) of a greedy verify step. +SHARD_ACCEPT = RowRange( + "m3_shard_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SHARD_ACCEPT_MAX_REQUESTS + 1)), + m3_shard_accept, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) + +# Rows of a fused-draft-loop all-reduce boundary (drafter hidden 6144) on the multicast all-reduce. +DRAFT_MC_AR = RowRange( + "m3_draft_mc_allreduce_norm", + _rows(1, DRAFT_MC_AR_MAX_ROWS), + m3_mc_allreduce_norm, + lambda spec: (spec["m"],) if spec.get("drafter") else (), +) + +# (draft length, requests) of a sampled verify step. +SAMPLED_ACCEPT = RowRange( + "m3_sampled_accept", + frozenset((SHARD_ACCEPT_DRAFT_LEN, n) for n in range(1, SAMPLED_VERIFY_MAX_REQUESTS + 1)), + m3_sampled_accept, + lambda spec: ((SHARD_ACCEPT_DRAFT_LEN, spec["nreq"]),), +) + + +# (query rows per request, requests) of a decode step whose dense-layer attention runs on PrimTS split-KV decode. +FULLCTX_DECODE = RowRange( + "m3_fullctx_decode", + frozenset(FULLCTX_DECODE_RECIPES) if FULLCTX_DECODE_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + +# (query rows per request, requests) of a generation-only drafter pass that runs on the full-context decode. +DRAFTER_FULLCTX = RowRange( + "m3_drafter_fullctx", + frozenset(FULLCTX_DECODE_RECIPES) if DRAFTER_FULLCTX_ENABLED else frozenset(), + m3_fullctx_decode, + lambda spec: ((spec["qlen"], len(spec["kv_lens"])),), +) + + +def _verify_rows(max_rows: int) -> FrozenSet[int]: + """The verify rows of 1..8 requests at the shard-accept draft length, up to ``max_rows``.""" + return frozenset( + n * (SHARD_ACCEPT_DRAFT_LEN + 1) + for n in range(1, 9) + if n * (SHARD_ACCEPT_DRAFT_LEN + 1) <= max_rows + ) + + +# Requests of a fused-verify step whose first drafter inputs are the fused ones; certified by the chunk-norm cells +# at those requests' verify rows (the norm is the only kernel of its own; the gather rows and the fc are in-tree). +FUSED_FIRST_INPUTS = RowRange( + "m3_fused_first_inputs", + _rows(1, FUSED_FIRST_INPUTS_MAX_REQUESTS), + m3_chunk_norm, + lambda spec: ( + (spec["m"] // (SHARD_ACCEPT_DRAFT_LEN + 1),) + if spec["m"] % (SHARD_ACCEPT_DRAFT_LEN + 1) == 0 + else () + ), +) +# Verify rows of a fused-verify step whose Eagle3 captures apply the fc_norm. +CAPTURE_NORM = RowRange( + "m3_capture_norm", _verify_rows(CAPTURE_NORM_MAX_ROWS), m3_capture_norm, _spec_rows("m") +) +# Verify rows of a first draft pass whose fc_norm runs as the chunk norm. +CHUNK_NORM = RowRange( + "m3_chunk_norm", _verify_rows(CHUNK_NORM_MAX_ROWS), m3_chunk_norm, _spec_rows("m") +) + +# ((N, K), rows) of a drafter BF16 linear on its low-M kernel. +DRAFT_LOWM = RowRange( + "m3_draft_lowm", + frozenset((nk, m) for nk, rows in DRAFT_LOWM_ROWS.items() for m in rows), + m3_draft_lowm, + lambda spec: ((spec["nk"], spec["m"]),), +) +# Verify rows of a step whose target logits come from the die-placed head (in-tree GEMM halves per die). + +# (draft length, requests) of a fused draft loop, and the rows of its passes' draft input: the first pass's +# requests x (draft length + 1), one per request on the later passes. +DRAFT_STEP = RowRange( + "m3_draft_step", + frozenset((DRAFT_STEP_DRAFT_LEN, n) for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1)), + m3_draft_step, + lambda spec: ((spec["draft_len"], spec["nreq"]),), +) +DRAFT_INPUT = RowRange( + "m3_draft_input", + frozenset( + rows + for n in range(1, DRAFT_STEP_MAX_REQUESTS + 1) + for rows in (n, n * (DRAFT_STEP_DRAFT_LEN + 1)) + ), + m3_draft_input, + _spec_rows("m"), +) + +# The dense layers' down projection takes the gate_up GEMV's output and applies the swigluoai activation in its +# quantize prologue (ops/mx_gemv x_swiglu, bit for bit the Triton silu_and_mul kernel), instead of a separate +# silu_and_mul launch between the two GEMVs; rows the bf16 GEMV serves. Kill constant: False runs the activation +# kernel as before. +# OFF: bitwise but SLOWER (Host TP2 same-job pair 782860 / 782935: +60 / +65 us/step fixed, both orders; every GEMV CTA +# re-activates gate | up in its quantize prologue, redundantly across the N tiles). Same lesson as the drafter fold. +DENSE_SWIGLU_FOLD = False +DENSE_SWIGLU_CPP_PATH = "m3_dense_swiglu" + +FAST_PATHS: Tuple[RowRange, ...] = ( + EMBED_FIRSTNORM, + FULLCTX_DECODE, + DRAFTER_FULLCTX, + SPARSE_DECODE_PRODUCER, + SPARSE_DECODE_LC8, + PRODUCER, + INDEX_SCORE, + LONG_SELECT, + LONG_SELECT_DSL, + MC_ALLREDUCE, + ATTN_PUSH, + DENSE_PUSH, + MOE_PUSH, + MOE_ENGINE, + ROUTER_TC_GEN, + MX_GEMV, + MX_GEMV_BF16, + SHARD_ACCEPT, + SAMPLED_ACCEPT, + DRAFT_MC_AR, + DRAFT_STEP, + DRAFT_INPUT, + DRAFT_LOWM, + FUSED_FIRST_INPUTS, + CAPTURE_NORM, + CHUNK_NORM, +) + +# -------------------------------------------------------------------------------------------------------------------- +# Selection: pure functions of the step shape +# -------------------------------------------------------------------------------------------------------------------- + + +def plan_for(num_tokens: int) -> DecodePlan: + for plan in DECODE_PLANS: + if num_tokens <= plan.max_tokens: + return plan + raise AssertionError(f"no decode plan covers {num_tokens} tokens") + + +def is_decode_step(num_tokens: int) -> bool: + """Whether a pure generation step of ``num_tokens`` rows goes to the decode target.""" + return num_tokens <= DECODE_MAX_TOKENS + + +def moe_path(num_tokens: int) -> Optional[str]: + """The MoE kernel of a decode step: the engine, or None (generic).""" + if MOE_ENGINE.covers(num_tokens): + return MOE_ENGINE.path + return None + + +MX_GEMV_PATH = MX_GEMV.path +MX_GEMV_BF16_PATH = MX_GEMV_BF16.path +# The dense layers' swigluoai activation as this target's C++ op (ops/gemm/m3_dense_swiglu, +# minimaxM3Sm107Tp8/swigluOai.cu, LongAttn c267035d2f: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, +# beta 1), PDL-launched with an early release so the down GEMV stages its weights during it; the down GEMV reads the +# activation only after its own grid wait. Indexer kdev17 783285 (TP2 chain): -0.69 us per dense layer vs the Triton +# kernel; LongAttn gate a 783511 (oracle 23/23). False runs the Triton kernel. +DENSE_SWIGLU_CPP_PATH = "m3_dense_swiglu" +# FlashInfer's CuTe DSL MXFP8 GEMM on the producer's MXFP8 activation (rows the GEMV does not serve). +PREQUANT_GEMM_PATH = "flashinfer_mxfp8_gemm" +# The MXFP8 quantize (tuned backend) and FlashInfer's CuTe DSL GEMM: every row the decode kernels do not serve. +GENERIC_GEMM_PATH = "mxfp8_gemm" + + +def linear_path(n: int, k: int, num_tokens: int, prequant: bool, low_m_gemm: bool) -> str: + """The decode-graph MXFP8 GEMM of an [N, K] projection on ``num_tokens`` rows. + + ``prequant``: the producer published the activation in MXFP8. ``low_m_gemm`` is accepted for the call sites' + symmetry: no low-M GEMM configuration is certified at TP2, so it selects nothing. + """ + if prequant: + if num_tokens in MX_GEMV_ROWS.get((n, k), ()): + return MX_GEMV_PATH + return PREQUANT_GEMM_PATH + if num_tokens <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0): + return MX_GEMV_BF16_PATH + return GENERIC_GEMM_PATH + + +# -------------------------------------------------------------------------------------------------------------------- +# Fallback policy +# -------------------------------------------------------------------------------------------------------------------- + + +class FastPathUnavailable(RuntimeError): + """A covered decode shape could not take its fast path under ``modeling_v2='require'``.""" + + +_lock = threading.Lock() +_require = [False] +_logged: set = set() + + +def set_require(require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + _require[0] = bool(require) + + +def require_mode() -> bool: + return _require[0] + + +def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: + key = (path, reason) + with _lock: + if key in _logged: + return + _logged.add(key) + level(message) + + +def missed(path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if _require[0]: + raise FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " + f"{reason}" + ) + _log_once( + logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def taken(path: str, rows: Optional[int], what: str = "") -> None: + """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line + every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" + at = "" if rows is None else f" at {rows} rows" + detail = f" ({what})" if what else "" + _log_once( + logger.info, + path, + f"taken: {rows}" + (f": {what}" if what else ""), + f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", + ) + + +def long_select_covers(num_blocks: int) -> bool: + """Whether ``LONG_SELECT`` takes rows of ``num_blocks`` KV blocks (the engine's per-row block width).""" + return LONG_SELECT_MIN_BLOCKS <= num_blocks <= LONG_SELECT_MAX_BLOCKS + + +def long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: + """The (PDL launch, early trigger, trigger phase) keys of the 384-thread select a decode step can launch, compiled + before warmup; none when ``LONG_SELECT_DSL`` covers nothing.""" + if not LONG_SELECT_DSL.keys: + return () + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + pdl = get_env_enable_pdl() + keys = [(pdl, False, SELECT_LONG_TRIGGER_PHASE)] + if SELECT_EARLY_MAX_ROWS: + keys.append((pdl, True, 0)) + return tuple(keys) + + +def long_select_fallback(num_blocks: int, rows: int) -> None: + """The in-tree selector is about to run instead of the long-row select: a miss when ``LONG_SELECT`` covers + ``num_blocks`` (the kernel declined operands its cells certify; raises under require), uncovered otherwise (rows of + at most 128 blocks, or wider than ``LONG_SELECT_MAX_BLOCKS``: the in-tree selector serves them in every mode).""" + if long_select_covers(num_blocks): + missed( + LONG_SELECT.path, + rows, + f"the long-row select declined {num_blocks}-block scores it covers", + ) + else: + span = f"{LONG_SELECT_MIN_BLOCKS}..{LONG_SELECT_MAX_BLOCKS}" + uncovered( + LONG_SELECT.path, f"{num_blocks} KV blocks, outside {span}: the built-in selector" + ) + + +def uncovered(path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + + +def reset_logged() -> None: + """Forget which (path, reason) pairs were logged.""" + with _lock: + _logged.clear() + + +# Kernel modules report from inside their calls through ops/m3_report: the same policy. +m3_report.install( + on_taken=taken, on_missed=missed, on_uncovered=uncovered, on_capture_missed=capture_missed +) + + +# -------------------------------------------------------------------------------------------------------------------- +# Target constants +# -------------------------------------------------------------------------------------------------------------------- + +# The generic boundary takes the low-latency multicast all-reduce + residual + RMSNorm up to this many tokens; with +# the shared + routed add folded in, up to its own maximum: the decode boundaries' certified rows. +MC_AR_MAX_TOKENS = GENERIC_MC_AR_MAX_TOKENS +MC_AR_FOLDED_MAX_TOKENS = GENERIC_MC_AR_FOLDED_MAX_TOKENS +# The shared expert folds into the routed engine from one token (dedicated shared-expert CTAs). +MOE_FOLD_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# SMs left to the concurrent kernels while a persistent routed engine runs. +MOE_ENGINE_RESERVED_SMS = 24 # capacity: fixed: an SM split, not a size +# Id of the target's final hidden state in Eagle3's ``layers_to_capture``. +FINAL_HIDDEN_LAYER_ID = -1 # capacity: fixed: an Eagle3 layer id, not a size +# The trtllm-gen MoE runner's routing kind for MiniMax-M3 (sigmoid + correction bias, renormalized, scaled) and its +# gated activation (SwiGLU with the per-expert alpha / beta / clamp operands). +ROUTING_MINIMAX2 = int(RoutingMethodType.MiniMax2) +ACT_SWIGLU = int(ActType_TrtllmGen.SwiGlu) +INDEX_SCORE_PATH = INDEX_SCORE.path + +# -------------------------------------------------------------------------------------------------------------------- +# Streaming: also emit a response on decoding iterations 2..FIRST_CONTENT_EMIT_ITERS (RuntimeHooks), so a request's +# first content leaves one step after its first token. 1: the stock schedule. +FIRST_CONTENT_EMIT_ITERS = 2 + +# The runtime-value dump (TRTLLM_MODELING_V2_VALUE_DUMP=1): the metadata of every fast-path and entry decision a step +# takes, recorded on eager steps and captures (replay runs no Python), printed once at exit. +# -------------------------------------------------------------------------------------------------------------------- + +_VALUE_DUMP = os.environ.get("TRTLLM_MODELING_V2_VALUE_DUMP") == "1" +_VALUES: set = set() + + +def _record(*meta: Hashable) -> None: + if _VALUE_DUMP: + _VALUES.add(meta) + + +def _print_values() -> None: + if _VALUES: + rank = os.environ.get("OMPI_COMM_WORLD_RANK", os.environ.get("RANK", "0")) + for meta in sorted(_VALUES, key=repr): + print(f"[modeling_v2 value dump] rank {rank}: {meta}", flush=True) + + +if _VALUE_DUMP: + atexit.register(_print_values) + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing + sees only its KV cache algorithm; this check at construction refuses every other runtime. + + The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index + projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. + """ + sparse = model_config.sparse_attention_config + if ( + sparse is None + or getattr(sparse, "algorithm", None) != "minimax_m3" + or sparse.implementation != "msa" + ): + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs +# -------------------------------------------------------------------------------------------------------------------- + + +def _mxfp8_linear( + core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + _record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +_PINNED_RUNNERS: dict = {} + + +def _pinned_runner(cls, dtype): + runner = _PINNED_RUNNERS.get((cls, dtype)) + if runner is None: + runner = _PINNED_RUNNERS[(cls, dtype)] = cls(dtype) + return runner + + +def _pinned_mxfp8_gemm(act, act_scale, weight, weight_scale, out_dtype): + """The decode-graph MXFP8 GEMM at the target's fixed tactic (``MXFP8_DECODE_GEMM_TACTIC``), no AutoTuner: the + same choice on every rank and every start.""" + _record("pinned_mxfp8_gemm", tuple(weight.shape), act.shape[0], MXFP8_DECODE_GEMM_TACTIC) + runner = _pinned_runner(FlashInferMXFP8GemmRunner, out_dtype) + return runner.forward([act, act_scale, weight, weight_scale], tactic=MXFP8_DECODE_GEMM_TACTIC) + + +def _pinned_mxfp8_quantize(x): + """The decode-graph MXFP8 activation quantize at the target's fixed tactic (``MXFP8_DECODE_QUANT_TACTIC``).""" + _record("pinned_mxfp8_quantize", x.shape[0], MXFP8_DECODE_QUANT_TACTIC) + return _pinned_runner(MXFP8QuantizeRunner, x.dtype).forward( + [x], tactic=MXFP8_DECODE_QUANT_TACTIC + ) + + +_MX_TAC_ENGAGED: set = set() + + +def _note_mx_tac(kind: str, n: int, k: int, rows: int) -> None: + """MX-TAC's engagement line, once per (activation kind, projection, rows).""" + key = (kind, n, k, rows) + if key not in _MX_TAC_ENGAGED: + _MX_TAC_ENGAGED.add(key) + logger.info( + f"[M3 sm107 tp8] MX-TAC engaged: pinned MXFP8 decision (quantize tactic " + f"{MXFP8_DECODE_QUANT_TACTIC}, GEMM tactic {MXFP8_DECODE_GEMM_TACTIC}) " + f"for the [{n}, {k}] projection at {rows} rows, {kind} activation" + ) + + +def _gemv_variants(n: int, k: int, push: Optional[str], world: int): + """Every (rows, bf16 activation, push (world, bucket tokens, mode)) call the decode dispatch routes to an [N, K] + projection's GEMV: an MXFP8 activation on the rows ``MX_GEMV_ROWS`` gives it, a bf16 one up to + ``MX_GEMV_BF16_MAX_TOKENS``; each pushing when the projection's boundary push covers the rows. ``push``: "attn" + (the o_proj into the attention boundary), "dense" (the dense MLP's down projection into the post-MLP boundary) + or "moe" (the shared expert's down projection into the MoE boundary).""" + push_rows = {"attn": ATTN_PUSH.keys, "dense": DENSE_PUSH.keys, "moe": MOE_PUSH.keys}.get( + push, frozenset() + ) + variants = [] + for rows in range(1, DECODE_MAX_TOKENS + 1): + for x_bf16, served in ( + (False, rows in MX_GEMV_ROWS.get((n, k), ())), + (True, rows <= MX_GEMV_BF16_MAX_TOKENS.get((n, k), 0)), + ): + if not served: + continue + variants.append((rows, x_bf16, None)) + if rows in push_rows: + bucket = next(b for b in m3_mc_allreduce.BUCKETS if rows <= b) + variants.append((rows, x_bf16, (world, bucket, AR_PUSH_MODE))) + return variants + + +def _uncompiled_note() -> str: + key = mx_gemv.last_uncompiled + return "" if key is None else f" (its kernel {key} was not compiled before warmup)" + + +class _DeferredGemv: + """#117 (``QKV_PRODUCER_FOLD``): a sparse layer's QKV + index MXFP8 GEMV, recorded instead of launched. + + ``_linear(..., defer=True)`` returns one on the decode GEMV path of the 12-head packed projection at <= 8 rows with + an MXFP8 hand-off. The producer call site either runs the GEMV with the producer in its epilogue + (``_qkv_producer_fold``) or calls ``materialize`` (the plain GEMV) before anything reads the packed row; + ``placeholder`` (unwritten, the packed shape) stands for that row until then and is never read.""" + + def __init__(self, act, act_scale, weight, weight_scale, use_pdl): + self.placeholder = torch.empty( + act.shape[0], weight.shape[0], dtype=torch.bfloat16, device=act.device + ) + self.act, self.act_scale = act, act_scale + self.weight, self.weight_scale = weight, weight_scale + self.use_pdl = use_pdl + # Set once the record ran, fused (the producer GEMV) or plain (``materialize``). + self.done = False + + def materialize(self) -> torch.Tensor: + """Launch the plain GEMV this record stands for; returns the packed row the projection would have returned.""" + self.done = True + rows, (n, k) = self.act.shape[0], self.weight.shape + output = m3_mx_gemv( + self.act, self.act_scale, self.weight, self.weight_scale, True, None, self.use_pdl, None + ) + if output is None: + raise RuntimeError( + f"the deferred [{n}, {k}] GEMV declined at {rows} rows" + _uncompiled_note() + ) + taken(MX_GEMV_PATH, rows, f"[{n}, {k}] mxfp8 in") + return output + + +def _linear( + core, + x: torch.Tensor, + wk: str, + sk: str, + low_m_gemm: bool, + push: Optional[tuple] = None, + defer: bool = False, +) -> Tuple[Any, bool]: + """One MXFP8 projection ``x @ w[wk]^T``; returns (output, whether its GEMV pushed into ``push``). + + Outside decode CUDA graphs: the engine-tuned MXFP8 quantize + GEMM. Inside them the projection's (N, K) and the + step's rows pick one path (``linear_path``): + + 1. an activation its producer already quantized to MXFP8 (``m3_prequant.SLOT``, taken here): the tcgen05 decode + GEMV on the rows it serves, else the decode-graph MXFP8 GEMM on that activation; + 2. a bf16 activation on the rows the GEMV serves quantizing in its prologue; + 3. otherwise the decode-graph MXFP8 quantize + GEMM. + + MX-TAC: on a decode graph whose MXFP8 projections the engine tunes (``tune_decode_graph_backends``), the GEMM + (and, on a bf16 activation, the quantize) the GEMVs do not serve take one fixed decision + (``MXFP8_DECODE_QUANT_TACTIC`` / ``MXFP8_DECODE_GEMM_TACTIC``: FlashInfer's CuTe DSL quantize and CuTe DSL GEMM) + instead of the AutoTuner's per-start choice (whose L2-hot profiling left c2 bimodal at +-6%). On SM107 every TP2 + projection at 1..16 rows runs fastest on this pair, cold and L2-hot (graph us, cold: 4.7..10.3 against 5.0..11.5 + with the TRT-LLM quantize and 9.7..58 with CUTLASS), and all four candidates give the same bits (quantize bytes + and scales, GEMM outputs). A constant is rank-agreed and start-independent by construction. + + A kernel declining rows its path covers is a miss; the chain then continues down the list. ``push`` (an + ``m3_mx_gemv`` ``ar`` tuple) is the boundary all-reduce exchange the GEMV writes its rows into. ``low_m_gemm`` + marks the projections the fused low-M MXFP8 GEMM may take (``LOW_M_GEMM_MAX_TOKENS``: no TP2 configuration is + built, so ``linear_path`` never selects it). ``defer`` (``QKV_PRODUCER_FOLD``, the sparse layers' packed QKV + projection): on path 1's GEMV at <= 8 rows of the 12-head projection, a ``_DeferredGemv`` is returned in place of + the output and nothing is launched. + """ + w, s = core.w[wk], core.w[sk] + x = x.contiguous() + mode = _mx_mode(core.mxfp8) + if not mode.decode_graph: + return _mxfp8_linear(core, x, w, s, mode), False + use_pdl = get_env_enable_pdl() + rows = x.shape[0] + n, k = w.shape + prequant = m3_prequant.SLOT.take(x) + if prequant is None and x.dtype != torch.bfloat16: + path = GENERIC_GEMM_PATH + else: + path = linear_path(n, k, rows, prequant is not None, low_m_gemm) + _record("linear", (n, k), rows, path, push is not None) + if prequant is not None: + act, act_scale = prequant + if defer and path == MX_GEMV_PATH and n == mx_gemv.PROD_N and rows <= 8: + return _DeferredGemv(act, act_scale, w, s, use_pdl), False + if path == MX_GEMV_PATH: + output = m3_mx_gemv(act, act_scale, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] mxfp8 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, + rows, + f"the GEMV declined the [{n}, {k}] projection's MXFP8 activation" + + _uncompiled_note(), + ) + if path == PREQUANT_GEMM_PATH: + uncovered(MX_GEMV_PATH, f"the [{n}, {k}] projection at rows the GEMV does not serve") + if mode.tuned_quantize_gemm: + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("mxfp8", n, k, rows) + return output, False + return m3_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16, mode.autotune), False + if path == MX_GEMV_BF16_PATH: + output = m3_mx_gemv_bf16(x, w, s, True, None, use_pdl, push) + if output is not None: + taken(path, rows, f"[{n}, {k}] bf16 in") + return output, push is not None and mx_gemv.last_pushed + missed( + path, rows, f"the GEMV declined the [{n}, {k}] projection's input" + _uncompiled_note() + ) + if path == GENERIC_GEMM_PATH: + uncovered( + MX_GEMV_BF16_PATH, + f"the [{n}, {k}] projection's bf16 input beyond the decode GEMV's rows", + ) + if mode.tuned_quantize_gemm and x.dtype == torch.bfloat16: + act, act_scale = _pinned_mxfp8_quantize(x) + output = _pinned_mxfp8_gemm(act, act_scale, w, s, torch.bfloat16) + _note_mx_tac("bf16", n, k, rows) + return output, False + return _mxfp8_linear(core, x, w, s, mode), False + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer boundaries: all-reduce + residual + RMSNorm +# -------------------------------------------------------------------------------------------------------------------- + + +def _can_mc(x, folded: bool = False) -> bool: + """Whether the generic boundary takes the multicast all-reduce + residual + RMSNorm: the rows + ``MC_AR_MAX_TOKENS`` / ``MC_AR_FOLDED_MAX_TOKENS`` (with the shared + routed add folded in) give it.""" + return ( + isinstance(x, torch.Tensor) + and x.dim() == 2 + and x.shape[0] <= (MC_AR_FOLDED_MAX_TOKENS if folded else MC_AR_MAX_TOKENS) + and x.dtype == torch.bfloat16 + and x.is_contiguous() + ) + + +def _mc_allreduce_norm(core, x, x2, residual, norm_w, *, quant=False, prepushed=0, capture=None): + """The multicast all-reduce + optional second addend + residual + RMSNorm (``m3_mc_allreduce_norm``), or None + without a workspace or for operands the kernel does not take (the same on every rank). With ``quant`` the normed + rows are also published, MXFP8-quantized, to ``m3_prequant.SLOT`` for the next MXFP8 projection; ``capture`` = + (out, fc_norm weight, eps) also writes the layer's Eagle3 capture norm of the updated residual.""" + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + if workspace is None: + return None + _record( + "m3_mc_allreduce_norm", x.shape[0], x2 is not None, quant, prepushed, capture is not None + ) + return m3_mc_allreduce.m3_mc_allreduce_norm( + x, + residual.contiguous(), + norm_w, + core.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + x2=x2, + quant=quant, + prepushed=prepushed, + capture=capture, + ) + + +def _native_allreduce_norm(core, x, residual, norm_w): + """The built-in fused all-reduce + residual + RMSNorm.""" + _record("m3_allreduce", x.shape[0]) + return m3_allreduce( + x, + core.allreduce, + AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual, + norm_weight=norm_w, + eps=core.rms_eps, + trigger_completion_at_end=False, + ), + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The MoE: router, routed experts (trtllm-gen NVFP4), shared expert +# -------------------------------------------------------------------------------------------------------------------- + + +@dataclass +class _MoeScalars: + """The routed experts' kernel scalars of one layer, derived from the checkpoint's per-tensor scales once the + weights are real (``MiniMaxM3Tp8ModelingV2._derive_moe``).""" + + fc31_input_scale: torch.Tensor + fc2_input_scale: torch.Tensor + fc31_alpha: torch.Tensor + fc2_alpha: torch.Tensor + fc31_scale_c: torch.Tensor + act_alpha: torch.Tensor + act_beta: torch.Tensor + act_clamp: torch.Tensor + + +def _router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The fp32 router logits of MoE layer ``i`` when the MoE runs unfused (the decode steps' engines route + in-kernel): ``F.linear(float(h), w)``, the built-in gate's arithmetic. No router kernel serves TP2 rows + (``ROUTER_GEMV_MAX_TOKENS`` = ``ROUTER_TC_MAX_TOKENS`` = 0), so the experts quantize their bf16 input + themselves.""" + _record("m3_fp32_linear", h.shape[0] if h.dim() == 2 else 0) + return m3_fp32_linear(h, core.w[f"l{i}_router"]) + + +def _decode_router_logits(core, i: int, h: torch.Tensor) -> torch.Tensor: + """The router logits of a pure generation step on the generic layer: ``ROUTER_TC_GEN``'s rows take the TF32 + tensor-core router when the device co-schedules 16-CTA clusters (decided at prebuild); every other step the fp32 + router (``_router_logits``).""" + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + ROUTER_TC_GEN.covers(rows) + and core.router_tc_cluster16 + and h.dtype == torch.bfloat16 + and h.stride(1) == 1 + and (h.stride(0) * 2) % 16 == 0 + ): + taken(ROUTER_TC_GEN.path, rows) + _record("m3_router_tc", rows) + return m3_router_tc(h, core.w[f"l{i}_router"], None, get_env_enable_pdl())[0] + return _router_logits(core, i, h) + + +def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): + """The trtllm-gen NVFP4 MoE runner over this rank's tensor-parallel slice of every expert, routing in-kernel. + + ``x_sf`` None means ``x`` is the bf16 activation, quantized here with the experts' global input scale. Returns + the finalized rows, or the unfinalized ``[gemm2 out, expert scales, expanded->permuted index]``. + """ + moe = core.moe[i] + if x_sf is None: + x, x_sf = fp4_quantize(x, moe.fc31_input_scale, 16, False, False) + rows = x.shape[0] + x_sf = x_sf.view(rows, -1).flatten() + _record("fp4_block_scale_moe_runner", rows, do_finalize) + outs = fp4_block_scale_moe_runner( + router_logits, + core.w[f"l{i}_router_bias"], + x, + x_sf.view(torch.float8_e4m3fn), + core.w[f"l{i}_e_w31"], + core.w[f"l{i}_e_s31"].view(torch.float8_e4m3fn), + None, + moe.act_alpha, + moe.act_beta, + moe.act_clamp, + core.w[f"l{i}_e_w2"], + core.w[f"l{i}_e_s2"].view(torch.float8_e4m3fn), + None, + moe.fc31_scale_c, + moe.fc31_alpha, + moe.fc2_alpha, + core.num_experts, + core.top_k, + None, + None, + core.moe_inter, + 0, + core.num_experts, + core.routed_scaling_factor, + ROUTING_MINIMAX2, + do_finalize, + act_type=ACT_SWIGLU, + topk_weights=None, + topk_ids=None, + output=None, + tune_max_num_tokens=core.max_num_tokens, + use_dp=False, + ) + return outs[0] if do_finalize else outs + + +def _split_chunk(rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +# -------------------------------------------------------------------------------------------------------------------- +# Attention +# -------------------------------------------------------------------------------------------------------------------- + + +def _nvfp4_layer(manager, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _producer_covered(core) -> bool: + """Whether ``PRODUCER`` covers this target's sparse layers: an NVFP4 KV cache at the TP8 heads (8 query, one KV + and one index head). A covered layer that cannot take ``m3_producer`` is a miss.""" + return core.main_kv_is_nvfp4 and (core.heads, core.kv_heads, core.index_heads) == (8, 1, 1) + + +def _producer_layout(kv_data, kv_scale) -> bool: + """Whether a sparse layer's NVFP4 pools are the layout ``m3_producer`` is built for: [pages, 2, 1, 128, 64] with + its scale pool [pages, 2, 1, 128, 8].""" + return ( + kv_data is not None + and kv_scale is not None + and kv_data.dim() == 5 + and tuple(kv_data.shape[1:]) == (2, 1, 128, 64) + and kv_data.stride()[2:] == (128 * 64, 64, 1) + and kv_scale.dim() == 5 + and tuple(kv_scale.shape[1:]) == (2, 1, 128, 8) + and kv_scale.stride()[2:] == (128 * 8, 8, 1) + and kv_scale.shape[0] == kv_data.shape[0] + ) + + +def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or data.shape[0] != scales.shape[0] + or tuple(data.shape[1:]) != (2, num_heads, 128, 64) + or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) + or data.dtype not in (torch.int8, torch.uint8) + or scales.dtype != torch.uint8 + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or cache.stride(1) < num_heads * cache.stride(2) + or cache.stride(0) < 2 * cache.stride(1) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + + +#: The step phases ("decode" / "prefill") whose producer engagement has been logged, once per process. +_PRODUCER_ENGAGED: set = set() + + +def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: + if covered: + missed(PRODUCER.path, rows, reason) + return None + + +def _sparse_qkv_producer( + core, + i: int, + packed, + position_ids, + kv_cache_manager, + msa_idx_k_cache, + msa_out_cache_loc, + deferred: Optional[_DeferredGemv] = None, +): + """The horizontal producer of a sparse layer: Q / index-Q (FP8) out, K / V / index-K written into the paged + caches (NVFP4 or FP8 main cache), for every step type (token-major). An NVFP4 cache at the TP8 heads takes + ``m3_producer`` (``PRODUCER``); a covered layer that cannot is a miss. Returns (q, index_q), or None for caches + neither kernel takes. + + ``kv_cache_manager``, ``msa_idx_k_cache`` (the per-layer index-K cache getter) and ``msa_out_cache_loc`` are the + calling target's ``step_args`` entries. ``deferred``: the QKV GEMV ``_linear`` recorded instead of launching + (``QKV_PRODUCER_FOLD``; ``packed`` is then its unwritten placeholder). The ``m3_producer`` branch runs it with the + producer in its epilogue (``_qkv_producer_fold``); any other branch launches it first (``materialize``). A None + return leaves it unlaunched.""" + covered = _producer_covered(core) + rows = int(packed.shape[0]) + if packed.dtype != torch.bfloat16 or position_ids is None: + return _producer_miss( + covered, + rows, + "the packed QKV dtype, position ids, head dims, norm or RoPE configuration", + ) + manager = kv_cache_manager + if manager is None: + return _producer_miss(covered, rows, "the step has no KV cache manager") + buffers = manager.get_buffers(i, kv_layout="HND") + index_k_cache = msa_idx_k_cache(i) + out_cache_loc = msa_out_cache_loc + layer_uses_nvfp4 = _nvfp4_layer(manager, i) + scale_buffers = ( + manager.get_block_scale_buffers(i, kv_layout="HND") if layer_uses_nvfp4 else None + ) + kv_quant_scale = core.inv_kv_scales[i] if layer_uses_nvfp4 else None + if layer_uses_nvfp4: + supported_main = ( + buffers is not None + and scale_buffers is not None + and _nvfp4_fused_write_layout(buffers, scale_buffers, core.kv_heads) + and buffers.is_cuda + and scale_buffers.is_cuda + and scale_buffers.device == buffers.device + and kv_quant_scale.is_cuda + and kv_quant_scale.dtype == torch.float32 + and kv_quant_scale.numel() >= 3 + and kv_quant_scale.is_contiguous() + ) + else: + supported_main = ( + buffers is not None + and buffers.is_cuda + and buffers.dtype == torch.float8_e4m3fn + and buffers.dim() == 5 + and tuple(buffers.shape[1:]) == (2, core.kv_heads, 128, 128) + and buffers.stride(4) == 1 + and buffers.stride(3) == 128 + and buffers.stride(2) == 128 * 128 + and buffers.stride(1) >= core.kv_heads * buffers.stride(2) + and buffers.stride(0) >= 2 * buffers.stride(1) + and buffers.stride(0) % 4 == 0 + and buffers.stride(1) % 4 == 0 + ) + supported_index = ( + index_k_cache.is_cuda + and index_k_cache.dtype == torch.float8_e4m3fn + and index_k_cache.dim() == 4 + and tuple(index_k_cache.shape[1:]) == (1, 128, 128) + and index_k_cache.stride(3) == 1 + and index_k_cache.stride(2) == 128 + and index_k_cache.stride(1) >= 128 * 128 + and index_k_cache.stride(0) >= index_k_cache.stride(1) + and index_k_cache.stride(0) % 4 == 0 + and index_k_cache.stride(1) % 4 == 0 + and buffers is not None + and index_k_cache.shape[0] == buffers.shape[0] + ) + cos_sin = core.rotary_cos_sin + supported_rope = ( + cos_sin.is_cuda + and cos_sin.dtype == torch.float32 + and cos_sin.is_contiguous() + and cos_sin.dim() == 3 + and tuple(cos_sin.shape[1:]) == (2, 32) + ) + if ( + not supported_main + or not supported_index + or not supported_rope + or out_cache_loc is None + or not out_cache_loc.is_cuda + or out_cache_loc.dtype != torch.int32 + or not out_cache_loc.is_contiguous() + or out_cache_loc.numel() < rows + ): + return _producer_miss( + covered and layer_uses_nvfp4, + rows, + "the layer's NVFP4 / index caches, RoPE table, KV quant scale or cache slots are not the layout the " + "producer is built for", + ) + w = core.w + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + if ( + layer_uses_nvfp4 + and (core.heads, core.kv_heads, core.index_heads) == (8, 1, 1) + and _producer_layout(buffers, scale_buffers) + ): + if deferred is not None: + folded = _qkv_producer_fold( + core, + i, + deferred, + buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + cos_sin, + pos, + ) + if folded is not None: + return folded + packed = deferred.materialize() + _record("m3_producer", rows) + q, index_q = m3_producer( + packed.contiguous(), + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + core.heads, + get_env_enable_pdl(), + PRODUCER_EARLY_TRIGGER, + PRODUCER_PREWAIT_STEP_INPUTS, + ) + phase = "decode" if rows <= DECODE_MAX_TOKENS else "prefill" + if phase not in _PRODUCER_ENGAGED: + _PRODUCER_ENGAGED.add(phase) + logger.info( + f"[M3 sm107 tp8] fast path engaged: m3_producer ({phase}, T={rows}, " + f"early_trigger={PRODUCER_EARLY_TRIGGER}, prewait_step_inputs={PRODUCER_PREWAIT_STEP_INPUTS})" + ) + else: + if layer_uses_nvfp4 and covered: + missed( + PRODUCER.path, + rows, + f"the NVFP4 pool {tuple(buffers.shape)} / scale pool {tuple(scale_buffers.shape)} is not the " + "producer's layout", + ) + if deferred is not None: + packed = deferred.materialize() + _record("m3_qkv_insert", rows, layer_uses_nvfp4) + q, index_q = m3_qkv_insert( + packed.contiguous(), + buffers.view(torch.uint8) if layer_uses_nvfp4 else buffers, + scale_buffers, + index_k_cache, + out_cache_loc[:rows], + kv_quant_scale, + core.heads, + core.kv_heads, + core.index_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + ) + return q.flatten(1), index_q.flatten(1) + + +def _qkv_producer_fold( + core, + i: int, + deferred: _DeferredGemv, + buffers, + scale_buffers, + index_k_cache, + out_cache_loc, + kv_quant_scale, + cos_sin, + pos, +) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """#117: the deferred QKV GEMV with layer ``i``'s producer in its epilogue (``mx_gemv.m3_mx_gemv_producer``), the + same outputs and cache bytes as the GEMV + ``m3_producer``; None (nothing launched) when the kernel declines.""" + rows = int(deferred.act.shape[0]) + device = deferred.act.device + q = torch.empty(rows, core.heads, core.head_dim, dtype=torch.float8_e4m3fn, device=device) + index_q = torch.empty(rows, 1, core.head_dim, dtype=torch.float8_e4m3fn, device=device) + w = core.w + _record("m3_qkv_producer_fold", rows) + ok = mx_gemv.m3_mx_gemv_producer( + deferred.act, + deferred.act_scale, + deferred.weight, + deferred.weight_scale, + q, + index_q, + buffers.view(torch.uint8), + scale_buffers.view(torch.uint8), + index_k_cache, + out_cache_loc, + kv_quant_scale, + core.rms_eps, + w[f"l{i}_q_norm"], + w[f"l{i}_k_norm"], + w[f"l{i}_idx_q_norm"], + w[f"l{i}_idx_k_norm"], + cos_sin, + pos, + use_pdl=deferred.use_pdl, + ) + if not ok: + missed( + QKV_PRODUCER_FOLD_PATH, + rows, + "the producer GEMV declined the step" + _uncompiled_note(), + ) + return None + deferred.done = True + taken(QKV_PRODUCER_FOLD_PATH, rows, "[1536, 6144] mxfp8 in + producer") + return q.flatten(1), index_q.flatten(1) + + +def _qk_norm_rope(core, i: int, qkv: torch.Tensor, position_ids, out_fp8: bool) -> torch.Tensor: + """The per-head Gemma RMSNorm of Q and K and the partial NeoX RoPE of a dense layer's packed [Q | K | V], in one + kernel: in place, or an FP8 copy (Q / K normed and roped, V cast) when ``out_fp8``.""" + qkv = qkv.contiguous() + pos = position_ids.reshape(-1).contiguous().to(torch.int32) + args = ( + qkv, + core.heads, + core.kv_heads, + core.kv_heads, + core.head_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, + pos, + 1.0, # factor: no YARN + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + True, # use_gemma + False, # use_mrope + 0, + 0, + ) + _record("qk_norm_rope", qkv.shape[0], out_fp8) + if out_fp8: + return m3_qk_norm_rope_fp8(*args) + return fused_qk_norm_rope(*args) + + +def _dense_fp8_op_applies(qkv: torch.Tensor, position_ids) -> bool: + """Whether a dense layer's FP8 q/k/v come from this target's op (``minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8``, + else the BF16 fused_qk_norm_rope output cast to FP8).""" + return DENSE_QKV_FP8_VIA_BF16 and position_ids is not None and qkv.dtype == torch.bfloat16 + + +def _dense_qk_norm_rope_fp8( + core, i: int, qkv: torch.Tensor, position_ids, kv_dest: Optional[tuple] = None +) -> torch.Tensor: + """A dense layer's FP8 q/k/v: QK norm + RoPE with Q/K rounded to BF16 and then to E4M3 and V cast, the bits of the + BF16 fused_qk_norm_rope output cast to FP8 (K is stored in the FP8 cache), in one launch of this target's op + (``DENSE_QKV_FP8_VIA_BF16``); outside its supported inputs, the BF16 kernel and the cast. + + With ``kv_dest`` = (the layer's HND pool view, the step's int32 cache slots, live tokens, a [tokens, q heads, 128] + FP8 buffer) the op also stores the live tokens' K and V rows into the pool (the bytes ``write_msa_phase_kv`` + stores, one launch instead of about ten) and every row's Q into the buffer; ``kv_dest`` is passed only where the + op is known to run (``_dense_fp8_op_applies``).""" + if not _dense_fp8_op_applies(qkv, position_ids): + return fp8_e4m3_saturate(_qk_norm_rope(core, i, qkv, position_ids, False)) + _record("m3_dense_qk_norm_rope_fp8", qkv.shape[0], kv_dest is not None) + args = ( + qkv.contiguous(), + core.heads, + core.kv_heads, + core.kv_heads, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_q_norm"], + core.w[f"l{i}_k_norm"], + core.rope_theta, + True, # NeoX RoPE + position_ids.reshape(-1).contiguous().to(torch.int32), + True, # Gemma norms + ) + if kv_dest is not None: + return torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv(*args, *kv_dest) + return torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + + +def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: + """The dense MLP's SwiGLU-OAI activation: this target's C++ op (``m3_dense_swiglu``, minimaxM3Sm107Tp8/ + swigluOai.cu: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, beta 1, PDL-launched with an early + release so the down GEMV stages its weights during it) at the rows its cells certify, in any graph or eager step; + otherwise the Triton kernel (``m3_silu_and_mul``). Indexer kdev17 (TP2 chain): -0.69 us per dense layer vs the + Triton kernel. Every decline is logged once.""" + path = DENSE_SWIGLU_CPP_PATH + rows = gu.shape[0] + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + params = (core.swiglu_alpha, core.swiglu_limit, core.swiglu_beta) + act = None + if rows > max_rows: + uncovered(path, f"{rows} rows, past its certified {max_rows}") + elif params != (DENSE_ALPHA, DENSE_LIMIT, DENSE_BETA): + uncovered(path, f"activation parameters {params}") + else: + act = m3_dense_swiglu( + gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + ) + if act is None: + uncovered(path, "the gate_up layout") + else: + _record("m3_dense_swiglu", rows) + taken(path, rows, "dense swigluoai") + if act is None: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return act + + +def _fp8_index_qk_norm_rope( + core, i: int, idx_qk: torch.Tensor, position_ids, msa_idx_k_cache, msa_out_cache_loc +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " + f"got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + _record("m3_fp8_indexer_qk_norm_rope", rows) + return m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +# -------------------------------------------------------------------------------------------------------------------- +# Layer helpers both targets' bodies call: the Eagle3 capture hook and the norm after a layer +# -------------------------------------------------------------------------------------------------------------------- + + +def _plain_capture(spec_metadata, layer_id: int, hidden_states, residual) -> None: + """The in-tree Eagle3 hidden-state capture.""" + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + + +def _next_norm(core, i: int) -> torch.Tensor: + """The norm the boundary after layer ``i`` applies: the next layer's input norm, the final norm after the last.""" + return core.w[f"l{i + 1}_in_norm"] if i + 1 < core.num_layers else core.w["final_norm"] + + +# -------------------------------------------------------------------------------------------------------------------- +# Kernel prebuilds the core runs before the executor's warmup +# -------------------------------------------------------------------------------------------------------------------- + + +def _in_tree_scorer(): + """The in-tree CuTe DSL index scorer's runner (importing it registers ``cute_dsl_minimax_m3_index_decode_score``); + imported on first use, as the CuTe DSL ops are.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + CuteDSLMiniMaxM3IndexDecodeScoreRunner, + ) + + return CuteDSLMiniMaxM3IndexDecodeScoreRunner + + +def _dummy_scorer_operands(dtype, heads, qlen, dim, page, device): + """One request of two blocks: real-shaped operands that load a compiled scorer module.""" + rows, blocks = 1, 2 + tokens = rows * qlen + idx_q = torch.zeros(tokens, heads, dim, dtype=torch.bfloat16, device=device).to(dtype) + idx_k = torch.zeros(rows * blocks, 1, page, dim, dtype=torch.bfloat16, device=device).to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device=device).view(rows, blocks) + seq_lens = torch.full((rows,), blocks * page, dtype=torch.int32, device=device) + score = torch.empty(heads, blocks, tokens, dtype=torch.float32, device=device) + return idx_q, idx_k.squeeze(1), block_table, score.transpose(1, 2), seq_lens + + +#: Mixed (context + generation) steps run their generation rows on the built-in sparse decode +#: (``MsaDecodeFmha._run_sparse``). Its Triton producer and merge specialize on the split-K factor (16 / 8 / 4 / 2 / 1 +#: by the row count on SM107), on the row count's divisibility and on the 16-byte alignment of the request rows (a +#: mixed step's generation requests follow its context requests, so its block-table and length rows start mid-buffer), +#: so ``_prebuild_builtin_sparse_decode`` launches them per decode row count up to this many rows, both alignments. +BUILTIN_SPARSE_DECODE_PREBUILD_MAX_ROWS = ( # capacity: fixed: the prebuilt built-in decode rows + 64 # capacity: fixed: the mixed steps' generation rows prebuilt +) + + +def _prebuild_builtin_sparse_decode(query_lens, num_heads, num_kv_heads, device) -> list: + """Compile (and load) the built-in sparse decode's Triton producer and merge for every generation row count a mixed + step can have, so no step compiles one after the server is ready. Each launch mirrors the operands + ``MsaDecodeFmha._run_sparse`` passes for an NVFP4 layer (FP8 q, the HND pool's K / V planes, the scale pool's + planes, the head-major top-k view, BF16 output) on dummy values, on the current stream. Returns the row counts.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode, + ) + + page, head_dim, topk = 128, 128, MSA_REQUIRED_TOPK + pages = topk + pool = torch.zeros( + pages, 2, num_kv_heads, page, head_dim // 2, dtype=torch.uint8, device=device + ) + scales = torch.zeros( + pages, 2, num_kv_heads, page, head_dim // 16, dtype=torch.uint8, device=device + ) + global_scales = torch.ones(2, 4, dtype=torch.float32, device=device) + done = [] + for qlen in sorted(set(query_lens)): + for rows in range(qlen, BUILTIN_SPARSE_DECODE_PREBUILD_MAX_ROWS + 1, qlen): + requests = rows // qlen + q = torch.zeros(rows, num_heads, head_dim, dtype=torch.bfloat16, device=device).to( + torch.float8_e4m3fn + ) + indexes = ( + torch.arange(topk, dtype=torch.int32, device=device) + .repeat(rows, num_kv_heads, 1) + .contiguous() + ) + # One leading context request: its rows are skipped, so the generation rows start mid-buffer. + block_rows = torch.arange(topk, dtype=torch.int32, device=device).repeat( + requests + 1, 1 + ) + lengths = torch.full((requests + 1,), topk * page, dtype=torch.int32, device=device) + output = torch.empty(rows, num_heads, head_dim, dtype=torch.bfloat16, device=device) + for first in (0, 1): + minimax_m3_sparse_attn_decode( + q, + pool[:, 0], + pool[:, 1], + indexes.permute(1, 0, 2), + block_rows[first : first + requests], + lengths[first : first + requests], + sm_scale=head_dim**-0.5, + output=output, + decode_query_len=qlen, + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=global_scales[0, :1], + v_global_scale=global_scales[1, :1], + ) + done.append(rows) + return done + + +# -------------------------------------------------------------------------------------------------------------------- +# The core +# -------------------------------------------------------------------------------------------------------------------- + +# Engine states (compiled kernel + workspace) per (device, token count[, shared fold]), shared by every layer: a step +# runs the layers one after another on one stream. +_ENGINE_STATES: Dict[tuple, object] = {} + + +class MiniMaxM3Tp8ModelingV2(DecoderModel): + """The MiniMax-M3 text model of this target: its weights, its runtime resources, and the per-step dispatch.""" + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + _validate_runtime(model_config) + super().__init__(model_config) + cfg = model_config.pretrained_config + mapping = model_config.mapping + self.mapping = mapping + # Whether the 9..16-row generation router takes the TF32 tensor-core kernel (set by prebuild_kernels). + self.router_tc_cluster16 = False + self.vocab_size = cfg.vocab_size + d = _weights.MODEL_WEIGHTS.dims(self) + self.num_layers = d.num_layers + self.hidden = d.hidden + self.heads = d.heads + self.kv_heads = d.kv_heads + self.head_dim = d.head_dim + self.index_heads = d.index_heads + self.index_dim = d.index_dim + self.q_size = d.heads * d.head_dim + self.kv_size = d.kv_heads * d.head_dim + self.moe_inter = d.moe_inter + self.num_experts = d.num_experts + self.top_k = int(cfg.num_experts_per_tok) + self.routed_scaling_factor = float(getattr(cfg, "routed_scaling_factor", 1.0)) + self.rms_eps = float(cfg.rms_norm_eps) + self.swiglu_alpha = float(getattr(cfg, "swiglu_alpha", 1.702)) + self.swiglu_beta = 1.0 # the (up + 1) offset of SwiGLU-OAI + self.swiglu_limit = float(getattr(cfg, "swiglu_limit", 7.0)) + self.max_num_tokens = model_config.max_num_tokens + self.moe_max_num_tokens = model_config.moe_max_num_tokens + self.sparse_layers = frozenset(_sparse_layer_ids(cfg)) + self.moe_layers = frozenset(_moe_layer_ids(cfg)) + mode = model_config.quant_config.quant_mode + self.main_kv_is_fp8 = bool(mode is not None and mode.has_fp8_kv_cache()) + self.main_kv_is_nvfp4 = bool(mode is not None and mode.has_fp4_kv_cache()) + # The sparse layers' QKV and index projections: one five-way packed projection, or the QKV projection and a + # separate BF16 index projection (fuse_qkv_index_projection=False) run on two streams. + self.fused_projection = bool( + getattr(model_config.sparse_attention_config, "fuse_qkv_index_projection", False) + ) + + self.w = _weights.MODEL_WEIGHTS.declare(d) + + # The engine reads these two off the module tree: the MXFP8 GEMM state it tunes, and the opt-in to the + # automatic FlashInfer MXFP8 path on decode graphs. + self.mxfp8 = _Mxfp8GemmState() + self._use_flashinfer_mxfp8_decode_graph_default = True + + # The built-in fused all-reduce + residual + RMSNorm of the layer boundaries the multicast kernel does not take. + self.allreduce = AllReduce( + mapping=mapping, strategy=model_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + + # Streams: the shared / routed expert overlap and, when the MoE chunks its input, the chunking overlap. + self.aux_stream_moe = torch.cuda.Stream() + chunked = self.moe_max_num_tokens < model_config.max_num_tokens * mapping.dp_size + self.aux_stream_moe_chunk = torch.cuda.Stream() if chunked else None + self.moe_events = {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers} + self.moe_chunk_events = { + i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.moe_layers + } + # The unfused sparse layers' index branch overlaps the QKV branch on its own stream. + self.aux_stream_attn = None if self.fused_projection else torch.cuda.Stream() + self.attn_events = ( + {} + if self.fused_projection + else {i: (torch.cuda.Event(), torch.cuda.Event()) for i in self.sparse_layers} + ) + + # Partial NeoX RoPE (rotary_dim of head_dim): the cos/sin table every attention kernel reads. + rope = RopeParams.from_config(cfg) + self.rope_theta = rope.theta + self.rotary_dim = int(rope.dim) + self.rotary_cos_sin = rope.create_rope_const_params(interleave=False)[1].reshape( + rope.max_positions, 2, -1 + ) + + # One MSA attention backend per layer (dense and sparse; the step metadata is the backend's). + quant_config = model_config.get_quant_config() + self.attn: List[MiniMaxM3MsaSparseAttention] = [] + for i in range(self.num_layers): + backend = create_attention( + model_config.attn_backend, + i, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=None, + quant_config=quant_config, + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + q_scaling=1.0, + sparse_params=model_config.sparse_attention_config.to_sparse_params( + pretrained_config=cfg, layer_idx=i + ), + ) + backend.update_quant_config(quant_config) + if not isinstance(backend, MiniMaxM3MsaSparseAttention): + raise ValueError( + f"layer {i}: expected the MiniMax-M3 MSA backend, got {type(backend).__name__}" + ) + if i not in self.sparse_layers: + # The KV split trtllm-gen dense decode runs at, per request count. + backend.dense_decode_max_seq_len_by_batch = DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH + self.attn.append(backend) + + # Built once the weights are real (``post_load_weights``): the KV (Q, K, V) scales, the MoE scalars, the + # decode engines and workspaces, the two targets. + self.kv_scales: Dict[int, torch.Tensor] = {} + self.inv_kv_scales: Dict[int, torch.Tensor] = {} + self.moe: Dict[int, _MoeScalars] = {} + self.mx_alpha: Optional[torch.Tensor] = None + self.moe_engines: Dict[int, Dict[Tuple[int, bool, bool], Optional[object]]] = {} + self.embed_owners: Dict[tuple, DecodeEmbedding] = {} + # Each sparse layer's selected block ids per query row at its last long-context decode + # (``SPARSE_DECODE_LC8_PREFETCH_NEXT_KV``, up to ``N1_MAX_ROWS`` query rows), read by the previous sparse + # layer's decode to prefetch this layer's KV. + self.n1_ids: Dict[int, torch.Tensor] = ( + { + i: torch.full( + (N1_MAX_ROWS, MSA_REQUIRED_TOPK), + -1, + dtype=torch.int32, + device=torch.device("cuda", torch.cuda.current_device()), + ) + for i in sorted(self.sparse_layers) + } + if SPARSE_DECODE_LC8_PREFETCH_NEXT_KV + else {} + ) + self._built_for: Optional[tuple] = None + self._prefill: Optional[PrefillTarget] = None + self._decode: Optional[DecodeTarget] = None + + def sm_scale(self, i: int) -> float: + return (self.head_dim**-0.5) / float(self.attn[i].q_scaling) + + # ---------------------------------------------------------------------------------------------- after the load + + def post_load_weights(self) -> None: + """Derive the kernels' state from the loaded weights, then build both targets and every decode kernel. + + Idempotent per set of weight storages: the loader's post-load walk reaches this module directly and through + the shell, and a second call on the same storages does nothing. + """ + identity = ( + self.w["embed"].data_ptr(), + self.w["l0_in_norm"].data_ptr(), + self.w["final_norm"].data_ptr(), + ) + if self._built_for == identity: + return + device = self.w["embed"].device + derived = nn.Module() + derived.register_buffer( + "mx_alpha", torch.ones(1, dtype=torch.float32, device=device), persistent=False + ) + for i in range(self.num_layers): + # No KV calibration in the checkpoint: unit (Q, K, V) scales and their reciprocals. + derived.register_buffer( + f"l{i}_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + derived.register_buffer( + f"l{i}_inv_kv_scales", + torch.ones(3, dtype=torch.float32, device=device), + persistent=False, + ) + for i in sorted(self.moe_layers): + for name, value in dataclasses.asdict(self._derive_moe(i)).items(): + derived.register_buffer(f"l{i}_{name}", value, persistent=False) + self.derived = derived + self.mx_alpha = derived.mx_alpha + self.kv_scales = {i: getattr(derived, f"l{i}_kv_scales") for i in range(self.num_layers)} + self.inv_kv_scales = { + i: getattr(derived, f"l{i}_inv_kv_scales") for i in range(self.num_layers) + } + self.moe = { + i: _MoeScalars( + **{ + f.name: getattr(derived, f"l{i}_{f.name}") + for f in dataclasses.fields(_MoeScalars) + } + ) + for i in sorted(self.moe_layers) + } + self._build_embed_owners() + self.prebuild_kernels() + self._prefill = PrefillTarget(self) + self._decode = DecodeTarget(self) + self._built_for = identity + + def cache_derived_state(self) -> None: + """The staged loaders' entry to the same build (``post_load_weights``).""" + self.post_load_weights() + + def _derive_moe(self, i: int) -> _MoeScalars: + """The routed experts' kernel scalars from the checkpoint's per-tensor scales (trtllm-gen NVFP4 layout). + + Global input scales: the reciprocal of the largest per-expert input scale (gate and up must agree). Per expert: + ``alpha = 1 / (input_scale * (1 / weight_scale_2))`` for the gate (fc31_alpha), the up half (its GEMM2-input + rescale ``fc31_scale_c = fc2_input_scale * up_alpha``) and the down projection; SwiGLU-OAI's activation + constants per expert, with beta and the clamp in the pre-dequant GEMM output domain (divided by fc31_alpha). + """ + w = self.w + isc1, isc3, isc2 = w[f"l{i}_e_isc1"].data, w[f"l{i}_e_isc3"].data, w[f"l{i}_e_isc2"].data + if not torch.allclose(isc1, isc3): + raise ValueError(f"layer {i}: the routed experts' w1 and w3 input scales differ") + fc31_input_scale = isc1.max().reciprocal() + fc2_input_scale = isc2.max().reciprocal() + ws2_1, ws2_3, ws2_2 = ( + w[f"l{i}_e_ws2_1"].data, + w[f"l{i}_e_ws2_3"].data, + w[f"l{i}_e_ws2_2"].data, + ) + fc31_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_1)) + up_alpha = 1.0 / (fc31_input_scale * (1.0 / ws2_3)) + fc2_alpha = 1.0 / (fc2_input_scale * (1.0 / ws2_2)) + fc31_scale_c = fc2_input_scale * up_alpha + n = self.num_experts + device = fc31_alpha.device + act_alpha = torch.full((n,), self.swiglu_alpha, dtype=torch.float32, device=device) + act_beta = torch.full((n,), self.swiglu_beta, dtype=torch.float32, device=device) + act_clamp = torch.full((n,), self.swiglu_limit, dtype=torch.float32, device=device) + act_beta.div_(fc31_alpha) + act_clamp.div_(fc31_alpha) + return _MoeScalars( + fc31_input_scale=fc31_input_scale, + fc2_input_scale=fc2_input_scale, + fc31_alpha=fc31_alpha, + fc2_alpha=fc2_alpha, + fc31_scale_c=fc31_scale_c, + act_alpha=act_alpha, + act_beta=act_beta, + act_clamp=act_clamp, + ) + + def _build_embed_owners(self) -> None: + """The fused embedding + first RMSNorm's caller-owned buffers, one per (rows, id dtype) of a decode step.""" + table, weight = self.w["embed"], self.w["l0_in_norm"] + self.embed_owners = { + (rows, dtype): DecodeEmbedding(table, weight, self.rms_eps, rows, dtype) + for rows in range(1, EMBED_FIRSTNORM_MAX_TOKENS + 1) + for dtype in (torch.int32, torch.int64) + } + + def projections(self) -> List[Tuple[str, str, Optional[str]]]: + """Every MXFP8 projection in layer order: (weight key, scale key, the boundary its GEMV pushes into).""" + out = [] + for i in range(self.num_layers): + out.append((f"l{i}_qkv_w", f"l{i}_qkv_s", None)) + out.append((f"l{i}_o_w", f"l{i}_o_s", "attn")) + if i in self.moe_layers: + out.append((f"l{i}_sh_gu_w", f"l{i}_sh_gu_s", None)) + out.append((f"l{i}_sh_dn_w", f"l{i}_sh_dn_s", "moe")) + else: + out.append((f"l{i}_mlp_gu_w", f"l{i}_mlp_gu_s", None)) + out.append((f"l{i}_mlp_dn_w", f"l{i}_mlp_dn_s", "dense")) + return out + + def decode_query_lens(self) -> List[int]: + """The query rows per request a decode step can have: 1, and 1 + each draft length speculation runs.""" + lens = {1} + spec = getattr(self.model_config, "spec_config", None) + if spec is not None: + for drafts in ( + getattr(spec, "max_draft_len", None), + getattr(spec, "max_total_draft_tokens", None), + *(getattr(spec, "draft_len_schedule", None) or {}).values(), + ): + if drafts: + lens.add(1 + int(drafts)) + return sorted(lens) + + def prebuild_kernels(self) -> Dict[str, object]: + """Compile (and load) every kernel a decode forward dispatches to that would otherwise compile on its first + call: each projection's decode GEMVs (every row count, activation and boundary push), the MoE engines of + ``MOE_ENGINE``'s token counts, the index scorers at the decode query lengths (this target's and the in-tree + one wider steps take), the 384-thread long-row select, the dense layers' full-context decode plans, and the + built-in sparse decode a mixed step's generation rows run. Runs on the current (main) stream before the + executor's warmup, in the same order on every rank, then synchronizes; after it, GEMV compiles are frozen and + the forward only looks kernels up (a missing one is a miss). Returns the seconds each part took and the keys. + """ + times: Dict[str, object] = {} + # The 9..16-row generation router: decided once, from the device, never inside a CUDA-graph capture. + # The built-in gate's conditions, mirrored exactly: 128 experts, a hidden size in 1024s, 16-CTA clusters. + cfg = self.model_config.pretrained_config + self.router_tc_cluster16 = ( + int(cfg.num_local_experts) == 128 + and int(cfg.hidden_size) % 1024 == 0 + and bool(torch.ops.trtllm.minimax_m3_router_tc_cluster16(int(cfg.hidden_size))) + ) + logger.info( + f"[M3 prebuild] rank {mpi_rank()}: TF32 tensor-core router for {ROUTER_TC_GEN_MIN_ROWS}.." + f"{ROUTER_TC_GEN_MAX_ROWS}-row generation steps: " + + ( + "engaged (16-CTA clusters)" + if self.router_tc_cluster16 + else "not used (no 16-CTA clusters)" + ) + ) + if not self.mxfp8.quant_method.use_cutlass: + raise ValueError("the MiniMax-M3 decode dispatch needs the CUTLASS MXFP8 GEMM") + from .ops import m3_fullctx_decode as _fullctx_decode_kernel + from .ops import m3_select_long_dsl as _select_long_dsl_kernel + from .ops import m3_tokmap_guard + + device = self.w["embed"].device + # The one-request token-map guard of the stock generation preprocess (``PREPROCESS_TOKMAP_FIX``): every path + # that can run it (the drafter's full-context decode or its fallback, any trtllm-gen fallback), so it is + # installed whatever else is enabled. + m3_tokmap_guard.install( + torch.device("cuda", torch.cuda.current_device()), + DECODE_MAX_TOKENS, + lambda rows: taken(PREPROCESS_TOKMAP_PATH, rows, "one-request token map"), + lambda rows, reason: missed(PREPROCESS_TOKMAP_PATH, rows, reason), + ) + mx_gemv.freeze(False) + use_pdl = get_env_enable_pdl() + world = int(self.mapping.tp_size) + start = time.perf_counter() + for wk, sk, push in self.projections(): + weight, scale = self.w[wk], self.w[sk] + mx_gemv.ensure_compiled(weight, scale, use_pdl) + n, k = weight.shape + for rows, x_bf16, ar in _gemv_variants(n, k, push, world): + mx_gemv.precompile(weight, scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + # #117: the 12-head packed projection also takes its MXFP8 rows <= 8 with the producer epilogue. + if ( + QKV_PRODUCER_FOLD + and n == mx_gemv.PROD_N + and not x_bf16 + and ar is None + and rows <= 8 + ): + mx_gemv.precompile(weight, scale, rows, False, use_pdl=use_pdl, prod=True) + torch.cuda.synchronize() + times["gemv_s"] = time.perf_counter() - start + start = time.perf_counter() + for i in sorted(self.moe_layers): + self._prebuild_engines(i) + torch.cuda.synchronize() + times["moe_s"] = time.perf_counter() - start + start = time.perf_counter() + times["scorer_keys"] = self._prebuild_index_scorer(self.decode_query_lens(), device) + torch.cuda.synchronize() + times["scorer_s"] = time.perf_counter() - start + start = time.perf_counter() + times["select_keys"] = [] + if long_select_dsl_keys(): + # The 384-thread long-row select (``SELECT_LONG_DSL``): one kernel per launch key. + times["select_keys"] = _select_long_dsl_kernel.prebuild(long_select_dsl_keys(), device) + torch.cuda.synchronize() + times["select_s"] = time.perf_counter() - start + start = time.perf_counter() + times["fullctx_plans"] = 0 + if FULLCTX_DECODE.keys: + # The dense layers' full-context decode: one compiled PrimTS plan per covered (rows, requests). + times["fullctx_plans"] = _fullctx_decode_kernel.prebuild( + device, + {key: FULLCTX_DECODE_RECIPES[key] for key in FULLCTX_DECODE.keys}, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and _fullctx_decode_kernel.has_own_reducer(), + use_pdl, + ) + torch.cuda.synchronize() + times["fullctx_s"] = time.perf_counter() - start + start = time.perf_counter() + times["builtin_decode_rows"] = _prebuild_builtin_sparse_decode( + self.decode_query_lens(), self.heads, self.kv_heads, device + ) + torch.cuda.synchronize() + times["builtin_decode_s"] = time.perf_counter() - start + mx_gemv.freeze(True) + logger.info( + "[M3 prebuild] decode kernels built before warmup: GEMVs %.1f s, MoE engines %.1f s, index scorer " + "%.1f s (%d variants), long-row select %.1f s (%d variants), full-context decode %.1f s (%d plans), " + "built-in sparse decode %.1f s (%d row counts)" + % ( + times["gemv_s"], + times["moe_s"], + times["scorer_s"], + len(times["scorer_keys"]), + times["select_s"], + len(times["select_keys"]), + times["fullctx_s"], + times["fullctx_plans"], + times["builtin_decode_s"], + len(times["builtin_decode_rows"]), + ) + ) + return times + + # ------------------------------------------------------------------------------------------ the MoE engines + + def _expert_operands(self, i: int): + """The routed experts' operands the engines read (trtllm-gen NVFP4 SwiGLU-OAI layout), and the input scale.""" + w, m = self.w, self.moe[i] + tensors = [ + w[f"l{i}_e_w31"].data, + w[f"l{i}_e_s31"].data, + w[f"l{i}_e_w2"].data, + w[f"l{i}_e_s2"].data, + m.fc31_alpha, + m.fc31_scale_c, + m.act_alpha, + m.act_beta, + m.act_clamp, + m.fc2_alpha, + ] + return tensors, m.fc31_input_scale + + def _shared_weights(self, i: int): + w = self.w + return ( + w[f"l{i}_sh_gu_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_gu_s"].data, + w[f"l{i}_sh_dn_w"].data.view(torch.float8_e4m3fn), + w[f"l{i}_sh_dn_s"].data, + ) + + def engine(self, i: int, num_tokens: int, shared: bool) -> Optional[object]: + """Layer ``i``'s prebuilt routed-expert engine for ``num_tokens`` (1..8), router folded in (and the shared + expert with ``shared``). None after a miss: the build failed or was never made (a forward never builds one).""" + return self._lookup( + i, self.moe_engines[i], (num_tokens, True, shared), MOE_ENGINE.path, num_tokens + ) + + def _lookup(self, i: int, engines: dict, key, path: str, num_tokens: int) -> Optional[object]: + if key not in engines: + what = f"layer {i}'s engine" + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, what) + else: + missed(path, num_tokens, f"{what} was not built before warmup") + return None + engine = engines[key] + if engine is None: + missed(path, num_tokens, "its build failed (logged)") + return engine + + def _prebuild_engines(self, i: int) -> None: + """Build (compile, on first use of each shared state, and warm) every engine a decode step of layer ``i`` can + take: ``MOE_ENGINE``'s token counts with the shared expert folded in. Whether an engine is kept is agreed + across the TP group (one host collective per engine): a build that fails on one rank drops that engine on + every rank, so the ranks never mix an engine (and its boundary push) with the generic MoE.""" + engines = self.moe_engines.setdefault(i, {}) + for num_tokens in sorted(MOE_ENGINE.keys): + key = (num_tokens, True, True) + if key not in engines: + engines[key] = m3_mc_allreduce.tp_agree_built( + self.mapping, + self._try_build( + i, MOE_ENGINE.path, num_tokens, self._build_engine, shared=True + ), + f"[M3 dispatch] {MOE_ENGINE.path} at {num_tokens} rows: layer {i}'s engine", + ) + + def _try_build(self, i: int, path: str, num_tokens: int, build, **kwargs) -> Optional[object]: + try: + return build(i, num_tokens, **kwargs) + except (ImportError, RuntimeError, ValueError) as err: + logger.warning( + f"[M3 dispatch] {path} at {num_tokens} rows: layer {i}'s engine failed to build: {err}" + ) + return None + + def _build_engine(self, i: int, num_tokens: int, shared: bool) -> object: + """The 1..8-token engine (router, routed experts and the shared expert in one launch, at the TP8 width 384). + Rubin TP2 ablation chain (57 layers behind a 58.4 MB PDL producer, row = (P + MoE) - P, both orders): 31.3 / + 50.7 us/layer at 4 / 8 tokens against the built-in MoE's 42.4 / 55.4 (FlashInfer MXFP8 shared expert); bit for + bit the built-in block's shared.add_(routed). When the engine pushes its sum into the boundary all-reduce, + every CTA releases the PDL dependents right after its grid dependency wait instead of after the push + (``MOE_EARLY_AR_RELEASE``), so the all-reduce (which polls the pushed tagged words, not the grid) becomes + resident on the SMs the engine frees: scheduling only, the same bits (Rubin TP2 greedy early column, same-job + A,B + B,A: c1 -15 us).""" + tensors, scale = self._expert_operands(i) + device = tensors[0].device + sms = torch.cuda.get_device_properties(device).multi_processor_count + kwargs = {} + if shared: + kwargs["shared_weights"] = self._shared_weights(i) + state = _ENGINE_STATES.get((str(device), num_tokens, shared)) + if state is None: + state = _ENGINE_STATES[(str(device), num_tokens, shared)] = M3MoeState( + device, + num_tokens=num_tokens, + num_ctas=sms - MOE_ENGINE_RESERVED_SMS, + shared=shared, + early_release=MOE_EARLY_AR_RELEASE, + ) + engine = state.layer( + *tensors, + self.w[f"l{i}_router_bias"].data, + router_weight=self.w[f"l{i}_router"].data, + fc31_input_scale=scale.reshape(1).float(), + **kwargs, + ) + # Compile (first use of the state) and warm. + m3_moe_engine( + torch.zeros(num_tokens, tensors[0].shape[-1] * 2, dtype=torch.bfloat16, device=device), + engine, + ) + return engine + + # ------------------------------------------------------------------------------------------ the index scorers + + def _index_geometries(self) -> list: + """Every sparse layer's (index dtype, index heads per rank, index dim), deduplicated.""" + geoms = [] + for i in sorted(self.sparse_layers): + backend = self.attn[i] + config = backend.m3_config + dtype = torch.float8_e4m3fn if backend.indexer_kv_dtype == "fp8" else torch.bfloat16 + geom = (dtype, int(config.num_index_heads), int(config.sparse_index_dim)) + if geom not in geoms: + geoms.append(geom) + return geoms + + def _scorer_variants(self, query_lens) -> list: + """The scorer keys a decode step can take: every sparse layer's index geometry at each decode query length, + with each covered request count's schedule (``index_score_schedule``), where the kernel serves it.""" + runner = _index_score_kernel.IndexScoreRunner + keys = [] + for dtype, heads, dim in self._index_geometries(): + for qlen in sorted(set(query_lens)): + if not runner.is_supported( + q_dtype=dtype, + num_heads=heads, + head_dim=dim, + page_size=runner.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + for requests in sorted(INDEX_SCORE.keys): + split_k, early = index_score_schedule(requests) + key = (dtype, heads, int(qlen), dim, int(split_k), bool(early)) + if key not in keys: + keys.append(key) + return keys + + def _in_tree_scorer_variants(self, query_lens) -> list: + """The in-tree CuTe DSL scorer's keys (``CuteDSLMiniMaxM3IndexDecodeScoreRunner``: index dtype, index heads, + decode query length, index dim, its 256-CTA split) a decode step outside ``INDEX_SCORE`` can take: every query + length at which a step of more than ``INDEX_SCORE_MAX_REQUESTS`` requests fits in ``DECODE_MAX_TOKENS`` (e.g. + 3..8 requests decoding one row each), where that kernel serves the geometry.""" + in_tree = _in_tree_scorer() + widest = DECODE_MAX_TOKENS // (INDEX_SCORE_MAX_REQUESTS + 1) + keys = [] + for dtype, heads, dim in self._index_geometries(): + for qlen in sorted(set(query_lens)): + if qlen > widest or not in_tree.is_supported( + q_dtype=dtype, + num_heads=heads, + head_dim=dim, + page_size=in_tree.SUPPORTED_PAGE_SIZE, + max_decode_query_len=qlen, + ): + continue + key = (dtype, heads, qlen, dim, in_tree.SPLIT_K) + if key not in keys: + keys.append(key) + return keys + + def _prebuild_index_scorer(self, query_lens, device) -> list: + """Compile every key of ``_scorer_variants`` and ``_in_tree_scorer_variants`` and launch each once on + real-shaped dummy operands (loading its module), on the current stream, so no decode step compiles a scorer. + The kernels' batch, token, page and block counts are symbolic: their compiles need no step metadata. Returns + the keys.""" + runner = _index_score_kernel.IndexScoreRunner + keys = self._scorer_variants(query_lens) + for key in keys: + dtype, heads, qlen, dim, split_k, early = key + runner._compile(dtype, heads, qlen, dim, split_k, early) + q, k, bt, score, seq = _dummy_scorer_operands( + dtype, heads, qlen, dim, runner.SUPPORTED_PAGE_SIZE, device + ) + runner.kernel_cache[key](q, k, bt, score, seq) + in_tree = _in_tree_scorer() + in_tree_keys = self._in_tree_scorer_variants(query_lens) + for key in in_tree_keys: + dtype, heads, qlen, dim, split_k = key + in_tree._compile(dtype, heads, qlen, dim, split_k) + q, k, bt, score, seq = _dummy_scorer_operands( + dtype, heads, qlen, dim, in_tree.SUPPORTED_PAGE_SIZE, device + ) + in_tree.kernel_cache[key](q, k, bt, score, seq) + return keys + in_tree_keys + + # --------------------------------------------------------------------------------------------- the dispatch + + def _select_target(self, attn_metadata: AttentionMetadata): + """Which target runs this step: one host-side integer, fixed per CUDA-graph capture. A step with context rows + (a mixed step included) is prefill's; every pure generation step is decode's, whatever its width.""" + return self._decode if attn_metadata.num_contexts == 0 else self._prefill + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + capture: Callable[..., None] = _plain_capture, + capture_target: Optional[Callable[..., Optional[tuple]]] = None, + **kwargs, + ) -> torch.Tensor: + """``capture``: the Eagle3 hidden-state capture of a layer (the worker's, with its capture norm), and + ``capture_target``: where a decode boundary all-reduce writes a capture layer's capture norm itself + (``CAPTURE_NORM_IN_BOUNDARY``; the worker's, None without one).""" + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if kwargs.get("lora_params") is not None: + raise ValueError("this MiniMax-M3 target serves no LoRA adapters") + return self._select_target(attn_metadata).forward( + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ) + + +# -------------------------------------------------------------------------------------------------------------------- +# The targets +# -------------------------------------------------------------------------------------------------------------------- + + +class PrefillTarget: + """A step with context rows: prefill chunks, and the mixed steps that carry generation rows beside them. Every + layer runs in its generic form; the attention core is the MSA backend's for every row.""" + + def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + Counts and per-step host scalars are read here; device tensors are the metadata's persistent buffers, held + by reference; ``msa_idx_k_cache`` is the metadata's per-layer index-K cache getter. ``metadata`` is the + prepared object itself: the MSA backend's step (``m3_msa_attention``) takes it as its argument. + """ + return dict( + num_contexts=md.num_contexts, + num_ctx_tokens=md.num_ctx_tokens, + num_tokens=md.num_tokens, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + """``capture_target`` is the decode boundaries' (unused: no prefill boundary writes a capture norm).""" + core = self.core + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) prefill CUDA graphs" + ) + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + _record("prefill", int(step["num_tokens"]), int(step["num_contexts"])) + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + hidden, residual = self._layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + return hidden + + def _swiglu_mlp( + self, h: torch.Tensor, prefix: str, low_m_gemm: bool, dense: bool = False + ) -> torch.Tensor: + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI (the dense MLP's on this target's + C++ op, ``_dense_swiglu``), down, no reduction.""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + out, _ = _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm) + return out + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: + the MoE scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: the router logits and the experts quantizing their bf16 input + (chunked).""" + return self._routed_chunks(i, h, _router_logits(self.core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool): + """The MoE block of layer ``i`` without its reduction. Returns ``(out, routed)``: with ``return_parts`` the + shared and routed parts (routed None when the engine folded the shared expert in), else their sum in + ``out``.""" + core = self.core + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, False) + taken(MOE_ENGINE.path, rows) + folded = m3_moe_engine(h.contiguous(), engine, None, None) + return folded, None + routed_out, shared_out = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the + # same collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step): + """Attention of layer ``i`` up to its o_proj (no reduction). Returns (o_proj output, pushed): a prefill step + arms no boundary push, so pushed is always False.""" + core = self.core + if i not in core.sparse_layers: + # Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v for an FP8 cache, on this target's op; else in + # place in bf16), the MSA backend (it writes the layer's cache), o_proj. + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = ( + _dense_qk_norm_rope_fp8(core, i, qkv, position_ids) + if core.main_kv_is_fp8 + else _qk_norm_rope(core, i, qkv, position_ids, False) + ) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step) + packed, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection + and its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and + the FP8 indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the + o_proj. A layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): the MSA + backend (``m3_msa_attention``: the KV write unless a producer wrote the caches, the sparse layers' top-k block + selection, the FMHA over the context rows and any generation rows beside them).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + sparse = i in core.sparse_layers + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + +class DecodeTarget: + """A pure generation step (no context rows), of any width. + + Up to ``DECODE_MAX_TOKENS`` tokens with 1-D input ids (up to 2 requests verifying 1 + 3 rows, the GEN worker's + batch), per layer: attention (o_proj all-reduce deferred, its rows pushed into the boundary all-reduce) -> + attention boundary (all-reduce + residual + post-attention RMSNorm) -> MoE (the 1..8-row engine pushing its sum + into the boundary all-reduce) or dense MLP (its down projection pushing into the post-MLP boundary) -> MoE / dense + boundary (all-reduce + shared/routed add + residual + the next layer's RMSNorm, emitting the next QKV's MXFP8 + activation). An Eagle3 capture layer runs the same decode layer and captures its feed-forward boundary's reduced + residual sum (the boundary writing the capture norm itself when the worker gives it a target). The first layer + runs the generic layer when the fused embedding does not take the step. A wider step (or one carrying embeddings) + runs every layer in its generic form, with this target's attention (the dense layers' full-context decode, the + decode selection and sparse decode where the producer wrote the caches). + """ + + def __init__(self, core: MiniMaxM3Tp8ModelingV2) -> None: + self.core = core + self._moe_push_descs: Dict[tuple, Any] = {} + # Set by a feed-forward boundary that wrote the layer's Eagle3 capture (see _layer's ``capture``). + self._capture_written = False + # Off unless TRTLLM_MODELING_V2_VALIDATE asks for it; read once here rather than per forward. + self._contract_pending = step_contract_enabled() + + def _check_step_contract(self, md) -> None: + """Opt-in first-forward fail-fast (TRTLLM_MODELING_V2_VALIDATE): the metadata fields this target consumes + must exist (they are private trtllm surface). Calling ``step_args`` is the check; once per target instance.""" + if not self._contract_pending: + return + try: + self.step_args(md) + except AttributeError as exc: + raise AssertionError(f"attention metadata surface drifted: {exc}") from exc + self._contract_pending = False + + def step_args(self, md) -> dict: + """The one projection of the MSA metadata this target's forward reads, computed once per step. + + States the two context counts instead of reading them back: this target runs only when ``num_contexts`` is + zero, so reading them would ask the metadata to confirm what the routing already decided. + + Graph classes: counts, the generation span, the identity-selection flag, the worst-case score extent and the + cache length bound are host values fixed per capture (prepare() stages them before the forward); the block + table, lengths, cache slots, identity table and valid-block counts are the metadata's persistent buffers, held + by reference and refreshed in place every step; ``msa_idx_k_cache`` / ``msa_proxy_max_score_view`` are the + metadata's getters of the per-layer index-K cache and of the score scratch view. ``metadata`` is the prepared + object itself: the MSA backend's step (``m3_msa_attention``) and the full-context decode's K/V write + (``write_msa_phase_kv``) take it as their argument. + """ + return dict( + num_contexts=0, + num_ctx_tokens=0, + num_tokens=md.num_tokens, + num_seqs=md.num_seqs, + max_seq_len=md.max_seq_len, + kv_cache_manager=md.kv_cache_manager, + msa_idx_k_cache=md.msa_idx_k_cache, + msa_out_cache_loc=md.msa_out_cache_loc, + msa_decode_span=md.msa_decode_span, + msa_block_table=md.msa_block_table, + msa_seq_lens_cuda=md.msa_seq_lens_cuda, + msa_short_ok=md._msa_short_ok, + msa_short_table=md.msa_short_table, + msa_proxy_max_score_view=md.msa_proxy_max_score_view, + msa_worst_case_max_k_tiles=md.msa_worst_case_max_k_tiles, + msa_n_valid_blocks=md.msa_n_valid_blocks, + metadata=md, + ) + + def forward( + self, + attn_metadata, + input_ids, + position_ids, + inputs_embeds, + spec_metadata, + capture, + capture_target, + ): + core = self.core + self._check_step_contract(attn_metadata) + step = self.step_args(attn_metadata) + # The decode kernels' rows (a per-capture constant): up to DECODE_MAX_TOKENS tokens of 1-D input ids. + kernels = ( + inputs_embeds is None and input_ids.ndim == 1 and is_decode_step(input_ids.numel()) + ) + if kernels: + num_tokens = int(input_ids.numel()) + plan = plan_for(num_tokens) + _record("decode", num_tokens, plan.name) + hidden, residual = self._embed(input_ids, num_tokens) + else: + if is_in_breakable_cuda_graph(): + raise RuntimeError( + "this MiniMax-M3 target does not run under piecewise (breakable) CUDA graphs" + ) + plan = None + _record("decode", int(step["num_tokens"]), "generic") + hidden = ( + m3_embedding(input_ids, core.w["embed"]) if inputs_embeds is None else inputs_embeds + ) + residual = None + for i in range(core.num_layers): + captured = spec_metadata is not None and spec_metadata.is_layer_capture(i) + if kernels and residual is not None and captured: + # The capture layer on the decode layer (its boundaries push); the capture reads the MoE / dense + # boundary's reduced residual sum, as the generic layer's fused boundary capture does. Given a target + # (``CAPTURE_NORM_IN_BOUNDARY``), the boundary all-reduce writes the capture norm itself. + target = ( + capture_target(spec_metadata, i, num_tokens) + if capture_target is not None + else None + ) + self._capture_written = False + hidden, residual = self._layer( + i, plan, position_ids, hidden, residual, step, capture=target + ) + taken("capture_layer_fast", num_tokens, f"layer {i}") + if target is not None and self._capture_written: + taken("capture_norm_in_boundary", num_tokens, f"layer {i}") + capture(spec_metadata, i, residual, None, written=True) + else: + capture(spec_metadata, i, residual, None) + elif not kernels or residual is None or captured: + hidden, residual = self._generic_layer( + i, position_ids, hidden, residual, step, spec_metadata, capture + ) + else: + hidden, residual = self._layer(i, plan, position_ids, hidden, residual, step) + return hidden + + def _embed(self, input_ids: torch.Tensor, num_tokens: int): + """(hidden, residual) of the fused embedding + first RMSNorm; (embedding, None) when it does not run, so the + first layer runs its own input norm. SM107: 1.55..1.65 us against 3.4..6.0 us for the stock embedding + + RMSNorm at 1..8 rows (the gather bitwise, the norm one ulp off on ~5e-6 of the elements: summation order).""" + core = self.core + if not EMBED_FIRSTNORM.covers(num_tokens): + uncovered(EMBED_FIRSTNORM.path, f"more than {EMBED_FIRSTNORM_MAX_TOKENS} tokens") + return m3_embedding(input_ids, core.w["embed"]), None + owner = core.embed_owners.get((num_tokens, input_ids.dtype)) + if owner is None or not input_ids.is_contiguous(): + missed( + EMBED_FIRSTNORM.path, + num_tokens, + f"no embedding owner for {input_ids.dtype} ids (built for int32 / int64 after the load)" + if owner is None + else "the input ids are not contiguous", + ) + return m3_embedding(input_ids, core.w["embed"]), None + _record("m3_embedding_firstnorm", num_tokens, input_ids.dtype) + taken(EMBED_FIRSTNORM.path, num_tokens, str(input_ids.dtype)) + return m3_embedding_firstnorm( + input_ids, owner.table, owner.weight, owner.eps, owner.raw, owner.output, owner.use_pdl + ) + + def _bucket_flags(self, num_tokens: int): + """(bucket workspace, its 8-CTA call counters) of the boundary all-reduce for ``num_tokens``, or None.""" + workspace = m3_mc_allreduce.get_workspace(self.core.mapping, self.core.hidden) + workspace = None if workspace is None else m3_mc_allreduce._bucket(workspace, num_tokens) + flags = None if workspace is None else workspace.flags_for(8) + return None if flags is None else (workspace, flags) + + @staticmethod + def _no_bucket(path: str, num_tokens: int) -> None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, num_tokens, "the multicast all-reduce workspace") + else: + missed(path, num_tokens, "no multicast all-reduce bucket holds the step") + + def _attn_push_target(self, num_tokens: int): + """The push tuple into the attention boundary's all-reduce (the o_proj's; also the dense down projection's into + the post-MLP boundary, the same exchange), or None (after a miss). + + None comes only from the token count or from workspace state created collectively in the same eager call on + every rank, so all ranks agree. (A lone non-pushing rank would still be safe: its all-reduce pushes the same + bf16 rows into the same set A.)""" + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(ATTN_PUSH.path, num_tokens) + return None + workspace, flags = found + return ( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + AR_PUSH_MODE, + ) + + def _moe_push_targets(self, num_tokens: int): + """(engine push descriptor, shared down-projection push tuple) for the MoE boundary's exchange sets, or None + (after a miss). + + This decision must match across ranks: prepushed 3 reads set B of every rank, which a non-pushing rank never + fills. None comes only from the token count or from state created in the same eager call on every rank + (workspace, flags, the descriptor cache), and the engine it is armed for is agreed at build time + (``_prebuild_engines``), so it does.""" + hidden = self.core.hidden + found = self._bucket_flags(num_tokens) + if found is None: + self._no_bucket(MOE_PUSH.path, num_tokens) + return None + workspace, flags = found + key = (workspace.mc.data_ptr(), flags.data_ptr()) + desc = self._moe_push_descs.get(key) + if desc is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(MOE_PUSH.path, num_tokens, "the MoE engine's push descriptor") + return None + desc = self._moe_push_descs[key] = ar_push_desc( + workspace.mc.data_ptr(), + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + flags.device, + ) + set_b = workspace.mc.data_ptr() + 2 * 4 * workspace.max_tokens * workspace.world_size * ( + hidden // 2 + ) + return desc, ( + set_b, + flags.data_ptr(), + workspace.rank, + workspace.world_size, + workspace.max_tokens, + 2, + ) + + def _boundary(self, use_mc: bool, x, residual, norm_w, x2=None, quant=False, capture=None): + """All-reduce + (second addend) + residual + RMSNorm: the multicast kernel, else the built-in all-reduce. + ``capture``: the Eagle3 capture the multicast kernel also writes (``_capture_written`` records it).""" + core = self.core + if use_mc: + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w, quant=quant, capture=capture) + if fused is not None: + if capture is not None: + self._capture_written = True + taken( + MC_ALLREDUCE.path, + x.shape[0], + ("x + x2" if x2 is not None else "x") + (", MXFP8 copy" if quant else ""), + ) + return fused + workspace = m3_mc_allreduce.get_workspace(core.mapping, x.shape[-1]) + missed( + MC_ALLREDUCE.path, + x.shape[0], + m3_mc_allreduce.decline_reason(x, residual, workspace, x2=x2), + ) + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _prepushed( + self, i: int, x, residual, norm_w, quant, prepushed, source, x2=None, capture=None + ): + fused = _mc_allreduce_norm( + self.core, x, x2, residual, norm_w, quant=quant, prepushed=prepushed, capture=capture + ) + if fused is None: + raise RuntimeError(f"layer {i}: rows were pushed into an all-reduce that did not run") + if capture is not None: + self._capture_written = True + taken(MC_ALLREDUCE.path, x.shape[0], f"prepushed {prepushed}, {source}") + return fused + + def _layer(self, i: int, plan: DecodePlan, position_ids, hidden, residual, step, capture=None): + """Layer ``i`` on the decode kernels. ``capture``: (out, fc_norm weight, eps) for the feed-forward boundary to + write the layer's Eagle3 capture norm (``CAPTURE_NORM_IN_BOUNDARY``); a boundary that writes it sets + ``_capture_written``.""" + core = self.core + num_tokens = hidden.shape[0] + # Padded rows: the boundary all-reduce polls and re-arms rows t < M = residual.shape[0], CUDA-graph padding + # rows included. Every producer pushes all rows of the same tensor: the o_proj / shared down / dense down GEMV + # t < m = x.shape[0] (its row count comes from the host shape, not a live count), the MoE engine every row of + # the num_tokens it was built for (= num_tokens here). The all-reduce checks x.sizes() == residual.sizes(), so + # no padded row is left with a stale flag. + # The GEMV dispatch (and with it every push into a boundary all-reduce) is a decode-graph one. + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + if not decode_graph: + uncovered( + "decode_graph", + "an eager decode step: the decode-graph GEMV dispatch and the boundary pushes are off", + ) + push = plan.attn_mc_ar and decode_graph and ATTN_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if push else None + hidden, attn_pushed = self._attention(i, position_ids, hidden, step, armed) + post_norm = core.w[f"l{i}_post_norm"] + if armed is not None and attn_pushed: + hidden, residual = self._prepushed( + i, hidden, residual, post_norm, False, AR_PUSH_MODE, "o_proj push" + ) + else: + if armed is not None: + missed( + ATTN_PUSH.path, + num_tokens, + "the o_proj did not push its rows into the armed boundary all-reduce", + ) + hidden, residual = self._boundary(plan.attn_mc_ar, hidden, residual, post_norm) + + norm_w = _next_norm(core, i) + if i not in core.moe_layers: + # A dense layer: its MLP's down projection pushes into the post-MLP boundary's all-reduce, as the o_proj + # does into the attention one (the same exchange; the gate_up projection never pushes). SM107 TP2: the + # two non-capture dense layers' post-MLP hops end 5.5-5.8 us after the down GEMV unpushed vs 2.4-2.5 us + # pushed; bitwise equal to the unpushed boundary (the same bf16 words). + dense_push = plan.attn_mc_ar and decode_graph and DENSE_PUSH.covers(num_tokens) + armed = self._attn_push_target(num_tokens) if dense_push else None + out, pushed = self._swiglu_mlp(hidden, f"l{i}_mlp", True, armed, dense=True) + if armed is not None and pushed: + return self._prepushed( + i, + out, + residual, + norm_w, + False, + AR_PUSH_MODE, + "dense down push", + capture=capture, + ) + if armed is not None: + missed( + DENSE_PUSH.path, + num_tokens, + "the dense MLP's down projection did not push its rows into the armed boundary all-reduce", + ) + return self._boundary(plan.attn_mc_ar, out, residual, norm_w, capture=capture) + + quant = plan.qkv_prequant and i + 1 < core.num_layers + moe_push = None + if moe_path(num_tokens) == MOE_ENGINE.path: + # The engine runs inside the MoE; looked up (or found missing) here first so a push is armed only for an + # engine that will take it. + engine = core.engine(i, num_tokens, shared=True) + if ( + engine is not None + and plan.moe_mc_ar + and decode_graph + and MOE_PUSH.covers(num_tokens) + ): + moe_push = self._moe_push_targets(num_tokens) + shared, routed, engine_pushed, shared_pushed = self._moe( + i, hidden, return_parts=True, push=moe_push + ) + if moe_push is not None: + if routed is None: + # The engine folded the shared expert in and pushed routed + shared into set A. + if engine_pushed: + return self._prepushed( + i, shared, residual, norm_w, quant, 2, "MoE engine push", capture=capture + ) + else: + if engine_pushed and shared_pushed: + return self._prepushed( + i, + shared, + residual, + norm_w, + quant, + 3, + "MoE engine + shared push", + x2=routed, + capture=capture, + ) + if engine_pushed or shared_pushed: + raise RuntimeError("only one MoE part pushed into the boundary all-reduce") + missed( + MOE_PUSH.path, + num_tokens, + "the MoE did not push its output into the armed boundary all-reduce", + ) + return self._boundary( + plan.moe_mc_ar, shared, residual, norm_w, x2=routed, quant=quant, capture=capture + ) + + def _swiglu_mlp( + self, + h: torch.Tensor, + prefix: str, + low_m_gemm: bool, + push: Optional[tuple] = None, + dense: bool = False, + ): + """A gated MXFP8 MLP (the dense MLP, the shared expert): gate_up, SwiGLU-OAI (the dense MLP's on this target's + C++ op, ``_dense_swiglu``), down, no reduction. Returns (output, whether the down projection's GEMV pushed into + ``push``).""" + core = self.core + gu, _ = _linear(core, h, f"{prefix}_gu_w", f"{prefix}_gu_s", low_m_gemm) + if dense: + act = _dense_swiglu(core, gu) + else: + act = m3_silu_and_mul(gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta) + return _linear(core, act, f"{prefix}_dn_w", f"{prefix}_dn_s", low_m_gemm, push) + + def _routed_chunks(self, i: int, x, router_logits): + """The routed experts from the bf16 activation over at most ``moe_max_num_tokens`` rows per runner call: the MoE + scheduler's even split, even chunks on the MoE chunking stream.""" + core = self.core + rows = x.shape[0] + num_chunks = (rows + core.moe_max_num_tokens - 1) // core.moe_max_num_tokens + if num_chunks <= 1: + return _routed_experts(core, i, x, None, router_logits, True) + sizes = _split_chunk(rows, num_chunks) + xs, logits = x.split(sizes), router_logits.split(sizes) + aux = core.aux_stream_moe_chunk + ev_main, ev_overlap = core.moe_chunk_events[i] + if aux is not None: + ev_main.record() + with torch.cuda.stream(aux): + ev_main.wait() + outs: List[torch.Tensor] = [] + for c in range(num_chunks): + if aux is not None and c % 2 == 0: + with torch.cuda.stream(aux): + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + else: + outs.append(_routed_experts(core, i, xs[c], None, logits[c], True)) + if aux is not None: + with torch.cuda.stream(aux): + ev_overlap.record() + ev_overlap.wait() + return torch.cat(outs) + + def _routed(self, i: int, h: torch.Tensor): + """Router + routed experts, finalized: the router logits (the TF32 tensor-core router on ``ROUTER_TC_GEN``'s + rows, else the fp32 router) and the experts quantizing their bf16 input (chunked).""" + return self._routed_chunks(i, h, _decode_router_logits(self.core, i, h)) + + def _moe(self, i: int, h: torch.Tensor, *, return_parts: bool, push=None): + """The MoE block of layer ``i`` without its reduction. + + ``push`` = (the engine's ``arp`` descriptor, the shared down projection's ``ar`` tuple) the decode layer armed + for the MoE boundary's exchange sets, or None. Returns ``(out, routed, engine_pushed, shared_pushed)``: with + ``return_parts`` the shared and routed parts (routed None when the engine folded the shared expert in), else + their sum in ``out``. + """ + core = self.core + engine_arp, shared_ar = push if push is not None else (None, None) + rows = h.shape[0] if h.dim() == 2 else 0 + if ( + h.dim() == 2 + and h.dtype == torch.bfloat16 + and 1 <= rows <= MOE_ENGINE_MAX_TOKENS + and rows >= MOE_FOLD_MIN_TOKENS + ): + engine = core.engine(i, rows, shared=True) + if engine is not None: + _record("m3_moe_engine", rows, engine_arp is not None) + taken(MOE_ENGINE.path, rows) + folded = m3_moe_engine(h.contiguous(), engine, None, engine_arp) + return folded, None, engine_arp is not None, False + routed_out, (shared_out, shared_pushed) = maybe_execute_in_parallel( + lambda: self._routed(i, h), + lambda: self._swiglu_mlp(h, f"l{i}_sh", False, shared_ar), + core.moe_events[i][0], + core.moe_events[i][1], + # The warmup tuning pass profiles the shared expert's GEMMs: keep it on the main stream, between the same + # collectives on every rank (the capture pass and serving still overlap it on the aux stream). + None if flashinfer_mxfp8_autotune_active() else core.aux_stream_moe, + disable_on_compile=True, + ) + if return_parts: + return shared_out, routed_out, False, shared_pushed + # In-place add into the shared output to avoid allocating a temporary. + return shared_out.add_(routed_out), None, False, shared_pushed + + def _generic_boundary(self, x, x2, residual, norm_w): + """A generic layer boundary: the multicast all-reduce + residual + RMSNorm on the rows it takes (with the + shared + routed add folded in), else the built-in fused all-reduce.""" + core = self.core + if _can_mc(x, folded=x2 is not None): + fused = _mc_allreduce_norm(core, x, x2, residual, norm_w) + if fused is not None: + return fused + if x2 is not None: + x = x.add_(x2) + return _native_allreduce_norm(core, x, residual, norm_w) + + def _generic_layer(self, i: int, position_ids, hidden, residual, step, spec_metadata, capture): + """Layer ``i`` in its generic form: input norm (first layer without the fused embedding), attention, the + attention boundary, the MoE (or the dense MLP) and the boundary into the next layer's input norm, every + collective fused with the residual add and the norm. An Eagle3 capture layer captures the boundary's reduced + residual sum.""" + core = self.core + if residual is None: + residual = hidden + hidden = flashinfer_rmsnorm(hidden, core.w[f"l{i}_in_norm"], core.rms_eps) + hidden, _ = self._attention(i, position_ids, hidden, step) + if _can_mc(hidden): + fused = _mc_allreduce_norm(core, hidden, None, residual, core.w[f"l{i}_post_norm"]) + hidden, residual = ( + fused + if fused is not None + else _native_allreduce_norm(core, hidden, residual, core.w[f"l{i}_post_norm"]) + ) + else: + hidden, residual = _native_allreduce_norm( + core, hidden, residual, core.w[f"l{i}_post_norm"] + ) + if i in core.moe_layers: + out, routed, _, _ = self._moe(i, hidden, return_parts=_can_mc(hidden, folded=True)) + else: + out, _ = self._swiglu_mlp(hidden, f"l{i}_mlp", True, dense=True) + routed = None + normed, updated = self._generic_boundary(out, routed, residual, _next_norm(core, i)) + if spec_metadata is not None and spec_metadata.is_layer_capture(i): + capture(spec_metadata, i, updated, None) + return normed, updated + + def _attention(self, i: int, position_ids, hidden, step, push: Optional[tuple] = None): + """Attention of layer ``i`` up to its o_proj (no reduction). ``push`` is the attention boundary's exchange the + o_proj may write into. Returns (o_proj output, pushed).""" + core = self.core + if i not in core.sparse_layers: + return self._dense_attention(i, position_ids, hidden, step, push) + if not core.fused_projection: + return self._unfused_sparse_attention(i, position_ids, hidden, step, push) + # #117: the QKV GEMV may be deferred to the producer call site (``QKV_PRODUCER_FOLD``). + packed, _ = _linear( + core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True, defer=QKV_PRODUCER_FOLD + ) + deferred = None + if isinstance(packed, _DeferredGemv): + deferred, packed = packed, packed.placeholder + horizontal = _sparse_qkv_producer( + core, + i, + packed, + position_ids, + step["kv_cache_manager"], + step["msa_idx_k_cache"], + step["msa_out_cache_loc"], + deferred, + ) + if horizontal is None: + raise RuntimeError( + f"MiniMax-M3 sparse layer {i}: no horizontal producer takes this step's caches (an FP8 or NVFP4 HND " + "pool with an FP8 index-K cache)" + ) + q, idx_q = horizontal + o = self._attention_core(i, q, None, None, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _dense_attention(self, i: int, position_ids, hidden, step, push): + """Dense layers 0-2: QKV, the fused norm + RoPE (FP8 q/k/v: an FP8 cache, or the FP8 pool the dense layers sit + in within a hybrid NVFP4 cache on a decode step), the attention core, o_proj. + + On a covered pure generation step (``_fullctx_plan``) the attention is this target's PrimTS split-KV decode + over the layer's full causal context, and the norm/RoPE op also stores the step's new K/V rows into the + layer's paged pool and the query compact, as the decode reads it (``DENSE_KV_INSERT_FUSED``; the bytes + ``write_msa_phase_kv`` stores). Otherwise the MSA backend.""" + core = self.core + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fp8 = core.main_kv_is_fp8 or core.main_kv_is_nvfp4 + plan = self._fullctx_plan(i, qkv.device, step) + kv_dest = None + if ( + plan is not None + and fp8 + and DENSE_KV_INSERT_FUSED + and _dense_fp8_op_applies(qkv, position_ids) + ): + kv_dest = ( + step["kv_cache_manager"].get_buffers(i, kv_layout="HND"), + step["msa_out_cache_loc"], + plan[1], + torch.empty( + (qkv.shape[0], core.heads, core.head_dim), + dtype=torch.float8_e4m3fn, + device=qkv.device, + ), + ) + if fp8: + fused = _dense_qk_norm_rope_fp8(core, i, qkv, position_ids, kv_dest) + if kv_dest is not None: + taken(DENSE_KV_INSERT_PATH, kv_dest[2]) + else: + fused = _qk_norm_rope(core, i, qkv, position_ids, False) + q, k, v = fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + if kv_dest is not None: + q = kv_dest[3].view(q.shape[0], -1) + o = None + if plan is not None: + o = self._fullctx_decode(i, q, k, v, step, plan, kv_written=kv_dest is not None) + if o is None: + o = self._attention_core(i, q, k, v, None, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _fullctx_plan(self, i: int, device: torch.device, step): + """(the PrimTS split-KV decode plan, the step's query tokens) where dense layer ``i``'s decode attention runs + on ``ops/m3_fullctx_decode`` (``FULLCTX_DECODE``): a covered (query rows per request, requests) shape over + the layer's FP8 pool. None otherwise (the MSA backend runs). + + Rubin TP2, FP8 KV, 165k context, graph-us per call (stock trtllm-gen in brackets): (4, 1) 18.8 (29.3), (1, 1) + 15.7 (28.0), (4, 2) 24.5 (29.3), (1, 2) 20.2 (28.1); the stock split is capped at 16 cluster CTAs per (request, + KV head), i.e. 32 of 212 CTAs at one request. With the BF16 split-KV merge (``FULLCTX_DECODE_OWN_REDUCER``): + (1, 1) 14.5 vs 18.5, (4, 1) 17.3 vs 21.3, (1, 2) 18.5 vs 23.1, (4, 2) 23.3 vs 27.5.""" + path = FULLCTX_DECODE.path + span = step["msa_decode_span"] + if span is None or span.row_first != 0: + uncovered(path, "no decode span starting at row 0") + return None + nreq = int(step["num_seqs"]) + qlen = int(span.query_len) + if not FULLCTX_DECODE.covers((qlen, nreq)): + uncovered(path, f"{qlen} query rows x {nreq} requests is not a covered shape") + return None + if int(step["max_seq_len"]) > FULLCTX_DECODE_MAX_KV_LEN: + uncovered(path, "the cache length bound is beyond the plans' K/V capacity") + return None + num_tokens = nreq * qlen + pool, _ = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + if pool.dtype != torch.float8_e4m3fn: + missed(path, num_tokens, "the dense layer's KV pool is not FP8") + return None + from .ops import m3_fullctx_decode as _fullctx_decode_kernel + + own_reducer = FULLCTX_DECODE_OWN_REDUCER + if own_reducer and not _fullctx_decode_kernel.has_own_reducer(): + missed(path, num_tokens, "the BF16 split-KV merge op is not built") + return None + # Compiled before warmup (``prebuild_kernels``); the forward never compiles one. + plan = _fullctx_decode_kernel.plan_for(device, qlen, nreq, own_reducer) + if plan is None: + missed(path, num_tokens, "the PrimTS decode plan is not prebuilt") + return None + return plan, num_tokens + + def _fullctx_decode(self, i: int, q, k, v, step, plan_tokens, *, kv_written: bool): + """Dense layer ``i``'s decode attention on the PrimTS plan of ``_fullctx_plan``: the step's new K/V go into the + paged cache as the MSA backend writes them (``write_msa_phase_kv``; already there when ``kv_written``), then + every query row attends its request's full causal context.""" + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + write_msa_phase_kv, + ) + + from .ops import m3_fullctx_decode as _fullctx_decode_kernel + + core = self.core + plan, num_tokens = plan_tokens + nreq = int(step["num_seqs"]) + if not kv_written: + write_msa_phase_kv( + core.attn[i], + k[:num_tokens], + v[:num_tokens], + step["metadata"], + AttentionInputType.mixed, + token_offset=0, + ) + pool, subpages = step["kv_cache_manager"].get_kv_subpage_pool(i, "HND") + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + k_pages, v_pages = _fullctx_decode_kernel.kv_page_views(pool, subpages) + # The merge also writes the o_proj's MXFP8 hand-off (``FULLCTX_DECODE_MXFP8``) where its TP8 copy runs. + mxfp8 = None + if ( + FULLCTX_DECODE_MXFP8 + and output.shape[0] == num_tokens + and _fullctx_decode_kernel.emits_mxfp8() + ): + merge_mxfp8 = _sparse_decode_kernel.MERGE_MXFP8 + if _sparse_decode_kernel.MERGE_MXFP8_MIN_TOKENS <= num_tokens <= merge_mxfp8.MAX_ROWS: + mxfp8 = merge_mxfp8.buffers(output.shape[1], output.device) + _record("m3_fullctx_decode", num_tokens, nreq, kv_written, mxfp8 is not None) + emitted = plan.run( + fp8_e4m3_saturate(q[:num_tokens].reshape(num_tokens, core.heads, core.head_dim)), + k_pages, + v_pages, + step["msa_block_table"][:nreq], + step["msa_seq_lens_cuda"][:nreq], + core.sm_scale(i), + output[:num_tokens], + mxfp8=mxfp8, + ) + if emitted: + m3_prequant.SLOT.publish(output, mxfp8[0][:num_tokens], mxfp8[1]) + taken(FULLCTX_DECODE_MXFP8_PATH, num_tokens) + taken(FULLCTX_DECODE.path, num_tokens) + return output + + def _unfused_sparse_attention(self, i: int, position_ids, hidden, step, push): + """A sparse layer without the five-way projection (``fuse_qkv_index_projection=False``): the QKV projection and + its fused norm + RoPE (FP8 q/k/v into an FP8 cache) on the current stream, the BF16 index projection and the FP8 + indexer kernel on the attention aux stream, then the MSA backend (it writes the main KV cache) and the o_proj. A + layer the NVFP4 producer covers is a miss first: the producer needs the five-way projection.""" + core = self.core + if _producer_covered(core): + missed(PRODUCER.path, None, "fuse_qkv_index_projection is disabled") + + def _main(): + qkv, _ = _linear(core, hidden, f"l{i}_qkv_w", f"l{i}_qkv_s", True) + fused = _qk_norm_rope(core, i, qkv, position_ids, core.main_kv_is_fp8) + return fused.split([core.q_size, core.kv_size, core.kv_size], dim=-1) + + def _index(): + idx_qk = m3_bf16_linear(hidden, core.w[f"l{i}_idx_w"]) + return _fp8_index_qk_norm_rope( + core, i, idx_qk, position_ids, step["msa_idx_k_cache"], step["msa_out_cache_loc"] + ) + + (q, k, v), idx_q = maybe_execute_in_parallel( + _main, + _index, + core.attn_events[i][0], + core.attn_events[i][1], + core.aux_stream_attn, + disable_on_compile=True, + ) + o = self._attention_core(i, q, k, v, idx_q, step) + return _linear(core, o, f"l{i}_o_w", f"l{i}_o_s", True, push) + + def _attention_core(self, i: int, q, k, v, idx_q, step) -> torch.Tensor: + """The attention core over the step's live tokens (a padded step's pad rows are left as allocated): a sparse + layer whose caches the producer wrote runs this target's selection and sparse decode when the generation span + starts at row 0; everything else the MSA backend (``m3_msa_attention``: the KV write unless a producer wrote the + caches, the sparse layers' top-k block selection, the FMHA).""" + core = self.core + output = q.new_empty((q.shape[0], core.heads * core.head_dim), dtype=torch.bfloat16) + live = int(step["num_tokens"]) + q_live, out_live = q[:live], output[:live] + k_live = k[:live] if k is not None else None + v_live = v[:live] if v is not None else None + idx_q_live = idx_q[:live] if idx_q is not None else None + span = step["msa_decode_span"] + sparse = i in core.sparse_layers + if sparse and k_live is None and span is not None and span.row_first == 0: + self._decode_sparse_core(i, q_live, idx_q_live, step, out_live) + return output + nvfp4 = core.main_kv_is_nvfp4 + _record("m3_msa_attention", sparse, k_live is not None) + m3_msa_attention( + core.attn[i], + q_live, + k_live, + v_live, + idx_q_live, + step["metadata"], + out_live, + sparse, + core.inv_kv_scales[i] if nvfp4 else None, + core.kv_scales[i] if nvfp4 else None, + ) + return output + + def _index_scores(self, i: int, idx_q, step, num_tokens: int, qlen: int): + """The per-index-head block maxima of a non-identity decode step: this target's CuTe DSL scorer at the request + counts ``INDEX_SCORE`` covers (each count's CTAs per request and early issue, ``index_score_schedule``), the + in-tree CuTe DSL scorer beyond them. Rubin TP2 (2 index heads), 165,521 tokens, behind a 58 MB PDL producer, + per layer: 1 request 3.27 us at 1536 CTAs (in-tree 256-CTA scorer 6.07); 2 requests 4.97 us at 768 (in-tree + 6.71).""" + config = self.core.attn[i].m3_config + rows = int(step["num_seqs"]) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + idx_q_view = idx_q.view(num_tokens, config.num_index_heads, config.sparse_index_dim) + idx_k_cache = step["msa_idx_k_cache"](i) + max_score = step["msa_proxy_max_score_view"]( + config.num_index_heads, step["msa_worst_case_max_k_tiles"], num_tokens + ) + if INDEX_SCORE.covers(rows): + split_k, early = index_score_schedule(rows) + key = ( + idx_q_view.dtype, + int(config.num_index_heads), + int(qlen), + int(idx_q_view.shape[-1]), + int(split_k), + bool(early), + ) + if key not in _index_score_kernel.IndexScoreRunner.kernel_cache: + # Outside require mode the op below compiles it (the only scorer this step has). + missed( + INDEX_SCORE_PATH, + num_tokens, + f"the index scorer {key} was not compiled before warmup", + ) + _record("m3_index_score", key) + m3_index_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + split_k, + early, + ) + taken(INDEX_SCORE_PATH, rows) + else: + uncovered( + INDEX_SCORE_PATH, + f"{rows} requests, past {INDEX_SCORE_MAX_REQUESTS}: the in-tree CuTe DSL scorer", + ) + _in_tree_scorer() # registers the op + _record("cute_dsl_minimax_m3_index_decode_score", rows, qlen) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q_view, + idx_k_cache.squeeze(1), + block_table, + seq_lens, + max_score.transpose(1, 2), + qlen, + ) + return max_score + + def _select(self, i: int, max_score, step, num_tokens: int): + """The top-16 table of a non-identity decode step from the block maxima reduced to the KV head: the long-row + select (the 384-thread kernel, ``SELECT_LONG_DSL``, when compiled) on rows of ``LONG_SELECT_MIN_BLOCKS.. + LONG_SELECT_MAX_BLOCKS`` blocks, else the built-in selector (``m3_select_blocks``).""" + config = self.core.attn[i].m3_config + n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] + group = config.num_index_heads // config.num_kv_heads + max_score_kv = ( + max_score.view(config.num_kv_heads, group, max_score.shape[1], max_score.shape[2]).amax( + dim=1 + ) + if group > 1 + else max_score + ) + from .ops import m3_select_long_dsl as _select_long_dsl_kernel + + num_blocks = int(max_score_kv.shape[1]) + table = None + dsl = False + if long_select_covers(num_blocks): + early_trigger = num_tokens <= SELECT_EARLY_MAX_ROWS + if LONG_SELECT_DSL.keys: + # Compiled before warmup (prebuild_kernels); a missing kernel falls back to the 1024-thread one. + key = _select_long_dsl_kernel.kernel_key( + get_env_enable_pdl(), early_trigger, SELECT_LONG_TRIGGER_PHASE + ) + dsl = key in _select_long_dsl_kernel.kernel_cache + if not dsl: + missed( + LONG_SELECT_DSL.path, + num_tokens, + f"the 384-thread select {key} was not compiled before warmup", + ) + _record("long_select", num_blocks, num_tokens, dsl) + table = long_select_blocks( + max_score_kv, + n_valid_blocks=n_valid_blocks, + init_blocks=config.init_blocks, + local_blocks=config.local_blocks, + use_pdl=get_env_enable_pdl(), + early_trigger=early_trigger, + trigger_phase=SELECT_LONG_TRIGGER_PHASE, + dsl=dsl, + ) + if table is not None: + taken((LONG_SELECT_DSL if dsl else LONG_SELECT).path, num_tokens) + else: + long_select_fallback(num_blocks, num_tokens) + _record("m3_select_blocks", num_blocks, num_tokens) + table = m3_select_blocks( + max_score_kv, + MSA_REQUIRED_TOPK, + n_valid_blocks, + config.init_blocks, + config.local_blocks, + True, + ) + return table + + def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: + """Top-16 block selection and the block-sparse decode of one pure generation step of a sparse layer whose caches + the producer already wrote. + + The selection is the identity when the metadata staged it (every request within 16 blocks). Otherwise the + index scorer writes the per-index-head block maxima, reduced to the KV head, and the long-row selector takes the + top 16 (the built-in selector beyond its shapes). The long-context decode (``m3_sparse_decode`` on + ``ops/m3_attn_lc8``) prefetches into L2 this layer's o_proj weights (``SPARSE_DECODE_LC8_PREFETCH_OPROJ``) + and the next sparse layer's K/V pages of the blocks that layer selected one step back + (``SPARSE_DECODE_LC8_PREFETCH_NEXT_KV``, ``_next_kv_prefetch``); hints only. + """ + from .ops import m3_attn_lc8 as _attn_lc8_kernel + + core = self.core + backend = core.attn[i] + num_tokens = int(q.shape[0]) + rows = int(step["num_seqs"]) + qlen = int(step["msa_decode_span"].query_len) + block_table = step["msa_block_table"][:rows] + seq_lens = step["msa_seq_lens_cuda"][:rows] + num_kv_heads = int(backend.indexer.config.num_kv_heads) + # The index scorers are compiled and loaded by prebuild_kernels, so no step compiles one and the identity can + # be taken from the first decode step. + identity = step["msa_short_ok"] and step["msa_short_table"] is not None + if identity: + topk = step["msa_short_table"][: num_kv_heads * num_tokens * MSA_REQUIRED_TOPK].view( + num_kv_heads, num_tokens, MSA_REQUIRED_TOPK + ) + else: + max_score = self._index_scores(i, idx_q, step, num_tokens, qlen) + topk = self._select(i, max_score, step, num_tokens).permute(1, 0, 2) + manager = step["kv_cache_manager"] + k_paged, v_paged = msa_paged_kv(manager, i) + nvfp4 = {} + if manager.is_nvfp4_layer(i): + k_scale, v_scale = _aligned_nvfp4_dequant_scales(backend, core.kv_scales[i]) + scales = manager.get_block_scale_buffers(i, "HND") + nvfp4 = dict( + k_block_scale=scales[:, 0], + v_block_scale=scales[:, 1], + k_global_scale=k_scale, + v_global_scale=v_scale, + ) + # The long-context decode (if it runs) prefetches this layer's o_proj weights into L2; any other route leaves + # the slot unused, and it is cleared after the call either way. + if _attn_lc8_kernel.MERGE_ONCE and SPARSE_DECODE_LC8_PREFETCH_OPROJ: + _attn_lc8_kernel.PREFETCH.tensors = (core.w[f"l{i}_o_w"], core.w[f"l{i}_o_s"]) + if _attn_lc8_kernel.MERGE_ONCE and i in core.n1_ids and num_tokens <= N1_MAX_ROWS: + _attn_lc8_kernel.NEXT_KV.tensors = self._next_kv_prefetch(i, manager) + _record("m3_sparse_decode", num_tokens, qlen, identity) + try: + m3_sparse_decode( + q.view(num_tokens, core.heads, core.head_dim), + k_paged, + v_paged, + topk, + block_table, + seq_lens, + sm_scale=core.sm_scale(i), + output=output.view(num_tokens, core.heads, core.head_dim), + decode_query_len=qlen, + identity_blocks=identity, + **nvfp4, + ) + finally: + _attn_lc8_kernel.PREFETCH.tensors = None + _attn_lc8_kernel.NEXT_KV.tensors = None + + def _next_kv_prefetch(self, i: int, manager): + """(layer ``i``'s id buffer, the next sparse layer's, that layer's K / V pools and block scales) for the + long-context decode's next-layer prefetch; the last sparse layer's next is the first one (next step). None when + the next layer's cache is not NVFP4.""" + n1_ids = self.core.n1_ids + order = sorted(n1_ids) + nxt = next((j for j in order if j > i), order[0]) + if not manager.is_nvfp4_layer(nxt): + return None + k_next, v_next = msa_paged_kv(manager, nxt) + scales = manager.get_block_scale_buffers(nxt, "HND") + return (n1_ids[i], n1_ids[nxt], k_next, v_next, scales[:, 0], scales[:, 1]) + + +_ENGAGED: set = set() + + +def _log_engaged(path: str, mapping) -> None: + """One info line per fast path and process (every GEN rank logs its own): the run's engagement evidence.""" + if path not in _ENGAGED: + _ENGAGED.add(path) + rank = getattr(mapping, "tp_rank", "?") if mapping is not None else "?" + logger.info(f"[M3 sm_107 tp8] fast path engaged: {path} (tp rank {rank})") + + +class M3SpecStepFold: + """The target's speculative step-input staging (the in-tree ``SpecStepFold`` protocol). + + Between two decode graphs of a one-model Eagle3 step, the spec sampler stores the step's outputs into its per-slot + store and the engine gathers them back into its flat input buffers for the next step. This object does both in + one launch each (``m3_spec_tail.slot_scatter`` / ``step_gather``): integer copies only, the same values the + engine's and the sampler's own copies write. + + The forward of a decode-only step also does both at its tail (``store_in_forward``, one ``m3_fold_tail`` launch + inside the step graph, this target's C++ op): it stores the step's outputs for the rows ``begin_step`` armed and + stages the next step's inputs for those rows in that order. The arming is a device-side row count and slot list + staged before the replay, so a captured graph whose step was not armed writes nothing. The sampler then skips its + store for exactly the armed rows (``sampler_rows_stored``), and the next step skips its gathers when its rows are + those rows (``begin_step``). The sampler copies the store to the host on its D2H side stream and records + ``store_copies_done`` after the copies; the forward's tail waits on it before it rewrites the store. + + * ``store`` -- the spec sampler's per-slot store (``bind_sampler_store``). + * ``step_inputs`` -- the engine's (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) device buffers + (``bind_engine_buffers``). + """ + + # Read by the shell's forward (the tail runs only when the fold stores in the forward). + in_forward = True + + def __init__(self) -> None: + self.store: Optional[object] = None + self.step_inputs: Optional[tuple] = None + # Device arming read by the forward's tail: ctrl[0] = armed rows (0 = write nothing), slots[:rows]. + self._ctrl: Optional[torch.Tensor] = None + self._slots: Optional[torch.Tensor] = None + # The rows (sequence slots) the next forward stores and stages for, as begin_step armed them. + self._armed: Optional[tuple] = None + # The rows the previous forward staged the next step's inputs for (confirmed by sampler_rows_stored). + self._staged_for: Optional[tuple] = None + # Output buffers the tail ran on: the captured graphs' static outputs (by address) and the last eager step's. + self._graph_outputs: set = set() + self._eager_outputs: Optional[torch.Tensor] = None + self._logged = False + # The (external) event the sampler records on its D2H side stream after copying the store to the host, waited + # on by the forward's tail before it rewrites the store (see SpecSamplerBase._copy_outputs_to_host). + self.store_copies_done: Optional[torch.cuda.Event] = None + + def bind_engine_buffers( + self, input_ids, draft_tokens, position_id_offsets, kv_lens_offsets, max_batch_size + ): + self.step_inputs = (input_ids, draft_tokens, position_id_offsets, kv_lens_offsets) + # Recorded once now so a captured wait has a completed record to wait on before the first copy. + self.store_copies_done = torch.cuda.Event(external=True) + self.store_copies_done.record() + logger.info( + "[M3 host] spec store host copies on the D2H side stream, awaited by the forward's tail" + ) + logger.info("[M3 host] spec fold tail: one C++ launch (no Triton) at the forward's tail") + self._ctrl = torch.zeros(1, dtype=torch.int32, device=input_ids.device) + self._slots = torch.zeros(max(max_batch_size, 1), dtype=torch.long, device=input_ids.device) + + def bind_sampler_store(self, store: object) -> None: + self.store = store + + def begin_step(self, seq_slots: Optional[list], num_rows: int) -> bool: + if self._ctrl is None: + return False + rows = tuple(seq_slots) if seq_slots is not None else None + staged = rows is not None and rows == self._staged_for + self._staged_for = None + armed = rows if rows is not None and 0 < len(rows) <= self._slots.numel() else None + # Read only by this step's forward, so both uploads may wait for the step-input pack's flush. + if armed is not None: + slots = torch.tensor(armed, dtype=torch.long, pin_memory=prefer_pinned()) + copy_h2d_if_changed(self._slots[: len(armed)], slots, pack=True) + ctrl = torch.tensor( + [len(armed) if armed is not None else 0], dtype=torch.int32, pin_memory=prefer_pinned() + ) + copy_h2d_if_changed(self._ctrl, ctrl, pack=True) + self._armed = armed + return staged + + # The launches the store wait may precede. store_copies_done is an external event: under graph capture its wait + # is a node of its own, so the kernel launched right after it loses its programmatic (PDL) edge and with it the + # overlap of its pre-wait part. Only a launch that waits for its predecessor's completion anyway may follow it: + # the fold tail (a plain launch, no griddepcontrol). Waiting before the last draft_step instead cost +13 / +17 + # us/step. + STORE_WAIT_BEFORE = frozenset({"fold_tail"}) + + def _wait_store_copies(self, before: str) -> None: + """Order the store's rewrite after the sampler's host copies of it; ``before`` names the launch that follows. + Under graph capture a launch outside STORE_WAIT_BEFORE fails the capture.""" + if self.store_copies_done is None: + return + if before not in self.STORE_WAIT_BEFORE and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"M3SpecStepFold: the store wait would precede {before!r}, a launch outside " + f"{sorted(self.STORE_WAIT_BEFORE)}; an external-event wait node drops the next kernel's PDL edge" + ) + torch.cuda.current_stream().wait_event(self.store_copies_done) + + def store_in_forward(self, outputs: dict) -> None: + """At the forward's tail: store the armed rows' outputs and stage the next step's inputs (one launch).""" + if self._ctrl is None or self.store is None or self.step_inputs is None: + return + store = self.store + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + new_tokens = outputs.get("new_tokens") + next_new_tokens = outputs.get("next_new_tokens") + new_tokens_lens = outputs.get("new_tokens_lens") + next_draft_tokens = outputs.get("next_draft_tokens") + if any( + t is None or not isinstance(t, torch.Tensor) + for t in (new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens) + ): + return + rows = new_tokens.shape[0] + # The previous step's host copies of the store run on the sampler's side stream. + self._wait_store_copies(before="fold_tail") + _record("m3_fold_tail", rows) + if not m3_fold_tail.fold_tail( + self._ctrl, + self._slots, + min(rows, self._slots.numel()), + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + next_draft_tokens, + store.next_draft_tokens, + new_tokens_lens, + store.new_tokens_lens, + input_ids, + draft_tokens, + position_id_offsets, + kv_lens_offsets, + next_new_tokens.shape[1], + cpp=True, + ): + return + if torch.cuda.is_current_stream_capturing(): + self._graph_outputs.add(new_tokens.data_ptr()) + else: + self._eager_outputs = new_tokens + + def gather( + self, + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len: int, + tokens_per_gen_step: int, + draft_token_buffer_width: int, + num_tokens: int, + num_draft_tokens: int, + first_previous_row: int, + ) -> bool: + if self.step_inputs is None: + return False + input_ids, draft_tokens, position_id_offsets, kv_lens_offsets = self.step_inputs + return m3_spec_tail.step_gather( + new_tokens, + next_draft_tokens, + new_tokens_lens, + previous_slots, + previous_pos_indices, + previous_batch_len, + tokens_per_gen_step, + draft_token_buffer_width, + input_ids, + num_tokens, + draft_tokens, + num_draft_tokens, + position_id_offsets, + kv_lens_offsets, + first_previous_row, + ) + + def sampler_rows_stored(self, outputs: dict, num_skip: int, seq_slots: list) -> bool: + armed, self._armed = self._armed, None + new_tokens = outputs.get("new_tokens") + ran = isinstance(new_tokens, torch.Tensor) and ( + new_tokens is self._eager_outputs or new_tokens.data_ptr() in self._graph_outputs + ) + self._eager_outputs = None + stored = ( + ran + and armed is not None + and num_skip == 0 + and tuple(seq_slots) == armed + and new_tokens.shape[0] == len(armed) + ) + self._staged_for = armed if stored else None + if stored and not self._logged: + self._logged = True + logger.info("[M3 host] spec step stored and next step staged in the forward tail") + return stored + + def scatter( + self, seq_slots, new_tokens, next_new_tokens, new_tokens_lens, next_draft_tokens + ) -> bool: + store = self.store + if store is None: + return False + return m3_spec_tail.slot_scatter( + seq_slots, + new_tokens, + store.new_tokens.squeeze(-1).T, + next_new_tokens, + store.next_new_tokens.squeeze(-1).T, + new_tokens_lens, + store.new_tokens_lens, + next_draft_tokens, + store.next_draft_tokens, + ) + + +class _DraftForward: + """``draft_model.model`` for the in-tree worker: the drafter layer's forward, called as + ``model(input_ids=, position_ids=, hidden_states=, attn_metadata=, spec_metadata=, all_rank_num_tokens=)`` and + returning (hidden, hidden to carry to the next pass). ``d2t`` None: the draft vocabulary is the target's.""" + + d2t = None + + def __init__(self, draft: "M3Eagle3Draft") -> None: + self._draft = draft + + def __call__(self, **kwargs): + return self._draft.forward_layer(**kwargs) + + +# The drafter's projections by weight-table name, for the low-M linears' engagement and prebuild. +_DRAFT_LINEARS = ("fc", "qkv", "o", "gu", "dn") + + +class M3Eagle3Draft(nn.Module): + """The one-layer BF16 Eagle3 drafter (``speculative_model``): its weights (``weights.DRAFT_WEIGHTS``), its + vocab-parallel BF16 head and logits processor, the drafter layer's TRTLLM attention backend (layer index = the + target's layer count, FP8 KV) and its TP all-reduce. + + What the worker reads: ``model`` (the layer forward), ``apply_eagle3_fc`` (the per-layer fc_norms + fc over the + captured target hidden states), ``lm_head`` and ``logits_processor``; the fused draft loop reads the layer's parts + (``draft_input_weights``, ``attention``, ``mlp``, ``ar_norm``). The drafter embeds with the target's table. + + The BF16 projections (fc, qkv, o_proj, gate_up, down; not the head) run FlashInfer's fixed-tactic low-M kernels + (``ops/gemm/m3_draft_lowm``, the TP8 port) at the rows ``DRAFT_LOWM`` covers, compiled before the warmup + (``prebuild_linears``); every other call is the in-tree cuBLAS GEMM. Rubin TP2 (mclk locked, 4 rotating sets) at + 1 row: qkv 13.18 -> 9.71, o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. The + split-K ('s') tactics run this target's copy of FlashInfer's kernel (``ops/m3_lowm_splitk``, bit for bit the + same output); with ``DRAFT_AR_PUSH`` the o_proj and down projection of the fused draft loop also push their rows + into the drafter boundary all-reduce's exchange (``arm_push`` / ``take_pushed``); the direct ('d') tactics never + push. + """ + + def __init__( + self, + core: MiniMaxM3Tp8ModelingV2, + draft_config: ModelConfig, + start_layer_idx: int, + capture_layers: int, + ) -> None: + super().__init__() + cfg = draft_config.pretrained_config + _ensure_draft_vocab_size(cfg) + eagle = getattr(cfg, "eagle_config", None) or {} + mapping = draft_config.mapping + self.draft_config = draft_config + # The engine's draft weight mapper binds to these two. + self.model_config = draft_config + self.config = cfg + self.capture_layers = int(capture_layers) + self.layer_idx = int(start_layer_idx) + self.dtype = cfg.torch_dtype + self.hidden = int(cfg.hidden_size) + self.hidden_size_in = int(getattr(cfg, "target_hidden_size", cfg.hidden_size)) + self.rms_eps = float(cfg.rms_norm_eps) + self.return_hidden_post_norm = bool( + eagle.get("return_hidden_post_norm", False) or getattr(cfg, "norm_output", False) + ) + if ( + int(cfg.num_hidden_layers) != 1 + or eagle.get("eh_proj_before_attn", False) + or eagle.get("norm_before_fc", False) + or not getattr(cfg, "fc_norm", False) + or getattr(cfg, "kv_lora_rank", None) + or self.capture_layers <= 1 + or self.hidden_size_in != self.hidden + or cfg.draft_vocab_size != cfg.vocab_size + or getattr(cfg, "attention_bias", False) + or getattr(cfg, "mlp_bias", False) + or getattr(cfg, "bias", False) + or getattr(cfg, "tie_word_embeddings", False) + ): + raise ValueError( + "this MiniMax-M3 target drafts with a one-layer Llama-style Eagle3 drafter with per-layer " + "fc_norms over " + "several captured layers, the target's hidden size and vocabulary and no biases" + ) + tp = mapping.tp_size + self.heads = cfg.num_attention_heads // tp + self.kv_heads = (cfg.num_key_value_heads + tp - 1) // tp + self.head_dim = int( + getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads + ) + self.q_size = self.heads * self.head_dim + self.kv_size = self.kv_heads * self.head_dim + self.w = _weights.DRAFT_WEIGHTS.declare(_weights.DRAFT_WEIGHTS.dims(self)) + self.lm_head = LMHead( + cfg.draft_vocab_size, + cfg.hidden_size, + dtype=cfg.torch_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=draft_config.lm_head_gather_output, + reduce_output=False, + use_custom_cublas_mm=False, + quant_config=None, + ) + self.logits_processor = LogitsProcessor() + self.allreduce = AllReduce( + mapping=mapping, strategy=draft_config.allreduce_strategy, dtype=cfg.torch_dtype + ) + quant_config = draft_config.get_quant_config() + self.attn = create_attention( + draft_config.attn_backend, + self.layer_idx, + self.heads, + self.head_dim, + self.kv_heads, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, rope=RopeParams.from_config(cfg) + ), + quant_config=quant_config, + skip_create_weights_in_init=draft_config.skip_create_weights_in_init, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + if not self.attn.support_fused_rope() or not self.attn.support_fused_qkv(): + raise ValueError( + f"the drafter's attention backend {type(self.attn).__name__} must fuse RoPE and QKV" + ) + self.attn.update_quant_config(quant_config) + # The drafter embeds with the target's table, read off the core on every call (the engine materializes the + # core's weights after construction). + self._core = [core] + self.model = _DraftForward(self) + # The low-M linears: (N, K, M, PDL) of every call prebuild_linears compiled, and (N, K, M) of every call that + # logged its engagement line. + self._lowm_built: set = set() + self._lowm_engaged: set = set() + # DRAFT_AR_PUSH: the weight whose next low-M call also pushes its rows into a boundary all-reduce (data + # pointer, the bucket's multicast words, its flags, rank, world, max tokens), set by arm_push; and the data + # pointer of the output the last armed call pushed, 0 if none (read once by take_pushed). Host state set and + # read in one forward's issue order. + self._push_arm: Optional[tuple] = None + self._pushed = 0 + + @property + def embed_table(self) -> torch.Tensor: + return self._core[0].w["embed"] + + def load_weights_from_target_model(self, target_model: nn.Module) -> None: + """Nothing to alias: the drafter reads the target's embedding table through the core and owns its head.""" + + # ------------------------------------------------------------------------------------------- the projections + + def linear(self, x: torch.Tensor, key: str) -> torch.Tensor: + """``x @ w[key]^T``: the low-M kernel where ``DRAFT_LOWM`` covers the call, else the in-tree cuBLAS GEMM.""" + w = self.w[key] + out = self._lowm_linear(w, x) + return out if out is not None else m3_bf16_linear(x, w) + + def _lowm_linear(self, weight: torch.Tensor, x: torch.Tensor) -> Optional[torch.Tensor]: + """x [M, K] @ weight^T through the low-M kernel when ``DRAFT_LOWM`` covers it; None for the in-tree GEMM. + A projection none of whose row counts is covered never reaches the low-M checks.""" + n, k = weight.shape + if x.dim() != 2 or not any(nk == (n, k) for nk, _ in DRAFT_LOWM.keys): + return None + m = int(x.shape[0]) + if not DRAFT_LOWM.covers(((n, k), m)): + uncovered(DRAFT_LOWM.path, f"the drafter's {n}x{k} linear past its tactic table's rows") + return None + if not x.is_contiguous(): + # The first pass's fc may read a column window of the captured hidden states; the in-tree GEMM takes + # strides. + uncovered(DRAFT_LOWM.path, "a strided activation (the in-tree GEMM reads it in place)") + return None + use_pdl = get_env_enable_pdl() + if x.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16: + missed(DRAFT_LOWM.path, m, "the activation or weight is not BF16") + return None + if (n, k, m, use_pdl) not in self._lowm_built: + missed( + DRAFT_LOWM.path, m, f"the {n}x{k} kernel at {m} rows was not compiled before warmup" + ) + return None + kind, tactic = tactic_for(n, k, m) + if (n, k, m) not in self._lowm_engaged: + self._lowm_engaged.add((n, k, m)) + logger.info( + f"[M3 sm_107 tp8] fast path engaged: m3_draft_lowm {n}x{k} at {m} rows (device {x.device.index})" + ) + out = torch.empty((m, n), dtype=x.dtype, device=x.device) + ar = None + if ( + self._push_arm is not None + and self._push_arm[0] == weight.data_ptr() + and kind == "s" + and DRAFT_LOWM_OWN_SPLITK + ): + ar = self._push_arm[1:] + self._pushed = out.data_ptr() + _record("m3_draft_lowm", n, k, m, kind, ar is not None) + return m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, DRAFT_LOWM_OWN_SPLITK, ar) + + def prebuild_linears(self) -> int: + """Compile (and run once) every covered row count of every drafter projection, outside capture, on the + current stream, before the executor's warmup. Returns the number of calls compiled.""" + from .ops.m3_lowm_splitk import compile_push + + use_pdl = get_env_enable_pdl() + built = 0 + shapes = {tuple(self.w[key].shape): self.w[key] for key in _DRAFT_LINEARS} + for (n, k), weight in sorted(shapes.items()): + if weight.dtype != torch.bfloat16: + continue + for _, m in sorted(key for key in DRAFT_LOWM.keys if key[0] == (n, k)): + if (n, k, m, use_pdl) in self._lowm_built: + continue + kind, tactic = tactic_for(n, k, m) + x = torch.zeros(m, k, dtype=torch.bfloat16, device=weight.device) + out = torch.empty(m, n, dtype=torch.bfloat16, device=weight.device) + m3_draft_lowm(x, weight, out, kind, tactic, use_pdl, DRAFT_LOWM_OWN_SPLITK) + if DRAFT_AR_PUSH and DRAFT_LOWM_OWN_SPLITK and kind == "s": + # The push variant a boundary all-reduce may arm (o_proj / down): compiled only, never run here + # (its pushes need the all-reduce's exchange and a matching call). + compile_push(x, weight.t(), out, use_pdl, tactic) + self._lowm_built.add((n, k, m, use_pdl)) + built += 1 + return built + + def arm_push(self, key: str, workspace, rows: int) -> None: + """Arm ``w[key]``'s next covered low-M call to push into ``workspace``'s bucket for ``rows``.""" + bucket = m3_mc_allreduce._bucket(workspace, rows) + flags = None if bucket is None else bucket.flags_for(8) + self._pushed = 0 + self._push_arm = None + if bucket is not None and flags is not None and bucket.max_tokens == 8: + self._push_arm = ( + self.w[key].data_ptr(), + bucket.mc.data_ptr(), + flags.data_ptr(), + bucket.rank, + bucket.world_size, + bucket.max_tokens, + ) + + def take_pushed(self, x: torch.Tensor) -> bool: + """Whether the armed call pushed its rows; disarms. ``x`` is what the boundary all-reduce is about to reduce: + it must be the pushing call's own output (the prepushed all-reduce reads only the pushed words, never ``x``), + so an operation between that call and the all-reduce raises instead of silently reducing the wrong rows.""" + pushed, self._pushed, self._push_arm = self._pushed, 0, None + if pushed and x.data_ptr() != pushed: + raise RuntimeError( + "the drafter's pushed rows are not the boundary all-reduce's input (an operation ran in between)" + ) + return bool(pushed) + + # ------------------------------------------------------------------------------------------------- the layer + + def fc(self, hs: torch.Tensor) -> torch.Tensor: + return self.linear(hs, "fc") + + def apply_eagle3_fc(self, hidden_states: torch.Tensor) -> torch.Tensor: + """The captured target hidden states [T, layers x H] to the drafter's input width: each layer's fc_norm, + then the fc (the in-tree ``Eagle3ForCausalLM.apply_eagle3_fc``).""" + hidden_states = hidden_states.to(self.dtype) + if hidden_states.shape[-1] != self.hidden: + chunks = hidden_states.chunk(self.capture_layers, dim=-1) + hidden_states = torch.cat( + [ + flashinfer_rmsnorm(c, self.w[f"fc_norm{i}"], self.rms_eps) + for i, c in enumerate(chunks) + ], + dim=-1, + ) + hidden_states = self.fc(hidden_states) + return hidden_states + + def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: + """QKV, the backend's fused RoPE + KV append + causal attention, o_proj (all-reduced unless ``reduce`` is + False, where the caller fuses the reduction). The backend's generation phase runs this target's full-context + decode with the fused RoPE + FP8 K/V append when the shell installed it (``ops/m3_drafter_attn``).""" + qkv = self.linear(x, "qkv") + _record("m3_draft_attention", int(md.num_tokens)) + o = m3_draft_attention(self.attn, qkv, md) + out = self.linear(o, "o") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def mlp(self, h: torch.Tensor, reduce: bool, cpp_silu: bool = False) -> torch.Tensor: + """gate_up, the plain SwiGLU, down (all-reduced unless ``reduce`` is False). ``cpp_silu`` (the fused draft + loop): the SwiGLU on this target's C++ op (``ops/spec/m3_draft_silu``, minimaxM3Sm107Tp8/siluAndMul.cu), bit + for bit the in-tree Triton silu_and_mul in one launch either way (no Triton on the GEN decode path).""" + gu = self.linear(h, "gu") + act = None + if cpp_silu: + act = m3_draft_silu(gu) + if act is None: + uncovered( + DRAFT_SILU_PATH, "the drafter SwiGLU input layout (the Triton op runs it)" + ) + else: + taken(DRAFT_SILU_PATH, int(gu.shape[0]), "drafter MLP SwiGLU, one C++ launch") + if act is None: + act = m3_silu_and_mul(gu, None, None, None) + out = self.linear(act, "dn") + return m3_allreduce(out, self.allreduce, None) if reduce else out + + def ar_norm(self, x: torch.Tensor, residual: torch.Tensor, norm_key: str): + """One draft-layer boundary in the in-tree structure: the projection's plain all-reduce (BF16 out), then the + fused residual add + RMSNorm (FP32 add, norm of the FP32 sum). Returns (normed, residual sum).""" + return flashinfer_fused_add_rmsnorm( + m3_allreduce(x, self.allreduce, None), residual, self.w[norm_key], self.rms_eps + ) + + def forward_layer( + self, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + spec_metadata, + all_rank_num_tokens: Optional[List[int]] = None, + ): + """The drafter layer (the in-tree ``Eagle3DraftModel.forward`` for one layer): the embedded tokens and the + fc'd hidden states normed and concatenated, attention, the post-attention norm with the residual, the MLP and + the final norm. Returns (normed output, the hidden state the next pass carries).""" + previous = attn_metadata.all_rank_num_tokens + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = all_rank_num_tokens + try: + embeds = m3_embedding(input_ids, self.embed_table).to(self.dtype) + residual = hidden_states + h = flashinfer_rmsnorm(hidden_states, self.w["hidden_norm"], self.rms_eps) + e = flashinfer_rmsnorm(embeds, self.w["in_norm"], self.rms_eps) + h = torch.cat([e, h], dim=-1) + h = self.attention(h, attn_metadata, reduce=True) + h, residual = flashinfer_fused_add_rmsnorm( + h, residual, self.w["post_norm"], self.rms_eps + ) + h = self.mlp(h, reduce=True) + spec_metadata.maybe_capture_hidden_states(self.layer_idx, h, residual) + h, to_save = flashinfer_fused_add_rmsnorm( + h, residual, self.w["final_norm"], self.rms_eps + ) + return (h, h) if self.return_hidden_post_norm else (h, to_save) + finally: + if all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = previous + + +def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, capture_norm) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and getattr(lm_head, "tp_size", 1) > 1 + ) + + +class M3Eagle3Worker(Eagle3OneModelWorker): + """``Eagle3OneModelWorker`` with the MiniMax-M3 fused verify and fused draft loop, at the in-tree numerics. + + * Verify: on decode-only batches whose verify rows the fused kernels cover (``tail_eligible``), the target's BF16 + lm_head computes this rank's vocabulary shard of the logits (the in-tree head's own local GEMM) and one kernel + replaces the logits all-gather, the argmax (or the temperature + top-p draw) and the strict acceptance. + * Draft: the drafter's first pass may take the per-layer fc_norms in one launch (or folded into the hidden-state + capture) and run its fc beside the target head GEMM (forked right after the last captured target layer at most + ``FC_EARLY_MAX_ROWS`` verify rows; after a sampled verify's head at most ``FC_AFTER_HEAD_MAX_ROWS``). On a + decode-only step ``DRAFT_STEP`` covers, the draft loop is fused: each pass's draft input in one launch, the + drafter layer's attention and MLP with their all-reduces fused with the residual add and the next RMSNorm (the + multicast all-reduce at ``DRAFT_MC_AR``'s rows, fed by the projection's push), and one kernel for the BF16 head, + the cross-rank argmax and the pass's bookkeeping. + + Every other batch takes the in-tree paths. Methods whose docstring says "Mirrors" copy the named in-tree method + with the fused branches inserted. + """ + + def __init__(self, spec_config, mapping, use_separate_draft_kv_cache: bool = False) -> None: + super().__init__( + spec_config, mapping=mapping, use_separate_draft_kv_cache=use_separate_draft_kv_cache + ) + # Per-step state, reset at the top of every _forward_impl: the first draft pass's gather rows when this step + # takes the fused first drafter inputs, and the fc output a fork computed beside the target head. + self._m3_first_gather_ids: Optional[torch.Tensor] = None + self._m3_fc_hidden: Optional[torch.Tensor] = None + self._m3_target_lm_head: Optional[nn.Module] = None + # (fc_norm0 weight, stacked fc_norm weights [layers, H]), built eagerly on the first fused first pass. + self._m3_fc_norm_w: Optional[tuple] = None + # Armed per step by arm_capture_norm: (stacked fc_norm weights, eps) or None. + self._m3_capture_norm: Optional[tuple] = None + # Per step: the captured layers written with the capture norm (set by capture_hidden_states). + self._m3_capture_normed: set = set() + self._m3_aux_stream: Optional[torch.cuda.Stream] = None + # Armed per step by arm_capture_norm: the drafter whose first fc is forked at the last captured target layer, + # and the fork that did it, joined by _m3_target_accept. + self._m3_fc_early_drafter: Optional[M3Eagle3Draft] = None + self._m3_fc_early_stream: Optional[torch.cuda.Stream] = None + # DI_FC_FLAG: int64 [3] device counters (forks bumped, forks waited, wait overruns), allocated by ``prebuild``; + # whether the latest fork bumped them; the fork whose stream join is deferred past the first draft input. + self._m3_fc_counts: Optional[torch.Tensor] = None + self._m3_fc_flagged = False + self._m3_fc_pending: Optional[torch.cuda.Stream] = None + # The 12-bit coded copy of the target head shard the fused verify streams, built before the warmup. + self._m3_force_vec_warm = False + # The streamed copy of the drafter's BF16 head shard the fused draft loop reads, built before the warmup, the + # head weight version it was built from, and the replaced copies a captured graph may still read. + self._m3_draft_head = None + self._m3_draft_head_version: Optional[int] = None + self._m3_draft_head_graveyard: list = [] + # Per step: the next step's input tokens [nreq, D + 1], column 0 written by the fused verify (or the seed) and + # columns 1.. by the fused draft passes; set as _m3_next_new_tokens once the fused draft loop filled it. + self._m3_next_new: Optional[torch.Tensor] = None + self._m3_next_new_tokens: Optional[torch.Tensor] = None + # Per step: the first draft inputs are the fused ones (accepted tokens as they lie, fc output of every row). + self._m3_fused_first_inputs = False + + def capture_target(self, spec_metadata, layer_id: int, num_tokens: int) -> Optional[tuple]: + """Where this step's capture norm writes layer ``layer_id``'s capture, for a boundary all-reduce that writes it + itself (CAPTURE_NORM_IN_BOUNDARY): (the capture buffer's [num_tokens, H] column block, the layer's fc_norm + weight, eps), or None (no capture norm armed, not a captured layer, or more rows than the buffer). The + boundary writes every row it reduces; rows past num_capture_tokens are unread.""" + capture_norm = self._m3_capture_norm + if capture_norm is None or spec_metadata.hidden_states is None: + return None + if not m3_mc_allreduce.capture_supported(): + return None + if num_tokens > spec_metadata.hidden_states.shape[0]: + return None + hidden_size = spec_metadata.hidden_size + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + out = spec_metadata.hidden_states[ + :num_tokens, i * hidden_size : (i + 1) * hidden_size + ] + return out, capture_norm[0][i], capture_norm[1] + return None + + def capture_hidden_states( + self, spec_metadata, layer_id: int, hidden_states, residual=None, written: bool = False + ) -> None: + """The Eagle3 capture with this step's capture norm (``written``: the layer's boundary all-reduce already wrote + it, see ``capture_target``); on a step armed for the early fc fork (``arm_capture_norm``), the last captured + layer then forks the drafter's first fc onto the aux stream, where it overlaps the remaining target layers + and the head (joined by ``_m3_target_accept``).""" + if not written: + _eagle3_capture(spec_metadata, layer_id, hidden_states, residual, self._m3_capture_norm) + if self._m3_capture_norm is not None: + # Written with this step's capture norm: by the boundary all-reduce (written) or just above. + self._m3_capture_normed.add(layer_id) + drafter = self._m3_fc_early_drafter + layers = spec_metadata.layers_to_capture + if drafter is not None and layers and layer_id == layers[-1]: + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = self._m3_fc_fork( + spec_metadata, spec_metadata.num_capture_tokens, drafter + ) + + def prebuild(self, lm_head: nn.Module, draft: Optional[M3Eagle3Draft]) -> None: + """From the loaded weights, before the executor's warmup, on every rank: the fused draft loop's head copy and + the drafter's low-M linears.""" + start = time.perf_counter() + if self.prebuild_draft_head(draft): + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] fused draft head built before warmup: %.1f s" + % (time.perf_counter() - start) + ) + if DI_FC_FLAG and draft is not None and self._m3_fc_counts is None: + # The fc fork counters (DI_FC_FLAG), zeroed here, before warmup, in order before any fork. + self._m3_fc_counts = torch.zeros( + 3, dtype=torch.int64, device=draft.w["hidden_norm"].device + ) + start = time.perf_counter() + built = draft.prebuild_linears() if draft is not None else 0 + torch.cuda.synchronize() + logger.info( + "[M3 prebuild] drafter low-M GEMMs built before warmup: %d calls, %.1f s" + % (built, time.perf_counter() - start) + ) + + def _forward_impl( + self, + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + resource_manager=None, + ): + """Mirrors Eagle3OneModelWorker._forward_impl, with the fused verify branch.""" + runtime_draft_len = spec_metadata.runtime_draft_len + # skip the draft forward if the runtime draft length is 0 + if runtime_draft_len == 0: + return self.skip_drafting( + input_ids, + position_ids, + hidden_states, + logits, + attn_metadata, + spec_metadata, + draft_model, + ) + + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + + raw_logits = logits + + self._m3_first_gather_ids = None + self._m3_next_new = None + self._m3_next_new_tokens = None + self._m3_fused_first_inputs = False + if logits is None: + # tail_eligible() skipped the gathered target logits: verify from this rank's logits shard. + accepted_tokens, num_accepted_tokens, self._m3_first_gather_ids = ( + self._m3_target_accept( + hidden_states, + spec_metadata, + batch_size, + runtime_draft_len, + draft_model, + kv_lens=getattr(attn_metadata, "kv_lens_cuda", None), + ) + ) + else: + self._execute_guided_decoder_if_present(logits) + + # Sample and accept tokens. ``input_ids`` is required by the relaxed-acceptance path (scans for + # thinking-phase tokens); ignored otherwise. + accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( + input_ids, logits, attn_metadata, spec_metadata + ) + if self.draft_only_eligible(attn_metadata, spec_metadata, draft_model): + self._m3_seed_first_pass(num_accepted_tokens, spec_metadata, batch_size) + + # Mamba hybrid models need state updates after token acceptance (a no-op on this target's cache manager). + if self._is_mamba_hybrid_cache is None: + self._is_mamba_hybrid_cache = isinstance( + attn_metadata.kv_cache_manager, MambaHybridCacheManager + ) + if num_gens > 0 and self._is_mamba_hybrid_cache: + attn_metadata.kv_cache_manager.update_mamba_states( + attn_metadata=attn_metadata, + num_accepted_tokens=num_accepted_tokens, + state_indices=attn_metadata.mamba_metadata.state_indices, + ) + + sa_manager = getattr(spec_metadata.spec_resource_manager, "sa_manager", None) + if self.sa_enhancer is not None and sa_manager is not None: + self.sa_enhancer.extend_and_prepare( + sa_manager=sa_manager, + request_ids=spec_metadata.request_ids, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_gens=num_gens, + num_contexts=num_contexts, + max_draft_len=runtime_draft_len, + ) + + if self._uses_external_shared_target_kv: + next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + accepted_tokens=accepted_tokens, + num_accepted_tokens=num_accepted_tokens, + num_contexts=num_contexts, + batch_size=batch_size, + ) + else: + # Save the old attn_metadata and spec_metadata + self._prepare_attn_metadata_for_spec_dec(attn_metadata) + + # Prepare inputs for the 1st draft model forward + position_ids = position_ids.squeeze(0) + inputs = self.prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + # Predict draft tokens. ``original_all_rank_num_tokens`` is saved here so the post-loop restore (below) + # can put attn_metadata back into a state the target model expects. + original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + # Get the draft KV cache manager if using separate layouts + draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) + + next_draft_tokens = self._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + # restore attn_metadata to support cuda graph + self._restore_attn_metadata_from_spec_dec(attn_metadata) + # restore all_rank_num_tokens for attention DP + if original_all_rank_num_tokens is not None: + attn_metadata.all_rank_num_tokens = original_all_rank_num_tokens + + self._m3_join_pending() + # prepare next new tokens to support overlap scheduler + if self._m3_next_new_tokens is not None: + # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). + next_new_tokens = self._m3_next_new_tokens + else: + next_new_tokens = self._prepare_next_new_tokens( + accepted_tokens, + next_draft_tokens, + spec_metadata.batch_indices_cuda, + batch_size, + num_accepted_tokens, + ) + + attn_metadata.use_spec_decoding = True + + # Commit model-owned target state only after draft preparation has succeeded; SpecWorkerBase aborts the + # still-pending snapshots on any exception above this point. + if self._auxiliary_state_handlers: + self.commit_auxiliary_speculative_states( + num_accepted_tokens, + attn_metadata.mamba_metadata.state_indices[:batch_size], + num_contexts, + ) + + return { + "logits": raw_logits, + "new_tokens": accepted_tokens, + "new_tokens_lens": num_accepted_tokens, + "next_draft_tokens": next_draft_tokens, + "next_new_tokens": next_new_tokens, + } + + def prepare_1st_drafter_inputs( + self, + input_ids, + position_ids, + hidden_states, + accepted_tokens, + attn_metadata, + spec_metadata, + draft_model, + ): + """The fused first draft inputs on a step that verified with the fused kernels (or was seeded for them), + else the in-tree ones.""" + fused = self._m3_first_drafter_inputs( + position_ids, + accepted_tokens, + input_ids.shape[0], + attn_metadata.num_contexts, + spec_metadata, + draft_model, + ) + self._m3_fused_first_inputs = fused is not None + if fused is not None: + fused.update(attn_metadata=attn_metadata, spec_metadata=spec_metadata) + return fused + return super().prepare_1st_drafter_inputs( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + accepted_tokens=accepted_tokens, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=draft_model, + ) + + def _m3_structural_ok(self, attn_metadata, spec_metadata) -> bool: + """The batch conditions shared by the fused verify and the fused first drafter inputs: a decode-only plain + Eagle3 step without attention DP, guided decoding, relaxed acceptance or a suffix-automaton enhancer, of at + most DECODE_MAX_TOKENS verify rows. Host state only, so the TP ranks agree.""" + if self.is_mtp_eagle or self.use_dynamic_tree: + return False + if self._uses_external_shared_target_kv or self.sa_enhancer is not None: + return False + if ( + self.guided_decoder is not None + or self.mapping is None + or self.mapping.enable_attention_dp + ): + return False + if getattr(self.spec_config, "use_relaxed_acceptance_for_thinking", False): + return False + if attn_metadata.num_contexts != 0 or attn_metadata.num_seqs == 0: + return False + draft_len = spec_metadata.runtime_draft_len + if draft_len == 0 or attn_metadata.num_seqs * (draft_len + 1) > DECODE_MAX_TOKENS: + return False + return not spec_metadata.wants_advanced_draft_sampling + + def tail_eligible(self, attn_metadata, spec_metadata, lm_head: nn.Module) -> bool: + """Whether this step verifies with the fused kernels instead of the gathered target logits. + + Decode-only Eagle3 batches (``_m3_structural_ok``) over an unpadded vocab-sharded bias-free target head, all + greedy at a request count ``SHARD_ACCEPT`` covers, or sampled under the temperature + top-p contract (see + _m3_sampled_verify_ok) at a request count ``SAMPLED_ACCEPT`` covers, with nothing else reading the target + logits. Every input is host state shared by all TP ranks, so the ranks agree. + """ + if not self._m3_structural_ok(attn_metadata, spec_metadata): + return False + if getattr(spec_metadata, "enable_penalty", False): + return False + if not _sharded_head(lm_head) or getattr(lm_head, "bias", None) is not None: + return False + draft_len = spec_metadata.runtime_draft_len + if spec_metadata.is_all_greedy_sample: + key = (draft_len, attn_metadata.num_seqs) + if not SHARD_ACCEPT.covers(key): + uncovered( + SHARD_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + else: + if not self._m3_sampled_verify_ok(spec_metadata): + return False + key = (draft_len, attn_metadata.num_seqs) + if not SAMPLED_ACCEPT.covers(key): + uncovered( + SAMPLED_ACCEPT.path, + f"draft length {draft_len} at {attn_metadata.num_seqs} requests", + ) + return False + self._m3_target_lm_head = lm_head + _log_engaged( + SHARD_ACCEPT.path if spec_metadata.is_all_greedy_sample else SAMPLED_ACCEPT.path, + self.mapping, + ) + return True + + def draft_only_eligible(self, attn_metadata, spec_metadata, draft_model) -> bool: + """Whether a step verified by the in-tree path can still take the fused first drafter inputs: up to + FIRST_DRAFTER_INPUTS_MAX_REQUESTS requests, the structural conditions of tail_eligible, without the + target-sampling ones.""" + if attn_metadata.num_seqs > FIRST_DRAFTER_INPUTS_MAX_REQUESTS: + uncovered( + "m3_first_drafter_inputs", f"in-tree verify at {attn_metadata.num_seqs} requests" + ) + return False + return self._m3_structural_ok(attn_metadata, spec_metadata) and draft_model is not None + + def _m3_seed_first_pass(self, num_accepted_tokens, spec_metadata, batch_size: int) -> None: + """The first draft pass's gather rows (the last accepted row of each request) from the in-tree acceptance, as + _m3_target_accept leaves them.""" + k1 = spec_metadata.runtime_draft_len + 1 + n = num_accepted_tokens[:batch_size].to(torch.int64) + rows = spec_metadata.batch_indices_cuda[:batch_size].to(torch.int64) + self._m3_first_gather_ids = rows * k1 + n - 1 + + def _m3_first_drafter_inputs( + self, + position_ids, + accepted_tokens, + num_tokens: int, + num_contexts: int, + spec_metadata, + draft_model: Optional[M3Eagle3Draft], + ): + """The in-tree first draft inputs of a decode-only step, with the per-layer fc_norm RMSNorms of + ``apply_eagle3_fc`` in one kernel (unless the capture already applied them) before the fc; the input ids are + the accepted tokens as they lie. None to use the in-tree path.""" + if ( + self._m3_first_gather_ids is None + or self.is_mtp_eagle + or num_contexts != 0 + or draft_model is None + ): + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + ): + return None + if not self._m3_first_inputs_covered(num_tokens, spec_metadata): + return None + cache = self._m3_fc_norm_w + if cache is None or cache[0] is not draft_model.w["fc_norm0"]: + if torch.cuda.is_current_stream_capturing(): + capture_missed("spec_tail", None, "the stacked Eagle3 fc_norm weights") + return None + cache = ( + draft_model.w["fc_norm0"], + torch.stack( + [draft_model.w[f"fc_norm{i}"].detach() for i in range(layers)] + ).contiguous(), + ) + self._m3_fc_norm_w = cache + hidden = self._m3_fc_hidden + self._m3_fc_hidden = None + if hidden is None or hidden.shape[0] != num_tokens: + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + _record("m3_chunk_norm", num_tokens) + hs = m3_spec_tail.chunk_norm(hs, cache[1], draft_model.rms_eps) + _log_engaged(CHUNK_NORM.path, self.mapping) + else: + hs = self._m3_renorm_raw_captures(hs, spec_metadata, draft_model) + hidden = draft_model.fc(hs) + input_ids = accepted_tokens[:, : spec_metadata.runtime_draft_len + 1].reshape(-1) + _log_engaged(FUSED_FIRST_INPUTS.path, self.mapping) + return {"input_ids": input_ids, "position_ids": position_ids, "hidden_states": hidden} + + def _m3_first_inputs_covered(self, num_tokens: int, spec_metadata) -> bool: + """Whether the fused first drafter inputs of ``num_tokens`` verify rows are covered: their request count + (``FUSED_FIRST_INPUTS``), and the chunk norm's rows (``CHUNK_NORM``) unless the captures already applied the + fc_norm.""" + k1 = spec_metadata.runtime_draft_len + 1 + if num_tokens % k1 or not FUSED_FIRST_INPUTS.covers(num_tokens // k1): + uncovered( + FUSED_FIRST_INPUTS.path, + f"{num_tokens} verify rows at draft length {k1 - 1} (the in-tree first drafter inputs run)", + ) + return False + if self._m3_capture_norm is None and not CHUNK_NORM.covers(num_tokens): + uncovered( + CHUNK_NORM.path, f"{num_tokens} verify rows (the in-tree first drafter inputs run)" + ) + return False + return True + + def _m3_fc_fork( + self, + spec_metadata, + num_tokens: int, + draft_model: Optional[M3Eagle3Draft], + after: Optional[torch.cuda.Event] = None, + ) -> Optional[torch.cuda.Stream]: + """Launch the first draft pass's fc_norm + fc on a forked stream beside the target head GEMM (after the event + ``after`` when given, else after everything queued on the current stream); the result is picked up by + _m3_first_drafter_inputs. Returns the stream to join, or None when the fc path does not apply (the first + drafter inputs then compute it).""" + self._m3_fc_hidden = None + if draft_model is None: + return None + layers = len(spec_metadata.layers_to_capture) + if ( + draft_model.hidden_size_in != spec_metadata.hidden_size + or layers != draft_model.capture_layers + or self._m3_fc_norm_w is None + or not self._m3_first_inputs_covered(num_tokens, spec_metadata) + ): + return None + aux = self._m3_aux_stream + if aux is None: + aux = self._m3_aux_stream = torch.cuda.Stream() + if after is not None: + aux.wait_event(after) + else: + aux.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(aux): + hs = spec_metadata.hidden_states[:num_tokens, : spec_metadata.hidden_size * layers] + if self._m3_capture_norm is None: + _record("m3_chunk_norm", num_tokens) + hs = m3_spec_tail.chunk_norm( + hs, self._m3_fc_norm_w[1], draft_model.rms_eps, use_pdl=False + ) + _log_engaged(CHUNK_NORM.path, self.mapping) + else: + hs = self._m3_renorm_raw_captures(hs, spec_metadata, draft_model) + self._m3_fc_hidden = draft_model.fc(hs) + self._m3_fc_flagged = self._m3_fc_bump(hs.device) + return aux + + def _m3_renorm_raw_captures( + self, hs: torch.Tensor, spec_metadata, draft_model: M3Eagle3Draft + ) -> torch.Tensor: + """The captured rows ``hs`` [rows, layers * H] for a reader that skips the fc_norm (the capture norm is armed): + unchanged when every captured layer was written with this step's norm (capture_hidden_states records them), + else -- not expected on any path -- a copy with each layer captured raw normalized by its own fc_norm here, + logged as a miss, so a raw row never reaches the drafter's fc.""" + missing = [ + i + for i, layer_id in enumerate(spec_metadata.layers_to_capture) + if layer_id not in self._m3_capture_normed + ] + if not missing: + return hs + missed( + CAPTURE_NORM.path, + int(hs.shape[0]), + f"captured layer(s) {[spec_metadata.layers_to_capture[i] for i in missing]} were written without the " + "capture norm; their fc_norm is applied before the drafter's fc", + ) + hs = hs.clone() + hidden = spec_metadata.hidden_size + for i in missing: + block = hs[:, i * hidden : (i + 1) * hidden] + block.copy_( + flashinfer_rmsnorm(block, draft_model.w[f"fc_norm{i}"], draft_model.rms_eps) + ) + return hs + + def _m3_fc_bump(self, device: torch.device) -> bool: + """DI_FC_FLAG: count the fc just queued on the aux stream (the wait kernel's release). The counters are zeroed + by ``prebuild``; without them the fork keeps the stream join.""" + if not DI_FC_FLAG or self._m3_fc_counts is None or self._m3_fc_counts.device != device: + return False + self._m3_fc_counts[:1].add_(1) + return True + + def _m3_fc_join(self, fork: torch.cuda.Stream) -> None: + """Join the fc fork onto the current stream and keep its output's memory alive for that stream. + + _m3_fc_hidden is allocated in fork (aux) stream order and read on the current stream. Without the record, + dropping it returns its block to the aux stream's pool, where later aux work can take it while the current + stream's reads are still queued (eager warm-up / mixed steps). Nothing is recorded under graph capture, as in + maybe_execute_in_parallel: a captured graph relies on the next fork's aux.wait_stream(main) to order the + reuse.""" + main = torch.cuda.current_stream() + if self._m3_fc_flagged: + # DI_FC_FLAG: the first draft input waits on the fork's counter; the stream join follows that input. + self._m3_fc_flagged = False + self._m3_join_pending() + m3_fc_wait.fc_wait(self._m3_fc_counts) + self._m3_fc_pending = fork + _log_engaged(DI_FC_FLAG_PATH, self.mapping) + else: + main.wait_stream(fork) + hidden = self._m3_fc_hidden + if ( + hidden is not None + and hidden.is_cuda + and not (torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing()) + ): + hidden.record_stream(main) + + def _m3_join_pending(self) -> None: + """Join the fc fork whose join DI_FC_FLAG deferred (its output's readiness already went through the wait + kernel): after the first draft input, at the end of the step, and before any newer fork is waited for.""" + fork = self._m3_fc_pending + if fork is not None: + self._m3_fc_pending = None + torch.cuda.current_stream().wait_stream(fork) + + def arm_capture_norm( + self, attn_metadata, spec_metadata, lm_head: nn.Module, draft_model: Optional[M3Eagle3Draft] + ) -> None: + """Before every target forward: whether this step's Eagle3 captures apply the drafter's fc_norm themselves. + + On a fused-verify step of ``CAPTURE_NORM``'s rows whose first drafter inputs are the fused ones and whose + fc_norm weights are already stacked (by an earlier fused first pass), the capture writes + rmsnorm(h + r) * fc_norm[i] (one kernel per captured layer instead of add + copy, or the layer's boundary + all-reduce itself, ``capture_target``) and the fused first drafter inputs read it without their own + chunk_norm. A fused-verify step of at most ``FC_EARLY_MAX_ROWS`` verify rows is also armed to fork the + drafter's first fc at the last captured layer (``capture_hidden_states``). + """ + if self._m3_fc_early_stream is not None: + # A fork the last step made but did not reach the verify: its wait keeps the counters paired. + self._m3_fc_join(self._m3_fc_early_stream) + self._m3_join_pending() + self._m3_capture_norm = None + self._m3_fc_early_drafter = None + self._m3_fc_early_stream = None + # The captured layers this step wrote with its capture norm (capture_hidden_states). + self._m3_capture_normed = set() + # A fork result no first draft pass consumed must not reach a later step's (same-size) first inputs. + self._m3_fc_hidden = None + if draft_model is None: + return + if not self.tail_eligible(attn_metadata, spec_metadata, lm_head): + return + rows = attn_metadata.num_seqs * (spec_metadata.runtime_draft_len + 1) + # The fork reads every captured column at the last captured layer, so a final-state capture, which is written + # after the model returns, keeps the fc at its stock place. + if ( + rows <= FC_EARLY_MAX_ROWS + and FINAL_HIDDEN_LAYER_ID not in spec_metadata.layers_to_capture + ): + self._m3_fc_early_drafter = draft_model + if not CAPTURE_NORM.covers(rows): + uncovered(CAPTURE_NORM.path, f"{rows} verify rows (the in-tree capture runs)") + return + # The normed captures are read only by the fused first drafter inputs; the in-tree ones norm again. + if not FUSED_FIRST_INPUTS.covers(attn_metadata.num_seqs): + return + fc_norm_w = self._m3_fc_norm_w + if ( + fc_norm_w is not None + and fc_norm_w[0] is draft_model.w["fc_norm0"] + and len(spec_metadata.layers_to_capture) == fc_norm_w[1].shape[0] + ): + self._m3_capture_norm = (fc_norm_w[1], draft_model.rms_eps) + _log_engaged(CAPTURE_NORM.path, self.mapping) + + @staticmethod + def _m3_sampled_verify_ok(spec_metadata) -> bool: + """Whether sampled verify rows can use the fused sampled op: only under advanced_sampling_mode=no_topk, the + deploy-time guarantee that no row uses top-k or min-p (that op applies neither).""" + mode = getattr(spec_metadata, "advanced_sampling_mode", None) + return mode is not None and mode.skips_top_k and not mode.skips_top_p + + def _m3_target_accept( + self, + hidden_states, + spec_metadata, + batch_size: int, + draft_len: int, + draft_model: Optional[M3Eagle3Draft] = None, + kv_lens: Optional[torch.Tensor] = None, + ): + """Fused verification (see tail_eligible); same outputs as the strict acceptance path. ``kv_lens``: the + requests' KV lengths at this verify (the per-request forced-acceptance key, FORCE_AL_PER_REQUEST).""" + force_total = 0 + force_vec = None + if self.force_num_accepted_tokens != 0.0: + is_warmup = spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + # The draw pool is built by an H2D copy and the draw kernels JIT-compile per block width, so both + # happen here, before the first captured step (a scratch counter keeps the real draw sequence). + self._ensure_force_accept_rng_state(hidden_states.device) + if not self._m3_force_vec_warm: + scratch = torch.zeros(1, dtype=torch.int64, device=hidden_states.device) + dummy = torch.empty(64, dtype=torch.int32, device=hidden_states.device) + for block in (16, 32, 64): + m3_force_vec( + self._force_accept_rng_pool, + scratch, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + m3_force_vec_req( + self._force_accept_rng_pool, + dummy, + dummy, + block, + 1, + 2, + 0.5, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + self._m3_force_vec_warm = True + if not is_warmup: + force_total = min(int(self.force_num_accepted_tokens) + 1, draft_len + 1) + force_vec = self._m3_force_vec( + batch_size, force_total, draft_len + 1, hidden_states.device, kv_lens + ) + draft_tokens = spec_metadata.draft_tokens.view(batch_size, draft_len) + if draft_tokens.dtype != torch.int32: + draft_tokens = draft_tokens.to(torch.int32) + # Column 0 receives the last accepted token, the next step's first input (the fused draft loop adds the rest). + self._m3_next_new = torch.empty( + (batch_size, draft_len + 1), dtype=torch.int32, device=hidden_states.device + ) + rows = hidden_states[: batch_size * (draft_len + 1)] + lm_head = self._m3_target_lm_head + batch_idx = spec_metadata.batch_indices_cuda[:batch_size] + # The Eagle3 fc on the aux stream: the fork made at the last captured layer, else one beside the head. A + # sampled verify's head (the in-tree GEMM up to FC_AFTER_HEAD_MAX_ROWS verify rows) holds + # every SM while it streams: its fc is forked after it instead (overlapping the sampled acceptance), where its + # many split-K CTAs do not delay the head's start. Rubin TP2 c1: the head started 17.4 us after its input was + # ready with the fc beside it. + fork = self._m3_fc_early_stream + self._m3_fc_early_stream = None + head_first = ( + fork is None + and rows.shape[0] <= FC_AFTER_HEAD_MAX_ROWS + and not spec_metadata.is_all_greedy_sample + ) + if fork is None and not head_first: + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model) + # This rank's vocabulary shard of the target logits: the head's own local GEMM, whose all-gather the fused + # kernels replace (they exchange per-row records across ranks, not logits). + logits = lm_head.apply_linear(rows, None) + head_done = None + if head_first: + head_done = torch.cuda.Event() + head_done.record() + vocab_offset = self.mapping.tp_rank * lm_head.weight.shape[0] + if spec_metadata.is_all_greedy_sample: + result = m3_shard_accept_kernel.shard_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = SHARD_ACCEPT.path + else: + n_rows = rows.shape[0] + seeds, offsets = self._rng_state_per_token(spec_metadata, n_rows) + result = m3_sampled_accept_kernel.sampled_accept( + logits, + vocab_offset, + self.mapping, + draft_tokens, + batch_idx, + spec_metadata.temperatures[:n_rows], + spec_metadata.top_ks[:n_rows], + spec_metadata.top_ps[:n_rows], + None, + seeds, + offsets, + force_total, + self._m3_next_new, + force_vec=force_vec, + ) + path = "spec_tail" + if head_done is not None: + # Enqueued after the sampled acceptance (which starts first) but waits on the head only. + fork = self._m3_fc_fork(spec_metadata, rows.shape[0], draft_model, after=head_done) + if fork is not None: + self._m3_fc_join(fork) + if result is None: + if torch.cuda.is_current_stream_capturing(): + capture_missed(path, None, "the fused verify's exchange workspace") + raise RuntimeError( + "M3 fused verify unavailable after tail_eligible() accepted the step" + ) + return result + + def _m3_force_vec( + self, + batch_size: int, + base_total: int, + max_total: int, + device, + kv_lens: Optional[torch.Tensor] = None, + ) -> Optional[torch.Tensor]: + """Per-request forced accepted counts for a fractional forced acceptance, else None. + + With ``kv_lens`` (FORCE_AL_PER_REQUEST): each request draws from the fixed-seed pool at its own key, its KV + length at this verify (its prompt plus what it has accepted so far), so its draws do not depend on the + server's history or on the other requests; every TP rank holds the same lengths, so the ranks agree. Without + lengths: exactly as _apply_force_accepted_tokens draws for a decode-only batch (one shared counter bump per + step, the (counter, slot) hash), the generic path's sequence. + """ + frac = self.force_num_accepted_tokens - int(self.force_num_accepted_tokens) + if frac <= 0.0 or base_total >= max_total: + return None + self._ensure_force_accept_rng_state(device) + out = torch.empty(batch_size, dtype=torch.int32, device=device) + block = max(16, 1 << (batch_size - 1).bit_length()) + if FORCE_AL_PER_REQUEST and kv_lens is not None: + keys = kv_lens[:batch_size] + if keys.dtype != torch.int32: + keys = keys.to(torch.int32) + return m3_force_vec_req( + self._force_accept_rng_pool, + keys, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + # One launch: bump the shared counter and draw, with _apply_force_accepted_tokens' exact arithmetic. + return m3_force_vec( + self._force_accept_rng_pool, + self._force_accept_rng_counter, + out, + batch_size, + base_total, + max_total, + float(frac), + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + block, + ) + + # ---------------------------------------------------------------------------------------------- fused draft loop + + def prebuild_draft_head(self, draft_model: Optional[M3Eagle3Draft]) -> bool: + """Build the copy of the drafter's BF16 head shard that the fused draft loop streams (rebuilt only when the + head weight tensor or its version changed; a replaced copy is kept alive), on every rank. False when the + drafter is not one the loop serves.""" + old = self._m3_draft_head + if draft_model is None or not self._m3_drafter_supported(draft_model): + self._m3_draft_head = None + return False + weight = draft_model.lm_head.weight + # Idempotent: the same weight tensor at the same version keeps the built copy, so a repeated + # cache_derived_state never re-places it under a captured CUDA graph. + if ( + old is not None + and old.weight is weight + and self._m3_draft_head_version == weight._version + ): + return True + if old is not None: + # A rebuilt copy (new weight values): keep the old one alive, since a graph captured earlier may still + # read its (placed) memory until the graphs are recaptured. + self._m3_draft_head_graveyard.append(old) + self._m3_draft_head = m3_draft_step_kernel.prepare_head( + draft_model.lm_head, self.mapping.tp_rank + ) + self._m3_draft_head_version = weight._version + return True + + def _m3_drafter_supported(self, draft_model: M3Eagle3Draft) -> bool: + """The drafter the fused draft loop runs (``M3Eagle3Draft`` validates its shape at construction) over this + worker's tensor-parallel group with a bias-free vocab-sharded BF16 head without padding.""" + if self.mapping is None or self.mapping.tp_size <= 1 or self.mapping.enable_attention_dp: + return False + if self._d2t is not None: + return False + lm_head = draft_model.lm_head + return _sharded_head(lm_head) and m3_draft_step_kernel.head_supported(lm_head) + + def _m3_draft_ar_norm( + self, + draft_model: M3Eagle3Draft, + x: torch.Tensor, + residual: torch.Tensor, + norm_key: str, + carveout: bool = True, + pushed: bool = False, + ): + """One drafter boundary of the fused draft loop: the multicast all-reduce + residual + RMSNorm with the + drafter's stock rounding (``DRAFT_MC_AR``: opts bit 16, normalizing the unrounded FP32 sum as FlashInfer + fused_add_rmsnorm does; Rubin TP2 10.0..11.0 -> 4.6..5.3 us per boundary), else the in-tree structure + (``M3Eagle3Draft.ar_norm``). ``carveout`` False drops the reduce's max-smem carveout (``DRAFT_MC_AR_CARVEOUT``). + ``pushed``: the boundary's projection already pushed its rows into the exchange (prepushed 2). The choice + depends on shapes and a collectively created workspace only, so the TP ranks agree. Returns (normed, residual + sum).""" + rows = x.shape[0] + path = DRAFT_MC_AR.path + if DRAFT_MC_AR.covers(rows): + workspace = m3_mc_allreduce.get_workspace(self.mapping, x.shape[-1]) + out = None + if workspace is not None: + _record("m3_draft_mc_ar", rows, pushed) + out = m3_mc_allreduce.m3_mc_allreduce_norm( + x.contiguous(), + residual.contiguous(), + draft_model.w[norm_key], + draft_model.rms_eps, + workspace, + use_pdl=get_env_enable_pdl(), + drafter_norm=True, + carveout=carveout, + prepushed=2 if pushed else 0, + ) + if out is not None: + _log_engaged(path, self.mapping) + if pushed: + taken(DRAFT_AR_PUSH_PATH, rows, "the projection pushed its rows (prepushed 2)") + elif DRAFT_AR_PUSH: + uncovered( + DRAFT_AR_PUSH_PATH, + "the boundary's projection did not push (the AR pushes x)", + ) + return out + if pushed: + raise RuntimeError( + "the drafter's rows were pushed into a boundary all-reduce that did not run" + ) + missed(path, rows, "the multicast all-reduce workspace or operand layout") + return draft_model.ar_norm(x, residual, norm_key) + + def _m3_arm_draft_push(self, draft_model: M3Eagle3Draft, key: str, rows: int) -> None: + """DRAFT_AR_PUSH: let projection ``key``'s next low-M call push its rows into the boundary all-reduce that + ``_m3_draft_ar_norm`` then runs (the multicast path only; shapes and a collectively created workspace decide + it, so the TP ranks agree).""" + if not (DRAFT_AR_PUSH and DRAFT_MC_AR.covers(rows)): + return + workspace = m3_mc_allreduce.get_workspace(self.mapping, draft_model.w[key].shape[0]) + if workspace is not None: + draft_model.arm_push(key, workspace, rows) + + def _m3_fused_loop_ok( + self, attn_metadata, spec_metadata, batch_size: int, draft_model: Optional[M3Eagle3Draft] + ) -> bool: + """Whether this step drafts with the fused loop: a decode-only step the fused verify or the seed prepared + (``_m3_first_gather_ids``) with the fused first draft inputs, a supported drafter, a (draft length, requests) + ``DRAFT_STEP`` covers, the prebuilt head and the exchange workspace. Host state only, so the TP ranks + agree.""" + if ( + self._m3_first_gather_ids is None + or not self._m3_fused_first_inputs + or draft_model is None + ): + return False + path = DRAFT_STEP.path + if not self._m3_drafter_supported(draft_model) or not hasattr( + attn_metadata, "kv_lens_cuda" + ): + uncovered( + path, "the drafter or its attention metadata is not one the fused loop serves" + ) + return False + draft_len = spec_metadata.runtime_draft_len + if not DRAFT_STEP.covers((draft_len, batch_size)): + uncovered(path, f"draft length {draft_len} at {batch_size} requests") + return False + head = self._m3_draft_head + if head is None or head.weight is not draft_model.lm_head.weight: + missed( + path, + batch_size, + "the drafter head's streamed copy is not built for its current weight (prebuild)", + ) + return False + if m3_draft_step_kernel.workspace(self.mapping, head.wf.device) is None: + capture_missed(path, None, "the draft-step exchange workspace") + return False + _log_engaged(path, self.mapping) + return True + + def _forward_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ): + """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop.""" + if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): + return self._m3_fused_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + batch_size, + num_accepted_tokens, + ) + return super()._forward_draft_loop( + inputs, + attn_metadata, + spec_metadata, + draft_model, + draft_kv_cache_manager, + num_contexts, + num_gens, + batch_size, + num_accepted_tokens, + original_all_rank_num_tokens, + resource_manager, + ) + + def _m3_fused_draft_loop( + self, + inputs, + attn_metadata, + spec_metadata, + draft_model: M3Eagle3Draft, + draft_kv_cache_manager, + batch_size: int, + num_accepted_tokens, + ): + """Mirrors Eagle3OneModelWorker._forward_linear_draft_loop for a decode-only Eagle3 step, with each pass's + glue fused (see ``_m3_fused_loop_ok``): + + * the draft input -- the embedding lookup, the layer's input norm and hidden norm and their concat -- in one + launch (``m3_spec_tail.draft_input``; the second pass gathers the last accepted rows there); + * the drafter layer's attention and MLP (its SwiGLU on this target's C++ op), each followed by the all-reduce + with the residual add and the next RMSNorm (the post-attention norm, then the drafter's final norm): the + multicast boundary fed by the o_proj / down projection's push where covered (``_m3_draft_ar_norm``); + * one kernel for the BF16 head over the head rows (pass 0: the last accepted row of each request), the + cross-rank argmax and the bookkeeping the in-tree loop does between passes: the draft token, the next pass's + input id and position, the draft KV lengths (and pass 0's sequence lengths) (``m3_draft_step``). + + The host metadata transitions are the in-tree loop's. The drafts land in columns 1.. of the step's next-input + tokens (column 0 is the last accepted token), which ``_forward_impl`` returns as they lie. + """ + draft_len = spec_metadata.runtime_draft_len + table = draft_model.embed_table + dev = table.device + head = self._m3_draft_head + gather_ids = self._m3_first_gather_ids + ids = inputs["input_ids"] + hidden = inputs["hidden_states"] + pos = inputs["position_ids"] + if ids.dtype != torch.int32: + ids = ids.to(torch.int32) + if pos.dtype != torch.int32: + pos = pos.to(torch.int32) + num_acc = num_accepted_tokens[:batch_size] + if num_acc.dtype != torch.int32: + num_acc = num_acc.to(torch.int32) + next_new = self._m3_next_new + if next_new is None: + # Seeded from the in-tree acceptance: the last accepted token is the first-pass input at the gather row. + next_new = torch.empty((batch_size, draft_len + 1), dtype=torch.int32, device=dev) + next_new[:, 0] = ids[gather_ids] + draft_tokens = next_new[:, 1:] + w = draft_model.w + eps = draft_model.rms_eps + + with self.draft_kv_cache_context(attn_metadata, draft_kv_cache_manager) as md: + for i in range(draft_len): + x, residual = m3_spec_tail.draft_input( + ids, + table, + hidden, + gather_ids if i == 1 else None, + w["in_norm"], + eps, + w["hidden_norm"], + eps, + ) + if i == 0: + self._m3_join_pending() + rows = x.shape[0] + self._m3_arm_draft_push(draft_model, "o", rows) + h = draft_model.attention(x, md, reduce=False) + h, residual = self._m3_draft_ar_norm( + draft_model, h, residual, "post_norm", pushed=draft_model.take_pushed(h) + ) + self._m3_arm_draft_push(draft_model, "dn", rows) + h = draft_model.mlp(h, reduce=False, cpp_silu=True) + h, to_save = self._m3_draft_ar_norm( + draft_model, + h, + residual, + "final_norm", + carveout=DRAFT_MC_AR_CARVEOUT, + pushed=draft_model.take_pushed(h), + ) + if draft_model.return_hidden_post_norm: + to_save = h + ids_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + pos_next = torch.empty(batch_size, dtype=torch.int32, device=dev) + if not m3_draft_step_kernel.draft_step( + head, + h, + batch_size, + gather_ids if i == 0 else None, + self.mapping, + draft_tokens, + i, + ids_next, + pos, + pos_next, + md.kv_lens_cuda, + num_acc if i == 0 else None, + md._seq_lens_cuda if i == 0 else None, + draft_len, + ): + raise RuntimeError( + "M3 fused draft step unavailable after _m3_fused_loop_ok() accepted the step" + ) + ids, pos, hidden = ids_next, pos_next, to_save + if i == 0: + # The draft step set the device sequence lengths to 1 and rolled the KV lengths back. + md._seq_lens[:batch_size].fill_(1) + md.on_update() + has_kv_cache = md.kv_cache_manager is not None + if has_kv_cache and hasattr(md, "host_request_types"): + md.host_request_types[: md.num_contexts].fill_(1) + if has_kv_cache: + md.num_contexts = 0 + self._prepare_flash_mla_generation_layout(md, 0, batch_size) + md.update_for_spec_dec() + md.use_spec_decoding = False + else: + md.update_for_spec_dec() + self._m3_next_new_tokens = next_new + return draft_tokens + + +def _build_speculation(shell, model_config: ModelConfig): + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise ValueError( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=_get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + _set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = should_use_separate_draft_kv_cache(spec) + draft = M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return draft_config, draft, worker, use_separate + + +# -------------------------------------------------------------------------------------------------------------------- +# Runtime hooks and the shell +# -------------------------------------------------------------------------------------------------------------------- + + +def _runtime_hooks(spec_fold: Optional[M3SpecStepFold]) -> RuntimeHooks: + """The step-loop hooks this target opts into. + + * ``pack_step_inputs``: a decode step's small input copies go out as one upload and one scatter at the graph + replay instead of a chain of copies the next graph waits on; the scatter is parallel (a program per copy), issued + by the decode graph's own prologue (this target's one-launch C++ op, ``m3_step_prologue``) from a window copied + on a side stream while the previous graph runs, so no operation is left between the step graphs and the replay + does not start with a host copy (its copy-engine -> compute turnaround was 10.3 us at c1). Sampled speculative + batches stage their Philox seed/offset uploads, and the accepted-draft counts, into the same pack; the sampler's + slot ids stay in a device buffer re-uploaded only when they change. + * ``pack_step_epoch``: the step epoch (step_inputs.step_epoch), one int64 device counter the decode graph's + prologue increments once per replay (program 0 of its first scatter launch, no extra kernel) and never resets. + The per-layer megakernel pieces tag their per-step device flags with it (contract v1: tag = epoch << 8 | layer). + * ``skip_unchanged_h2d``: steady decode re-uploads slot indices, query lengths and page tables that rarely change. + * ``h2d_changed_span``: a changed packed 2-D metadata upload stages only each row's changed column span. On a step + where a request crosses into a new page the MSA block table changed by one entry, but the whole padded table + (32 KB a request) was staged, overflowing the 16 KB in-graph window into an eager upload and scatter between the + step graphs (TP4 57-69 us at the ~1 in 19 block-crossing steps). + * ``kv_lens_update_once``: a speculative overlap step re-derives the MSA KV-length metadata once, after the + forward's KV-length correction (the first pass is rewritten in full by the second). + * ``spec_offsets_op``: the overlap scheduler's speculative step correction (position_ids += offsets, kv_lens += + offsets) as this target's one plain launch (``m3_spec_offsets``) instead of two ATen adds at the head of the + decode graph: the same int32 adds. + * ``loop_on_execution_stream``: the overlap loop's own stream work and the forward on one stream, so step N+1's + inputs wait on graph N in program order rather than through two cross-stream event hops. + * ``merge_tuned_tactics_across_ranks``: tensor-parallel peers capture the same GEMM tactics (rank-local profiling + noise otherwise splits them, and the slowest rank sets every step). + * ``skip_greedy_rng_upload``: all-greedy speculative batches never read the Philox seeds. + * ``spec_step_fold``: the spec sampler's slot stores and the engine's next-step input gathers; the decode forward + stores its outputs and stages the next step's inputs at its own tail (``M3SpecStepFold``). + * ``sampler_d2h_side_stream``: the spec sampler copies its per-slot store to the host on a side stream, overlapping + the next step's input staging. + * ``iteration_timing_one_event``: the iteration log times the device step with one event per boundary on a side + stream (the reported step times are the same). + * MSA: identity selections for short sequences (and short/long graph families), the one-launch KV-length update, + unchanged block-table copies skipped, pool views built once, block-id tables copied from the page-index views + (the stock builder makes a Python list of the whole padded row: on the Vera CPU 282 us per call at one request + and 561 us at two, against 8 and 13 us). + """ + return RuntimeHooks( + first_content_emit_iters=FIRST_CONTENT_EMIT_ITERS, + mxfp8_pinned_tactics=True, + late_compile_detect=True, + pack_step_inputs=True, + pack_parallel_scatter=True, + pack_in_graph=True, + pack_early_copy=True, + pack_step_epoch=True, + pack_prologue_op=m3_step_prologue, + skip_unchanged_h2d=True, + kv_lens_update_once=True, + spec_offsets_op=m3_spec_offsets, + h2d_changed_span=True, + loop_on_execution_stream=True, + merge_tuned_tactics_across_ranks=True, + skip_greedy_rng_upload=True, + pack_spec_rng_uploads=True, + spec_small_h2d=True, + spec_step_fold=spec_fold, + sampler_d2h_side_stream=True, + iteration_timing_one_event=True, + msa=MsaRuntimePolicy( + short_sequence_identity=True, + kv_update=_msa_kv_update_kernel.msa_kv_update, + skip_unchanged_block_offsets=True, + reuse_pool_views=True, + block_ids_from_views=True, + ), + ) + + +class ModelingV2MiniMaxM3Nvfp4Sm107Tp8( + DecoderModelForCausalLM[MiniMaxM3Tp8ModelingV2, PretrainedConfig] +): + """MiniMax-M3 NVFP4 / sm_107 / tp8, text only, with one-model Eagle3 or no speculation. + + The shell owns the vocab-parallel BF16 ``lm_head`` and the logits processor (``DecoderModelForCausalLM``); with + speculation, the drafter (``M3Eagle3Draft``, its weights from the speculative checkpoint through + ``load_draft_weights``) and the worker (``M3Eagle3Worker``) that verifies and drafts each step. + """ + + def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None: + assert torch.cuda.get_device_capability() == _SM, ( + f"this target is built for sm_{_SM[0]}{_SM[1]}, the device is sm_" + f"{''.join(map(str, torch.cuda.get_device_capability()))}" + ) + # Under modeling_v2='require' a covered decode shape that cannot take its fast path raises. + set_require(resolved_mode(type(self)) is ModelingV2Mode.REQUIRE) + model_config = _text_model_config(model_config) + cfg = model_config.pretrained_config + spec_config = getattr(model_config, "spec_config", None) + if model_config.quant_config.kv_cache_quant_algo == QuantAlgo.NVFP4: + # The one-model Eagle3 layer has no NVFP4 trtllm-gen decode kernel: keep its modules and the shared + # draft cache on FP8 while the target stays NVFP4. + model_config.extra_attrs["draft_kv_cache_quant_algo_override"] = QuantAlgo.FP8 + super().__init__( + MiniMaxM3Tp8ModelingV2(model_config), + config=model_config, + hidden_size=cfg.hidden_size, + vocab_size=cfg.vocab_size, + ) + self.spec_config = spec_config + self.draft_model: Optional[M3Eagle3Draft] = None + self.draft_config: Optional[ModelConfig] = None + self.spec_worker: Optional[M3Eagle3Worker] = None + self.use_separate_draft_kv_cache = False + self.layer_idx = FINAL_HIDDEN_LAYER_ID + self._spec_fold: Optional[M3SpecStepFold] = None + if spec_config is not None: + ( + self.draft_config, + self.draft_model, + self.spec_worker, + self.use_separate_draft_kv_cache, + ) = _build_speculation(self, model_config) + self._spec_fold = M3SpecStepFold() + self.runtime_hooks = _runtime_hooks(self._spec_fold) + + def load_weights(self, weights: Dict, *args, **kwargs) -> None: + _weights.MODEL_WEIGHTS.load(self, weights) + + def load_draft_weights(self, weights: Dict, weight_mapper=None) -> None: + """The drafter's checkpoint (``speculative_model``), loaded by the engine after the target's.""" + _weights.DRAFT_WEIGHTS.load(self.draft_model, weights) + + def set_guided_decoder(self, guided_decoder) -> bool: + """Mirrors ``SpecDecOneEngineForCausalLM.set_guided_decoder``: guided decoding with speculation is installed + on the worker.""" + if hasattr(self.spec_worker, "set_guided_decoder"): + return self.spec_worker.set_guided_decoder(guided_decoder) + return False + + def cache_derived_state(self) -> None: + """Build the targets and every decode kernel once the weights are real, before the executor's warmup; with + speculation also the drafter's full-context decode, the fused verify's coded target head, the die-placed + target head copy, the fused draft loop's head copy and the drafter linears' low-M kernels.""" + super().cache_derived_state() + self.model.post_load_weights() + if self.draft_model is not None and DRAFTER_FULLCTX.keys: + path = DRAFTER_FULLCTX.path + recipes = {k: FULLCTX_DECODE_RECIPES[k] for k in DRAFTER_FULLCTX.keys} + installed = m3_drafter_attn.install( + self.draft_model, + recipes, + FULLCTX_DECODE_MAX_KV_LEN, + FULLCTX_DECODE_OWN_REDUCER and m3_fullctx_decode_kernel.has_own_reducer(), + get_env_enable_pdl(), + fused_rope_mode=(DRAFTER_ROPE_MODE if DRAFTER_ROPE_FUSED else None), + on_rope_taken=lambda rows, what: taken(DRAFTER_ROPE_PATH, rows, what), + on_taken=lambda rows: taken(path, rows), + on_uncovered=lambda reason: uncovered(path, reason), + on_missed=lambda rows, reason: missed(path, rows, reason), + ) + logger.info(f"[M3 prebuild] {path} installed on {installed} drafter attention layer(s)") + if self.spec_worker is not None: + self.spec_worker.prebuild(self.lm_head, self.draft_model) + + def post_load_weights(self) -> None: + super().post_load_weights() + self.model.post_load_weights() + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + return_context_logits: bool = False, + spec_metadata: Optional[SpecMetadata] = None, + resource_manager: Optional[Any] = None, + **kwargs, + ) -> Any: + """One engine step: the target forward, then the logits (or, with speculation, the worker's verify and + draft, and on a decode step the spec fold's store of its outputs). Text only.""" + multimodal_params = kwargs.pop("multimodal_params", None) or [] + if any(p is not None and getattr(p, "multimodal_data", None) for p in multimodal_params): + raise ValueError( + "this MiniMax-M3 target is text only; the request carries multimodal input" + ) + worker = self.spec_worker + capture = _plain_capture + capture_target = None + if worker is not None: + capture = worker.capture_hidden_states + capture_target = worker.capture_target + if spec_metadata is not None: + worker.arm_capture_norm( + attn_metadata, spec_metadata, self.lm_head, self.draft_model + ) + hidden_states = self.model( + attn_metadata=attn_metadata, + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + spec_metadata=spec_metadata, + capture=capture, + capture_target=capture_target, + **kwargs, + ) + if getattr(attn_metadata, "csa2_remote_tail_mode", None) == "source": + # A context worker stopped at the encoder/decoder boundary: a disposable token for the context-only + # completion, and no draft worker. + return self.logits_processor.forward(hidden_states, self.lm_head, attn_metadata, False) + if spec_metadata is not None and spec_metadata.is_layer_capture(FINAL_HIDDEN_LAYER_ID): + capture(spec_metadata, FINAL_HIDDEN_LAYER_ID, hidden_states, None) + if attn_metadata.padded_num_tokens is not None: + hidden_states = hidden_states[: attn_metadata.num_tokens] + if worker is None: + return self.logits_processor.forward( + hidden_states, self.lm_head, attn_metadata, return_context_logits + ) + # The fused verify computes this rank's logits shard itself and needs no gathered logits. + if worker.tail_eligible(attn_metadata, spec_metadata, self.lm_head): + logits = None + else: + logits = self.logits_processor.forward( + hidden_states[spec_metadata.gather_ids], self.lm_head, attn_metadata, True + ) + spec_input_ids = input_ids if input_ids is not None else kwargs.get("orig_input_ids") + spec_position_ids = position_ids + if attn_metadata.padded_num_tokens is not None: + if spec_input_ids is not None: + spec_input_ids = spec_input_ids[: attn_metadata.num_tokens] + if position_ids is not None: + spec_position_ids = _slice_spec_position_ids(position_ids, attn_metadata.num_tokens) + outputs = worker( + input_ids=spec_input_ids, + position_ids=spec_position_ids, + hidden_states=hidden_states, + logits=logits, + attn_metadata=attn_metadata, + spec_metadata=spec_metadata, + draft_model=self.draft_model, + resource_manager=resource_manager, + ) + spec_fold = self._spec_fold + if ( + spec_fold is not None + and spec_fold.in_forward + and attn_metadata.num_contexts == 0 + and isinstance(outputs, dict) + ): + # The decode forward stores its outputs into the spec sampler's store and stages the next step's inputs + # at its tail (one launch inside the step graph). + spec_fold.store_in_forward(outputs) + return outputs diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/__init__.py new file mode 100644 index 000000000000..1059445198db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ops used only by this target. Shared ops stay in the catalog.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/__init__.py new file mode 100644 index 000000000000..b8c7426e3933 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Activation ops of this target: the gated MLPs' fused SwiGLU.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py new file mode 100644 index 000000000000..e7ceefc5bdab --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SiluAndMul(OpWrapper): + """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with + ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain + SwiGLU ``silu(gate) * up``. + + The MXFP8 dense MLPs (layers 0-2, intermediate 6144 per rank) and the MoE's shared expert (intermediate 1536 per + rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 9216 per + rank) runs it plain. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=384, oai=True)), + Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=384, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=1536, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=1536, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=2304, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=2304, oai=False), + ), + ) + + note = """ + The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two + differ only by the fast sigmoid's last bits (the test's band). + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def reference( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> torch.Tensor: + gate, up = x.float().chunk(2, dim=-1) + if limit is not None: + gate = gate.clamp(max=limit) + up = up.clamp(min=-limit, max=limit) + a = 1.0 if alpha is None else alpha + b = 0.0 if beta is None else beta + return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) + + def is_valid( + self, + x: torch.Tensor, + limit: Optional[float], + alpha: Optional[float], + beta: Optional[float], + ) -> None: + assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" + + +m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/__init__.py new file mode 100644 index 000000000000..a7f7fee7a234 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Attention ops of this target: MSA block selection, the sparse decode and the KV-length update.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py new file mode 100644 index 000000000000..789ae79db88e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per +request count (modeling's ``DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 + minimax_m3_trtllm_gen_dense_decode, +) + +__all__ = ["m3_dense_decode"] + +#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. +HEADS = 16 +KV_HEADS = 1 +HEAD_DIM = 128 +#: The split this target hands trtllm-gen at one request (modeling's DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model +#: maximum the in-tree step passes otherwise. +CAPPED = 16384 +MODEL_MAX = 1 << 20 + + +def dense_attention_reference( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first + ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's + tokens in order). Returns fp32 [total_q, heads, head_dim].""" + batch, kv_heads, max_len, _ = keys.shape + heads = q.shape[1] + group = heads // kv_heads + out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) + pos = torch.arange(max_len, device=q.device) + with true_fp32_matmul(): + for b in range(batch): + full = int(seq_lens[b]) + rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] + limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) + valid = pos[None, :] < limit[:, None] # [qlen, max_len] + for h in range(kv_heads): + qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] + logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale + logits = logits.masked_fill(~valid[:, None, :], float("-inf")) + probs = torch.softmax(logits, dim=-1) + out[ + b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group + ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) + return out + + +def gather_paged_kv( + kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool + (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + page = pool.shape[2] + pages = (int(seq_lens.max()) + page - 1) // page + slots = block_table[:, :pages].long() + k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] + v = pool[slots * subpages + 1] + batch, _, kv_heads, _, d = k.shape + + def order(t): + return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) + + return order(k), order(v) + + +class M3DenseDecode(OpWrapper): + """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len + ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. + + Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into + ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: + every CTA still covers its share of each request's actual length, which may exceed it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", + spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", + spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", + spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), + ), + Cell( + why="two requests keep the model maximum (no cap at two requests)", + spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), + ), + ) + + note = """ + The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped + one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV + band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + minimax_m3_trtllm_gen_dense_decode( + q, + kv_cache_manager, + layer_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + max_seq_len=max_seq_len, + max_num_requests=max_num_requests, + ) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). + + ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + if keys.dtype == torch.float8_e4m3fn: + q = fp8_e4m3_saturate(q) + return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + max_seq_len: int, + max_num_requests: int, + ) -> None: + # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. + assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" + # Each verify row attends a causal prefix of at least one key. + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + + +m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py new file mode 100644 index 000000000000..c2b23aa18750 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention on the TRTLLM backend: RoPE, KV-cache append and causal attention.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class M3DraftAttention(OpWrapper): + """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the + backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's + NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` + ``num_tokens``); returns [num_tokens, q heads x 128] bf16. + + ``backend`` (``TrtllmAttention``) and ``metadata`` are the runtime's objects for the drafter layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="first draft pass of a 2-request verify (2 x 4 rows)", spec=dict(reqs=2, qlen=4)), + Cell(why="a later draft pass, one row per request", spec=dict(reqs=2, qlen=1)), + Cell(why="a context chunk's draft rows", spec=dict(reqs=1, context=512)), + ) + + note = """ + `reference` is exact (float64) causal softmax attention: Q rotated by the backend's NeoX RoPE at each token's + position, over the drafter layer's K / V rows read back from its FP8 cache (dequantized by the backend's + ``kv_scale_quant_orig``). What remains between the two is the FMHA's own arithmetic (Q's rounding to the + kernel's input precision, its P rounding and accumulation order), which the test bounds. + """ + + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors grouped KV heads (query head h reads KV head h // (heads / kv heads); two per rank at tp2) and + un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + # [kv heads, tokens, dim]: each KV head's rows in token order across the pages. + keys = ( + pool[pages, 0] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + values = ( + pool[pages, 1] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + keys = keys.repeat_interleave(group, dim=0) # [heads, tokens, dim] + values = values.repeat_interleave(group, dim=0) + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] + rows = q[first : first + n] + x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,hkd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.einsum( + "thk,hkd->thd", torch.softmax(logits, dim=-1), values + ) + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) + + +m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..2fecad0e9afa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q +out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3Fp8IndexerQkNormRope(OpWrapper): + """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and + writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, + then RoPE on its leading ``rotary_dim`` channels (NeoX). + + The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; + the fused layers' producer writes the index cache itself. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, two index heads per rank", spec=dict(m=1, heads=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, heads=2)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=2)), + ) + + note = """ + `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 + before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of + values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. + """ + + def __call__( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( + qk, + index_k_cache, + out_cache_loc, + num_heads_q, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + position_ids, + ) + + def reference( + self, + qk: torch.Tensor, + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + num_heads_q: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + base: float, + position_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, + then the E4M3 cast.""" + out = fused_qk_norm_rope( + qk.clone(), + num_heads_q, + 1, + 0, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + True, + position_ids, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) + return ( + q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), + k.to(torch.float8_e4m3fn), + ) + + +m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py new file mode 100644 index 000000000000..88ce9728f40b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' full-context decode attention at TP8 on Rubin (``ops/m3_fullctx_decode``): PrimTS split-KV +decode at the recipe modeling's ``FULLCTX_DECODE_RECIPES`` pins per (query rows per request, requests).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) + +from .. import m3_fullctx_decode as kernel +from .m3_dense_decode import dense_attention_reference, gather_paged_kv + +__all__ = ["m3_fullctx_decode"] + +HEADS = kernel.HEADS +KV_HEADS = kernel.KV_HEADS +HEAD_DIM = kernel.HEAD_DIM +#: The plan capacity (the per-request K/V length bound PrimTS sizes its policy for): dispatch's +#: FULLCTX_DECODE_MAX_KV_LEN, which the forward's prebuilt plans use. +MAX_KV_LEN = 1 << 21 + + +class M3FullCtxDecode(OpWrapper): + """Decode attention of every query row over its request's full causal context in the paged FP8 pool, written + into ``output`` (BF16 [T, heads, head_dim]) in place. + + ``recipe`` is ``modeling's FULLCTX_DECODE_RECIPES[(qlen, requests)]``: the PrimTS Q tile and the split-KV fanout. + ``own_reducer`` merges the splits on the BF16 op (modeling's ``FULLCTX_DECODE_OWN_REDUCER``) rather than PrimTS's + FP16 reducer and a cast. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 1 + 3 rows over an AgentX-length (165k) context", + spec=dict(kv_lens=(165525,), qlen=4), + ), + Cell( + why="one request, one draft-pass row, over a 165k context", + spec=dict(kv_lens=(165522,), qlen=1), + ), + Cell( + why="two requests verifying 1 + 3 rows at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165525, 120011), qlen=4), + ), + Cell( + why="two requests, one row each, at unequal lengths (165k, 120k)", + spec=dict(kv_lens=(165522, 120008), qlen=1), + ), + Cell( + why="one request verifying 1 + 3 rows over a short (4.1k) context: most splits are empty", + spec=dict(kv_lens=(4104,), qlen=4), + ), + ) + + note = """ + Same operands and arithmetic as the stock trtllm-gen call (E4M3 q, E4M3 K/V, P quantized to E4M3 per split, + fp32 accumulation): the split differs, so the result differs from stock as a stock split change does (the P + quantization is per split). Deterministic, and the same bits eager and in a CUDA graph. Held to the fp32 + reference within the FP8-KV band of the in-tree dense decode test. + """ + + def __call__( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + shape = (decode_query_len, int(seq_lens.shape[0])) + kernel.prebuild(q.device, {shape: recipe}, max_kv_len, own_reducer, get_env_enable_pdl()) + plan = kernel.plan_for(q.device, *shape, own_reducer) + pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + k_pages, v_pages = kernel.kv_page_views(pool, subpages) + plan.run(fp8_e4m3_saturate(q), k_pages, v_pages, block_table, seq_lens, sm_scale, output) + return output + + def reference( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> torch.Tensor: + """fp32 attention over the pool's K / V, q rounded to E4M3 first as the kernel runs it.""" + keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) + return dense_attention_reference( + fp8_e4m3_saturate(q), keys, values, seq_lens, decode_query_len, sm_scale + ) + + def is_valid( + self, + q: torch.Tensor, + kv_cache_manager, + layer_idx: int, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + recipe, + max_kv_len: int = MAX_KV_LEN, + own_reducer: bool = True, + ) -> None: + pool, _ = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") + assert pool.dtype == torch.float8_e4m3fn, "an FP8 E4M3 pool" + assert tuple(pool.shape[1:]) == (KV_HEADS, kernel.PAGE, HEAD_DIM), ( + "the TP8 dense-layer pool geometry" + ) + assert q.shape[1:] == (HEADS, HEAD_DIM), "8 query heads of 128 per rank" + assert int(q.shape[0]) == int(seq_lens.shape[0]) * decode_query_len, ( + "q covers exactly the batch" + ) + if torch.cuda.is_current_stream_capturing(): + return # the length checks read the device: eager calls only + assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" + assert bool((seq_lens <= max_kv_len).all()), "a request longer than the plan capacity" + + +m3_fullctx_decode = M3FullCtxDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..5eb42546612e --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# PyTorch oracle ported from vLLM (Apache-2.0), _reference_decode_index_score: +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/tests/kernels/attention/test_minimax_m3.py#L188 +"""The MSA index scorer of a sparse decode step: per (index head, query token, KV block) the causal max of Q . K over +the block's 128 index-K positions, with the pre-wait schedule of ``ops/m3_index_score``.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_index_score"] + +PAGE = 128 +HEAD_DIM = 128 + + +def compiled_scorer(key: tuple): + """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss + (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" + from ..m3_index_score import IndexScoreRunner + + if key not in IndexScoreRunner.kernel_cache: + IndexScoreRunner._compile(*key) + return IndexScoreRunner.kernel_cache[key] + + +class M3IndexScore(OpWrapper): + """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the + positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of + ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is + written). FP8 or BF16 q and index-K, fp32 products and sums. + + ``early`` (the target's kernel at every request count it covers) issues the Q and first K tiles right after the + grid wait from a pre-wait length hint and prefetches one older block per CTA into L2 before the wait; otherwise + every older block of the CTA is prefetched. Neither, nor ``split_k``, changes a score. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request, 1 + 3 rows over 165521 tokens (AgentX ISL), 1 index head (TP8): 1536 CTAs, early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request of 129 tokens: the early tiles of most of the 1536 CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request of 400000 tokens: 3125 blocks, so the 1536 CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request decoding one row over 65536 tokens", + spec=dict(seq_lens=(65536,), qlen=1, heads=1, split_k=1536, early=True), + ), + Cell( + why="one request, 1 + 3 rows over 165521 tokens, 1 index head (TP8): 1272 CTAs (one wave), early issue", + spec=dict(seq_lens=(165521,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="one request of 129 tokens at 1272 CTAs: the early tiles of most CTAs lie past the length and drain", + spec=dict(seq_lens=(129,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="one request of 400000 tokens at 1272 CTAs: 3125 blocks, so the CTAs loop", + spec=dict(seq_lens=(400000,), qlen=4, heads=1, split_k=1272, early=True), + ), + Cell( + why="two requests verifying 4 rows each over 165521 and 20000 tokens: 768 CTAs each, early issue", + spec=dict(seq_lens=(165521, 20000), qlen=4, heads=1, split_k=768, early=True), + ), + Cell( + why="two requests, one of 300 tokens and one of 262144 (2048 blocks: the 768 CTAs loop)", + spec=dict(seq_lens=(300, 262144), qlen=4, heads=1, split_k=768, early=True), + ), + ) + + note = """ + The products of FP8 (or BF16) values are exact in fp32 and the kernel sums 128 of them in its k-split order, so it + differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. + """ + + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) + + +m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py new file mode 100644 index 000000000000..f1a7f210a590 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MiniMax-M3 MSA attention backend's step: KV write, sparse top-k block selection, FMHA.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + +from .m3_sparse_decode_lc8 import dequant_nvfp4_pages + + +class M3MsaAttention(OpWrapper): + """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` + [T, heads * 128] (bf16, the step's live rows). + + * ``k`` / ``v`` given (the dense layers): written into the layer's paged main cache first (NVFP4 with + ``kv_scale_orig_quant``, else FP8 / BF16 as the pool is). None: a producer already wrote the caches. + * ``sparse``: the indexer scores the layer's index-K cache (already written) with ``idx_q`` and publishes the + top-k blocks; the FMHA attends those blocks. Otherwise the FMHA attends the whole page table. + + The FMHA dispatches on the step's metadata (context rows: the MSA prefill FMHA; generation rows: the MSA decode + FMHA, the dense layers' trtllm-gen decode at their capped KV split). ``backend`` and ``metadata`` are the runtime's + objects for this layer and step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="dense layer, a prefill chunk, FP8 pool", + spec=dict(layer="dense", phase="context", tokens=1024), + ), + Cell( + why="dense layer, a decode step of 2 requests", + spec=dict(layer="dense", phase="generation", reqs=2), + ), + Cell( + why="sparse layer, a prefill chunk beyond 16 blocks, NVFP4 pool", + spec=dict(layer="sparse", phase="context", tokens=4096), + ), + Cell( + why="sparse layer, a mixed step (context + generation rows)", + spec=dict(layer="sparse", phase="mixed", tokens=512, reqs=2), + ), + ) + + note = """ + `reference` is exact (float64) softmax attention over the cache rows the step attends, read back from the paged + cache the call wrote (or a producer wrote before it) and dequantized as the FMHA dequantizes them: FP8 rows as + E4M3, NVFP4 rows as E2M1 x E4M3 block scale x the per-tensor ``kv_scale_quant_orig``. A sparse layer's top-k blocks + are re-selected from ``idx_q`` and the index-K cache with the backend's rule (per-block max over the causal + keys, the query's own block forced in, ties to the lower block). On an FP8 pool the FMHA takes Q as a plain E4M3 + cast, which the reference applies too. What remains between the two is the FMHA's own arithmetic (its P + rounding and accumulation order), which the test bounds. + """ + + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors grouped KV heads (query + head h reads KV head h // (heads / kv heads); two per rank at tp2), each KV head selecting its own top-k + blocks from its own index heads against the rank's single index-K head. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = backend.num_heads, backend.head_dim + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + page = int(manager.tokens_per_block) + sm_scale = dim**-0.5 / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + assert config.num_index_heads % kv_heads == 0 + idx_group = config.num_index_heads // kv_heads + + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = qlens[row], cached[row] + qlens[row] + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + qpos = cached[row] + torch.arange(n, device=q.device) + causal = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + for g in range(kv_heads): + if nvfp4: + keys = dequant_nvfp4_pages( + pool[:, 0].view(torch.uint8), + scale_pool[:, 0], + pages, + quad_scales=False, + kv_head=g, + ) + values = dequant_nvfp4_pages( + pool[:, 1].view(torch.uint8), + scale_pool[:, 1], + pages, + quad_scales=True, + kv_head=g, + ) + keys, values = keys * k_global, values * v_global + else: + keys = pool[pages, 0, g].to(torch.float64) + values = pool[pages, 1, g].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + allowed = causal + if sparse: + # One index-K head per rank, shared by every index head (the cache is [pages, 1, page, dim]). + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = _selected_blocks( + idx64[first : first + n, g * idx_group : (g + 1) * idx_group], + idx_keys, + causal, + qpos, + page, + config, + ) + allowed = causal & chosen.repeat_interleave(page, dim=1)[:, :length] + h0, h1 = g * group, (g + 1) * group + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = ( + torch.einsum("thd,kd->thk", q64[first + lo : first + hi, h0:h1], keys) + * sm_scale + ) + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi, h0:h1] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + (ids >= n_valid - config.local_blocks) & (ids < n_valid) + ) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen + + +m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py new file mode 100644 index 000000000000..ff34b8eac631 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer, with its constant loads staged +before the PDL grid-dependency wait and an early dependents trigger (``kernels/minimaxM3/producerPw.cu``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +__all__ = ["m3_producer"] + +HEAD_DIM = 128 +ROTARY_DIM = 64 +PAGE = 128 +#: Query heads per rank at TP8 (MiniMax-M3: 64 query heads; 1 KV head and 1 index head per rank). +HEADS_Q = 8 +#: NVFP4 block: 16 values share one E4M3 scale. +FP4_BLOCK = 16 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) + + +def _norm_rope( + x: torch.Tensor, weight: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, eps: float +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(FP4_BLOCK) + + +def scale_index(within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return (within_page // 4) * 32 + cols * 4 + within_page % 4 + + +class M3Producer(OpWrapper): + """The NVFP4 QKV + indexer producer of a sparse layer: ``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert`` for one + KV head, one index head, head dim 128 and rotary dim 64, with the same arguments and outputs bit for bit. + + ``packed`` [T, (heads_q + 4) * 128] BF16 holds each token's q heads, k, v, index-q and index-k. Every head but v + gets the Gemma RMSNorm with its weight and NeoX RoPE on its first 64 dims (``rotary_cos_sin[position]``). Returns + (q FP8 [T, heads_q, 128], index-q FP8 [T, 1, 128]); writes the token's slot ``out_cache_loc[t]`` (skipped when + negative or past the pool): k and v as NVFP4 (16-value blocks, E4M3 scales against ``kv_quant_scale[1]`` / + ``[2]``, the V scales in the 4-token swizzle) and index-k as FP8. ``early`` triggers the dependent kernel on + entry (its outputs are read only after the dependent's own grid wait). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token (c1 without speculation)", spec=dict(tokens=1, scale=1.0)), + Cell(why="a c1 verify: 4 tokens", spec=dict(tokens=4, scale=3.0)), + Cell(why="a c4 verify: 16 tokens, small activations", spec=dict(tokens=16, scale=1e-3)), + Cell( + why="a c8 verify: 32 tokens, large activations, a skipped and an out-of-pool slot", + spec=dict(tokens=32, scale=40.0, bad_slots=True), + ), + Cell(why="the widest decode step: 48 tokens", spec=dict(tokens=48, scale=1.0)), + Cell( + why="a prefill chunk of 1000 tokens, early trigger off", + spec=dict(tokens=1000, scale=1.0, early=False), + ), + ) + + note = """ + The kernel is the native producer's arithmetic, compiled with the library's --use_fast_math, so it is held to that + kernel bit for bit on every output and cache byte. The torch reference states the math in fp32: the FP8 outputs are + held to it within one E4M3 step, and the NVFP4 K/V, decoded, within one E2M1 step of their block scale. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if prewait_step_inputs: + # This target's copy (producerPw8): the same outputs bit for bit, with the position ids, slots and RoPE + # rows read before its grid wait (none of them may be written by a PDL predecessor). + return tuple( + torch.ops.trtllm.minimax_m3_sm107_tp8_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + pdl, + early, + True, + ) + ) + q, index_q = torch.ops.trtllm.minimax_m3_producer_pw( + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + ) + return q, index_q + + def reference( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", + "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change + them.""" + t = packed.shape[0] + x = packed.float().view(t, heads_q + 4, HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = rope[:, 0], rope[:, 1] + q = torch.stack( + [_norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + ) + return { + "q": q, + "k": _norm_rope(x[:, heads_q], k_weight, cos, sin, eps), + "v": x[:, heads_q + 1], + "index_q": _norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), + "index_k": _norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), + } + + def is_valid( + self, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, + ) -> None: + t = packed.shape[0] + # The RoPE row is read at the position with no bound. + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + # Two tokens writing one slot race. + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" + + +m3_producer = M3Producer() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..88180c2f1a1f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( + fused_qk_norm_rope, +) + + +class M3QkNormRopeFp8(OpWrapper): + """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded + once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place + entry. + + The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps + of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one decode token, tp8 heads (8 q, 1 kv)", spec=dict(m=1, q=8, kv=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=8, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=8, kv=1)), + ) + + note = """ + `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from + fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. + """ + + def __call__( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + + def reference( + self, + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) -> torch.Tensor: + """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" + out = fused_qk_norm_rope( + qkv.clone(), + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + eps, + q_weight, + k_weight, + base, + is_neox, + position_ids, + factor, + low, + high, + attention_factor, + is_qk_norm, + use_gemma, + use_mrope, + mrope_section1, + mrope_section2, + ) + return out.to(torch.float8_e4m3fn) + + +m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py new file mode 100644 index 000000000000..c6f4be426ba3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into +the paged caches.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3QkvInsert(OpWrapper): + """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the + partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and + index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. + + ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool + [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 + [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp8 + heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), + Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), + ) + + note = """ + The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` + returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. + """ + + def __call__( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + common = ( + num_heads, + num_kv_heads, + num_index_heads, + head_dim, + rotary_dim, + eps, + q_norm, + k_norm, + index_q_norm, + index_k_norm, + rotary_cos_sin, + position_ids, + ) + if kv_scale is not None: + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common + ) + return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( + packed, kv_cache, index_k_cache, out_cache_loc, *common + ) + + def reference( + self, + packed: torch.Tensor, + kv_cache: torch.Tensor, + kv_scale: Optional[torch.Tensor], + index_k_cache: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: Optional[torch.Tensor], + num_heads: int, + num_kv_heads: int, + num_index_heads: int, + head_dim: int, + rotary_dim: int, + eps: float, + q_norm: torch.Tensor, + k_norm: torch.Tensor, + index_q_norm: torch.Tensor, + index_k_norm: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + ): + """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of + Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as + packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" + m = packed.shape[0] + sizes = [ + num_heads * head_dim, + num_kv_heads * head_dim, + num_kv_heads * head_dim, + num_index_heads * head_dim, + head_dim, + ] + q, k, v, iq, ik = packed.float().split(sizes, dim=-1) + cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) + half = rotary_dim // 2 + + def norm_rope(x, w, heads): + x = x.reshape(m, heads, head_dim) + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) + x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] + c, s = cos.unsqueeze(1), sin.unsqueeze(1) + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) + + return { + "q": norm_rope(q, q_norm, num_heads), + "index_q": norm_rope(iq, index_q_norm, num_index_heads), + "k": norm_rope(k, k_norm, num_kv_heads), + "v": v.reshape(m, num_kv_heads, head_dim), + "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], + } + + +m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py new file mode 100644 index 000000000000..76841cca732d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3SelectBlocks(OpWrapper): + """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), + with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the + row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` + backs the result head-major so ``permute(1, 0, 2)`` is contiguous. + + The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="eight decode rows of 17..128 blocks", + spec=dict(total_q=8, n_blocks=128, mixed=True), + ), + Cell( + why="one decode row at the 128-block edge", + spec=dict(total_q=1, n_blocks=128, mixed=False), + ), + Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), + ) + + note = """ + Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. + """ + + def __call__( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() + return torch.ops.trtllm.minimax_m3_select_blocks( + max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output + ) + + def reference( + self, + max_score_kv: torch.Tensor, + topk: int, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + head_major_output: bool, + ) -> torch.Tensor: + """Plain torch selection; ``head_major_output`` changes the backing only.""" + heads, n_blocks, total_q = max_score_kv.shape + out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) + blocks = torch.arange(n_blocks, device=max_score_kv.device) + for t in range(total_q): + nv = int(n_valid_blocks[t]) + for h in range(heads): + s = max_score_kv[h, :, t].clone().double() + s[blocks >= nv] = float("-inf") + forced = (blocks < min(init_blocks, nv)) | ( + (blocks >= nv - local_blocks) & (blocks < nv) + ) + s[forced] = float("inf") + order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] + ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) + out[t, h, : ids.numel()] = ids + return out + + +m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py new file mode 100644 index 000000000000..9a92050c6abb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). +MAX_BLOCKS = 9216 +#: Blocks the selection keeps per (query, KV head). +TOPK = 16 +#: Scores the forcing substitutes: initial blocks rank above local ones, and both above any finite score. +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 + +# M3's checkpoint at TP8 (one KV head per rank: 4 KV heads over 8 ranks, each group on a rank pair): no forced initial +# blocks, the one block holding the query forced in. +_M3 = dict(num_kv_heads=1, init_blocks=0, local_blocks=1) + + +class M3SelectLong(OpWrapper): + """Fill ``out`` [total_q, num_kv_heads, 16] with each query's top-16 blocks. + + ``scores`` is the per-KV-head block maxima [num_kv_heads, num_blocks, + total_q] the CuTe DSL index scorer wrote (amax over the KV head's index + heads). Per (query, KV head) row the op keeps the ``n_valid_blocks[query]`` + leading blocks (clamped to ``[0, num_blocks]``), forces the last + ``local_blocks`` of them (counted from the *unclamped* valid count) to + ``LOCAL_SCORE`` and the first ``init_blocks`` to ``INIT_SCORE`` -- local wins + where both apply -- and selects the 16 largest, ties to the smaller block. + A pick whose forced score is -inf is emitted as -1. The ids are written + ascending with a -1 tail. + + Same contract as ``torch.ops.trtllm.minimax_m3_select_blocks`` (the in-tree + ``select_blocks_from_maxscore``) and bit-identical to it, except on a row + where more than 2048 valid blocks tie at the 16th-largest score: there the + in-tree op fills the tied slots in atomic order (its staging holds 2048 + candidates), while this op breaks the tie to the smallest blocks as the + contract says. Up to 2048 tied blocks both take the smallest. The test checks + both against ``reference``. Returns ``out``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the TP2 engine's 5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + note = """ + Certified at one KV head per rank (TP8, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`). The kernel takes any head count and + forcing; other forcing is not driven here. + + `scores` may be strided (the op reads it by stride); the target hands over + the scorer's contiguous view, which is what the cells drive. Rows of at + most 128 blocks are within the kernel's range, but the target keeps those on + the in-tree op, so they are not certified here. + + Ordering is by the descending-score key `orderedKey`: a positive NaN ranks + above +inf, +0 above -0, and a negative NaN below -inf -- which, unlike + -inf, is emitted as a block id rather than -1. The cells drive positive + NaN only. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_sm107_tp2_select_long( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and + ``trigger_phase`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() # [total_q, num_blocks] + forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where( + in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max + ) + k = min(TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) + ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> None: + """The inputs this op takes and answers wrongly. + + Negative forcing counts are accepted and push the forced ranges past the + row instead of failing; the in-tree op this one stands in for refuses + them, so a caller relying on this one would silently diverge from it. + An ``out`` overlapping ``scores`` is overwritten while it is still read. + """ + assert init_blocks >= 0 and local_blocks >= 0, ( + f"init_blocks ({init_blocks}) and local_blocks ({local_blocks}) must be non-negative; the " + "op accepts negative counts and selects as if the forced ranges were elsewhere" + ) + storage = scores.untyped_storage() + assert not (storage.data_ptr() <= out.data_ptr() < storage.data_ptr() + storage.nbytes()), ( + "out must not alias scores: the op writes it while scores are still read" + ) + + +m3_select_long = M3SelectLong() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py new file mode 100644 index 000000000000..7d3627bba5f3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long_dsl.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 MSA block selection for decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_select_long import _M3, MAX_BLOCKS, M3SelectLong + + +class M3SelectLongDsl(OpWrapper): + """``m3_select_long``'s contract, cells and reference, run by ``ops/m3_select_long_dsl``: the same selection + bit for bit as the 1024-thread ``minimax_m3_sm107_tp2_select_long`` (also on rows where more than 2048 blocks tie + at the 16th-largest score), from one 384-thread CTA per row. Returns ``out``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request verifying 4 rows over 165521 tokens (1294 blocks, the AgentX ISL): the 2-key path", + spec=dict(num_blocks=1294, total_q=4, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the TP2 engine's 5888-block rows at two requests x 4 rows: the 9-key path", + spec=dict(num_blocks=5888, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="the longest covered rows (9216 blocks, 1.18M tokens), 2 requests x 4 rows: the 9-key path at its end", + spec=dict(num_blocks=9216, total_q=8, valid="near_full", scores="random", **_M3), + ), + Cell( + why="~1500 blocks tied at the 16th score: the in-tree op also takes the smaller blocks (bit for bit)", + spec=dict(num_blocks=2000, total_q=8, valid="near_full", scores="ties1500", **_M3), + ), + Cell( + why="two requests x 4 rows of mixed lengths, 4500-block engine: every key width, short and padding rows", + spec=dict(num_blocks=4500, total_q=8, valid="mixed", scores="random", **_M3), + ), + Cell( + why="duplicated scores: ties broken to the smaller block, and > 2048 keys at the bound (re-read path)", + spec=dict(num_blocks=4500, total_q=4, valid="near_full", scores="ties", **_M3), + ), + Cell( + why="non-finite scores: +inf, NaN ranked above +inf, and rows whose picks run into -inf (emitted -1)", + spec=dict(num_blocks=600, total_q=8, valid="mixed", scores="nonfinite", **_M3), + ), + ) + + def reference( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + """``m3_select_long``'s reference: the selection is the same bit for bit.""" + return M3SelectLong.reference( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + + def is_valid( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> None: + """``m3_select_long``'s domain.""" + M3SelectLong.is_valid( + self, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + + note = """ + Certified at one KV head per rank (TP8, Rubin) with M3's forcing + (`init_blocks=0`, `local_blocks=1`), and bit for bit against the + 1024-thread `m3_select_long` on every cell. The kernel takes any head count + and forcing; other forcing is not driven here. + + The kernel is compiled per (PDL launch, early trigger, trigger phase) with + symbolic shapes and strides; the target compiles its keys before warmup + (`ops/m3_select_long_dsl.prebuild`), and a call outside them compiles + on first use. + """ + + def __call__( + self, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + ) -> torch.Tensor: + from ..m3_select_long_dsl import select_long_dsl + + if scores.shape[1] > MAX_BLOCKS: + raise RuntimeError( + f"m3_select_long_dsl supports at most {MAX_BLOCKS} blocks, got {scores.shape[1]}" + ) + select_long_dsl( + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) + return out + + +m3_select_long_dsl = M3SelectLongDsl() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py new file mode 100644 index 000000000000..fc770d010643 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""This target's block-sparse decode attention (Triton, with the one-launch long-context route).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..triton_sparse_decode import minimax_m3_sparse_attn_decode +from .m3_sparse_decode_lc8 import PAGE, sparse_attention_reference + + +class M3SparseDecode(OpWrapper): + """``minimax_m3_sparse_attn_decode`` as ``M3Attention._decode_sparse_core`` calls it. + + Attention of each query token over the blocks ``topk_idx`` selected for it + (``[num_kv_heads, T, topk]``, ascending ids with a -1 tail), each block + truncated at the token's causal extent, written into ``output`` in place. + The target passes FP8 q, the packed NVFP4 sparse-layer cache (one KV head) with its block and per-tensor + scales, and BF16 ``output`` [T, 8, 128]. Returns ``output``. + + Which kernel runs is decided inside, per step: a selected (non-identity) table on the TP8 geometry goes to the + one-launch ``m3_sparse_decode_lc8`` kernel while modeling's ``SPARSE_DECODE_LC8`` covers it; an identity table + (every request within 16 blocks, staged by the metadata), or a selected table with that range killed, runs the + Triton per-row producer (the multi-query producer for identity steps of at least 17 rows), then the LSE merge. + Every route also publishes the MXFP8 o_proj hand-off to ``m3_prequant.SLOT``, bit for bit + ``torch.ops.trtllm.mxfp8_quantize(output, True)``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="identity selection, decode without speculation at batch 8: the Triton producer + merge", + spec=dict( + seq_lens=(1, 128, 129, 700, 2048, 1500, 64, 0), decode_query_len=1, identity=True + ), + ), + Cell( + why="identity selection, Eagle3 verify at one request (T=4): the Triton producer + merge", + spec=dict(seq_lens=(1900,), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 2 (T=8), a padding request", + spec=dict(seq_lens=(2048, 0), decode_query_len=4, identity=True), + ), + Cell( + why="identity selection, Eagle3 verify at batch 8 (T=32): the multi-query producer + merge", + spec=dict( + seq_lens=(1900, 2048, 700, 1500, 1200, 2000, 300, 1024), + decode_query_len=4, + identity=True, + ), + ), + Cell( + why="selected blocks, Eagle3 verify at one request (T=4): the one-launch long-context kernel", + spec=dict(seq_lens=(40000,), decode_query_len=4, identity=False), + ), + Cell( + why="selected blocks, Eagle3 verify at batch 2 (T=8): the one-launch long-context kernel", + spec=dict(seq_lens=(9000, 40000), decode_query_len=4, identity=False), + ), + ) + + note = """ + Certified only as the target calls it: one KV head of 8 query heads, head_dim 128, FP8 q, packed NVFP4 + cache, BF16 output, `num_topk_chunks` left to the dispatcher. The function also takes BF16 / FP8 caches, other + head geometries and an explicit split-K factor (the in-tree tests drive those); none is certified here. + + The MXFP8 hand-off of the long-context route is certified on its own entry (`m3_sparse_decode_lc8`); the Triton + merge's (`_store_mxfp8_head`) is held to `torch.ops.trtllm.mxfp8_quantize` by this entry's test. + """ + + def __call__( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + minimax_m3_sparse_attn_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + kv_scale=kv_scale, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + num_topk_chunks=num_topk_chunks, + identity_blocks=identity_blocks, + ) + return output + + def reference( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> torch.Tensor: + """``output`` as the call is supposed to write it: float64 block-sparse attention over the + dequantized NVFP4 blocks (``sparse_attention_reference``), rounded once to ``output``'s dtype. + + Only the packed NVFP4 cache is mirrored -- the target's only configuration -- one KV head at a time: KV head + ``h`` serves query heads ``[h * group, (h + 1) * group)`` with its own top-k row. ``kv_scale``, + ``num_topk_chunks`` and ``identity_blocks`` do not change what is computed. + """ + assert k_block_scale is not None, "mirrors the NVFP4 cache only" + kv_heads = k_paged.shape[1] + assert kv_heads == 1, "mirrors the TP8 cache (one KV head per rank) only" + group = q.shape[1] // kv_heads + heads = [] + for h in range(kv_heads): + heads.append( + sparse_attention_reference( + q[:, h * group : (h + 1) * group], + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[h], + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + ) + return torch.cat(heads, dim=1).to(output.dtype) + + def is_valid( + self, + q: torch.Tensor, + k_paged: torch.Tensor, + v_paged: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, + ) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) + + +m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..a845ac68f257 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One-launch long-context sparse decode over the selected NVFP4 blocks at 8 query heads (16-CTA clusters). + +Also the plain-torch block-sparse attention and MXFP8 hand-off mirrors the +Triton sparse decode entry (``m3_sparse_decode``) reuses: both compute the same +thing over the same cache. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The TP8 geometry the kernel is compiled for: 8 query heads on one KV head. +NUM_HEADS = 8 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +#: NVFP4 elements per E4M3 block scale. +SF_VEC = 16 +#: MXFP8 elements per UE8M0 scale, and the o_proj width the hand-off is laid out for. +MX_BLOCK = 32 +WIDTH = NUM_HEADS * HEAD_DIM + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, + block_scale: torch.Tensor, + pages: torch.Tensor, + quad_scales: bool, + kv_head: int = 0, +) -> torch.Tensor: + """Dequantize KV head ``kv_head`` of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, kv_head].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, kv_head].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, +) -> torch.Tensor: + """Block-sparse decode attention over one KV head of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages(k, k_block_scale, pages, quad_scales=False).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True).reshape(-1, HEAD_DIM) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, + sf [num_rows_alloc * WIDTH // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) + num_sf_k_tiles = WIDTH // 128 + sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +class M3SparseDecodeLc8(OpWrapper): + """Sparse decode attention, split-K merge and MXFP8 hand-off of one TP8 decode step in one launch. + + One cluster of 16 CTAs per query token: CTA ``r`` attends the token's + ``r``-th selected block (``topk[tok, r]``) for all 8 heads (transposed MMAs, + the heads on the n dimension), and the cluster merges the 16 partials over + distributed shared memory into the BF16 row ``out[tok]`` [1024]. With ``q8`` / + ``sf`` it also writes the row's MXFP8 copy and swizzled UE8M0 scales for the + o_proj GEMM. All operands are byte views (E4M3 ``q``, packed E2M1 pages, E4M3 + block scales). The kernel releases its PDL dependents right after its grid + wait. Returns ``out``. + + ``prewait`` stages the request length, the query rows (widened to FP16) and + the request's block-table row (up to 8192 entries, longer rows are read from + global memory) before the grid wait; the arithmetic is the same, so both + variants write the same bits. + + Per block: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16 + (exact), QK and PV on FP16 tensor cores with FP32 accumulation, the softmax + probabilities rounded to FP16 for PV, log2-domain softmax with ``ex2.approx``. + ``linear_softmax`` picks the denominator recurrence (the two agree to + rounding); the target picks it per batch shape. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation, one request: T=1, log-domain recurrence, MXFP8 hand-off, prewait", + spec=dict( + seq_lens=(9000,), decode_query_len=1, linear_softmax=False, mxfp8=True, prewait=True + ), + ), + Cell( + why="Eagle3 verify at one request (the TP8 operating point): T=4, log-domain recurrence, prewait", + spec=dict( + seq_lens=(20000,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="Eagle3 verify at two requests: T=8 on the linear denominator, prewait at its row limit", + spec=dict( + seq_lens=(5000, 12345), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="a 1M-token request: a block-table row longer than the staged 8192 entries, read from global", + spec=dict( + seq_lens=(1048576,), + decode_query_len=4, + linear_softmax=False, + mxfp8=True, + prewait=True, + ), + ), + Cell( + why="Eagle3 verify at batch 8 (T=32): short requests with -1 tails and zero-length padding rows", + spec=dict( + seq_lens=(16384, 2049, 130, 0, 7777, 4, 30000, 0), + decode_query_len=4, + linear_softmax=True, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="decode without speculation at batch 8: T=8 rows of distinct requests, staged after the wait", + spec=dict( + seq_lens=(3000, 2100, 9999, 2048, 2049, 4500, 16000, 2200), + decode_query_len=1, + linear_softmax=False, + mxfp8=True, + prewait=False, + ), + ), + Cell( + why="no hand-off buffers (the variant the call site takes past 128 rows), driven at T=8", + spec=dict( + seq_lens=(5000, 12345), + decode_query_len=4, + linear_softmax=True, + mxfp8=False, + prewait=False, + ), + ), + ) + + note = """ + Certified: `linear_softmax` x hand-off on / off x `prewait`, the variants + the cells drive. The hand-off cells are certified bit-exactly against + `mxfp8_handoff_reference` applied to the op's own BF16 rows; the BF16 rows + themselves within the band the test derives, and the two `prewait` variants + against each other bit for bit. + + `topk` is read for the first `min(16, ceil(kv_len / 128))` entries of each + token; a -1 among them skips that CTA's block, entries past them are never + read. `use_pdl` only changes the launch attribute. The hand-off assumes + `num_sf_k_tiles = 8` (a 1024-wide o_proj input) and buffers with at least + as many rows as the step; `MERGE_MXFP8` sizes them at 128 rows. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + ) -> torch.Tensor: + # merge_once: this target's copy of the op (minimaxM3Sm107Tp8/sparseDecodeLc8m.cu), whose split-K merge forms + # the 16 weights once per CTA (same values, same order: bitwise the image op's output). + lc8 = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_sparse_decode_lc8 + if merge_once + else torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8 + ) + # prefetch: up to two read-only tensors this target's op pulls into L2 (a hint; the image op takes none). + # next_kv: (this layer's id buffer, the next sparse layer's, its K / V pools and block scales), target op only. + extra = () + if merge_once and (prefetch or next_kv): + extra = tuple(prefetch or ()) + (None,) * (2 - len(prefetch or ())) + if next_kv: + extra += tuple(next_kv) + lc8( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + decode_query_len, + sm_scale, + linear_softmax, + use_pdl, + prewait, + *extra, + **({"identity": True} if identity else {}), + **({"identity_step": True} if identity_step else {}), + **({"one_div": True} if one_div else {}), + ) + return out + + def reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + ) -> torch.Tensor: + """The BF16 ``out`` [T, 1024] the op is supposed to write: ``sparse_attention_reference`` in float64, + rounded once. ``linear_softmax``, ``use_pdl``, ``prewait`` and ``one_div`` do not change what is computed; + the hand-off buffers are checked separately with ``mxfp8_handoff_reference``. ``identity`` attends each token's + first min(16, blocks) blocks (its own causal length), whatever ``topk`` holds. + """ + if identity: + topk = _identity_table(seq_lens, decode_query_len, int(q.shape[0]), topk.device) + attn = sparse_attention_reference( + q.view(-1, NUM_HEADS, HEAD_DIM), + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, + ) + return attn.reshape(out.shape).to(out.dtype) + + def is_valid( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + out: torch.Tensor, + q8: Optional[torch.Tensor], + sf: Optional[torch.Tensor], + num_sf_k_tiles: int, + decode_query_len: int, + sm_scale: float, + linear_softmax: bool, + use_pdl: bool, + prewait: bool, + merge_once: bool = False, + prefetch: Optional[tuple] = None, + next_kv: Optional[tuple] = None, + identity: bool = False, + identity_step: bool = False, + one_div: bool = False, + ) -> None: + """The inputs this op takes and answers wrongly (``m3_attn_lc8.lc8_applies`` is the target's form of the + layout half). + + The kernel addresses a page's token rows at a fixed 64-byte (data) and + 8-byte (scale) pitch and loads them 16 / 8 bytes at a time, so only the + page stride is read from the views. It reads the selected block ids, + the block table and the hand-off buffers unchecked. Reads ``topk`` and + ``seq_lens`` back to the host: a test-only cost. + """ + for name, t in (("k", k), ("v", v)): + assert t.dim() == 4 and t.shape[1] == 1 and t.shape[2:] == (PAGE, HEAD_DIM // 2), ( + f"{name} must be the [pages, 1, 128, 64] byte view of one KV head" + ) + assert t.stride(3) == 1 and t.stride(2) == HEAD_DIM // 2, ( + f"{name} rows must be dense 64-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 16 == 0 and t.stride(0) % 16 == 0, ( + f"{name} pages must be 16-byte aligned" + ) + for name, t in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + assert t.dim() == 4 and t.stride(3) == 1 and t.stride(2) == HEAD_DIM // SF_VEC, ( + f"{name} rows must be dense 8-byte rows: the kernel ignores the row stride" + ) + assert t.data_ptr() % 8 == 0 and t.stride(0) % 8 == 0, ( + f"{name} pages must be 8-byte aligned" + ) + assert q.data_ptr() % 16 == 0 and q.stride(0) % 16 == 0 and q.stride(1) % 16 == 0, ( + "q rows are loaded 16 bytes at a time and must be 16-byte aligned" + ) + assert k_global_scale.numel() == 1 and v_global_scale.numel() == 1, ( + "global scales are single values" + ) + total_q = q.shape[0] + assert ( + total_q % decode_query_len == 0 and seq_lens.numel() >= total_q // decode_query_len + ), ( + "every query token must belong to a request: seq_lens is read at tok // decode_query_len unchecked" + ) + seq = seq_lens.tolist() + num_cols = block_table.shape[1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + PAGE - 1) // PAGE + ids = topk[tok, : min(TOPK, nblocks)] + assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( + f"token {tok}: selected block ids must be -1 or below the request's block count; the block " + "table is indexed with them unchecked" + ) + if q8 is not None: + assert num_sf_k_tiles == WIDTH // 128, ( + f"num_sf_k_tiles must be {WIDTH // 128}: the hand-off rows are {WIDTH} wide and a smaller " + "value overlaps them silently" + ) + assert ( + q8.numel() >= total_q * WIDTH + and sf is not None + and sf.numel() >= ((total_q + 127) // 128 * 128 * (WIDTH // MX_BLOCK)) + ), "the hand-off buffers must hold every row of the step: they are written unchecked" + + +def _identity_table(seq_lens, decode_query_len, total_q, device): + """[T, 16] ids (i < min(16, blocks of the token's causal length) ? i : -1), the selection identity mode reads.""" + ids = torch.full((total_q, 16), -1, dtype=torch.int32, device=device) + for tok in range(total_q): + req, qoff = divmod(tok, decode_query_len) + kv_len = max(int(seq_lens[req]) - decode_query_len + qoff + 1, 0) + n = min(16, (kv_len + PAGE - 1) // PAGE) + ids[tok, :n] = torch.arange(n, dtype=torch.int32, device=device) + return ids + + +m3_sparse_decode_lc8 = M3SparseDecodeLc8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py new file mode 100644 index 000000000000..43166b472eca --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: The MSA pool's page: one sparse block is one 128-token KV page. +TOKENS_PER_BLOCK = 128 + + +class MsaKvUpdate(OpWrapper): + """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. + + ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its + own buffers at the head of every captured speculative decode step. One + single-CTA launch writes, for ``batch`` requests and ``total_q`` query + tokens (all int32): + + * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, + ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); + * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` + attends up to, through ``block_table``; + * ``n_valid_blocks[t]``, when given, the block count top-k selection is + bounded by. + + It returns None; the cells read the four buffers. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="decode without speculation at one request: qo = 1, every buffer written", + spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="decode without speculation at the captured batch 8", + spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), + ), + Cell( + why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", + spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), + ), + Cell( + why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", + spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), + ), + Cell( + why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", + spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), + ), + Cell( + why="positions crossing a page boundary: the slot moves to the previous page of the request", + spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), + ), + ) + + note = """ + The kernel covers any batch up to the block size's loop (it strides over + 256 threads), and any page size; only `tokens_per_block = 128`, the MSA + pool's page, is certified. + + The op does not bounds-check the request index a token reads + (`q_batch_row`) or the block-table column it derives: both come from the + metadata's own staging, which keeps them in range by construction, and an + out-of-range value reads whatever memory sits there. `is_valid` states both. + """ + + def __call__( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + torch.ops.trtllm.minimax_m3_msa_kv_update( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + + def reference( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: + """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], + n_valid_blocks[:total_q])`` as the op is supposed to write them. + + The in-tree metadata's own torch re-derivation + (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook + call), statement for statement, returning new tensors rather than + writing the buffers. The buffer arguments only say which outputs exist: + a ``None`` buffer gives a ``None`` result. + """ + page = int(tokens_per_block) + kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) + cu_kv = None + if cu_kv_lens is not None: + cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) + cu_kv[1:] = torch.cumsum(kv_true, 0) + qbr = q_batch_row[:total_q].to(torch.long) + qo_dev = seq_lens[:batch] + pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] + pos_slot = pos.clamp_min(0) + block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( + max=int(block_table.shape[1]) - 1 + ) + slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) + n_valid = None + if n_valid_blocks is not None: + n_valid = torch.div( + (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" + ).to(torch.int32) + return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid + + def is_valid( + self, + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], + ) -> None: + """The inputs this op takes and answers wrongly. + + Every operand but the block table is read as a dense run from its first + element, so a strided view reads the wrong elements. The request index of + each token and the table row it selects are not bounds-checked. Reads + ``q_batch_row`` back to the host: a test-only cost. + """ + for name, t in ( + ("kv_lens", kv_lens), + ("kv_lens_staged", kv_lens_staged), + ("seq_lens", seq_lens), + ("q_batch_row", q_batch_row), + ("q_intra", q_intra), + ("msa_seq_lens", msa_seq_lens), + ("out_cache_loc", out_cache_loc), + ("cu_kv_lens", cu_kv_lens), + ("n_valid_blocks", n_valid_blocks), + ): + assert t is None or t.is_contiguous(), ( + f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" + ) + rows = q_batch_row[:total_q] + assert total_q == 0 or ( + int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) + ), ( + "q_batch_row must index a request below batch (and a block-table row): the op reads the " + "request's lengths and table row unchecked" + ) + + +msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/__init__.py new file mode 100644 index 000000000000..5da86fe24bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective-communication entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py new file mode 100644 index 000000000000..8daa2d6b9540 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams + + +class M3AllReduce(OpWrapper): + """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the + plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: + ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; + returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the + communicator's per call. + + The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps + beyond the multicast kernel's rows. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # No cell names a world size: the result does not depend on it, and the certifying matrix runs at every world + # size one node can form (the target's two-tray group of 8 cannot be formed on one node). + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, hidden=6144)), + Cell(why="a 64-row generation step", spec=dict(m=64, hidden=6144)), + Cell(why="plain sum of eight drafter rows", spec=dict(m=8, hidden=6144, plain=True)), + ) + + note = """ + The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then + adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum + to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. + """ + + def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): + return comm(x, all_reduce_params=params) + + def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): + """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" + total = torch.stack([x.float() for x in xs]).sum(0) + if params is None: + return total.to(xs[0].dtype) + updated = (params.residual.float() + total).to(xs[0].dtype) + h = updated.float() + normed = ( + h + * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) + * params.norm_weight.float() + ) + return normed.to(xs[0].dtype), updated + + +m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py new file mode 100644 index 000000000000..cf0a4be970d9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py @@ -0,0 +1,300 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode layer boundary, 2 ranks on SM107.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer + fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; + only the stored residual is rounded to BF16. + + The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are + summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic + and the normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + pre = total.to(torch.bfloat16).float() + residual.float() + updated = pre.to(torch.bfloat16) + u = pre.double() if drafter else updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +class M3McAllreduceNorm(OpWrapper): + """One TP group's boundary: ``updated = bf16(residual + bf16(sum_r bf16(x_r + x2_r)))``, ``normed = + RMSNorm(updated)``: bitwise the stock all-reduce + fused residual RMSNorm (+ mxfp8_quantize) of SM107 TP2. + + Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls + its own view ``ws_uc`` (re-arming the words it read) and rotates the per-CTA call parities in ``flags``; the + workspace is the ``McAllReduceWorkspace`` bucket of ``maxt`` tokens. All ranks must make the same calls in the + same stream order. ``prepushed`` 2 means a producer already pushed ``x`` (and launched this grid after its + pushes); 3 that two producers pushed ``x`` into exchange set A and ``x2`` into set B. Returns + ``[normed, updated]`` or, with ``quant``, also the MXFP8 copy of ``normed`` (E4M3 rows, swizzled UE8M0). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The tp2 decode boundaries (hidden 6144, 2 ranks, eps 1e-6): every decode step fits the 8-token bucket, epilogue + # option 7 at 1, 2, 7, 8 tokens and 15 at 3..6 (ops/m3_mc_allreduce.AR_OPTS_BY_M; every option is + # bit-identical). + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token at the attention boundary, x only, option 7", + spec=dict(m=1, maxt=8, opts=7, x2=False, quant=False, prepushed=0), + ), + Cell( + why="MoE boundary at 4 rows (one request verifying 1 + 3): shared + routed summed in-kernel, MXFP8 copy", + spec=dict(m=4, maxt=8, opts=15, x2=True, quant=True, prepushed=0), + ), + Cell( + why="MoE boundary at 8 rows (two requests), MXFP8 copy for the next QKV", + spec=dict(m=8, maxt=8, opts=7, x2=True, quant=True, prepushed=0), + ), + Cell( + why="attention boundary whose 8 rows the o_proj GEMV already pushed (prepushed 2)", + spec=dict(m=8, maxt=8, opts=7, x2=False, quant=False, prepushed=2), + ), + # Every other decode row count (modeling's MC_ALLREDUCE), the second addend and the MXFP8 copy alternating. + *( + Cell( + why=f"{m} rows: a boundary row count between the ones above", + spec=dict( + m=m, + maxt=8, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + x2=m % 2 == 0, + quant=m % 3 == 0, + prepushed=0, + ), + ) + for m in range(1, 9) + if m not in (1, 4, 8) + ), + # The Eagle3 drafter boundaries (modeling's DRAFT_MC_AR, SpecTail): option 16 (the drafter's stock rounding) on + # 23 / 31 at every draft-pass row count 1..8 (pass 0: 4 x requests, later passes: requests), x only. + *( + Cell( + why=f"Eagle3 drafter boundary at {m} rows (opts {23 if m in (1, 2, 7, 8) else 31})", + spec=dict( + m=m, + maxt=8, + opts=23 if m in (1, 2, 7, 8) else 31, + x2=False, + quant=False, + prepushed=0, + drafter=True, + ), + ) + for m in range(1, 9) + ), + # Every other row count the o_proj pushes (modeling's ATTN_PUSH). + *( + Cell( + why=f"{m} rows prepushed by the o_proj GEMV", + spec=dict( + m=m, + maxt=8, + opts=(7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + x2=False, + quant=False, + prepushed=2, + ), + ) + for m in range(1, 8) + ), + ) + + note = """ + The epilogue option bits (1 early residual / weight load, 2 st.async cluster exchange, 4 early parity rotation, + 8 max carveout) are bit-identical by construction; only the four (bucket, option) pairs of the sm_103 kernel are + built. SM107 TP2 against the stock AUTO all-reduce + fused residual RMSNorm + mxfp8_quantize, rows 1..24, x alone + and x + x2: normed, updated and the MXFP8 copy (data and the scale bytes of valid rows) bitwise equal. `prepushed` + 1 (pushed, grid wait kept) is accepted by the op but no caller passes it. `pdl` changes only launch overlap and is + driven off in the cells. + + `normed` is gated to one BF16 rounding step: the kernel's RMS statistic is a warp / cluster tree of FP32 sums + and its `rsqrtf` is approximate, so a row whose scaled value lands next to a BF16 midpoint can round either + way. `updated` and the MXFP8 copy of the op's own `normed` are bit-exact. + """ + + def __call__( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + own_op: bool = False, + ) -> list[torch.Tensor]: + # The capture operands only with a capture: an op without the capture epilogue (an older image) takes the + # rest unchanged. own_op: this target's copy (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm, the same + # kernels with the capture epilogue; modeling's MC_AR_OWN_OP). + op = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_mc_allreduce_norm + if own_op + else torch.ops.trtllm.minimax_m3_sm107_tp2_mc_allreduce_norm + ) + return op( + x, + x2, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + None, + None, + maxt, + quant, + prepushed, + opts, + *((cap_out, cap_w, cap_eps) if cap_out is not None else ()), + ) + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant, drafter) + + def is_valid( + self, + x: torch.Tensor, + x2: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + prepushed: int, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + own_op: bool = False, + ) -> None: + h = residual.shape[-1] + # The op checks norm_w's element count, not its layout: a strided view is read as packed memory. + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + # The workspace is addressed as [set][parity][maxt][world][H / 2] words and never bounds-checked: a smaller + # buffer (another bucket's, or one built for fewer ranks) is polled and re-armed out of range. + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + # Flags are per (token, cluster CTA) of the bucket; the op checks the count, not that they belong to the + # workspace. Flags of another bucket hold another call history, and the parity they select is not the one the + # peers push into: the poll then waits on the other half forever or reads last call's rows. + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + + +m3_mc_allreduce_norm = M3McAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..8b74c9974397 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MoE finalize + multicast all-reduce + residual + RMSNorm (+ MXFP8 copy) of the decode MoE boundary, 4 ranks.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_mc_allreduce_norm import boundary + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class M3McFinalizeAllreduceNorm(OpWrapper): + """The MoE boundary with the routed experts' finalize folded in. + + Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, + ``scales`` fp32 or bf16), then the group's all-reduce, residual add and RMSNorm exactly as + ``m3_mc_allreduce_norm``: same workspace, flags and call-order contract. Returns ``[normed, updated]`` or, with + ``quant``, also the MXFP8 copy of ``normed``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Not taken at tp2 (modeling's MC_FINALIZE is empty: the decode step is at most 8 rows); the cells are the tp4 + # target's, on the SM107 kernel. The tp4 target takes this path from 9 tokens (top-4 of 128 experts, fp32 routing + # weights, the shared expert always present): the 16- and 48-token buckets, epilogue option 15. + CELLS: tuple[Cell, ...] = ( + Cell( + why="the smallest finalize call (9 tokens, 16-token bucket) with the next QKV's MXFP8 copy", + spec=dict(m=9, maxt=16, opts=15, topk=4, quant=True, dropped=0), + ), + Cell( + why="a full 16-token bucket with dropped expert slots", + spec=dict(m=16, maxt=16, opts=15, topk=4, quant=False, dropped=5), + ), + Cell( + why="the 48-token bucket at a partial fill", + spec=dict(m=24, maxt=48, opts=15, topk=4, quant=False, dropped=0), + ), + Cell( + why="the 48-token bucket full, MXFP8 copy", + spec=dict(m=48, maxt=48, opts=15, topk=4, quant=True, dropped=3), + ), + # Every other row count of the unfinalized MoE (MC_FINALIZE): the bucket and options the target picks. + *( + Cell( + why=f"{m} tokens: a finalize row count between the ones above", + spec=dict( + m=m, + maxt=16 if m <= 16 else 48, + opts=15, + topk=4, + quant=m % 2 == 0, + dropped=m % 4, + ), + ) + for m in range(9, 49) + if m not in (9, 16, 24, 48) + ), + ) + + note = """ + BF16 routing weights and the 8-token bucket are accepted by the op (option 7 there) but the target only enters + this path at 9 tokens and up, with the built-in MoE's fp32 weights; neither is a cell. The rest is as in + `m3_mc_allreduce_norm`'s note. + """ + + def __call__( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + own_op: bool = False, + ) -> list[torch.Tensor]: + # The capture operands only with a capture: an op without the capture epilogue (an older image) takes the + # rest unchanged. own_op: this target's copy (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm, the same + # kernels with the capture epilogue; modeling's MC_AR_OWN_OP). + op = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_mc_allreduce_norm + if own_op + else torch.ops.trtllm.minimax_m3_sm107_tp2_mc_allreduce_norm + ) + return op( + fc2, + shared, + residual, + norm_w, + eps, + ws_uc, + ws_mc_ptr, + flags, + rank, + world, + cluster, + pdl, + perm, + scales, + maxt, + quant, + 0, + opts, + *((cap_out, cap_w, cap_eps) if cap_out is not None else ()), + ) + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) + + def is_valid( + self, + fc2: torch.Tensor, + perm: torch.Tensor, + scales: torch.Tensor, + shared: Optional[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + flags: torch.Tensor, + rank: int, + world: int, + cluster: int, + pdl: bool, + maxt: int, + quant: bool, + opts: int, + cap_out: Optional[torch.Tensor] = None, + cap_w: Optional[torch.Tensor] = None, + cap_eps: float = 0.0, + own_op: bool = False, + ) -> None: + h = residual.shape[-1] + assert norm_w.is_contiguous(), ( + "norm_w must be contiguous; the kernel reads it as packed memory" + ) + need = 4 * maxt * world * (h // 2) + assert ws_uc.is_contiguous() and ws_uc.numel() >= need, ( + f"ws_uc holds {ws_uc.numel()} words, the {maxt}-token bucket of {world} ranks needs {need}" + ) + assert flags.numel() == maxt * cluster, ( + f"flags has {flags.numel()} counters; the {maxt}-token bucket's are exactly {maxt * cluster}" + ) + # The kernel gathers fc2[perm] without a bound: an index past the FC2 rows reads whatever follows them. + lo, hi = int(perm.min()), int(perm.max()) + assert lo >= -1 and hi < fc2.shape[0], ( + f"perm spans [{lo}, {hi}]; entries must be -1 (dropped) or a row of fc2 (< {fc2.shape[0]})" + ) + + +m3_mc_finalize_allreduce_norm = M3McFinalizeAllreduceNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/__init__.py new file mode 100644 index 000000000000..e9f4d8e20b94 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Embedding ops of this target: the generic token-embedding gather.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py new file mode 100644 index 000000000000..71fce6b253cb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="an 8-row decode verify, int64 ids", + spec=dict(rows=8, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ) + + note = """ + An exact gather: `reference` must match bit for bit. + """ + + def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) + + def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] + return rows * valid.unsqueeze(-1).to(table.dtype) + + +m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/__init__.py new file mode 100644 index 000000000000..4c6fc92cd49c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 decode embedding gather + first-layer RMSNorm as one CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py new file mode 100644 index 000000000000..78563b4db40d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. + +128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original +K3 wrapper this target does not promise FlashInfer's bit-identical reduction. +""" + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +NORM_THREADS = 128 +NORM_VEC = 8 +NORM_WARPS = 4 # capacity: fixed: the CTA's warps + + +@cute.kernel +def m3_embed_norm_kernel( + ids: cute.Tensor, # int32 or int64 [N] + table: cute.Tensor, # bf16 [V, H] + weight: cute.Tensor, # bf16 [H] + raw: cute.Tensor, # bf16 [N, H], out: the rows + out: cute.Tensor, # bf16 [N, H], out: the normed rows + vocab: cutlass.Int64, + eps: cutlass.Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + hidden: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 + ) + + tok = cutlass.Int64(ids[bidx]) + valid = (tok >= cutlass.Int64(0)) & (tok < vocab) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + + gX = cute.local_tile(table, tiler_mn, (row, 0)) + gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) + gY = cute.local_tile(out, tiler_mn, (bidx, 0)) + w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 + ) + thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + tXgR = thr_copy_O.partition_D(gR) + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. + tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) + if valid: + cute.copy(copy_atom_load, tXgX, tXrX) + cute.copy(copy_atom_load, tWgW, tWrW) + cute.copy(copy_atom_store, tXrX, tXgR) + + # flashinfer's RMSNormKernel arithmetic, in its order (see the header). + x = tXrX.load().to(cutlass.Float32) + x_sq = x * x + sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) + for i in cutlass.range_constexpr(5): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + if lane == 0: + reduction_buffer[0, warp] = sum_sq + cute.arch.barrier() + total = cutlass.Float32(0.0) + if lane < NORM_WARPS: + total = reduction_buffer[0, lane] + for i in cutlass.range_constexpr(5): + total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) + mean_sq = total / cutlass.Float32(hidden) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + w = tXrW.load().to(cutlass.Float32) + y = x * rstd * (w + cutlass.Float32(0.0)) + tXrO.store(y.to(out.element_type)) + cute.copy(copy_atom_store, tXrO, tXgO) + + +@cute.jit +def m3_embed_norm( + ids: cute.Tensor, + table: cute.Tensor, + weight: cute.Tensor, + raw: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int64, + eps: cutlass.Float32, + n_tokens: cutlass.Constexpr[int], + hidden: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. + blocks = hidden // (NORM_VEC * NORM_THREADS) + tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), + stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip + m3_embed_norm_kernel( + ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden + ).launch( + grid=[n_tokens, 1, 1], + block=[NORM_THREADS, 1, 1], + smem=NORM_WARPS * 4, + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py new file mode 100644 index 000000000000..276f8c32023b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host side of the MiniMax-M3 decode embedding + first-layer RMSNorm (m3_embedding_firstnorm_kernel.py): H 6144, +1..8 rows, into caller-owned buffers. + +Compiled on the first call per (rows, id dtype, table shape, PDL), which must happen outside CUDA-graph capture +(:class:`DecodeEmbedding` compiles at construction). +""" + +import threading +from typing import Dict, Tuple + +import cuda.bindings.driver as cuda_driver +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +MAX_ROWS = 8 # capacity: fixed: the kernel's built rows (decode steps) + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def arg(tensor): + return from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _check(ids, table, weight, raw, out) -> None: + rows = ids.shape[0] if ids.dim() == 1 else 0 + if not ( + 1 <= rows <= MAX_ROWS + and ids.dtype in (torch.int32, torch.int64) + and ids.is_contiguous() + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN,) + and all( + t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() + for t in (raw, out) + ) + ): + raise ValueError( + "m3_embed_norm: expects ids int32/int64 [1..8], a contiguous bf16 table [V, 6144], bf16 weight [6144] " + "and contiguous bf16 raw / out [rows, 6144]" + ) + + +def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) + fn = _compiled[key] = cute.compile( + _m3_embed_norm_jit, + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + _check(ids, table, weight, raw, out) + fn = _compile(ids, table, weight, eps, raw, out, use_pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return out, raw + + +class DecodeEmbedding: + """Caller-owned buffers of one (rows, id dtype) decode embedding: ``raw`` (the gathered rows, the first residual) + and ``output`` (their RMSNorm), over ``table`` and the first layer's plain (non-Gemma) input norm ``weight`` / + ``eps``. Built (and compiled) before CUDA-graph capture.""" + + def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): + assert not torch.cuda.is_current_stream_capturing() + assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN + assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) + assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( + table, + weight, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(eps) + self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) + + def __call__(self, ids): + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return m3_embed_norm( + ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/__init__.py new file mode 100644 index 000000000000..9cc374bdf1cc --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMM entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py new file mode 100644 index 000000000000..b973b175d8b4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3Bf16Linear(OpWrapper): + """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. + + The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), + o_proj (this rank's input columns), gate_up and down projections run it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, nine draft rows", spec=dict(m=9, n=1280, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=1024)), + Cell(why="drafter gate_up, 16 rows", spec=dict(m=16, n=4608, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=2304)), + ) + + note = """ + cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x, w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return (x.float() @ w.float().t()).to(x.dtype) + + +m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py new file mode 100644 index 000000000000..a0013e3ded3b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The dense MLP's swigluoai activation as one C++ launch with PDL and an early release (DENSE_SWIGLU_CPP).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +# The dense layers' swigluoai parameters, which the kernel hard-codes (modeling_minimaxm3._build_swiglu_oai_dense_mlp). +LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 + + +class M3DenseSwiglu(OpWrapper): + """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, + rounded once to bf16: trtllm::minimax_m3_sm107_tp8_swiglu_oai (minimaxM3Sm107Tp8/swigluOai.cu), bit for bit the + Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the + grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights + during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even + half width; ``__call__`` returns None otherwise or for other activation parameters (the caller runs the Triton + op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP8 dense MLP shard: gate | up 3072 wide (d = 1536), at the decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP8 dense MLP at 1 row", spec=dict(rows=1, d=1536)), + Cell(why="the TP8 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=1536)), + Cell(why="the TP8 dense MLP at 8 rows", spec=dict(rows=8, d=1536)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the + alpha-scaled ex2.approx sigmoid, div.full), checked over every bf16 gate and every bf16 up pattern. + """ + + def __call__( + self, + x: torch.Tensor, + limit: Optional[float] = LIMIT, + alpha: Optional[float] = ALPHA, + beta: Optional[float] = BETA, + use_pdl: bool = True, + ) -> Optional[torch.Tensor]: + if ( + (limit, alpha, beta) != (LIMIT, ALPHA, BETA) + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 4 != 0 + or x.data_ptr() % 4 != 0 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp8_swiglu_oai(x, out, bool(use_pdl)) + return out + + def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta + ) + + def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_dense_swiglu = M3DenseSwiglu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py new file mode 100644 index 000000000000..8d4a56e3b437 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""BF16 x BF16 -> BF16 low-M GEMM of the Eagle3 drafter's linears at fixed FlashInfer tactics (SM10x).""" + +from __future__ import annotations + +from typing import Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: (N, K) of the TP8 drafter shards -> ((max rows, kind, tactic), ...), rows ascending. Kind "d" is FlashInfer's direct +#: SIMT kernel with tactic (block, outputs per block, rows per block); kind "s" is its split-K tensor-core kernel with +#: tactic (mma_m, mma_n, split_k, ab_stages). A row count is served by the first entry whose max rows holds it; the +#: rows modeling's DRAFT_LOWM covers are only the ones measured (the fastest tactic at no larger error vs fp64 than +#: cuBLAS). TP8 shards (64 query / 4 KV heads of 128, KV heads replicated across rank pairs; intermediate 18432): +#: qkv 1280 x 12288, o_proj 6144 x 1024, gate_up 4608 x 6144, down 6144 x 2304. Rubin SM107 TP8 sweep 777890 (mclk +#: locked, rotating weights >= 3x L2, CUDA-graph timing), us cuBLAS -> low-M at 1 / 4 rows: qkv 6.15 -> 3.94 / 6.26 -> +#: 5.69, o 4.97 -> 2.42 / 4.48 -> 2.68, gate_up 8.09 -> 5.78 / 8.11 -> 6.99, down 5.62 -> 4.00 / 5.74 -> 3.94. The fc is +#: replicated (the TP2 shape): its TP2 tactic and row. +TACTICS: Mapping[Tuple[int, int], tuple] = { + # qkv + (1280, 12288): ((1, "d", (256, 2, 1)), (4, "d", (128, 4, 4))), + # o_proj + (6144, 1024): ((1, "s", (64, 8, 1, 10)), (4, "s", (64, 8, 2, 6))), + # gate_up + (4608, 6144): ((1, "d", (128, 2, 1)), (4, "d", (64, 4, 4))), + # down + (6144, 2304): ((1, "s", (64, 8, 2, 6)), (4, "s", (64, 8, 1, 10))), + # the Eagle3 fc over the three captured hidden states + (6144, 18432): ((1, "d", (384, 2, 1)),), +} +#: Rows each (N, K) is certified at (a subset of what its table would serve). +ROWS: Mapping[Tuple[int, int], Tuple[int, ...]] = { + (1280, 12288): (1, 4), + (6144, 1024): (1, 4), + (4608, 6144): (1, 4), + (6144, 2304): (1, 4), + (6144, 18432): (1,), +} + + +def tactic_for(n: int, k: int, m: int) -> Optional[Tuple[str, tuple]]: + """(kind, tactic) of an [m, k] x [n, k]^T call, or None when the table does not serve it.""" + for rows, kind, tactic in TACTICS.get((n, k), ()): + if m <= rows: + return kind, tactic + return None + + +class M3DraftLowM(OpWrapper): + """``out`` [M, N] BF16 = ``x`` [M, K] BF16 @ ``weight`` [N, K]^T BF16, FP32 accumulation. + + ``kind`` / ``tactic`` pick FlashInfer's kernel: "d" runs ``run_direct_dense`` (SIMT, one dot product per output in + K order), "s" runs ``run_splitk_dense`` (tcgen05 MMA, K split over ``split_k`` CTAs and summed in FP32). Both + round once to BF16. The kernel is compiled on its first call per (M, N, K, tactic, PDL); the target compiles every + covered call before warmup (``M3Eagle3Draft.prebuild_linears``). + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS: tuple[Cell, ...] = ( + Cell(why="drafter qkv, 1 rows: d tactic", spec=dict(nk=(1280, 12288), m=1)), + Cell(why="drafter qkv, 4 rows: d tactic", spec=dict(nk=(1280, 12288), m=4)), + Cell(why="drafter o_proj, 1 rows: s tactic", spec=dict(nk=(6144, 1024), m=1)), + Cell(why="drafter o_proj, 4 rows: s tactic", spec=dict(nk=(6144, 1024), m=4)), + Cell(why="drafter gate_up, 1 rows: d tactic", spec=dict(nk=(4608, 6144), m=1)), + Cell(why="drafter gate_up, 4 rows: d tactic", spec=dict(nk=(4608, 6144), m=4)), + Cell(why="drafter down, 1 rows: s tactic", spec=dict(nk=(6144, 2304), m=1)), + Cell(why="drafter down, 4 rows: s tactic", spec=dict(nk=(6144, 2304), m=4)), + Cell(why="drafter fc, 1 rows: d tactic", spec=dict(nk=(6144, 18432), m=1)), + ) + + note = """ + The table is the TP8 drafter's: its shapes are the per-rank shards of the Eagle3 layer at 8-way tensor + parallelism (64 query and 4 KV heads, the KV heads replicated across rank pairs; intermediate 18432) and the + replicated fc, at the b1 decode rows (1 row in draft passes 1-2, 4 in pass 0). Certified at those shapes and + rows only; FlashInfer's kernels take other shapes, which are not measured here. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + if kind == "d": + from flashinfer.gemm.kernels.dense_bf16_gemm_direct import ( + DirectTactic, + run_direct_dense, + ) + + run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) + elif own_splitk: + # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # it also pushes the rows into a drafter boundary all-reduce (modeling's DRAFT_AR_PUSH). + from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push + + if ar is not None: + run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) + else: + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + else: + from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import ( + SplitKTactic, + run_splitk_dense, + ) + + run_splitk_dense(x, weight.t(), None, out, use_pdl, SplitKTactic(*tactic)) + return out + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """bf16(x @ weight^T) summed in fp64 and rounded once; ``out``, ``kind``, ``tactic`` and ``use_pdl`` only + choose where and how the op writes.""" + return (x.double() @ weight.double().t()).to(torch.bfloat16) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + out: torch.Tensor, + kind: str, + tactic: tuple, + use_pdl: bool, + own_splitk: bool = False, + ar: Optional[tuple] = None, + ) -> None: + """The kernels index ``x`` and ``out`` densely: a strided view is read as if it were packed.""" + assert x.dim() == 2 and x.is_contiguous(), "x must be a contiguous [M, K] matrix" + assert weight.dim() == 2 and weight.is_contiguous(), ( + "weight must be a contiguous [N, K] matrix" + ) + assert out.is_contiguous() and tuple(out.shape) == (x.shape[0], weight.shape[0]), ( + "out must be a contiguous [M, N] matrix" + ) + assert ar is None or (own_splitk and kind == "s" and len(ar) == 5), ( + "ar (the boundary all-reduce push) runs only on this target's split-K copy" + ) + + +m3_draft_lowm = M3DraftLowM() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py new file mode 100644 index 000000000000..9702cc0778a4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + + +class M3Fp32Linear(OpWrapper): + """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the + fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies + (the engine leaves torch's default).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), + Cell( + why="a generation step above the tensor-core router's 48 rows", + spec=dict(m=64, e=128, k=6144), + ), + ) + + note = """ + Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against + the true fp32 reference. + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + return F.linear(x.to(torch.float32), w) + + def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + with true_fp32_matmul(): + return x.float() @ w.t() + + +m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py new file mode 100644 index 000000000000..a16702f2cc1a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 x MXFP8 GEMV (tcgen05, block-scaled, CuTe DSL) on an activation its producer already quantized.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_dequant +from ..mx_gemv import op as _gemv + +__all__ = ["PROJECTIONS", "gemv_is_valid", "m3_mx_gemv", "served_cells"] + +# (N, K) of this target's MXFP8 projections at tp2, the GEMV's tuned shapes. +QKV_SPARSE = (1536, 6144) # the fused QKV + index projection (index values disabled) +QKV = (1280, 6144) +O_PROJ = (6144, 1024) +SHARED_GATE_UP = (768, 6144) +SHARED_DOWN = (6144, 384) +DENSE_GATE_UP = (3072, 6144) +DENSE_DOWN = (6144, 1536) +PROJECTIONS = { + "sparse-layer QKV": QKV_SPARSE, + "dense-attention QKV": QKV, + "o_proj": O_PROJ, + "shared gate_up": SHARED_GATE_UP, + "shared down": SHARED_DOWN, + "dense gate_up": DENSE_GATE_UP, + "dense down": DENSE_DOWN, +} + + +def served_cells(named: tuple[Cell, ...], x_bf16: bool, max_m: int) -> tuple[Cell, ...]: + """A cell for every (projection, rows) the GEMV serves up to ``max_m`` rows that ``named`` does not list. + + Certification is per projection: a projection's split-K, ring and MMA width differ, so each (N, K, M) the + decode dispatch can send here is its own cell. + """ + listed = {(cell.spec["nk"], cell.spec["m"]) for cell in named} + return tuple( + Cell(why=f"{name}, {m} rows", spec=dict(nk=nk, m=m)) + for name, nk in PROJECTIONS.items() + for m in range(1, max_m + 1) + if (nk, m) not in listed and _gemv.serves(nk[0], nk[1], m, x_bf16) + ) + + +def gemv_reference( + x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor +) -> torch.Tensor: + """bf16(x @ dequant(weight)^T) for fp32 activation values x [M, K], summed in fp64 and rounded once.""" + return (x.double() @ mxfp8_dequant(weight, weight_scale).double().t()).to(torch.bfloat16) + + +def gemv_is_valid(weight: torch.Tensor, weight_scale: torch.Tensor) -> None: + n, k = weight.shape + # The kernel reads both from data_ptr() as dense buffers at the layout it was compiled for. + assert weight.dtype == torch.float8_e4m3fn and weight.is_contiguous(), ( + "weight must be contiguous e4m3 [N, K]" + ) + assert weight_scale.numel() == n * k // 32, ( + "weight_scale must be the 128x4-swizzled UE8M0 buffer [N * K / 32]" + ) + + +class M3MxGemv(OpWrapper): + """bf16 [M, N] = xq [M, K] (e4m3, 128x4-swizzled UE8M0 scales ``xsf``) x weight [N, K]^T (e4m3, swizzled + ``weight_scale``), M <= 48; None when the shape is not served (the caller falls back). + + ``trigger_early``: release the PDL dependents before the epilogue. ``cfg``: a (split, ring[, contig, nacc, a_tmem, + st_async]) tuning override; default the table's. ``use_pdl``: default the environment's. ``ar`` = (multicast + workspace address, flags address, rank, world, bucket tokens, mode): also push the rows into the next boundary + all-reduce's exchange (then ``op.last_pushed`` is True). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The projections' (N, K) at the token counts each build differs in: the MMA width (8 columns up to the TP2 decode + # step's 8 rows) is one, so one cell per projection pins the shape, and the rest are every served row count. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token (split 8)", spec=dict(nk=QKV_SPARSE, m=1)), + Cell( + why="sparse-layer QKV, 8 rows: two requests verifying 1 + 3", + spec=dict(nk=QKV_SPARSE, m=8), + ), + Cell( + why="dense-attention QKV, 4 rows: one request verifying 1 + 3", spec=dict(nk=QKV, m=4) + ), + Cell(why="o_proj, one token (split 2, K 1024)", spec=dict(nk=O_PROJ, m=1)), + Cell(why="o_proj, 8 rows", spec=dict(nk=O_PROJ, m=8)), + Cell(why="shared gate_up, 4 rows (split 8)", spec=dict(nk=SHARED_GATE_UP, m=4)), + Cell(why="shared down, one token (split 1, K 384)", spec=dict(nk=SHARED_DOWN, m=1)), + Cell(why="dense gate_up, 8 rows (N 3072)", spec=dict(nk=DENSE_GATE_UP, m=8)), + Cell(why="dense down, 4 rows (K 1536)", spec=dict(nk=DENSE_DOWN, m=4)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling's MX_GEMV_ROWS). + CELLS = CELLS + served_cells(CELLS, x_bf16=False, max_m=8) + + note = """ + The kernel sums in fp32 in split-K order where `reference` sums in fp64: the bf16 outputs agree within the band the + test derives (SM107: at most as many misrounded elements as the stock FlashInfer CuTe DSL GEMM on the same + operands, rand / wide-scale / cancellation / tie / zero inputs). Not covered: the all-reduce push (`ar`), which + needs a multicast exchange; `cfg` overrides beyond the table (the target passes none). Shapes outside the table + and M > 48 return None; the target sends at most 8 rows. + """ + + def __call__( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv(xq, xsf, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The dequantized product; the launch options are ignored.""" + return gemv_reference(mxfp8_dequant(xq, xsf), weight, weight_scale) + + def is_valid( + self, + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + m, k = xq.shape + assert xq.dtype == torch.float8_e4m3fn and xq.is_contiguous(), ( + "xq must be contiguous e4m3 [M, K]" + ) + # A linear-layout [M, K / 32] scale buffer is read as the swizzled one: wrong scales, no error. + assert xsf.numel() >= (m + 127) // 128 * 128 * (k // 32), ( + "xsf must be the 128x4-swizzled scale buffer (rows padded to 128), not the linear one" + ) + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv = M3MxGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..159b49ff69af --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mx_gemv_bf16.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode MXFP8 GEMV (CuTe DSL) on a bf16 activation, quantized to MXFP8 in the kernel's prologue.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..moe.m3_moe_engine import mxfp8_quant_dequant +from ..mx_gemv import op as _gemv +from .m3_mx_gemv import ( + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + SHARED_DOWN, + SHARED_GATE_UP, + gemv_is_valid, + gemv_reference, + served_cells, +) + +__all__ = ["m3_mx_gemv_bf16"] + + +class M3MxGemvBf16(OpWrapper): + """bf16 [M, N] = mxfp8_quantize(x) x weight [N, K]^T for a contiguous, 16-byte aligned bf16 ``x`` [M, K]; the + quantization is mxfp8_quantize's (UE8M0 scale per 32 = the smallest power of two >= amax / 448). None when the + shape is not served. The other arguments as ``m3_mx_gemv``'s. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The prologue quantizes one 32-element block per thread and round: the projections at the token counts the + # target sends here (no producer-quantized activation), every TP2 decode row count. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, one token", spec=dict(nk=QKV_SPARSE, m=1)), + Cell(why="dense-attention QKV, 8 rows", spec=dict(nk=QKV, m=8)), + Cell(why="o_proj, 4 rows (the attention output in bf16)", spec=dict(nk=O_PROJ, m=4)), + Cell(why="shared gate_up, 8 rows", spec=dict(nk=SHARED_GATE_UP, m=8)), + Cell(why="shared down, 4 rows (K 1536)", spec=dict(nk=SHARED_DOWN, m=4)), + Cell( + why="dense gate_up, 8 rows: 24 K tiles per split-2 CTA", + spec=dict(nk=DENSE_GATE_UP, m=8), + ), + Cell(why="dense down, one token", spec=dict(nk=DENSE_DOWN, m=1)), + ) + # Every other (projection, rows) the decode dispatch can send here (modeling's MX_GEMV_BF16_MAX_TOKENS). + CELLS = CELLS + served_cells(CELLS, x_bf16=True, max_m=8) + + note = """ + As `m3_mx_gemv`. The in-kernel quantization is mxfp8_quantize's bit for bit except where the kernel's + rcp.approx(448) and an exact 1/448 put a block's amax on different sides of a power of two (not observed on the + cells). Not covered: `ar`, `cfg` overrides. Returns None for rows beyond 8 rounds of the prologue per CTA. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> Optional[torch.Tensor]: + return _gemv.m3_mx_gemv_bf16(x, weight, weight_scale, trigger_early, cfg, use_pdl, ar) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> torch.Tensor: + """The MXFP8 round trip of x times the dequantized weight; the launch options are ignored.""" + return gemv_reference(mxfp8_quant_dequant(x), weight, weight_scale) + + def is_valid( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, + ) -> None: + gemv_is_valid(weight, weight_scale) + + +m3_mx_gemv_bf16 = M3MxGemvBf16() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py new file mode 100644 index 000000000000..2bfc6ee495fb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned + +from .m3_mxfp8_linear import mxfp8_reference + + +class M3Mxfp8Gemm(OpWrapper): + """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight + [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key + and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The decode dispatch hands it the boundary / merge MXFP8 activation when the GEMV declines rows it serves: the + # sparse-layer QKV and the o_proj at the tp8 shards. + CELLS: tuple[Cell, ...] = ( + Cell(why="sparse-layer QKV, 4 rows", spec=dict(m=4, n=1536, k=6144)), + Cell(why="o_proj, 32 rows", spec=dict(m=32, n=6144, k=1024)), + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=384)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=384)), + ) + + note = """ + Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and + multiplies in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) + + def reference( + self, + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + output_dtype: torch.dtype, + tune: bool, + ) -> torch.Tensor: + """The dequantized product; ``tune`` is ignored.""" + return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) + + +m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py new file mode 100644 index 000000000000..3f07fe2fb256 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned + +_BLOCK = 32 + + +def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" + pad_r = (rows + 127) // 128 * 128 + pad_c = (cols + 3) // 4 * 4 + s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) + return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] + + +def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + rows, k = x.shape + s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) + factor = torch.pow(2.0, (s - 127).to(torch.float32)) + return ( + x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) + ).view(rows, k) + + +def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: + """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" + with true_fp32_matmul(): + return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=1536, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=1024, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=1536, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=768, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=1024, path="tuned")), + ) + + note = """ + Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` + quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and + multiplies dequantized operands in fp32; the test's band covers the accumulation order. + """ + + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + if tuned_quantize_gemm: + return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) + act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + if flashinfer: + return torch.ops.trtllm.flashinfer_mm_mxfp8( + act, act_scale, weight, weight_scale, output_dtype + ) + gemm = ( + torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned + if native_tuning + else torch.ops.trtllm.mxfp8_mxfp8_gemm + ) + return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) + + def reference( + self, + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + tuned_quantize_gemm: bool, + tune: bool, + flashinfer: bool, + native_tuning: bool, + ) -> torch.Tensor: + """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are + ignored.""" + m, k = x.shape + blocks = x.float().view(m, k // _BLOCK, _BLOCK) + amax = blocks.abs().amax(-1, keepdim=True) + exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) + scale = torch.pow(2.0, exp) + q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale + with true_fp32_matmul(): + out = q.view(m, k) @ _dequant(weight, weight_scale).t() + return (out * alpha.float()).to(output_dtype) + + +m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py new file mode 100644 index 000000000000..f8deb858c843 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, +and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). + +The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values +``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 +marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a +per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: +per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 +code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved +it decodes to the BF16 weight bit for bit. +""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +__all__ = ["m3_target_head"] + +#: Verify rows per call (requests x (draft length + 1)). +MAX_ROWS = 8 +#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). +CTA_WARPS = 16 +#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. +TARGET_SHARD = 50016 +HIDDEN = 6144 +#: Codes per window: 0..14 decode, 15 marks an exception. +WINDOW = 15 +#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. +MAX_BASE = 254 - WINDOW +#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). +ROWS_PER_CHUNK = 4096 + + +def _pad_rows(weight: torch.Tensor) -> torch.Tensor: + n = weight.shape[0] + tiles = (n + 15) // 16 + if tiles * 16 == n: + return weight + return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) + + +def _fragment_view(w: torch.Tensor) -> torch.Tensor: + """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" + n, k = w.shape + # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) + return ( + w.view(n // 16, 2, 8, k // 16, 2, 4, 2) + .permute(0, 3, 2, 5, 4, 1, 6) + .reshape(n // 16, k // 16, 32, 8) + ) + + +def _exponents(weight: torch.Tensor) -> torch.Tensor: + return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF + + +def _window_base(hist: torch.Tensor) -> int: + """The base of the 15-exponent window holding the most weights (the lowest such base).""" + window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) + return int(window.argmax()) + + +def encode(weight: torch.Tensor, base: Optional[int] = None): + """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's + own window. + + Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile + << 16 | k), exc_w fp32 [E, 2] ({w, 0})). + """ + w = _pad_rows(weight) + n, k = w.shape + tiles, ks = n // 16, k // 16 + assert ks % 4 == 0, "K must be a multiple of 64" + bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF + expo = (bits >> 7) & 0xFF + if base is None: + base = _window_base(torch.bincount(expo.flatten(), minlength=256)) + code = expo - base + exc = (code < 0) | (code >= WINDOW) + code = torch.where(exc, torch.full_like(code, WINDOW), code) + sm = ((bits >> 8) & 0x80) | (bits & 0x7F) + sm = torch.where(exc, torch.zeros_like(sm), sm) + # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. + rows, cols = exc.nonzero(as_tuple=True) + order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) + rows, cols = rows[order], cols[order] + exc_w = torch.stack( + [ + w[rows, cols].float(), + torch.zeros((rows.numel(),), device=w.device), + ], + dim=1, + ) + exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) + exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) + exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( + torch.int32 + ) + # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. + smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) + cf = ( + _fragment_view(code.to(torch.uint8).view(n, k)) + .view(tiles, ks // 4, 4, 32, 8) + .to(torch.int32) + ) + # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. + cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) + for p in range(4): + cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) + cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 + # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. + seg01 = ( + smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) + ) # [.., 2, 32, 16] + seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) + packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] + return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() + + +def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): + """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" + n, _ = weight.shape + assert rows_per_chunk % 16 == 0 + hist = torch.zeros(256, dtype=torch.int64, device=weight.device) + for r in range(0, n, rows_per_chunk): + hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) + base = _window_base(hist) + packed, offs, poss, ws = [], [], [], [] + n_exc = 0 + for r in range(0, n, rows_per_chunk): + p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) + packed.append(p) + offs.append(off[:-1] + n_exc) + poss.append(pos) + ws.append(ew) + n_exc += pos.numel() + exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) + return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) + + +def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: + """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" + tiles, ks = (n + 15) // 16, k // 16 + p = packed.view(tiles, ks // 4, 3, 32, 16) + sm = ( + p[:, :, :2] + .permute(0, 1, 3, 2, 4) + .reshape(tiles, ks // 4, 32, 4, 8) + .permute(0, 1, 3, 2, 4) + .to(torch.int32) + ) + cw = ( + p[:, :, 2] + .contiguous() + .view(torch.int32) + .view(tiles, ks // 4, 32, 4) + .permute(0, 1, 3, 2) + .long() + & 0xFFFFFFFF + ) + codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( + torch.int32 + ) + bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) + bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) + frag = bits.reshape(tiles, ks, 32, 8) + # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). + w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) + return w.to(torch.int16).view(torch.bfloat16)[:n] + + +def decode( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + n: int, + k: int, + row0: int = 0, + row1: Optional[int] = None, +) -> torch.Tensor: + """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). + ``row0`` must be a multiple of 16.""" + row1 = n if row1 is None else row1 + assert row0 % 16 == 0 + t0, t1 = row0 // 16, (row1 + 15) // 16 + unit_bytes = (k // 64) * 3 * 32 * 16 + dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() + counts = (exc_off[1:] - exc_off[:-1]).long() + tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) + rows = tile_of * 16 + (exc_pos >> 16).long() + cols = (exc_pos & 0xFFFF).long() + sel = (rows >= row0) & (rows < row1) + dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) + return dec + + +def round_trip_ok( + packed: torch.Tensor, + base: int, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + weight: torch.Tensor, + rows_per_chunk: int = ROWS_PER_CHUNK, +) -> bool: + """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent + window below the BF16 infinity exponent.""" + n, k = weight.shape + if not 0 <= base <= MAX_BASE or k % 64: + return False + if ( + packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 + or exc_off.numel() != (n + 15) // 16 + 1 + ): + return False + for r0 in range(0, n, rows_per_chunk): + r1 = min(n, r0 + rows_per_chunk) + dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) + if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): + return False + return True + + +class M3TargetHead(OpWrapper): + """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit + the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in + another summation order). + + ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` + (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are + the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and + # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with + # a shard whose rows are not a multiple of 16. + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"{m} verify row(s) of the TP4 target head shard", + spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), + ) + for m in range(1, MAX_ROWS + 1) + ), + Cell( + why="4 rows with outliers and exact zeros in many tiles (the exception lists)", + spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), + ), + Cell( + why="2 rows over a shard of 4000 rows (250 tiles), K 1024", + spec=dict(m=2, n=4000, k=1024, exceptions=True), + ), + ) + + note = """ + The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception + adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs + from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); + integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp + order: the result is deterministic. + """ + + def __call__( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_target_head( + w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl + ) + return out + + def reference( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> torch.Tensor: + """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. + + ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" + weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) + with true_fp32_matmul(): + logits = x.float() @ weight.float().t() + return logits.to(torch.bfloat16) + + def is_valid( + self, + w: torch.Tensor, + x: torch.Tensor, + out: torch.Tensor, + part: torch.Tensor, + counter: torch.Tensor, + exc_off: torch.Tensor, + exc_pos: torch.Tensor, + exc_w: torch.Tensor, + exp_base: int, + n: int, + grid: int, + pdl: bool = True, + ) -> None: + # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count + # stores a partial sum. + assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + assert out.data_ptr() != x.data_ptr(), "out must not alias x" + + +m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/long_select.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/long_select.py new file mode 100644 index 000000000000..79aaf6674300 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/long_select.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of 129 to 9216 blocks. + +Same contract and bit-identical output as ``torch.ops.trtllm.minimax_m3_select_blocks`` on per-KV-head scores +[num_kv_heads, n_blocks, total_q]: init / local forcing, per-query valid-block count, ties to the smaller block index, +forced -inf picks emitted as -1, ascending block ids with -1 tail padding. + +One 1024-thread CTA per (query, KV head) row holds the row's ordered score keys in registers (2, 4 or 9 per +thread, by the row's valid count). The weakest of the 16 maxima of the interleaved groups ``block % 16`` is a lower +bound on the 16th-largest key, so the top 16 is among the keys at or above it; those are compacted to shared memory +and ranked exactly by (key, block). More than 1024 keys tied at the bound (degenerate rows) rank every candidate +against the whole row (re-read). + +The kernel is ``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long`` (this target's copy of ``minimax_m3_select_long`` +with a choice of where it releases the sparse decode), called through this target's ``m3_select_long`` entry, or with +``dsl`` the same selection from one 384-thread CTA per row (``m3_select_long_dsl``, ops/m3_select_long_dsl). +""" + +from typing import Optional + +import torch + +from ..ops.attention.m3_select_long import MAX_BLOCKS, m3_select_long +from ..ops.attention.m3_select_long_dsl import m3_select_long_dsl + + +def long_select_blocks( + max_score_kv: torch.Tensor, + *, + n_valid_blocks: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, + dsl: bool = False, +) -> Optional[torch.Tensor]: + """Top-16 select over per-KV-head scores [num_kv_heads, n_blocks, total_q] with 129..9216 blocks. + + Returns the head-major backed [total_q, num_kv_heads, 16] int32 table, or None outside this kernel's shapes + (callers keep the reference op). ``early_trigger`` releases the dependent sparse decode right after the select's + grid wait (modeling's ``SELECT_EARLY_MAX_ROWS``); otherwise ``trigger_phase`` 1 releases it once every score of the + row is loaded and 0 at the select's exit (modeling's ``SELECT_LONG_TRIGGER_PHASE``). ``dsl`` runs the 384-thread + kernel (modeling's ``SELECT_LONG_DSL``; the same table). + """ + if ( + max_score_kv.dtype != torch.float32 + or not max_score_kv.is_cuda + or max_score_kv.dim() != 3 + or not 128 < max_score_kv.shape[1] <= MAX_BLOCKS + or n_valid_blocks.dtype != torch.int32 + or not n_valid_blocks.is_cuda + or n_valid_blocks.shape[0] < max_score_kv.shape[2] + ): + return None + num_kv_heads, _, total_q = max_score_kv.shape + out = torch.empty( + (num_kv_heads, total_q, 16), dtype=torch.int32, device=max_score_kv.device + ).permute(1, 0, 2) + return (m3_select_long_dsl if dsl else m3_select_long)( + max_score_kv, + n_valid_blocks.contiguous(), + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_attn_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_attn_lc8.py new file mode 100644 index 000000000000..2d446f2b45d3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_attn_lc8.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 long-context sparse decode attention at 8 query heads in one launch. + +Replaces the generic chain's split-K producer (``_gqa_sparse_decode_kernel``, one block per chunk) and the LSE +merge (``_merge_topk_attn_out_kernel``, with its MXFP8 copy for the o_proj) once the top-k block table exists. + +Grid: one cluster of 16 CTAs per query token; CTA ``r`` attends the token's ``r``-th selected block for all 8 heads, +keeps the partial in shared memory, and the cluster merges the 16 partials over DSMEM into the BF16 row and, +optionally, its MXFP8 copy and swizzled UE8M0 scales exactly as ``_store_mxfp8_head`` lays them out. + +Per-block math follows the Triton producer: NVFP4 K/V dequantized to FP16 (exact), FP8 q widened to FP16, QK and PV +on FP16 tensor cores with FP32 accumulation, log2-domain softmax with either the linear denominator or the +log-domain recurrence. Accumulation order inside the MMAs and the reductions differs, so results match to rounding, +not bit for bit. + +Schedule: the kernel releases its PDL dependents right after its grid wait, so the o_proj launches and streams its +weights while attention runs. Up to ``PREWAIT_MAX_ROWS`` query rows it also stages the request length, +the query rows and the request's block-table row before that wait, overlapping them with the select ahead when it +releases this grid early (modeling's ``SELECT_EARLY_MAX_ROWS``). Neither changes a value. + +Geometry: 8 query heads per KV head (TP8), one KV head, head_dim 128, page 128, top-k 16. The kernel is +``torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8``, called through this target's ``m3_sparse_decode_lc8`` entry. +""" + +import torch + +from ..ops.attention.m3_sparse_decode_lc8 import NUM_HEADS, WIDTH, m3_sparse_decode_lc8 + +# Up to this many query rows the kernel stages the request length, the query rows and the block-table row before its +# grid wait. +PREWAIT_MAX_ROWS = 8 # capacity: fixed: measured occupancy crossover +# The split-K merge forms its weights once per CTA (MERGE_ONCE: also enables the identity route and the L2 prefetches) +# and divides once per output (ONE_DIV_MERGE: sum_c w_c acc_c / sum_c w_c l_c in chunk order; deterministic, in the +# float64 reference's band). +MERGE_ONCE = True +ONE_DIV_MERGE = True + + +class PrefetchSlot: + """The read-only tensors the next long-context sparse decode prefetches into L2 (the attention layer's o_proj weight + and scales), set by the layer around its decode call and taken by the decode; empty means no prefetch.""" + + def __init__(self) -> None: + self.tensors = None + + def take(self): + tensors, self.tensors = self.tensors, None + return tensors + + +PREFETCH = PrefetchSlot() +# (this layer's id buffer, the next sparse layer's id buffer, its K / V pools, its K / V block scales) +NEXT_KV = PrefetchSlot() + + +def lc8_applies(q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output) -> bool: + """Whether the step's operands fit the kernel: TP8 geometry, packed NVFP4 pages, FP8 q, BF16 [T, 8, 128] out.""" + if q.dim() != 3 or tuple(q.shape[1:]) != (NUM_HEADS, 128) or q.dtype != torch.float8_e4m3fn: + return False + if q.stride(2) != 1 or q.stride(1) % 16 or q.stride(0) % 16 or q.data_ptr() % 16: + return False + if int(k_paged.shape[1]) != 1 or int(topk_idx.shape[0]) != 1 or int(topk_idx.shape[-1]) != 16: + return False + for t in (k_paged, v_paged): + if ( + t.stride(3) != 1 + or t.stride(2) != 64 + or int(t.shape[3]) != 64 + or t.data_ptr() % 16 + or t.stride(0) % 16 + ): + return False + for t in (k_block_scale, v_block_scale): + if t is None or t.stride(3) != 1 or t.stride(2) != 8 or t.data_ptr() % 8 or t.stride(0) % 8: + return False + return ( + output.dtype == torch.bfloat16 + and output.dim() == 3 + and output.is_contiguous() + and tuple(output.shape) == tuple(q.shape) + ) + + +def m3_lc8_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + linear_softmax=True, + mxfp8=None, + use_pdl=True, + identity=False, + identity_step=False, + prefetch=None, + next_kv=None, +): + """Drop-in for ``minimax_m3_sparse_attn_decode`` on the TP8 NVFP4 shape; callers check ``lc8_applies`` first. + + ``topk_idx`` is the head-major [1, total_q, 16] view the Triton producer takes. ``mxfp8`` is ``(q8, sf)`` from + ``MERGE_MXFP8.buffers(1024, dev)`` to also emit the o_proj hand-off. With ``MERGE_ONCE``, + ``prefetch`` (up to two contiguous tensors) is pulled into L2 by the op's CTAs before their grid dependency, and + ``next_kv`` drives the next-layer KV prefetch. + """ + merge_once = MERGE_ONCE + total_q = int(q.shape[0]) + q8, sf = ( + (mxfp8[0].view(torch.uint8), mxfp8[1].view(torch.uint8)) + if mxfp8 is not None + else (None, None) + ) + m3_sparse_decode_lc8( + q.view(torch.uint8), + k_paged.view(torch.uint8), + v_paged.view(torch.uint8), + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + k_global_scale, + v_global_scale, + topk_idx[0], + block_table, + seq_lens, + output.view(total_q, WIDTH), + q8, + sf, + WIDTH // 128, + int(decode_query_len), + float(sm_scale), + bool(linear_softmax), + bool(use_pdl), + total_q <= PREWAIT_MAX_ROWS, + merge_once=merge_once, + identity=bool(identity) and merge_once, + identity_step=bool(identity_step) and merge_once, + one_div=ONE_DIV_MERGE and merge_once, + prefetch=prefetch if merge_once else None, + next_kv=( + tuple(next_kv[:2]) + tuple(t.view(torch.uint8) for t in next_kv[2:]) + if merge_once and next_kv is not None + else None + ), + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py new file mode 100644 index 000000000000..c977362e92b8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 draft pass for the MiniMax-M3 drafter's vocab-sharded BF16 head (TP decode, <= 8 requests). + +One kernel replaces a draft pass's head and the draft loop's per-pass bookkeeping, at the in-tree precision: + +* The drafter's BF16 lm_head shard is kept a second time, built by ``prepare_head`` before the executor's warmup: in + the lossless 12-bit coding of the target head (``ops.gemm.m3_target_head``; 75% of the bytes, kept only when it + decodes to the BF16 shard bit for bit), else in ``mma.m16n8k16`` A-fragment order (``bf16_fragments``, the same + values). The kernel streams it with BF16 MMAs and FP32 accumulation over the BF16 head input rows; a coded + exception decodes to +0 and its exact product x * w is added at its tile's flush. +* Each logit is rounded to BF16 (the in-tree logits are the BF16 head GEMM's output) and keyed + ``ordered(value) << 32 | ~global_index`` with -0 as +0, so the larger value wins and ties take the lowest index, as + the in-tree ``torch.max`` / ``torch.argmax`` over the gathered logits (``greedy_sample_draft_with_tp_gather``). +* The last CTA exchanges each row's key across the TP ranks through a multicast mapping (Lamport slots, see + ``m3_shard_accept``), then stores the draft token, the next pass's input id and position, and advances the draft KV + lengths as the in-tree draft loop does between passes. + +The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/draftStep.cu) and called through the +catalog entry ``ops.spec.m3_draft_step``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from tensorrt_llm.logger import logger + +from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok +from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments, max_tile_pieces +from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op +from . import m3_report +from .m3_shard_accept import AcceptWorkspace, create_workspace + +_ENGAGED: set = set() + + +# The fused draft loop streams the 12-bit coded copy of the draft head when it decodes bit for bit; off (the coded +# head is slower on sm_107: the BF16 fragments stream). +HEAD_CODED = False +# The head copy placed per locality domain, each pass as two launches: off at TP8 (one launch over the whole device). +UGPU = False +# The draft step's CTA count (at most one per SM; 0 = one per SM). +GRID = 208 # capacity: fixed: a launch-grid choice, not a size +# The draft head sums a tile shared by more than 2 warp pieces in warp order instead of with atomics (the same sum +# every run; the atomics are deterministic only where a tile takes at most 2 pieces). +FIXED_ORDER = True +FIXED_ORDER_PATH = "m3_draft_step_fixed_order" + + +@dataclass(eq=False) +class DraftHead: + """The streamed copy of one rank's draft lm_head shard and its first global id: ``wf`` is the 12-bit coded stream + with its exception lists (``coded``), or the fragment-ordered BF16 copy (``bf16_fragments``).""" + + weight: torch.Tensor + wf: torch.Tensor + vocab_offset: int + coded: bool = False + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + exp_base: int = 0 + # BF16 fragments placed per locality domain (UGPU below): tiles [0, split_tile) live in domain 0 and + # the rest in domain 1; 0 = not placed (one launch over all SMs). + split_tile: int = 0 + + +def head_supported(lm_head) -> bool: + """A bias-free BF16 head whose local shard the kernel can own: whole 16-row tiles, K a multiple of 16, at most + MAX_TILES_PER_CTA tiles per CTA of a one-CTA-per-SM grid.""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + if getattr(lm_head, "bias", None) is not None: + return False + n, k = weight.shape + if n % 16 or k % 16: + return False + return (n // 16 + _grid() - 1) // _grid() <= MAX_TILES_PER_CTA + + +def prepare_head(lm_head, tp_rank: int) -> DraftHead: + """The streamed copy of ``lm_head``'s local shard (global ids from tp_rank * rows): its 12-bit coding when + ``HEAD_CODED``, K is a multiple of 64 and the coding decodes to the shard bit for bit, else (fail + closed) the fragment-ordered BF16 copy. + Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_draft_head``).""" + weight = lm_head.weight + offset = tp_rank * weight.shape[0] + with torch.no_grad(): + w = weight.detach() + if HEAD_CODED and w.shape[1] % 64 == 0: + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return DraftHead( + weight=weight, + wf=packed, + vocab_offset=offset, + coded=True, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + exp_base=base, + ) + del packed, exc_off, exc_pos, exc_w + wf = bf16_fragments(w) + wf, split_tile = _place_per_domain(wf) + return DraftHead(weight=weight, wf=wf, vocab_offset=offset, split_tile=split_tile) + + +def _place_per_domain(wf: torch.Tensor): + """(wf, split tile): the BF16 fragments split-mapped over the two locality domains at the granule-aligned tile + nearest half (locality.LocalityMap.place_split; its domain streams built here, before capture), or (wf, 0) when + UGPU is off. Placement only: the streamed values are the same.""" + if not UGPU: + return wf, 0 + from .m3_locality import LocalityMap + + lm = LocalityMap.get(wf.device) + tile_bytes = wf[0].numel() * wf.element_size() + step = 1 + while (step * tile_bytes) % lm.granularity: + step += 1 + n_tiles = wf.shape[0] + split_tile = max(step, round(n_tiles / 2 / step) * step) + per = lm.sms_per_domain + if split_tile >= n_tiles or max(split_tile, n_tiles - split_tile) > per * MAX_TILES_PER_CTA: + return wf, 0 + lm.domain_streams() + return lm.place_split(wf, split_tile), split_tile + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _grid() -> int: + """The draft step's CTA count: GRID when set (at most one per SM), else one per SM.""" + return min(GRID, _num_sms()) if GRID > 0 else _num_sms() + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace ([2][MAX_TOKENS][world] slots), created collectively on the first eager call; a + CUDA-graph capture that finds it missing is a miss.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("m3_draft_step", None, "the draft-step exchange workspace") + return None + _WS = create_workspace(mapping, device, MAX_TOKENS) + return _WS + + +def _fixed_order( + rows: int, head: DraftHead, x: torch.Tensor, num_ctas: int, split: bool = False +) -> bool: + """FIXED_ORDER where the launch lets a tile take more than 2 warp pieces (max_tile_pieces; a + split or equal-chunk launch always), with its engagement line (D-FIX); False keeps the atomics, deterministic on + this shape.""" + if not FIXED_ORDER: + return False + eq = getattr(head, "bnd", None) is not None + if not (split or eq): + chunks = x.shape[1] // (16 if head.exc_off is None else 64) + if ( + max_tile_pieces( + head.wf.shape[0] if head.exc_off is None else head.exc_off.numel() - 1, + num_ctas, + chunks, + ) + <= 2 + ): + return False + m3_report.taken(FIXED_ORDER_PATH, rows, "the draft head sums shared tiles in warp order") + return True + + +def draft_step( + head: DraftHead, + x: torch.Tensor, + rows: int, + x_rows: Optional[torch.Tensor], + mapping, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + fold: Optional[tuple] = None, +) -> bool: + """One fused greedy draft pass over ``rows`` requests (see the module docstring); False when the exchange + workspace is missing during a capture. Collective over the TP group. + + Args: + head: the prepared head of this rank (``prepare_head``), coded or BF16 fragments. + x: [*, K] BF16 head input rows (the drafter's final-norm output); row ``x_rows[b]`` (int64) or ``b`` is + request b's. + draft_out: int32 [rows, >= step + 1], unit column stride: column ``step`` receives the tokens. + ids_next / pos_next: int32 [rows]: the tokens, and ``pos_src[x_rows[b] or b] + 1`` (pos_src int32). + kv_lens: int32 draft KV lengths: ``-= draft_len - num_acc`` with ``seq_lens = 1`` on pass 0 when ``num_acc`` + is given, ``+= 1`` on later passes. + fold: on the last pass (modeling's DRAFT_STEP_FOLD_TAIL), (the 14 fold-tail tensors, rows, tokens per step) from + M3SpecStepFold.fused_tail_operands: the step's forward-tail store and next-step staging run by this pass's + last CTA after its tokens, instead of the standalone fold tail. + """ + ws = workspace(mapping, x.device) + if ws is None: + return False + if head.split_tile: + _draft_step_split( + head, + x, + rows, + x_rows, + ws, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + fold, + ) + return True + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + _grid(), + True, + None, + head.exc_off, + head.exc_pos, + head.exc_w, + head.exp_base, + fold=fold, + fixed_order=_fixed_order(rows, head, x, _grid()), + ) + return True + + +def _draft_step_split( + head, + x, + rows, + x_rows, + ws, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + fold=None, +) -> None: + """One pass as two launches, one per locality domain on that domain's green-context stream, each streaming its own + domain's tiles of the split-mapped fragments; the keys' max and the last-CTA exchange and bookkeeping run once over + both (the last of the 2 x SMs-per-domain CTAs). Forked from and joined back to the current stream; no PDL on the + forks.""" + from .m3_locality import LocalityMap + + lm = LocalityMap.get(x.device) + per = lm.sms_per_domain + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + streams = lm.domain_streams() + n_tiles = head.wf.shape[0] + if x.device not in _ENGAGED: + _ENGAGED.add(x.device) + logger.info( + f"[M3 sm_107 tp2] fast path engaged: m3_draft_step_ugpu (split tile {head.split_tile} of {n_tiles}, " + f"{per} CTAs per domain, on {x.device})" + ) + for st, lo, hi in ((streams[0], 0, head.split_tile), (streams[1], head.split_tile, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + _draft_step_op( + x, + x_rows, + rows, + head.wf, + head.vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + per, + False, + None, + None, + None, + None, + 0, + lo, + hi, + 2 * per, + fold=fold, # run by whichever launch holds the pass's last CTA + fixed_order=_fixed_order(rows, head, x, 2 * per, split=True), + ) + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py new file mode 100644 index 000000000000..e582e2309e61 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py @@ -0,0 +1,404 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's decode attention on Rubin: ``M3DrafterFullCtxFmha``. + +The drafter layer is a plain ``TrtllmAttention`` over its own FP8 paged cache (the full context, every draft pass). +On SM 10.7 its generation phase runs ``thop.attention``'s in-tree trtllm-gen runner, whose split is capped at 8 CTAs +per (request, KV head): 45-46 us per pass at 165k tokens on 212 SMs, for 85 MB. + +This library keeps the in-tree preprocessing -- ``thop.trtllm_gen_generation_preprocess``: RoPE, the FP8 K/V append +into the paged cache and the FP8 query, as ``FlashInferTrtllmGenFmha`` runs it -- and attends with the full-context +PrimTS split-KV decode of ``ops/m3_fullctx_decode`` at the same pinned recipes as the dense layers. + +With ``fused_rope_mode`` set, the preprocess and the query cast run as one launch instead +(``minimax_m3_sm107tp8_drafter_rope_fp8_kv``, the same bytes into the query and the cache), the layer's sub-page +pool and block table taken from the KV cache pool pointers and mapping the preprocess reads. + +It serves generation-only steps of the covered (query rows per request, requests) shapes; the target installs it +first in the drafter layer's FMHA list (``install``), so every other step falls through to the stock libraries. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Mapping, Optional, Tuple + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + FlashInferTrtllmGenFmha, +) +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode + +from . import m3_fullctx_decode as fullctx + +# With PDL, the fused RoPE + FP8 K/V launch takes the maximum shared-memory carveout so the drafter decode's large-smem +# CTAs co-reside with it (timing only). +ROPE_CARVEOUT = True + + +def fwd_has_no_mrope(fwd) -> bool: + return getattr(fwd, "mrope_position_deltas", None) is None + + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + +class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): + """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + + def __init__( + self, + attn: "TrtllmAttention", + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + pdl: bool = False, + on_rope_taken=None, + ) -> None: + super().__init__(attn) + self._recipes = recipes + self._max_kv_len = max_kv_len + self._own_reducer = own_reducer + self._fused_rope_mode = fused_rope_mode + self._pdl = pdl + self._on_rope_taken = on_rope_taken + self._on_taken = on_taken + self._on_uncovered = on_uncovered + self._on_missed = on_missed + + def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: + nreq = int(metadata.num_generations) + if nreq <= 0 or int(q.shape[0]) % nreq: + return None + return int(q.shape[0]) // nreq, nreq + + def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: + if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: + return False + if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: + return False + attn = self.attn + if ( + attn.num_heads != fullctx.HEADS + or attn.num_kv_heads != fullctx.KV_HEADS + or attn.head_dim != fullctx.HEAD_DIM + or not QuantMode(attn.quant_mode).has_fp8_kv_cache() + ): + self._on_uncovered("not the TP8 FP8 drafter geometry") + return False + # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's + # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. + if ( + forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant + or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig + ): + self._on_uncovered("per-step KV cache scales") + return False + if int(metadata.max_seq_len) > self._max_kv_len: + self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") + return False + shape = self._shape(q, metadata) + if shape is None or shape not in self._recipes: + self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") + return False + # Compiled before warmup (``install``); the forward never compiles one. + if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: + self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") + return False + return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) + + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + attn = params.attn + width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim + rotary_dim = int(attn.rope_params.dim) + return ( + self._fused_rope_mode is not None + and params.qkv_input is not None + and attention_input.dim() == 2 + and attention_input.dtype == torch.bfloat16 + and attention_input.is_contiguous() + and int(attention_input.shape[1]) == width + and not params.is_cross + and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) + and attn.rotary_cos_sin is not None + and attn.rotary_cos_sin.dtype == torch.float32 + and fwd_has_no_mrope(params.fwd) + and 0 < rotary_dim <= attn.head_dim + and rotary_dim % 16 == 0 + and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) + # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo + # runs for several requests only) and reads the workspace's, so it keeps that path. + and not (self._multi_token(params) and int(params.batch_size) == 1) + ) + + @staticmethod + def _multi_token(params: FmhaParams) -> bool: + return ( + params.spec_decoding_generation_lengths is not None + and int(params.attn.predicted_tokens_per_seq) > 1 + ) + + def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" + if not self._fused_rope_applies(params, attention_input): + return False + attn = params.attn + meta = params.meta + qlen, nreq = int(params.input_seq_length), int(params.batch_size) + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + if plan is None: + return False + # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary + # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the + # block table is the pool's [requests, 2, blocks] K / V sub-page indices. + pool_index, layer_in_pool = ( + int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] + ) + tpb = int(params.tokens_per_block) + sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim + base = ( + int(meta.host_kv_cache_pool_pointers[pool_index, 0]) + + layer_in_pool * int(params.kv_factor) * sub_page_bytes + ) + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) + ) + tables = meta.kv_cache_block_offsets[ + pool_index, params.seq_offset : params.seq_offset + nreq + ] + multi = self._multi_token(params) + if not torch.cuda.is_current_stream_capturing(): + # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess + # path checks; the op itself writes V at the table's V sub-page. + tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(tables.shape[-1], device=tables.device) + in_use = cols[None, :] < tpb_held[:, None] + offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + q = torch.empty( + int(params.num_tokens), + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=attention_input.device, + ) + torch.ops.trtllm.minimax_m3_sm107tp8_drafter_rope_fp8_kv( + attention_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + int(attn.rope_params.dim), + tpb, + self._fused_rope_mode, + self._pdl, + ROPE_CARVEOUT, + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + tables[:, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + if self._on_rope_taken is not None: + self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") + self._on_taken(params.num_tokens) + return True + + def run_generation(self, params: FmhaParams) -> None: + attention_input = params.qkv_input if params.qkv_input is not None else params.query_input + if attention_input is None: + raise RuntimeError("the drafter decode requires QKV or query input.") + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + if self._run_fused_rope(params, attention_input): + return + # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 + # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). + fp8_context_fmha = False + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + _fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + _window_left, + is_multi_token_gen, + ) = thop.trtllm_gen_generation_preprocess( + attention_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + fp8_context_fmha, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + ) + qlen, nreq = params.input_seq_length, params.batch_size + plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the + # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row + # pass is never multi-token. + if plan is None or (is_multi_token_gen and qlen == 1): + raise RuntimeError( + f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" + ) + rows = params.num_tokens * attn.num_heads * attn.head_dim + q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) + q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) + # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 + # of the pool (the two views have the same page count, as the kernel requires). + if not torch.cuda.is_current_stream_capturing(): + # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. + lengths = params.spec_decoding_generation_lengths + if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): + raise RuntimeError( + f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" + ) + # Only the blocks a request holds are laid out; the table's tail entries are unset. + tpb = int(params.tokens_per_block) + held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(block_tables.shape[-1], device=block_tables.device) + in_use = cols[None, :] < held[:, None] + offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. + # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + block_tables[:nreq, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + self._on_taken(params.num_tokens) + + +def install( + draft_model: torch.nn.Module, + recipes, + max_kv_len: int, + own_reducer: bool, + pdl: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + on_rope_taken=None, +) -> int: + """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention + backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it + was installed on.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + count = 0 + device = next(draft_model.parameters()).device + # The backend is not a module: each attention module of the drafter holds it as ``attn``. + for module in draft_model.modules(): + backend = getattr(module, "attn", None) + if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): + continue + if not bool((backend.kv_cache_scaling_factor == 1.0).all()): + on_uncovered("a scaled FP8 KV cache") + continue + fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) + libs = backend._fmha_manager.fmha_libs + if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): + libs.insert( + 0, + M3DrafterFullCtxFmha( + backend, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode=fused_rope_mode, + pdl=pdl, + on_rope_taken=on_rope_taken, + ), + ) + count += 1 + return count diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fc_wait.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fc_wait.py new file mode 100644 index 000000000000..a545b43857ec --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fc_wait.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter fc fork's release of the first draft input by a device counter instead of a stream join (modeling's +``DI_FC_FLAG``). + +Triton. ``fc_wait_kernel``: one thread, queued on the main stream after the acceptance. It lets the next kernel (the +first draft input) launch, waits for the acceptance before it, then spins until the fork's counter reaches the number +of forks waited for. Ordering only: no value changes. +""" + +import torch +import triton +import triton.language as tl + +# Spins of fc_wait_kernel before it gives up on the fc fork (~ms): a broken fork / wait pairing then counts an overrun +# in counts[2] instead of hanging the step (the draft input may read a stale fc output: draft tokens only, the verify +# stays exact). +MAX_SPINS = 1 << 22 # capacity: fixed: a spin bound, not a size + + +@triton.jit +def fc_wait_kernel(counts, MAX_SPINS: tl.constexpr): + # counts[0]: fc forks finished (bumped on the aux stream after each fc); counts[1]: forks waited for (here). + # Lets the next kernel (the first draft input) launch now; waits for the acceptance before it, then for the fork. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + target = tl.load(counts + 1) + 1 + done = tl.load(counts, volatile=True) + spins = 0 + while (done < target) & (spins < MAX_SPINS): + done = tl.load(counts, volatile=True) + spins += 1 + tl.store(counts + 1, target) + if done < target: + tl.atomic_add(counts + 2, 1) + + +def fc_wait(counts: torch.Tensor) -> None: + """Queue the wait on the current stream for the fork ``counts`` (int64 [3]) last counted.""" + fc_wait_kernel[(1,)](counts, MAX_SPINS=MAX_SPINS, num_warps=1, launch_pdl=True) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py new file mode 100644 index 000000000000..292d1df264b1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The speculative step's store and next-step staging at the tail of the forward (one Triton launch). + +For each of the first ``n`` rows (``n`` read from device memory, so a captured graph whose step was not armed writes +nothing), with ``s = slots[i]``: + +* the spec sampler's per-slot store, exactly as ``m3_spec_tail.slot_scatter`` writes it (zero-padded to each store + width): ``new[s][c]``, ``next_new[s][c]``, ``draft[s][c]`` and ``lens[s]``; +* the engine's next-step inputs, exactly as ``m3_spec_tail.step_gather`` writes them from that store for a step whose + rows are these rows in this order (no new tokens, first previous row 0): ``ids[i*T + t] = next_new[s][t]``, + ``draft_ids[i*D + d] = draft[s][d]``, ``pos_off[i*T + t] = lens[s]``, ``kv_off[i] = lens[s] - T``. + +Integer copies only: the values are the ones the sampler's and the engine's own copies write. + +With ``cpp`` the same writes run as this target's C++ op (trtllm::minimax_m3_sm107_tp8_fold_tail, +minimaxM3Sm107Tp8/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fold_tail_kernel( + ctrl, + slots, + src_new, + sn_s0, + sn_s1, + SN_W: tl.constexpr, + dst_new, + dn_s0, + dn_s1, + DN_W: tl.constexpr, + src_next, + sx_s0, + sx_s1, + SX_W: tl.constexpr, + dst_next, + dx_s0, + dx_s1, + DX_W: tl.constexpr, + src_draft, + sd_s0, + sd_s1, + SD_W: tl.constexpr, + dst_draft, + dd_s0, + dd_s1, + DD_W: tl.constexpr, + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + T: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(ctrl) + if i < n: + s = tl.load(slots + i).to(tl.int64) + c = tl.arange(0, BLOCK) + new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) + tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) + nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) + tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) + drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) + tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) + ln = tl.load(src_lens + i) + tl.store(dst_lens + s, ln) + # The next step's inputs read the store's padded rows: zero past each source width. + tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) + tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) + tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) + tl.store(kv_off + i, ln - T) + + +def fold_tail( + ctrl: torch.Tensor, + slots: torch.Tensor, + rows: int, + src_new: torch.Tensor, + dst_new: torch.Tensor, + src_next: torch.Tensor, + dst_next: torch.Tensor, + src_draft: torch.Tensor, + dst_draft: torch.Tensor, + src_lens: torch.Tensor, + dst_lens: torch.Tensor, + ids: torch.Tensor, + draft_ids: torch.Tensor, + pos_off: torch.Tensor, + kv_off: torch.Tensor, + tokens_per_step: int, + cpp: bool = False, + one_warp: bool = False, +) -> bool: + """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs + are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width]. + ``one_warp`` (with ``cpp``): the in-kernel form (one warp looping over the rows, foldTailWarp in + minimaxM3Sm107Tp8/foldTail.h) that a caller folding the tail into its own kernel runs; the same writes.""" + two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) + one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) + if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): + return False + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) + or slots.dtype != torch.long + ): + return False + t_step = tokens_per_step + d_step = t_step - 1 + if ( + rows <= 0 + or slots.numel() < rows + or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) + or ids.numel() < rows * t_step + or pos_off.numel() < rows * t_step + or draft_ids.numel() < rows * d_step + or kv_off.numel() < rows + ): + return False + widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) + widths += (src_draft.shape[1], dst_draft.shape[1], t_step) + if cpp and max(widths) <= 32: + op = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_fold_tail_one_warp + if one_warp + else torch.ops.trtllm.minimax_m3_sm107_tp8_fold_tail + ) + op( + ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, + src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, + ) # fmt: skip + return True + block = triton.next_power_of_2(max(widths)) + _fold_tail_kernel[(rows,)]( + ctrl, + slots, + src_new, + src_new.stride(0), + src_new.stride(1), + src_new.shape[1], + dst_new, + dst_new.stride(0), + dst_new.stride(1), + dst_new.shape[1], + src_next, + src_next.stride(0), + src_next.stride(1), + src_next.shape[1], + dst_next, + dst_next.stride(0), + dst_next.stride(1), + dst_next.shape[1], + src_draft, + src_draft.stride(0), + src_draft.stride(1), + src_draft.shape[1], + dst_draft, + dst_draft.stride(0), + dst_draft.stride(1), + dst_draft.shape[1], + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + t_step, + d_step, + block, + num_warps=1, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py new file mode 100644 index 000000000000..cc0ce28d9639 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's +``force_num_accepted_tokens``), in one launch. Reached through the ``ops/spec`` entries ``m3_force_vec`` and +``m3_force_vec_req``. + +Triton. ``force_vec_kernel``: one CTA bumps the shared step counter and draws each request's count from the fixed-seed +pool with the in-tree (counter, slot) hash, exactly as ``_apply_force_accepted_tokens`` does. ``force_vec_req_kernel``: +one draw per request from its own key (its KV length at this verify), no shared counter. +""" + +import triton +import triton.language as tl + + +# One compiled variant for every batch and accept count: specializing them (== 1, divisible by 16) would compile +# mid-serving the first time a step reaches a new one, past the worker's warm-up. +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_kernel( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK: tl.constexpr, +): + c = tl.load(counter) + 1 + slots = tl.arange(0, BLOCK) + tl.store(counter, c) + idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask + live = slots < batch_size + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) + + +@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) +def force_vec_req_kernel( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK: tl.constexpr, +): + # One draw per request from its own key (its KV length at this verify): no shared counter, so a request's forced + # counts depend only on its own prompt and progress, not on what else the server ran before or beside it. + slots = tl.arange(0, BLOCK) + live = slots < batch_size + k = tl.load(keys + slots, mask=live, other=0).to(tl.int64) + idx = (k * key_stride) & pool_mask + r = tl.load(pool + idx, mask=live, other=1.0) + v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) + tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py new file mode 100644 index 000000000000..de9f6d6666f3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py @@ -0,0 +1,328 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 full-context FP8 decode attention on Rubin: ``fullctx_decode``. + +The dense layers (0-2) attend every token of their FP8 E4M3 (unscaled) paged KV cache. On SM 10.7 the stock call +(flashinfer trtllm-gen, ``minimax_m3_trtllm_gen_dense_decode``) caps its split at 16 cluster CTAs per (request, KV +head): at TP8 (8 query heads and 1 KV head per rank) one request runs on 16 CTAs of 212. + +This runs the same arithmetic -- FP8 Q (the same E4M3 rounding of q), FP8 K/V, P quantized to E4M3 per split -- on +the vendored PrimTS decode kernel (``attention/backends/prims_ts``, SM 10.7 through the sm_100f family) at a pinned +split-KV recipe, which fans one request out over many more CTAs. The split partials are merged by +``torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce`` (``cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/ +fullCtxReduce.cu``, head-count generic, used as is): unnormalized fp32 accumulation over the active splits, one final +division, one BF16 rounding, the merge form of stock trtllm-gen. Only the split changes against stock +(modeling's ``FULLCTX_DECODE_RECIPES``). + +A plan is compiled per (device, query rows per request, requests) by ``prebuild`` before the executor's warmup; the +forward only reads plans (``plan_for``). The plan is built from PrimTS's own config, compile and workspace helpers at +an explicit recipe; the PrimTS wrapper's planning policy is not used, so no PrimTS state is replaced. With +``own_reducer`` off the PrimTS reducer merges into FP16 and the result is cast to BF16. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Mapping, Optional, Tuple + +import torch + +# The TP8 per-rank geometry: 8 query heads over 1 KV head. +HEADS = 8 +KV_HEADS = 1 # capacity: fixed: the tp8 KV heads per rank +HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim +PAGE = 128 +# The per-request K/V length PrimTS's config selection sizes its policy for (Q/KV tiles, launch mode); the pinned +# recipes were measured with it. The compiled capacity (``max_kv_len``) may be larger. +POLICY_KV_LEN = ( + 1 << 20 +) # capacity: fixed: the PrimTS policy's K/V length bound (the plans' compile shape) +HEADS_PER_KV = HEADS // KV_HEADS # capacity: fixed: the tp2 query heads per KV head + +_plans: Dict[Tuple[int, int, int, bool], "FullCtxPlan"] = {} +_lock = threading.Lock() + + +class FullCtxPlan: + """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" + + def __init__( + self, + device: torch.device, + qlen: int, + nreq: int, + max_kv_len: int, + recipe: Mapping, + own_reducer: bool, + pdl: bool, + ): + from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec + from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( + get_max_active_clusters_for_cluster_size, + make_decode_config, + ) + + self._tsdec = tsdec + self.qlen = qlen + self.nreq = nreq + self.pdl = pdl + self.recipe = dict(recipe) + splits = int(self.recipe["splits"]) + fp8 = tsdec._cutlass_dtype("float8_e4m3fn") + # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer + # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents + # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an + # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the + # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own + # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge + # (minimax_m3_sm107_tp2_fullctx_reduce) waits as its first instruction. + cfg = make_decode_config( + headdim=HEAD_DIM, + args={"use_pdl": pdl, **self.recipe.get("config", {})}, + seq_len_q=qlen, + seq_len_kv=min(max_kv_len, POLICY_KV_LEN), + batch_size=nreq, + num_heads_q=HEADS, + num_heads_kv=KV_HEADS, + q_dtype=fp8, + k_dtype=fp8, + v_dtype=fp8, + o_dtype=tsdec._cutlass_dtype("float16"), + qkv_layout="pagedKv", + num_tokens_per_page=PAGE, + storage_tokens_per_page=PAGE, + split_kv_mode="gmem_reduction_with_separate_kernel", + splits_kv=splits, + max_splits_kv=splits, + min_loop_iters_per_split=1, + mask_type="causal" if qlen > 1 else "dense", + auto_tuner=True, + split_kv=True, + ) + _check_layout(cfg, qlen, splits) + with torch.cuda.device(device): + max_active = get_max_active_clusters_for_cluster_size(1) + spec = tsdec._decode_launch_spec_from_config( + cfg, + batch_size=nreq, + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + seq_len_q=qlen, + max_active_clusters=int(max_active), + ) + compile_spec = tsdec._make_decode_compile_spec( + spec, + device_index=_index(device), + num_qo_heads=HEADS, + num_kv_heads=KV_HEADS, + head_dim=HEAD_DIM, + page_size=PAGE, + max_kv_len=max_kv_len, + seq_len_q=qlen, + q_dtype_key="float8_e4m3fn", + k_dtype_key="float8_e4m3fn", + v_dtype_key="float8_e4m3fn", + output_dtype_key="float16", + use_packed_q=False, + kv_prefix_mode="dynamic", + kv_lengths_mode="dynamic", + ) + self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) + layout = tsdec._make_decode_workspace_layout( + spec.scratch_shapes, + torch.float16, + use_separate_reduction_kernel=True, + use_split_kv=True, + ) + self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) + self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) + self.workspace.split_kv_counter.zero_() + self.workspace.cu_seqlens_q.zero_() + self.workspace.attention_sinks.zero_() + rows = HEADS_PER_KV * qlen + if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): + raise RuntimeError( + f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " + f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" + ) + self.own_reducer = own_reducer + self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) + self.tile_kv = int(cfg.tile_size_kv) + self.insts_kv = int(cfg.num_insts_kv) + self.causal = qlen > 1 + shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) + self.staging = torch.empty(shape, dtype=torch.float16, device=device) + + def run( + self, + q: torch.Tensor, + k_pages: torch.Tensor, + v_pages: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + sm_scale: float, + output: torch.Tensor, + mxfp8: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> bool: + """Decodes into output; with mxfp8 (E4M3 rows, UE8M0 scales) and the TP8 merge, the merge also writes + the MXFP8 copy of output. Returns whether it did.""" + # The kernel takes a compact query; a column view of the fused QKV projection is copied once. + q = q.reshape(self.staging.shape).contiguous() + runtime = self._tsdec._prepare_decode_runtime_unchecked( + q, + (k_pages, v_pages), + None, + output_dtype=torch.float16, + bmm1_scale=sm_scale, + bmm2_scale=1.0, + out=self.staging, + ) + ws = self.workspace + self._main( + runtime.q, + runtime.k_cache, + runtime.v_cache, + runtime.k_sf_cache, + runtime.v_sf_cache, + runtime.out, + seq_lens, + ws.cu_seqlens_q, + block_table, + block_table, # Dense decode constexpr-elides the membership table. + ws.partial_o, + ws.partial_stats, + ws.split_kv_counter, + ws.attention_sinks, + runtime.num_physical_pages, + runtime.k_page_stride, + runtime.k_head_stride, + runtime.k_token_stride, + runtime.v_page_stride, + runtime.v_head_stride, + runtime.v_token_stride, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + if self.own_reducer: + reduce = _reduce_op() + args = ( + ws.partial_o, + ws.partial_stats, + seq_lens, + output, + self.qlen, + self.q_tokens_per_cta, + self.tile_kv, + self.insts_kv, + self.causal, + self.pdl, + ) + if mxfp8 is not None and reduce is _tp8_reduce_op(): + reduce(*args, mxfp8[0], mxfp8[1]) + return True + reduce(*args) + return False + self._primts_reducer( + runtime.out, + seq_lens, + ws.cu_seqlens_q, + ws.partial_o, + ws.partial_stats, + ws.attention_sinks, + runtime.bmm1_scale, + runtime.bmm2_scale, + (), + ) + output.view(self.staging.shape).copy_(self.staging) + return False + + +def _check_layout(cfg, qlen: int, splits: int) -> None: + """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. + + Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is + ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding + window, attention sinks or packed Q. + """ + problems = [] + if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): + problems.append("not split-KV with the separate GMEM reducer") + if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: + problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") + if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: + problems.append("cluster-SMEM reduction or a persistent scheduler") + if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: + problems.append("packed Q, a sliding window or attention sinks") + if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): + problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") + if qlen > 1: + if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): + problems.append("multi-token Q without the grouped [token, head] tile layout") + elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): + problems.append( + f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" + ) + if problems: + raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + + +# Select the TP8 copy of the BF16 reducer; both copies use the same per-row arithmetic. +TP8_REDUCER = True + + +def _reduce_op(): + """The BF16 split-KV merge selected by TP8_REDUCER.""" + if TP8_REDUCER: + return torch.ops.trtllm.minimax_m3_sm107tp8_fullctx_reduce + return torch.ops.trtllm.minimax_m3_sm107_tp2_fullctx_reduce + + +def _tp8_reduce_op(): + return torch.ops.trtllm.minimax_m3_sm107tp8_fullctx_reduce + + +def emits_mxfp8() -> bool: + """Whether the selected merge writes the TP8 o_proj MXFP8 hand-off.""" + return _reduce_op() is _tp8_reduce_op() + + +def has_own_reducer() -> bool: + """Require the selected BF16 split-KV merge op supplied by this target build.""" + _reduce_op() + return True + + +def plan_for( + device: torch.device, qlen: int, nreq: int, own_reducer: bool = True +) -> Optional[FullCtxPlan]: + """The prebuilt plan for this shape, or None.""" + return _plans.get((_index(device), qlen, nreq, own_reducer)) + + +def prebuild( + device: torch.device, + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + pdl: bool, +) -> int: + """Compile the plan of every (query rows per request, requests) in ``recipes``. Call before warmup.""" + built = 0 + for (qlen, nreq), recipe in recipes.items(): + key = (_index(device), qlen, nreq, own_reducer) + with _lock: + if key not in _plans: + _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) + built += 1 + return built + + +def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: + """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" + return pool[0::subpages_per_slot], pool[1::subpages_per_slot] + + +def _index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py new file mode 100644 index 000000000000..1467852d67b0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py @@ -0,0 +1,692 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""MiniMax-M3 index decode block scorer (CuTe DSL, Blackwell) with its pre-wait schedule. + +The kernel of ``tensorrt_llm/_torch/attention/kernels/blackwell/minimax_m3_index_decode_score.py`` (the same scores, bit +for bit) with four schedule options this target compiles in: + +* ``prefetch_old``: before the grid-dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that hold only + tokens older than this step (the producer ahead writes only this step's tokens), up to ``prefetch_depth`` blocks per + CTA (0 = all of them). +* ``early_issue``: the TMA warp reads a ``seq_lens`` hint and its first page id before the wait and issues the Q tile + and the first K tile right after it. The tiles are read after the wait, so their contents are final; the block table + is only written between graph replays, so the early page id is the right one whenever the block is valid. A CTA + whose early tiles lie past the request's true length drains them before it exits. +* ``prestage_old`` (with ``early_issue``, ``INDEX_SCORE_PRESTAGE_OLD``): when the CTA's first block holds only + tokens older than this step (block < (seq_lens hint - decode query length) / 128), its K tile is TMA-loaded into its + stage before the wait; Q (and a first block that is not all old) is still loaded after it. +* ``early_trigger`` (``INDEX_SCORE_SET``): every CTA releases the dependent block selector right after its + setup, so the selector launches and makes its pre-wait reads while this grid scores; the selector still waits for + this whole grid before it reads a score, so the scores and the selection are unchanged. + +Every kernel is compiled before the executor's warmup (``layers.prebuild_index_scorer``) through ``IndexScoreRunner`` +and called through this target's ``m3_index_score`` entry. +""" + +import cutlass +import torch +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import ( + TRTLLM_ENABLE_PDL, + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) +from tensorrt_llm.logger import logger + +# With the early issue, a CTA whose first block holds only tokens older than this step also TMA-loads that block's +# index-K tile before the grid wait (bit-identical scores). +INDEX_SCORE_PRESTAGE_OLD = True +# Every CTA releases the dependent block selector right after its setup instead of at its exit (the selector's own +# grid wait still returns only once the whole scoring grid is complete, so the picks are unchanged), at every key. +INDEX_SCORE_SET = True + +__all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + prestage_old: bool = False, + early_trigger: bool = False, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K + # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's + # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. + self.prestage_old = prestage_old and early_issue + # early_trigger: release the dependent block selector right after this CTA's setup (see the module docstring). + self.early_trigger = early_trigger + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # Without early_trigger the dependent block selector is released as the CTAs exit. + early_q = False + early_k = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + seqlen_hint = seq_lens[batch_id] + hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + if cutlass.const_expr(self.prestage_old): + # Stage the first K tile now when every token of the block predates this step. + if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: + early_k = True + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + if not early_k: + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # The length that bounds this grid's work is read after the wait (the + # pre-wait reads above are prefetch / early-issue hints only, clamped to + # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens + # patches every step, and num_blocks bounds an unmasked block_table load + # and an unmasked score store, so a stale length is an out-of-bounds write + # rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA; ``early_trigger`` is ``INDEX_SCORE_SET`` for every key of the process.""" + + kernel_cache: dict = {} + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + prestage_old=early and INDEX_SCORE_PRESTAGE_OLD, + early_trigger=INDEX_SCORE_SET, + ) + if INDEX_SCORE_SET and not cls.kernel_cache: + logger.info( + "tp8 index scorer: the block selector is released after the scorer's setup (INDEX_SCORE_SET)" + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/__init__.py @@ -0,0 +1,14 @@ +# 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. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py new file mode 100644 index 000000000000..c8dfa3eee847 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py @@ -0,0 +1,487 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode producer as a phase of the TP8 per-layer kernel (m3_layer): q / k / index-q / index-k RMSNorm + NeoX +partial RoPE, NVFP4 K / V and FP8 index-K cache insert, FP8 q / index-q out. + +Bit for bit ``torch.ops.trtllm.minimax_m3_producer_pw`` (cpp/tensorrt_llm/kernels/minimaxM3/producerPw.cu, +m3_producer_kernel<8, *>) at the TP8 geometry: 8 query heads, 1 KV head, 1 index head, head_dim 128, rotary_dim 64, +page 128. One work item = one decode token; the CTA's warp w (< 12) owns packed head w of it, lane l holds elements +[4l, 4l + 4) exactly as the C++ kernel, so every reduction (the RMS sum of squares, the RoPE pair exchange, the +NVFP4 block max) runs the same shuffle tree in the same order. Every floating-point operation is written as its PTX +instruction (mul.rn / add.rn / fma.rn / rsqrt.approx.ftz / rcp.approx.ftz / cvt) so the DSL compiler cannot contract +or reassociate differently from the library build (--use_fast_math). The places where that build contracted are +marked CONTRACT and follow the shipped SASS (msa/results/sassdump-*). + +Packed row layout per token: [8 q | k | v | index_q | index_k] x 128 BF16 (the QKV GEMV output rounded to BF16). +""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op + +HEADS_Q = 8 # capacity: fixed: the tp8 query heads per rank +PACKED_HEADS = HEADS_Q + 4 # q x 8, k, v, index_q, index_k +HEAD_DIM = 128 +ROTARY_DIM = 64 # capacity: fixed: the checkpoint's rotary dim +PAGE = 128 +ELEMS = HEAD_DIM // 32 # capacity: fixed: elements per lane (per lane) +K_HEAD, V_HEAD, IQ_HEAD, IK_HEAD = HEADS_Q, HEADS_Q + 1, HEADS_Q + 2, HEADS_Q + 3 + + +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- exact float arithmetic (no contraction) --------------------------------------------------------------------- +@dsl_user_op +def f_mul(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_add(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_fma(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_max(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "max.f32 $0, $1, $2;", + "=f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def f_abs(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "abs.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_neg(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "neg.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_rsqrt(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rsqrt.approx.ftz.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def f_rcp(a, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", loc=loc, ip=ip) + ) + + +@dsl_user_op +def bf16_bits_to_f32(bits16, *, loc=None, ip=None): + """The float value of a BF16 given in the low 16 bits of an Int32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(bits16)], + "{\n\t.reg .b32 t;\n\tshl.b32 t, $1, 16;\n\tmov.b32 $0, t;\n\t}", + "=f,r", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def round_bf16(a, *, loc=None, ip=None): + """__bfloat162float(__float2bfloat16_rn(a)).""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a)], + "{\n\t.reg .b16 h;\n\t.reg .b32 t;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b32 t, {0, h};\n\tmov.b32 $0, t;\n\t}", + "=f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e4m3x4_word(e0, e1, e2, e3, *, loc=None, ip=None): + """store_fp8x4: bytes [e0, e1, e2, e3] (cvt.rn.satfinite.e4m3x2.f32 puts its first source in the high byte).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(e0), cutlass.Float32(e1), cutlass.Float32(e2), cutlass.Float32(e3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", + "=r,f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def e4m3_byte_and_value(x, *, loc=None, ip=None): + """(byte, float value) of __nv_fp8_e4m3(x): the encoded scale and static_cast of it.""" + byte = cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n\t.reg .b16 p;\n\t.reg .b32 w;\n\tcvt.rn.satfinite.e4m3x2.f32 p, 0f00000000, $1;\n\t" + "cvt.u32.u16 w, p;\n\tand.b32 $0, w, 0xFF;\n\t}", + "=r,f", + loc=loc, + ip=ip, + ) + ) + value = cutlass.Float32( + _asm( + _T.f32(), + [byte], + "{\n\t.reg .b16 p;\n\t.reg .b32 h2;\n\t.reg .b16 h;\n\tcvt.u16.u32 p, $1;\n\t" + "cvt.rn.f16x2.e4m3x2 h2, p;\n\tmov.b32 {h, _}, h2;\n\tcvt.f32.f16 $0, h;\n\t}", + "=f,r", + loc=loc, + ip=ip, + ) + ) + return byte, value + + +@dsl_user_op +def e2m1x4_halfword(q0, q1, q2, q3, *, loc=None, ip=None): + """The low 16 bits of fp32x8_to_e2m1 for q[0..3] (q[4..7] = 0): byte0 = (q1, q0), byte1 = (q3, q2).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(q0), cutlass.Float32(q1), cutlass.Float32(q2), cutlass.Float32(q3)], + "{\n\t.reg .b8 b0, b1;\n\t.reg .b32 w;\n\tcvt.rn.satfinite.e2m1x2.f32 b0, $2, $1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 b1, $4, $3;\n\tmov.b32 w, {b0, b1, 0, 0};\n\tmov.b32 $0, w;\n\t}", + "=r,f,f,f,f", + loc=loc, + ip=ip, + ) + ) + + +# ---- global memory ----------------------------------------------------------------------------------------------- +@dsl_user_op +def ld_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm(_T.i64(), [cutlass.Int64(addr)], "ld.global.u64 $0, [$1];", "=l,l", True, loc, ip) + ) + + +@dsl_user_op +def ld_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm(_T.i32(), [cutlass.Int64(addr)], "ld.global.s32 $0, [$1];", "=r,l", True, loc, ip) + ) + + +@dsl_user_op +def ld_f32(addr, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Int64(addr)], "ld.global.f32 $0, [$1];", "=f,l", True, loc, ip) + ) + + +@dsl_user_op +def st_u32(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "st.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_u16(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "{\n\t.reg .b16 h;\n\tcvt.u16.u32 h, $1;\n\tst.global.u16 [$0], h;\n\t}", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def st_u8(addr, v, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(v)], + "{\n\t.reg .b16 h;\n\tcvt.u16.u32 h, $1;\n\tst.global.u8 [$0], h;\n\t}", + "l,r", + True, + loc, + ip, + ) + + +# ---- CONTRACT points (follow the shipped SASS; see the module docstring) ------------------------------------------- +def sumsq_pair(acc, x, y): + """sumSquares += x * x + y * y (C++: one pair of the packed row). CONTRACT: fma(y, y, x * x) then add.""" + return f_add(acc, f_fma(y, y, f_mul(x, x))) + + +def rms_reciprocal(sum_squares, eps): + """rsqrtf(sumSquares / 128 + eps). CONTRACT: the / 128 is an exact * 2^-7; fused with + eps into one fma.""" + return f_rsqrt(f_fma(sum_squares, cutlass.Float32(1.0 / HEAD_DIM), eps)) + + +class ProducerArgs: + """Raw device addresses (int64) + strides, as the C++ Params (producerPw.h) carries them.""" + + def __init__( + self, + packed, + q_out, + index_q_out, + index_k, + kv_data, + kv_scale, + out_cache_loc, + kv_quant_scale, + position_ids, + rotary_cos_sin, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + eps, + num_pages, + index_page_stride, + index_token_stride, + data_page_stride, + data_plane_stride, + data_token_stride, + scale_page_stride, + scale_plane_stride, + ): + self.__dict__.update(locals()) + del self.__dict__["self"] + + +@cute.jit +def phase_producer(tok, a: ProducerArgs): + """Producer work item ``tok``: warp w (< 12) of the calling CTA does packed head w of decode token tok. + The caller guarantees tok < rows and that warps >= 12 skip. Stores only (no CTA barrier inside).""" + tid = cute.arch.thread_idx()[0] + lane = tid % 32 + head = tid // 32 + in_off = (cutlass.Int64(tok) * PACKED_HEADS + cutlass.Int64(head)) * HEAD_DIM + cutlass.Int64( + lane * ELEMS + ) + pk = ld_u64(a.packed + in_off * 2) + x0 = bf16_bits_to_f32(cutlass.Int32(pk & cutlass.Int64(0xFFFF))) + x1 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(16)) & cutlass.Int64(0xFFFF))) + x2 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(32)) & cutlass.Int64(0xFFFF))) + x3 = bf16_bits_to_f32(cutlass.Int32((pk >> cutlass.Int64(48)) & cutlass.Int64(0xFFFF))) + producer_head(tok, head, lane, x0, x1, x2, x3, a) + + +@cute.jit +def producer_head(tok, head, lane, x0, x1, x2, x3, a: ProducerArgs): + """Everything after the packed-row load for packed head ``head`` (warp-uniform; dynamic Int32 or Constexpr) of + decode token ``tok``, on the calling warp: lane ``lane`` holds the head's elements [4 lane, 4 lane + 4) as the + float values x0..x3 of the BF16 packed row (the QKV GEMV output rounded to BF16). Norm weights, KV quant scale, + position and RoPE loads, the RMS tree, RoPE, and the q / index-q / index-K / K-V NVFP4 stores. Shared by + phase_producer and the QKV GEMV owner epilogue (MK-A), so the producer math has one implementation. + + Data flow is branch-free (every load is in bounds for every head: V reads the q norm weight, a non-rotating lane + reads its token's RoPE row, and the results are then discarded with select); only the stores branch, on + warp-uniform conditions, so every shuffle runs on the full warp as in the C++ kernel.""" + l4 = cutlass.Int64(lane * ELEMS) + is_v = head == V_HEAD + zero = cutlass.Float32(0.0) + + # Norm weights (BF16): q / k / index-q / index-k weight by head (V: q's, unused). + w_base = cutlass.select_( + head == K_HEAD, + a.k_weight, + cutlass.select_( + head == IQ_HEAD, + a.index_q_weight, + cutlass.select_(head == IK_HEAD, a.index_k_weight, a.q_weight), + ), + ) + ww = ld_u64(w_base + l4 * 2) + w0 = bf16_bits_to_f32(cutlass.Int32(ww & cutlass.Int64(0xFFFF))) + w1 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(16)) & cutlass.Int64(0xFFFF))) + w2 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(32)) & cutlass.Int64(0xFFFF))) + w3 = bf16_bits_to_f32(cutlass.Int32((ww >> cutlass.Int64(48)) & cutlass.Int64(0xFFFF))) + quant = ld_f32(a.kv_quant_scale + cutlass.select_(is_v, cutlass.Int64(8), cutlass.Int64(4))) + + # Position id, RoPE row. + pos = ld_s32(a.position_ids + cutlass.Int64(tok) * 4) + row = cutlass.Int64(pos) * ROTARY_DIM + coef = cutlass.Int64((lane * ELEMS) % (ROTARY_DIM // 2)) + c0 = ld_f32(a.rotary_cos_sin + (row + coef) * 4) + c1 = ld_f32(a.rotary_cos_sin + (row + coef + 1) * 4) + c2 = ld_f32(a.rotary_cos_sin + (row + coef + 2) * 4) + c3 = ld_f32(a.rotary_cos_sin + (row + coef + 3) * 4) + s0 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef) * 4) + s1 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 1) * 4) + s2 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 2) * 4) + s3 = ld_f32(a.rotary_cos_sin + (row + ROTARY_DIM // 2 + coef + 3) * 4) + + # RMSNorm (all heads compute it; V keeps x). + ss = sumsq_pair(zero, x0, x1) + ss = sumsq_pair(ss, x2, x3) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=16)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=8)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=4)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=2)) + ss = f_add(ss, cute.arch.shuffle_sync_bfly(ss, offset=1)) + r = rms_reciprocal(ss, a.eps) + one = cutlass.Float32(1.0) + n0 = f_mul(x0, f_mul(r, f_add(one, w0))) + n1 = f_mul(x1, f_mul(r, f_add(one, w1))) + n2 = f_mul(x2, f_mul(r, f_add(one, w2))) + n3 = f_mul(x3, f_mul(r, f_add(one, w3))) + # NeoX partial RoPE: lane l's dim 4l + i pairs with lane l ^ 8 (dim +-32) within the first 64 dims. + pair_off = (ROTARY_DIM // 2) // ELEMS + p0 = cute.arch.shuffle_sync_bfly(n0, offset=pair_off) + p1 = cute.arch.shuffle_sync_bfly(n1, offset=pair_off) + p2 = cute.arch.shuffle_sync_bfly(n2, offset=pair_off) + p3 = cute.arch.shuffle_sync_bfly(n3, offset=pair_off) + rot = lane * ELEMS < ROTARY_DIM + low = lane * ELEMS < ROTARY_DIM // 2 + p0 = cutlass.select_(low, f_neg(p0), p0) + p1 = cutlass.select_(low, f_neg(p1), p1) + p2 = cutlass.select_(low, f_neg(p2), p2) + p3 = cutlass.select_(low, f_neg(p3), p3) + r0 = cutlass.select_(rot, f_fma(p0, s0, f_mul(n0, c0)), n0) + r1 = cutlass.select_(rot, f_fma(p1, s1, f_mul(n1, c1)), n1) + r2 = cutlass.select_(rot, f_fma(p2, s2, f_mul(n2, c2)), n2) + r3 = cutlass.select_(rot, f_fma(p3, s3, f_mul(n3, c3)), n3) + e0 = cutlass.select_(is_v, x0, r0) + e1 = cutlass.select_(is_v, x1, r1) + e2 = cutlass.select_(is_v, x2, r2) + e3 = cutlass.select_(is_v, x3, r3) + + if head < HEADS_Q: + st_u32( + a.q_out + (cutlass.Int64(tok) * HEADS_Q + cutlass.Int64(head)) * HEAD_DIM + l4, + e4m3x4_word(e0, e1, e2, e3), + ) + if head == IQ_HEAD: + st_u32(a.index_q_out + cutlass.Int64(tok) * HEAD_DIM + l4, e4m3x4_word(e0, e1, e2, e3)) + if head >= K_HEAD: + if head != IQ_HEAD: + slot = ld_s32(a.out_cache_loc + cutlass.Int64(tok) * 4) + page = slot // PAGE + if slot >= 0: + if page < a.num_pages: + within = slot & (PAGE - 1) + if head == IK_HEAD: + st_u32( + a.index_k + + cutlass.Int64(page) * a.index_page_stride + + cutlass.Int64(within) * a.index_token_stride + + l4, + e4m3x4_word(e0, e1, e2, e3), + ) + else: + plane = cutlass.select_(is_v, cutlass.Int64(1), cutlass.Int64(0)) + col = lane // 4 + sw = cutlass.select_( + is_v, + cutlass.Int64((within // 4) * 32 + col * 4 + within % 4), + cutlass.Int64(within * 8 + col), + ) + s_off = ( + cutlass.Int64(page) * a.scale_page_stride + + plane * a.scale_plane_stride + + sw + ) + # quantize_fp4x4: BF16-rounded values, 16-element block max over the 4-lane group. + q0 = round_bf16(e0) + q1 = round_bf16(e1) + q2 = round_bf16(e2) + q3 = round_bf16(e3) + lmax = f_max( + f_max(f_max(f_max(zero, f_abs(q0)), f_abs(q1)), f_abs(q2)), f_abs(q3) + ) + lmax = f_max(lmax, cute.arch.shuffle_sync_bfly(lmax, offset=1)) + lmax = f_max(lmax, cute.arch.shuffle_sync_bfly(lmax, offset=2)) + sbyte, sval = e4m3_byte_and_value( + f_mul(quant, f_mul(lmax, f_rcp(cutlass.Float32(6.0)))) + ) + if lane % 4 == 0: + st_u8(a.kv_scale + s_off, sbyte) + outs = cutlass.select_(lmax != zero, f_rcp(f_mul(sval, f_rcp(quant))), zero) + hw = e2m1x4_halfword( + f_mul(q0, outs), f_mul(q1, outs), f_mul(q2, outs), f_mul(q3, outs) + ) + d_off = ( + cutlass.Int64(page) * a.data_page_stride + + plane * a.data_plane_stride + + cutlass.Int64(within) * a.data_token_stride + + cutlass.Int64(lane * 2) + ) + st_u16(a.kv_data + d_off, hw) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py new file mode 100644 index 000000000000..3e88c7fb1ce2 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py @@ -0,0 +1,448 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Locality domains ("uGPU") of a multi-die GPU (Rubin VR200: 2 domains of 100 SMs + 12 remainder SMs) for the +streaming decode kernels of this target: the SM -> domain map, per-domain weight placement, and the CTA -> SM +assignment a persistent kernel uses to read only its own die's memory. + +A kernel launched with one CTA per SM reads ``%smid`` once and takes its logical block index from a host-built +permutation (:meth:`LocalityMap.logical_map`), so every per-CTA structure (work plan row, epochs, roles) follows its +die. Weights that are split by output rows into per-domain halves (:meth:`LocalityMap.split_rows`) live in one virtual +range whose halves are physically in the two domains. Placement only: the arithmetic is unchanged, so results are +bitwise those of the unplaced weights. + +The SM -> domain map is attested by the driver: the SMs are split by locality domain +(``cuDevSmResourceSplit`` with ``CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID``), and a ``%smid`` probe runs on a green +context of each split. Allocations use driver VMM with ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN`` and are checked +with ``CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL``; any failure raises (fail closed).""" + +from __future__ import annotations + +import threading +from typing import Dict, List, Sequence, Tuple + +import torch + +_LOCK = threading.Lock() +_MAPS: Dict[int, "LocalityMap"] = {} +_PROBE_SRC = r""" +extern "C" __global__ void smid_probe(int* out) { + if (threadIdx.x == 0) { + unsigned s; + asm volatile("mov.u32 %0, %%smid;" : "=r"(s)); + out[blockIdx.x] = (int)s; + } +} +""" + + +def _drv(): + from cuda.bindings import driver + + return driver + + +def _ck(res, what: str): + """cuda-python returns (CUresult, *values): raise on an error, return the values.""" + err = res[0] if isinstance(res, tuple) else res + if int(err) != 0: + raise RuntimeError(f"locality: {what} failed: {err}") + if isinstance(res, tuple): + return res[1] if len(res) == 2 else res[1:] + return None + + +class _Mapping: + """One virtual range backed by per-domain physical allocations, exposed to torch through + ``__cuda_array_interface__``. Lives as long as a tensor that views it (the tensor holds this object).""" + + def __init__(self, device: int, pieces: Sequence[Tuple[int, int]]): + """``pieces``: (bytes, domain) in address order; each a multiple of the allocation granularity.""" + d = _drv() + self.device = device + self.size = sum(n for n, _ in pieces) + self.va = _ck(d.cuMemAddressReserve(self.size, 0, 0, 0), "cuMemAddressReserve") + self.handles = [] + off = 0 + try: + for nbytes, dom in pieces: + prop = _alloc_prop(device, dom) + h = _ck(d.cuMemCreate(nbytes, prop, 0), f"cuMemCreate(domain {dom})") + self.handles.append((h, nbytes)) + _ck(d.cuMemMap(int(self.va) + off, nbytes, 0, h, 0), "cuMemMap") + off += nbytes + acc = d.CUmemAccessDesc() + acc.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + acc.location.id = device + acc.flags = d.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _ck(d.cuMemSetAccess(self.va, self.size, [acc], 1), "cuMemSetAccess") + except Exception: + self.close() + raise + self.__cuda_array_interface__ = { + "shape": (self.size,), + "typestr": "|u1", + "data": (int(self.va), False), + "version": 3, + "strides": None, + "stream": None, + } + + def close(self): + d = _drv() + off = 0 + for h, n in self.handles: + d.cuMemUnmap(int(self.va) + off, n) + d.cuMemRelease(h) + off += n + self.handles = [] + if self.va is not None: + d.cuMemAddressFree(self.va, self.size) + self.va = None + + def __del__(self): + try: + self.close() + except Exception: + pass + + +_LIBCUDA = None + + +def _libcuda(): + global _LIBCUDA + if _LIBCUDA is None: + import ctypes + + _LIBCUDA = ctypes.CDLL("libcuda.so.1") + return _LIBCUDA + + +def _alloc_prop(device: int, domain: int): + d = _drv() + prop = d.CUmemAllocationProp() + prop.type = d.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.requestedHandleTypes = d.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE + prop.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN + prop.location.localized.deviceId = device + prop.location.localized.localityDomainId = domain + return prop + + +class LocalityMap: + """The locality domains of one device. Build with :meth:`get` before CUDA-graph capture.""" + + def __init__(self, device: int): + d = _drv() + _ck(d.cuInit(0), "cuInit") + dev = _ck(d.cuDeviceGet(device), "cuDeviceGet") + attr = d.CUdevice_attribute + self.device = device + self.n_domains = int( + _ck( + d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, dev), + "domain count", + ) + ) + self.sms_per_domain = int( + _ck( + d.cuDeviceGetAttribute( + attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_MULTIPROCESSOR_COUNT, dev + ), + "SMs per domain", + ) + ) + self.num_sms = int( + _ck(d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev), "SMs") + ) + if self.n_domains != 2: + raise RuntimeError( + f"locality: {self.n_domains} locality domains (this layer handles 2)" + ) + self._domain_sms = self._probe(d, dev) + dom = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for s in sms: + if dom[s] != -1: + raise RuntimeError(f"locality: SM {s} probed in two domains") + dom[s] = k + self.sm_domain: List[int] = dom + self.remainder: List[int] = [s for s in range(self.num_sms) if dom[s] < 0] + rt = _alloc_prop(device, 0) + self.granularity = int( + _ck( + d.cuMemGetAllocationGranularity( + rt, d.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM + ), + "granularity", + ) + ) + + @classmethod + def get(cls, device) -> "LocalityMap": + idx = torch.device(device).index if not isinstance(device, int) else device + idx = torch.cuda.current_device() if idx is None else idx + with _LOCK: + m = _MAPS.get(idx) + if m is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "LocalityMap probes the device: build it before CUDA-graph capture" + ) + m = _MAPS[idx] = cls(idx) + return m + + def sms(self, domain: int) -> List[int]: + return list(self._domain_sms[domain]) + + @property + def sm_table(self) -> torch.Tensor: + """int32 [num_sms] on the device: (domain << 16) | rank within the domain's sorted SMs; -1 = remainder SM. + For kernels that pick work by %smid themselves.""" + t = getattr(self, "_sm_table", None) + if t is None: + v = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for i, s in enumerate(sms): + v[s] = (k << 16) | i + t = self._sm_table = torch.tensor(v, dtype=torch.int32, device=f"cuda:{self.device}") + return t + + def _probe(self, d, dev) -> List[List[int]]: + """Split the SMs by locality domain (strict, no backfill) and record %smid of a probe kernel launched on a + green context of each part: the driver's own SM -> domain assignment.""" + from cuda.bindings import nvrtc + + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + # the probe module, compiled for this device + major = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev + ), + "cc", + ) + minor = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev + ), + "cc", + ) + prog = _ck( + nvrtc.nvrtcCreateProgram(_PROBE_SRC.encode(), b"smid_probe.cu", 0, [], []), + "nvrtcCreateProgram", + ) + opts = [f"--gpu-architecture=sm_{major}{minor}".encode()] + err = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)[0] + if int(err) != 0: + raise RuntimeError(f"locality: probe compile failed: {err}") + n = _ck(nvrtc.nvrtcGetCUBINSize(prog), "cubin size") + cubin = b" " * n + _ck(nvrtc.nvrtcGetCUBIN(prog, cubin), "cubin") + out: List[List[int]] = [] + nb = 4 * self.sms_per_domain # more blocks than SMs: every SM of the part shows up + torch.cuda.set_device(self.device) + buf = torch.full((nb,), -1, dtype=torch.int32, device=f"cuda:{self.device}") + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate(desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM), + "green ctx", + ) + st = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + ctx = _ck(d.cuCtxFromGreenCtx(g), "ctx from green ctx") + _ck(d.cuCtxPushCurrent(ctx), "push ctx") + try: + mod = _ck(d.cuModuleLoadData(cubin), "probe module") + fn = _ck(d.cuModuleGetFunction(mod, b"smid_probe"), "probe fn") + buf.fill_(-1) + torch.cuda.synchronize(self.device) + import ctypes + + arg = ctypes.c_void_p(buf.data_ptr()) + _ck( + d.cuLaunchKernel( + fn, nb, 1, 1, 32, 1, 1, 0, st, ((arg,), (ctypes.c_void_p,)), 0 + ), + "probe launch", + ) + _ck(d.cuStreamSynchronize(st), "probe sync") + sms = sorted(set(int(x) for x in buf.tolist() if x >= 0)) + d.cuModuleUnload(mod) + finally: + d.cuCtxPopCurrent() + d.cuStreamDestroy(st) + d.cuGreenCtxDestroy(g) + if len(sms) != self.sms_per_domain: + raise RuntimeError( + f"locality: domain {k} probed {len(sms)} SMs, expected {self.sms_per_domain}" + ) + out.append(sms) + return out + + def logical_map(self, per_domain: int, extra: int) -> Tuple[torch.Tensor, List[int]]: + """The SM -> logical block index permutation for a kernel of ``2 * per_domain + extra`` CTAs, one per SM: + logical blocks [d * per_domain, (d + 1) * per_domain) run on domain d's first ``per_domain`` SMs; logical + blocks 2 * per_domain .. run on the rest (the leftover domain SMs, alternating domains, then the remainder + SMs). Returns (int32 [num_sms] on the device, -1 = no block; the domain of each extra block, -1 remainder).""" + if 2 * per_domain + extra != self.num_sms: + raise ValueError( + f"locality: {2 * per_domain} + {extra} CTAs must cover the {self.num_sms} SMs" + ) + if per_domain > self.sms_per_domain: + raise ValueError(f"locality: {per_domain} CTAs per domain > {self.sms_per_domain} SMs") + lmap = [-1] * self.num_sms + for k in range(2): + for i, s in enumerate(self._domain_sms[k][:per_domain]): + lmap[s] = k * per_domain + i + left = [self._domain_sms[k][per_domain:] for k in range(2)] + rest = [x for pair in zip(*left) for x in pair] + self.remainder + ext_dom = [] + for i, s in enumerate(rest): + lmap[s] = 2 * per_domain + i + ext_dom.append(self.sm_domain[s]) + assert sorted(lmap) == list(range(self.num_sms)) + return torch.tensor(lmap, dtype=torch.int32, device=f"cuda:{self.device}"), ext_dom + + def domain_streams(self) -> List[torch.cuda.ExternalStream]: + """One non-blocking stream per domain on a green context holding exactly that domain's SMs (kept alive with + this map): work launched on stream d runs only on domain d's SMs. Fork / join with events from the caller's + stream; graph capture follows the forks. Built once, before capture.""" + st = getattr(self, "_streams", None) + if st is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.domain_streams: build before CUDA-graph capture") + d = _drv() + dev = _ck(d.cuDeviceGet(self.device), "cuDeviceGet") + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + self._green = [] + st = [] + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate( + desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM + ), + "green", + ) + s_ = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + self._green.append((g, s_)) + st.append(torch.cuda.ExternalStream(int(s_), device=f"cuda:{self.device}")) + self._streams = st + return st + + def domain_of(self, ptr: int) -> int: + """The locality domain ordinal of a device address (-1: not placed in one domain). Through libcuda + directly: cuda-python's attribute helper does not know this attribute.""" + import ctypes + + d = _drv() + lib = _libcuda() + out = ctypes.c_int(-9) + attr = int(d.CUpointer_attribute.CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL) + err = lib.cuPointerGetAttribute(ctypes.byref(out), ctypes.c_int(attr), ctypes.c_uint64(ptr)) + if err != 0: + raise RuntimeError( + f"locality: cuPointerGetAttribute(LOCALITY_DOMAIN_ORDINAL) failed: {err}" + ) + return int(out.value) + + def place_split(self, t: torch.Tensor, split_row: int) -> torch.Tensor: + """Split-mapped placement: a copy of ``t`` (contiguous, rows along dim 0) with the SAME shape and layout, in one + virtual range whose rows [0, split_row) are physically in domain 0 and rows [split_row, N) in domain 1. The + split must fall on an allocation granule (rows [0, split_row) a multiple of 2 MiB on VR200); the tail is padded + to a granule in domain 1. Stock consumers read it as a plain tensor; no extra memory once ``t`` is dropped.""" + if not t.is_contiguous(): + raise ValueError("locality.place_split: needs a contiguous tensor") + row_bytes = t[0].numel() * t.element_size() if t.dim() > 1 else t.element_size() + head = split_row * row_bytes + total = t.numel() * t.element_size() + if not 0 < head < total or head % self.granularity: + raise ValueError( + f"locality.place_split: rows [0, {split_row}) = {head} B must be a granule multiple " + f"({self.granularity} B) inside the {total} B tensor" + ) + tail = -(-(total - head) // self.granularity) * self.granularity + m = _Mapping(self.device, [(head, 0), (tail, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + out = flat[:total].view(t.dtype).view(t.shape) + out.copy_(t) + out._locality_mapping = m + for q, k in ( + (out.data_ptr(), 0), + (out.data_ptr() + head - 1, 0), + (out.data_ptr() + head, 1), + (out.data_ptr() + total - 1, 1), + ): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.place_split: byte {q:#x} is not in domain {k}") + return out + + def replace_param_(self, param: torch.nn.Parameter, split_row: int) -> torch.nn.Parameter: + """Re-place a parameter in place with :meth:`place_split` (swaps ``param.data``; the old block returns to the + caching allocator, so call ``torch.cuda.empty_cache()`` between large re-placements).""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.replace_param_: not during CUDA-graph capture") + param.data = self.place_split(param.data, split_row) + return param + + def split_rows(self, t: torch.Tensor, dim: int = 1) -> torch.Tensor: + """Place ``t`` (contiguous, 1-byte elements, rows along ``dim`` split into two equal halves) as + [2, *t.shape[:dim], t.shape[dim] // 2, *t.shape[dim + 1:]]: half k (rows [k h, (k + 1) h) of every leading + index) in domain k's memory, one virtual range. Each half's byte size must be a multiple of the allocation + granularity (2 MiB on VR200). The copy is exact; the caller drops ``t`` afterwards.""" + if t.element_size() != 1 or not t.is_contiguous(): + raise ValueError("locality.split_rows: needs a contiguous 1-byte tensor") + n = t.shape[dim] + if n % 2: + raise ValueError(f"locality.split_rows: {n} rows don't split in halves") + half_bytes = t.numel() // 2 + if half_bytes % self.granularity: + raise ValueError( + f"locality.split_rows: a half is {half_bytes} B, not a multiple of the {self.granularity} B granularity" + ) + m = _Mapping(self.device, [(half_bytes, 0), (half_bytes, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + shape = (2,) + tuple(t.shape[:dim]) + (n // 2,) + tuple(t.shape[dim + 1 :]) + out = flat.view(torch.uint8).view(shape) + src = t.view(torch.uint8) + for k in range(2): + out[k].copy_(src.narrow(dim, k * (n // 2), n // 2)) + out._locality_mapping = m # the tensor keeps its mapping alive + for k in range(2): + p = out[k].data_ptr() + for q in (p, p + half_bytes - 1): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.split_rows: byte {q:#x} is not in domain {k}") + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py new file mode 100644 index 000000000000..2f1a599c45aa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py @@ -0,0 +1,1592 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. + +This target's copy of FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the Eagle3 +drafter's low-M linears, modeling's ``DRAFT_LOWM_OWN_SPLITK``): the same kernel and the same split order at every +tactic, bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's +final cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, +the SwiGLU prologue that follow. + +Each cluster rank accumulates an exact K slice in FP32. Peers publish partials +to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public +``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below +use kernel coordinates: kernel-M carries public N and kernel-N carries public M. +""" + +from __future__ import annotations + +import dataclasses +import functools + +import cuda.bindings.driver as _cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass import Int32 +from cutlass._mlir.dialects import llvm +from cutlass.cute import experimental as cute_ext +from cutlass.cute.nvgpu import tcgen05 +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. +_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") + +#: K extent of one CTA tile. +_CTA_K = 128 # capacity: fixed: a kernel tile constant + +#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. +_SUPPORTED_MMA_M = (64, 128) + +#: Kernel-N carries public M, which is limited to 32. +_SUPPORTED_MMA_N = (8, 16, 32) + +#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. +_SUPPORTED_SPLIT_K = (1, 2, 3, 4) + +#: Largest public M this low-M policy serves. +_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound + +#: Bytes per FP32 partial exchanged through DSMEM. +_FP32_BYTES = 4 + +#: DSMEM mailbox base alignment, in bytes. +_MAILBOX_ALIGN_BYTES = 128 + +#: Size and alignment of one mbarrier, in bytes. +_MBARRIER_BYTES = 8 + +#: Size and alignment of the TMEM base pointer slot. +_TMEM_POINTER_BYTES = 4 + +#: Bytes per BF16/FP16 element. +_AB_ELEMENT_BYTES = 2 + +#: Alignment of the A/B shared-memory buffers. +_AB_BUFFER_ALIGN_BYTES = 1024 # capacity: fixed: an alignment + +#: A/B pipeline stage bounds. One stage is useful only for a single K tile. +_MIN_AB_STAGES = 1 # capacity: fixed: pipeline stages +_DEFAULT_MIN_AB_STAGES = 2 +_MAX_AB_STAGES = 12 # capacity: fixed: pipeline stages + + +@dataclasses.dataclass(frozen=True) +class SplitKTactic: + """One specialization; mma_m carries public N and mma_n carries public M.""" + + mma_m: int + mma_n: int + split_k: int + ab_stages: int + + +def _align_up(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _smem_bytes( + tactic: SplitKTactic, + ab_stages: int, +) -> int: + """Mirror the device allocator's shared-memory layout.""" + cursor = ( + _align_up( + tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, + _AB_BUFFER_ALIGN_BYTES, + ) + + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages + ) + + cursor = _align_up(cursor, _MBARRIER_BYTES) + cursor += 2 * ab_stages * _MBARRIER_BYTES + cursor += 3 * _MBARRIER_BYTES + cursor = _align_up(cursor, _TMEM_POINTER_BYTES) + cursor += _TMEM_POINTER_BYTES + + if tactic.split_k == 1: + return cursor + + return ( + _align_up( + _align_up(cursor, _MAILBOX_ALIGN_BYTES) + + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, + _MBARRIER_BYTES, + ) + + _MBARRIER_BYTES + ) + + +def _max_ab_stages_for( + tactic: SplitKTactic, + smem_capacity: int, +) -> int: + return next( + ( + stages + for stages in range(_MAX_AB_STAGES, -1, -1) + if _smem_bytes(tactic, stages) <= smem_capacity + ), + 0, + ) + + +def validate_tactic( + tactic: SplitKTactic, + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> None: + """Reject a tactic that cannot serve ``(m, n, k)``.""" + if tactic.mma_m not in _SUPPORTED_MMA_M: + raise ValueError(f"unsupported mma_m={tactic.mma_m}") + if tactic.mma_n not in _SUPPORTED_MMA_N: + raise ValueError(f"unsupported mma_n={tactic.mma_n}") + if tactic.split_k not in _SUPPORTED_SPLIT_K: + raise ValueError(f"unsupported split_k={tactic.split_k}") + if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: + raise ValueError( + f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" + ) + if not 1 <= m <= _MAX_M: + raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: + raise ValueError( + f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" + ) + smem_bytes = _smem_bytes(tactic, tactic.ab_stages) + if smem_bytes > smem_capacity: + raise ValueError( + f"tactic {tactic} needs {smem_bytes} B of shared memory but only " + f"{smem_capacity} B are available; max ab_stages is " + f"{_max_ab_stages_for(tactic, smem_capacity)}" + ) + + +def autotune_tactics( + m: int, + n: int, + k: int, + *, + smem_capacity: int = _SMEM_CAPACITY_BYTES, +) -> list[SplitKTactic]: + """Return valid tactics in the shape-specific stage window.""" + tactics: list[SplitKTactic] = [] + for mma_m in _SUPPORTED_MMA_M: + for mma_n in _SUPPORTED_MMA_N: + for split_k in _SUPPORTED_SPLIT_K: + min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES + base = SplitKTactic(mma_m, mma_n, split_k, min_stages) + try: + validate_tactic(base, m, n, k, smem_capacity=smem_capacity) + except ValueError: + continue + max_stages = _max_ab_stages_for(base, smem_capacity) + # Short K favors shallow pipelines; long K stays near the cap. + tactics.extend( + dataclasses.replace(base, ab_stages=ab_stages) + for ab_stages in ( + range(min_stages, min(max_stages, 6) + 1) + if k <= 4 * _CTA_K + else range( + min(max(5, max_stages - 2), max_stages), + max_stages + 1, + ) + ) + ) + return tactics + + +def default_tactic(m: int, n: int, k: int) -> SplitKTactic: + """Choose the default occupancy-oriented tactic.""" + if n <= 512: + mma_m = 64 + mma_n = 16 if n == 512 and m > 24 else 8 + requested_split = 4 + else: + mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 + if n <= 3072: + mma_m = 128 if m <= 16 else 64 + requested_split = 4 if m <= 16 else 2 + elif n < 8192: + mma_m = 64 + requested_split = 2 + else: + mma_m = 128 if k <= 1024 and m <= 24 else 64 + requested_split = 1 + + if k <= 4 * _CTA_K: + requested_split = 1 + split_k = next( + split_k + for split_k in reversed(_SUPPORTED_SPLIT_K) + if split_k <= requested_split and (k // _CTA_K) % split_k == 0 + ) + tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) + max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) + tactic = dataclasses.replace( + tactic, + ab_stages=( + _MIN_AB_STAGES + if k == _CTA_K + else _DEFAULT_MIN_AB_STAGES + if k <= 2 * _CTA_K and m > 8 + else min(max_stages, 6) + if k <= 4 * _CTA_K + else max_stages + ), + ) + validate_tactic(tactic, m, n, k) + return tactic + + +__all__ = [ + "SplitKTactic", + "autotune_tactics", + "default_tactic", + "run_splitk_dense", + "run_splitk_dense_swiglu", + "compile_swiglu_push", +] + + +@dsl_user_op +def _map_shared_rank( + smem_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> "Int32": + """Map an SMEM pointer into a peer CTA's address space.""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + smem_ptr.toint(loc=loc, ip=ip).ir_value(), + peer_cta_rank_in_cluster.ir_value(), + ], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _store_shared_remote_v4( + value0, + value1, + value2, + value3, + smem_ptr: "cute.Pointer", + mbar_ptr: "cute.Pointer", + peer_cta_rank_in_cluster: "Int32", + *, + loc=None, + ip=None, +) -> None: + """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" + llvm.inline_asm( + None, + [ + _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + value0.bitcast(Int32).ir_value(loc=loc, ip=ip), + value1.bitcast(Int32).ir_value(loc=loc, ip=ip), + value2.bitcast(Int32).ir_value(loc=loc, ip=ip), + value3.bitcast(Int32).ir_value(loc=loc, ip=ip), + _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), + ], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[$0], {$1, $2, $3, $4}, [$5];", + "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _silu_mul_f32(g, u, *, loc=None, ip=None): + """silu(g) * (u + 0) in FP32 for one gate / up pair: the instruction sequence of the Triton silu_and_mul kernel at + its defaults (no limit, alpha 1, beta 0), so the bf16 rounding after it is bit for bit the Triton activation.""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [ + cutlass.Float32(g).ir_value(loc=loc, ip=ip), + cutlass.Float32(u).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .f32 t, e, s, p, q, one, zero;\n\t" + "mul.rn.f32 t, $1, 0fBFB8AA3B;\n\tex2.approx.f32 e, t;\n\t" + "mov.f32 one, 0f3F800000;\n\tadd.rn.f32 e, e, one;\n\tdiv.full.f32 s, one, e;\n\t" + "mul.rn.f32 p, $1, s;\n\tmov.f32 zero, 0f00000000;\n\tadd.rn.f32 q, $2, zero;\n\t" + "mul.rn.f32 $0, q, p;\n\t}", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _ld_bf16_f32(addr, *, loc=None, ip=None): + """The bf16 at global byte address ``addr``, widened to FP32 (exact).""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h;\n\tld.global.b16 h, [$1];\n\tcvt.f32.bf16 $0, h;\n\t}", + "=f,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_f32_bf16(addr, value, *, loc=None, ip=None) -> None: + """``value`` rounded to bf16 (round to nearest even) at global byte address ``addr``.""" + llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Float32(value).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b16 h;\n\tcvt.rn.bf16.f32 h, $1;\n\tst.global.b16 [$0], h;\n\t}", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None) -> None: + """Order this thread's generic global stores before later async-proxy (TMA) reads of them.""" + llvm.inline_asm( + None, + [], + "fence.proxy.async.global;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +#: Rank that gathers partials and stores the output. +OWNER_RANK = 0 # capacity: fixed: a split index, not a size + + +# The drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH): the exchange words of +# ops/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return Int32( + llvm.inline_asm( + T.i32(), + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], + "ld.volatile.global.s32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + llvm.inline_asm( + None, + [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Float32(hi).ir_value(loc=loc, ip=ip), + cutlass.Float32(lo).ir_value(loc=loc, ip=ip), + ], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +class SplitKDenseGemmKernel: + """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" + + def __init__( + self, + *, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + b_swiglu: bool = False, + ar_push: bool = False, + ) -> None: + self.ar_push = ar_push + self.acc_dtype = cutlass.Float32 + # The B operand (public A, the activation) is silu(gate) * up of a [M, 2K] gate | up tensor: the B DMA warp + # writes its CTA's K-slice of it to the B tensor (a scratch) before its unchanged TMA loop reads it back. + self.b_swiglu = b_swiglu + self.cta_m = tactic.mma_m + self.cta_n = tactic.mma_n + self.cta_k = _CTA_K + self.num_ab_stage = tactic.ab_stages + self.split_k = tactic.split_k + self.use_pdl = use_pdl + self.has_bias = has_bias + + self.threads_per_cta = 256 + self.epilog_threads = 128 + self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) + self.cta_group = tcgen05.CtaGroup.ONE + self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD + self.cluster_shape = (1, tactic.split_k, 1) + + values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads + if values_per_thread % 4: + raise ValueError( + f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " + f"{values_per_thread} " + "values per epilogue thread; remote stores require a multiple of 4" + ) + self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread + self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES + + @cute.experimental.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + g: cute.Tensor, + u: cute.Tensor, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64 = cutlass.Int64(0), + ar_flags: cutlass.Int64 = cutlass.Int64(0), + ar_rank: Int32 = Int32(0), + ar_world: Int32 = Int32(1), + ar_maxt: Int32 = Int32(8), + ar_tokens: Int32 = Int32(0), + ar_n: Int32 = Int32(0), + ): + # Grid-y packs output-N tile and cluster rank. + self.kernel( + a, + b, + c, + bias, + g, + u, + gu_ld, + b_ld, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ).launch( + grid=( + cute.ceil_div(c.layout.shape[0], self.cta_m), + cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, + c.layout.shape[2], + ), + block=(self.threads_per_cta, 1, 1), + cluster=self.cluster_shape, + smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.experimental.kernel + def kernel( + self, + mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major + mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major + mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major + mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False + mG: cute.Tensor, # b_swiglu: the gate half (Gemm_N, Gemm_K, Gemm_L); dead otherwise + mU: cute.Tensor, # b_swiglu: the up half, same layout; dead otherwise + gu_ld: cutlass.Int32, # b_swiglu: elements between gate (and up) rows + b_ld: cutlass.Int32, # b_swiglu: elements between B rows + ar_ws: cutlass.Int64, # ar_push: the boundary all-reduce's multicast exchange words (int32) + ar_flags: cutlass.Int64, # ar_push: its per-(token, cluster CTA) call flags (parity = flag & 1) + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, # the bucket's tokens + ar_tokens: Int32, # rows pushed (public M) + ar_n: Int32, # public N (the hidden size): the exchange row width + ): + """Allocate storage and dispatch the specialized warps.""" + stages = self.num_ab_stage + + ab_dtype = mA.element_type + tiled_mma = sm100_utils.make_trivial_tiled_mma( + ab_dtype, + ab_dtype, + utils.LayoutEnum.from_tensor(mA).mma_major_mode(), + utils.LayoutEnum.from_tensor(mB).mma_major_mode(), + self.acc_dtype, + self.cta_group, + self.mma_tiler_mn, + ) + + mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) + block_idx = cute.arch.block_idx() + bidx = block_idx[0] + split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + n_idx = block_idx[1] // self.split_k + l_idx = block_idx[2] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + sA = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + sB = cute_ext.allocate( + ab_dtype, + cute.AddressSpace.smem, + sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), + alignment=_AB_BUFFER_ALIGN_BYTES, + ) + + acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) + c_tiler_mn = (self.cta_m, self.cta_n) + + bar_full = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_empty = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(stages), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_mma_epilog = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + bar_tmem_alloc = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + tmem_base_ptr = cute_ext.allocate( + cutlass.Int32, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_TMEM_POINTER_BYTES, + ).iterator + # ar_push: the owner's BF16 tile (kernel-M = output column, kernel-N = token), M-major. + sD = cute_ext.allocate( + mC.element_type, + cute.AddressSpace.smem, + cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), + alignment=16, + ) + + if cutlass.const_expr(self.split_k > 1): + mailbox = cute_ext.allocate( + cutlass.Float32, + cute.AddressSpace.smem, + cute.make_layout(self.mailbox_elements), + alignment=_MAILBOX_ALIGN_BYTES, + ) + bar_reduce = cute_ext.allocate( + cutlass.Int64, + cute.AddressSpace.smem, + cute.make_layout(1), + alignment=_MBARRIER_BYTES, + ).iterator + else: + # Dummy operands for the compile-time-elided reduction. + mailbox = sA + bar_reduce = bar_mma_epilog + + if warp_idx == 0: + with cute.arch.elect_one(): + for i in range(stages): + cute.arch.mbarrier_init(bar_full + i, 2) + cute.arch.mbarrier_init(bar_empty + i, 1) + cute.arch.mbarrier_init(bar_tma_epilog, 32) + cute.arch.mbarrier_init(bar_mma_epilog, 1) + cute.arch.mbarrier_init(bar_tmem_alloc, 160) + + if cutlass.const_expr(self.split_k > 1): + # Owner arrival plus peer transaction-byte credits. + cute.arch.mbarrier_init(bar_reduce, 1) + + cute.arch.mbarrier_init_fence() + if cutlass.const_expr(self.split_k > 1): + # Publish peer barriers before cross-CTA stores. + cute.arch.cluster_arrive_relaxed() + else: + cute.arch.barrier() + + # Host validation guarantees an equal, tail-free K partition. + k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k + k_tile_start = split_rank * k_tile_count + + if cutlass.const_expr(self.split_k > 1): + cute.arch.cluster_wait() + + # Warp 3 is idle; warps 4-7 run the epilogue. + if warp_idx == 0: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), + sA, + cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), + k_tile_start, + k_tile_count, + True, + cutlass.Int64(0), + cutlass.Int64(0), + cutlass.Int64(0), + Int32(0), + Int32(0), + Int32(0), + Int32(0), + ) + elif warp_idx == 1: + self.dma_warp( + bar_full, + bar_empty, + bar_tma_epilog, + cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), + sB, + cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), + k_tile_start, + k_tile_count, + False, + mG.iterator.toint(), + mU.iterator.toint(), + mB.iterator.toint(), + gu_ld, + b_ld, + n_idx * Int32(self.cta_n), + Int32(cute.size(mB, mode=[0])) - n_idx * Int32(self.cta_n), + ) + elif warp_idx == 2: + self.mma_warp( + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma, + sA, + sB, + tmem_base_ptr, + acc_layout, + self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), + k_tile_count, + ) + elif warp_idx >= 4: + self.epilog_warp( + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout, + cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), + cute.arch.thread_idx()[0] - 128, + mC.element_type, + utils.LayoutEnum.from_tensor(mC), + mailbox, + bar_reduce, + split_rank, + sD, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + @cute.experimental.jit + def _swiglu_to_b( + self, + gate_base: cutlass.Int64, + up_base: cutlass.Int64, + out_base: cutlass.Int64, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + row0: cutlass.Int32, + rows_valid: cutlass.Int32, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + """This CTA's K-slice of the activation, silu(gate) * up rounded to bf16, into the B tensor's valid rows (TMA + fills the rest with zeros), through raw global addresses (the B tensor is in the TMA chain); then the stores + are made visible to the TMA reads that follow.""" + lane = cute.arch.lane_idx() + per_tile = cutlass.const_expr(self.cta_n * self.cta_k) + for k_tile in cutlass.range(k_tile_count, unroll=1): + k0 = (k_tile_start + k_tile) * Int32(self.cta_k) + for j in cutlass.range_constexpr(per_tile // 32): + idx = Int32(j * 32) + lane + row = idx // Int32(self.cta_k) + col = k0 + idx % Int32(self.cta_k) + if row < rows_valid: + src = ( + cutlass.Int64(row0 + row) * cutlass.Int64(gu_ld) + cutlass.Int64(col) + ) * cutlass.Int64(2) + dst = ( + cutlass.Int64(row0 + row) * cutlass.Int64(b_ld) + cutlass.Int64(col) + ) * cutlass.Int64(2) + act = _silu_mul_f32(_ld_bf16_f32(gate_base + src), _ld_bf16_f32(up_base + src)) + _st_f32_bf16(out_base + dst, act) + _fence_proxy_async_global() + cute.arch.sync_warp() + + @cute.experimental.jit + def dma_warp( + self, + bar_full, + bar_empty, + bar_tma_epilog, + g_tile: cute.Tensor, + s_tile: cute.Tensor, + cta_v_map: cute.Layout, + k_tile_start: cutlass.Int32, + k_tile_count: cutlass.Int32, + is_a: cutlass.Constexpr, + gate_base: cutlass.Int64, + up_base: cutlass.Int64, + out_base: cutlass.Int64, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + row0: cutlass.Int32, + rows_valid: cutlass.Int32, + ): + stages = self.num_ab_stage + if cutlass.const_expr(not is_a and self.use_pdl): + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not is_a and self.b_swiglu): + self._swiglu_to_b( + gate_base, + up_base, + out_base, + gu_ld, + b_ld, + row0, + rows_valid, + k_tile_start, + k_tile_count, + ) + + empty_phase = cutlass.Int32(1) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + bar_full + stage, + cute.size_in_bytes( + s_tile.element_type, + cute.slice_(s_tile.layout, (None, None, None, 0)), + ), + ) + cute_ext.tma_load( + g_tile[None, None, k_tile_start + k_tile], + s_tile[None, None, None, stage], + (bar_full + stage).value, + cta_v_map=cta_v_map, + tma_operation_type=self.tma_op, + update_expect_tx=False, + ) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + if cutlass.const_expr(is_a and self.use_pdl and not self.ar_push): + # ar_push: the dependent is a prepushed all-reduce, which skips its grid wait: it is released only after + # this grid's own grid wait and its pushes, in the owner's _store_tile (the prepushed >= 2 contract). + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(not is_a and self.has_bias): + cute.arch.mbarrier_arrive(bar_tma_epilog) + self._drain_producer(bar_empty, empty_phase, k_tile_count) + + @cute.experimental.jit + def _drain_producer( + self, + bar_empty, + empty_phase: cutlass.Int32, + k_tile_count: cutlass.Int32, + ): + stages = self.num_ab_stage + for tail in cutlass.range(stages, unroll=1): + stage = (tail + k_tile_count) % stages + cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) + if stage == stages - 1: + empty_phase = empty_phase ^ 1 + + @cute.experimental.jit + def mma_warp( + self, + bar_full, + bar_empty, + bar_mma_epilog, + bar_tmem_alloc, + tiled_mma: cute.TiledMma, + sA: cute.Tensor, + sB: cute.Tensor, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + mma_inst_tile_k: cutlass.Constexpr, + k_tile_count: cutlass.Int32, + ): + num_tmem_cols = 256 + 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) + + tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) + accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] + mma_atom = cute.make_mma_atom(tiled_mma.op) + full_phase = cutlass.Int32(0) + for k_tile in cutlass.range(k_tile_count, unroll=1): + stage = k_tile % self.num_ab_stage + cute.arch.mbarrier_wait(bar_full + stage, full_phase) + for k_block in range(mma_inst_tile_k): + if k_block == 0: + mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + else: + mma_atom.set(tcgen05.Field.ACCUMULATE, True) + cute_ext.dot( + mma_atom, + cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), + cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), + accumulator, + ) + with cute.arch.elect_one(): + tcgen05.commit(bar_empty + stage, None, self.cta_group) + if stage == self.num_ab_stage - 1: + full_phase = full_phase ^ 1 + + with cute.arch.elect_one(): + tcgen05.commit(bar_mma_epilog, None, self.cta_group) + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 1) + cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) + + @cute.experimental.jit + def epilog_warp( + self, + bar_tma_epilog, + bar_mma_epilog, + bar_tmem_alloc, + tmem_base_ptr, + acc_layout: cutlass.Constexpr, + gD_tile: cute.Tensor, + gBias_tile: cute.Tensor, + epi_tid: cutlass.Int32, + c_dtype: cutlass.Constexpr, + d_layout: cutlass.Constexpr, + mailbox, + bar_reduce, + split_rank: cutlass.Int32, + sD: cute.Tensor, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + # Wait until MMA publishes the TMEM base pointer. + cute.arch.mbarrier_arrive(bar_tmem_alloc) + cute.arch.mbarrier_wait(bar_tmem_alloc, 0) + + acc_view = cute.make_tensor( + cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), + acc_layout, + )[((None, None), 0, 0, 0)] + + epi_tile = (self.cta_m, self.cta_n) + tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( + sm100_utils.get_tmem_load_op( + (self.cta_m, self.cta_n, self.cta_k), + d_layout, + c_dtype, + self.acc_dtype, + epi_tile, + False, + ), + acc_view, + ) + gD_epi = cute.flat_divide(gD_tile, epi_tile) + + # Match each epilogue thread's TMEM partition in RMEM. + rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) + rAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rD = cute_ext.allocate( + c_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + thr_t2r = tiled_copy_t2r.get_slice(epi_tid) + + if cutlass.const_expr(self.has_bias): + bias_dtype = gBias_tile.element_type + rBias = cute_ext.allocate( + bias_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + rBiasAcc = cute_ext.allocate( + self.acc_dtype, + cute.AddressSpace.rmem, + rmem_layout, + alignment=32, + ) + if split_rank == OWNER_RANK: + cute.arch.mbarrier_wait(bar_tma_epilog, 0) + cute_ext.partition_and_copy( + cute.make_tiled_copy_D( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), + tiled_copy_t2r, + ).get_slice(epi_tid), + cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], + rBias, + ) + rBiasAcc.store(rBias.load().to(self.acc_dtype)) + + cute.arch.mbarrier_wait(bar_mma_epilog, 0) + cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) + # Make tcgen05.ld visible before TMEM release and RMEM use. + cute.arch.fence_view_async_tmem_load() + cute.arch.mbarrier_arrive(bar_tmem_alloc) + + # Peers publish FP32 partials; only rank 0 reduces and stores. + if cutlass.const_expr(self.split_k > 1): + assert cute.size(rmem_layout) == self.mailbox_elements // ( + (self.split_k - 1) * self.epilog_threads + ) + values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) + values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) + if split_rank != OWNER_RANK: + for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): + _store_shared_remote_v4( + rAcc[value_idx], + rAcc[value_idx + 1], + rAcc[value_idx + 2], + rAcc[value_idx + 3], + mailbox.iterator + + (split_rank - Int32(1)) * values_per_peer + + epi_tid * values_per_thread + + value_idx, + bar_reduce, + Int32(OWNER_RANK), + ) + else: + if epi_tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx( + bar_reduce, self.expected_transaction_bytes + ) + cute.arch.mbarrier_wait(bar_reduce, 0) + for peer in cutlass.range_constexpr(self.split_k - 1): + for value_idx in cutlass.range_constexpr(values_per_thread): + rAcc[value_idx] = ( + rAcc[value_idx] + + mailbox[ + peer * values_per_peer + epi_tid * values_per_thread + value_idx + ] + ) + + if split_rank == OWNER_RANK: + if cutlass.const_expr(self.has_bias): + rAcc.store(rAcc.load() + rBiasAcc.load()) + self._store_tile(thr_t2r, rAcc, rD, gD_epi, c_dtype) + if cutlass.const_expr(self.ar_push): + self._push_tile( + thr_t2r, + rD, + sD, + epi_tid, + bidx, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + # The reduction mbarrier covers remote stores; no cluster barrier needed. + + @cute.experimental.jit + def _store_tile(self, thr_t2r, rAcc, rD, gD_epi, c_dtype: cutlass.Constexpr): + """The owner rank's reduced FP32 tile: one cast to ``c_dtype`` and the store to global memory.""" + rD.store(rAcc.load().to(c_dtype)) + # Preserve TMEM coordinates; the copy predicates output tails. + cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) + + @cute.experimental.jit + def _push_tile( + self, + thr_t2r, + rD, + sD, + epi_tid: cutlass.Int32, + bidx: cutlass.Int32, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_world: cutlass.Int32, + ar_maxt: cutlass.Int32, + ar_tokens: cutlass.Int32, + ar_n: cutlass.Int32, + ): + """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output + would store them (ops/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) + cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" + cute_ext.partition_and_copy(thr_t2r, rD, sD) + cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + pairs = self.cta_m // 2 + col0 = bidx * self.cta_m + half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) + for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): + word = epi_tid + w0 + t = word // pairs + p = word % pairs + if t < ar_tokens: + n = col0 + 2 * p + par = _ld_volatile_s32( + ar_flags + cutlass.Int64(4) * cutlass.Int64(t * 8 + n // (ar_n // 8)) + ) & Int32(1) + off = ( + cutlass.Int64(par) * half + + cutlass.Int64(t * ar_world + ar_rank) * cutlass.Int64(ar_n // 2) + + cutlass.Int64(n // 2) + ) + _st_global_b32( + ar_ws + off * cutlass.Int64(4), + _bf16x2_word( + sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + ), + ) + if cutlass.const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + + +import torch as _torch # noqa: E402 (as in FlashInfer) + +_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) + + +@cute.experimental.jit +def _bmm_no_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + b_sel = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + b_sel, + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + b_sel, + b_sel, + Int32(0), + Int32(0), + stream, + ) + + +@cute.experimental.jit +def _bmm_swiglu( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + g: cute.Tensor, + u: cute.Tensor, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + stream: _cuda.CUstream, +): + """_bmm_no_bias with the B operand built in the kernel from g / u (laid out as b, rows gu_ld apart; B rows b_ld + apart).""" + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + cute.make_tensor(g.iterator, cute.select(g.layout, mode=[2, 1, 0])), + cute.make_tensor(u.iterator, cute.select(u.layout, mode=[2, 1, 0])), + gu_ld, + b_ld, + stream, + ) + + +@cute.experimental.jit +def _bmm_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ar_tokens: Int32, + ar_n: Int32, +): + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + Int32(0), + Int32(0), + stream, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + +@cute.experimental.jit +def _bmm_swiglu_push( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + g: cute.Tensor, + u: cute.Tensor, + gu_ld: cutlass.Int32, + b_ld: cutlass.Int32, + stream: _cuda.CUstream, + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: Int32, + ar_world: Int32, + ar_maxt: Int32, + ar_tokens: Int32, + ar_n: Int32, +): + """_bmm_swiglu that also pushes the output rows (_bmm_push's exchange).""" + c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), + c, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), + cute.make_tensor(g.iterator, cute.select(g.layout, mode=[2, 1, 0])), + cute.make_tensor(u.iterator, cute.select(u.layout, mode=[2, 1, 0])), + gu_ld, + b_ld, + stream, + ar_ws, + ar_flags, + ar_rank, + ar_world, + ar_maxt, + ar_tokens, + ar_n, + ) + + +@cute.experimental.jit +def _bmm_bias( + gemm_op: cutlass.Constexpr, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + bias: cute.Tensor, + stream: _cuda.CUstream, +): + b_sel = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])) + gemm_op( + cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), + b_sel, + cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), + cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), + b_sel, + b_sel, + Int32(0), + Int32(0), + stream, + ) + + +def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): + return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( + leading_dim=leading_dim + ) + + +def _detect_leading_dim(tensor: _torch.Tensor) -> int: + # Ignore synthetic batch stride, including 1x1. + for dim, stride in enumerate(tensor.stride()[1:], start=1): + if stride == 1: + return dim + raise ValueError("tensor has no stride-1 dimension") + + +def _make_layout_tensor( + shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int +) -> _torch.Tensor: + permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] + return _torch.empty( + tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" + ).permute([permutation.index(dim) for dim in range(len(shape))]) + + +def _make_compile_repr_tensors( + dtype: _torch.dtype, + has_bias: bool, + a_leading: int, + b_leading: int, + c_leading: int, +): + m, n, k, batch = 64, 8, _CTA_K, 1 + tensors = tuple( + _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) + for shape, leading_dim in ( + ((batch, n, k), a_leading), + ((batch, k, m), b_leading), + ((batch, n, m), c_leading), + ) + ) + if not has_bias: + return (*tensors, None) + return ( + *tensors, + _from_dlpack_dynamic( + _torch.empty((n,), dtype=dtype, device="cuda").as_strided( + size=(batch, n, m), stride=(0, 1, 0) + ), + 1, + 2, + ), + ) + + +def _to_cute_swap(a, b, out, bias): + a_swap = b.unsqueeze(0).transpose(-2, -1) + b_swap = a.unsqueeze(0).transpose(-2, -1) + c_swap = out.unsqueeze(0).transpose(-2, -1) + leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) + cute_tensors = tuple( + _from_dlpack_dynamic(tensor, leading_dim) + for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) + ) + if bias is None: + return (*cute_tensors, None, leading_dims) + return ( + *cute_tensors, + _from_dlpack_dynamic( + bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), + 1, + 2, + ), + leading_dims, + ) + + +@functools.cache +def _get_compiled_splitk_kernel( + dtype, + tactic: SplitKTactic, + use_pdl: bool, + has_bias: bool, + leading_dims: tuple[int, int, int], + b_swiglu: bool = False, +): + if dtype not in _SUPPORTED_TORCH_DTYPES: + raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") + + kernel = SplitKDenseGemmKernel( + tactic=tactic, + use_pdl=use_pdl, + has_bias=has_bias, + b_swiglu=b_swiglu, + ) + compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + if b_swiglu: + # The gate / up halves are laid out as the B operand (public A) is. + gu = _make_compile_repr_tensors(dtype, False, *leading_dims)[0] + compiled = cute_ext.compile( + _bmm_swiglu, kernel, *compile_tensors[:3], gu, gu, Int32(2), Int32(1), stream + ) + elif has_bias: + compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) + else: + compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) + return compiled + + +def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: + tensors = (a, b, out) + ((bias,) if bias is not None else ()) + if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): + raise ValueError("a, b, out, and bias must be torch tensors") + if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: + raise ValueError("split-K dense GEMM accepts only 2D 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, b, out, and bias must share BF16 or FP16 dtype") + if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): + raise ValueError("a, b, and out must be dense row-major or column-major matrices") + if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): + raise ValueError("a, b, and out must be 32-byte aligned") + + m, k = a.shape + if b.shape[0] != k: + raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") + n = b.shape[1] + if out.shape != (m, n): + raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") + if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): + raise ValueError( + f"bias must be contiguous with shape {(n,)}, " + f"got shape {tuple(bias.shape)} and stride {bias.stride()}" + ) + + return m, n, k + + +def run_splitk_dense( + a, + b, + bias, + out, + pdl: bool, + tactic: SplitKTactic, +): + """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) + has_bias = bias is not None + cute_tensors = _to_cute_swap(a, b, out, bias) + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=has_bias, + leading_dims=cute_tensors[4], + ) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + if has_bias: + compiled(*cute_tensors[:4], stream) + else: + compiled(*cute_tensors[:3], stream) + return out + + +@functools.cache +def _get_compiled_swiglu_push( + dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] +): + """The SwiGLU B load and the ar push in one kernel (the drafter down projection at a pushed boundary).""" + kernel = SplitKDenseGemmKernel( + tactic=tactic, use_pdl=use_pdl, has_bias=False, b_swiglu=True, ar_push=True + ) + compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) + gu = _make_compile_repr_tensors(dtype, False, *leading_dims)[0] + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + z64, z32 = cutlass.Int64(0), Int32(0) + return cute_ext.compile( + _bmm_swiglu_push, kernel, *compile_tensors[:3], gu, gu, Int32(2), Int32(1), stream, + z64, z64, z32, Int32(1), Int32(8), z32, Int32(6144), + ) # fmt: skip + + +def compile_swiglu_push(scratch, b, out, pdl: bool, tactic) -> None: + """Compile (only) the SwiGLU + push variant of ``run_splitk_dense_swiglu`` for these operands' layouts.""" + cute_tensors = _to_cute_swap(scratch, b, out, None) + _get_compiled_swiglu_push(scratch.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) + + +def run_splitk_dense_swiglu( + gate_up, b, out, pdl: bool, tactic: SplitKTactic, scratch, ar: tuple = None +): + """``out = silu_and_mul(gate_up) @ b`` in one launch: ``gate_up`` [M, 2K] (gate then up, contiguous rows) is + activated by the B DMA warp (the Triton silu_and_mul arithmetic at its defaults, bit for bit) into ``scratch`` + [>= M, K], which the unchanged GEMM then reads; the output is bit for bit ``run_splitk_dense(silu_and_mul(gate_up), + b, None, out, pdl, tactic)``. ``scratch`` must not be read or written by anything else while the kernel runs. With + ``ar`` (``run_splitk_dense_push``'s tuple) the rows are also pushed into the drafter boundary all-reduce's + exchange, from the same kernel.""" + if gate_up.ndim != 2 or gate_up.shape[1] % 2 or not gate_up.is_contiguous(): + raise ValueError("gate_up must be a contiguous [M, 2K] matrix") + m, k = gate_up.shape[0], gate_up.shape[1] // 2 + a = scratch[:m] + if a.shape != (m, k) or not a.is_contiguous(): + raise ValueError(f"scratch must be a contiguous [>= {m}, {k}] matrix") + gate, up = gate_up[:, :k], gate_up[:, k:] + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + if any(t.data_ptr() % 32 for t in (gate, up)) or gate.dtype != a.dtype: + raise ValueError("gate / up halves must be 32-byte aligned and of the GEMM dtype") + cute_tensors = _to_cute_swap(a, b, out, None) + g_swap = gate.unsqueeze(0).transpose(-2, -1) + u_swap = up.unsqueeze(0).transpose(-2, -1) + lead = cute_tensors[4][1] + if _detect_leading_dim(g_swap) != lead or _detect_leading_dim(u_swap) != lead: + raise ValueError("gate / up must share the activation's K-major layout") + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + operands = ( + *cute_tensors[:3], + _from_dlpack_dynamic(g_swap, lead), + _from_dlpack_dynamic(u_swap, lead), + Int32(gate_up.stride(0)), + Int32(a.stride(0)), + stream, + ) + if ar is None: + compiled = _get_compiled_splitk_kernel( + dtype=a.dtype, + tactic=tactic, + use_pdl=pdl, + has_bias=False, + leading_dims=cute_tensors[4], + b_swiglu=True, + ) + compiled(*operands) + return out + compiled = _get_compiled_swiglu_push(a.dtype, tactic, pdl, cute_tensors[4]) + ws, flags, rank, world, maxt = ar + compiled( + *operands, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + Int32(out.shape[0]), + Int32(out.shape[1]), + ) + return out + + +@functools.cache +def _get_compiled_push( + dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] +): + kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, ar_push=True) + compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) + stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) + z64, z32 = cutlass.Int64(0), Int32(0) + return cute_ext.compile( + _bmm_push, + kernel, + *compile_tensors[:3], + stream, + z64, + z64, + z32, + Int32(1), + Int32(8), + z32, + Int32(6144), + ) + + +def compile_push(a, b, out, pdl: bool, tactic) -> None: + """Compile (only) the push variant of ``run_splitk_dense_push`` for these operands' shapes and layouts.""" + cute_tensors = _to_cute_swap(a, b, out, None) + _get_compiled_push(a.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) + + +def run_splitk_dense_push(a, b, out, pdl: bool, tactic: SplitKTactic, ar: tuple): + """``run_splitk_dense`` without bias that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange; + ``ar`` = (multicast words address, flags address, rank, world, bucket tokens) of the 8-token bucket.""" + validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) + cute_tensors = _to_cute_swap(a, b, out, None) + compiled = _get_compiled_push(a.dtype, tactic, pdl, cute_tensors[4]) + stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) + ws, flags, rank, world, maxt = ar + compiled( + *cute_tensors[:3], + stream, + cutlass.Int64(ws), + cutlass.Int64(flags), + Int32(rank), + Int32(world), + Int32(maxt), + Int32(out.shape[0]), + Int32(out.shape[1]), + ) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py new file mode 100644 index 000000000000..eb025bc1694a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py @@ -0,0 +1,431 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-latency TP all-reduce + residual + RMSNorm for decode (M <= 8 tokens) over a multicast mapping. + + partial = bf16(x + x2) (optional second input, e.g. shared + routed experts) + updated = bf16(residual + sum_r partial_r) (ranks summed in FP32, rank order) + normed = bf16(updated * rsqrt(mean(updated^2) + eps) * w) + +One cluster of CLUSTER CTAs per token; CTA c owns columns [c * H / CLUSTER, ...), one 16-byte chunk (8 BF16) +per thread. Each CTA pushes its slice of this rank's partial through the multicast mapping (one store lands in +every rank's buffer), releases dependents, polls its own buffer until all ranks' slices arrived, sums them, adds +the residual and reduces the RMS statistic across the cluster through DSMEM. + +Exchange protocol as in the K3 sandwich: int32 words, 0x80000000 = empty (pushed -0.0 halves become +0.0), each +CTA keeps a call counter whose parity selects the buffer half, readers re-arm the words they read. All ranks of +the TP group must issue the same calls in the same stream order. + +The workspace holds two exchange sets. With prepushed 3 two producers pushed this rank's x (set A) and x2 (set B), +e.g. the routed experts' finalize and the shared expert's down projection, and the kernel forms bf16(x + x2) per +rank from the two slots, exactly as the unpushed path does. + +Producers that push (the o_proj / shared-expert GEMV ``ar`` tuple, the routed-expert engine's ``arp`` descriptor) are +pointed at a bucket's exchange by its ``mc`` address, ``flags_for(8)`` counters, ``rank``, ``world_size`` and +``max_tokens``; ``get_workspace`` and ``_bucket`` select it. + +The kernel is ``torch.ops.trtllm.minimax_m3_mc_allreduce_norm``, reached through the ``ops/comm`` entries +``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``; this module owns the workspaces, the bucket and +epilogue-option choice and the MXFP8 publication. +""" + +import functools +from dataclasses import dataclass +from typing import Any + +import torch + +from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op +from ..ops.comm.m3_mc_finalize_allreduce_norm import ( + m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, +) +from . import m3_report + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +BUCKETS = (8, 16, 48) # token capacity of each workspace (grid = bucket * cluster CTAs) +EMPTY_WORD = -2147483648 + +# Attention / MoE boundary epilogue options (each bit-identical): 1 residual + norm weight loaded at entry (before +# the poll); 2 cluster RMS exchange by st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated +# right after the poll (no trailing __syncthreads); 8 max-smem carveout so the next kernel's CTAs co-reside early. +# The 8-token bucket, by token count 1..8 (GB300 TP4 tuning). +AR_OPTS_BY_M = (7, 7, 15, 15, 15, 15, 7, 7) +# The 16- and 48-token buckets. +AR_OPTS_BY_BUCKET = {16: 15, 48: 15} +# The MoE-finalize all-reduce, per bucket. +AR_FIN_OPTS_BY_BUCKET = {8: 7, 16: 15, 48: 15} + + +# Option 2 is also the safe ordering of the cluster RMS exchange: its stores into peer CTAs' shared memory follow a +# cluster barrier arrive / wait pair. Without it they run before any cluster barrier, when a peer need not have +# started, so every call keeps it. +_OPT_CLUSTER_BARRIER = 2 + + +def _ar_opts(m: int, bucket: int, finalize: bool = False) -> int: + """Epilogue options for an m-row call on the given bucket; every option is bit-identical. A bucket missing from + the tables gets 0, and every result keeps _OPT_CLUSTER_BARRIER.""" + if finalize: + opts = AR_FIN_OPTS_BY_BUCKET.get(bucket, 0) + elif bucket == 8: + opts = AR_OPTS_BY_M[m - 1] + else: + opts = AR_OPTS_BY_BUCKET.get(bucket, 0) + return opts | _OPT_CLUSTER_BARRIER + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags + + +class McBuckets: + """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" + + def __init__(self, workspaces): + self.workspaces = workspaces + self.hidden = workspaces[0].hidden + + def for_tokens(self, m: int): + for ws in self.workspaces: + if m <= ws.max_tokens: + return ws + return None + + +def _bucket(workspace, m): + return workspace.for_tokens(m) if isinstance(workspace, McBuckets) else workspace + + +def decline_reason(x, residual, workspace, cluster=8, x2=None) -> str: + """Why ``m3_mc_allreduce_norm`` returned None for these operands (for the dispatch's miss message).""" + if workspace is None: + return "no multicast workspace: it is created collectively on the first eager decode step" + bucket = _bucket(workspace, x.shape[0]) + if bucket is None: + return f"no workspace bucket holds {x.shape[0]} tokens (buckets {BUCKETS})" + if bucket.flags.get(cluster) is None: + return f"no {cluster}-CTA call counters for the {bucket.max_tokens}-token bucket" + return ( + f"operands outside the kernel: x {tuple(x.shape)} {x.dtype} contiguous={x.is_contiguous()}, " + f"residual contiguous={residual.is_contiguous()}, hidden {bucket.hidden}" + + ("" if x2 is None else f", x2 {tuple(x2.shape)} contiguous={x2.is_contiguous()}") + ) + + +def m3_mc_allreduce_norm( + x, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + x2=None, + quant=False, + prepushed=0, + drafter_norm=False, + carveout=True, + capture=None, +): + """Return (normed, updated_residual) or None when unsupported. ``carveout`` False drops opts bit 8 (the max-smem + carveout) for this call. ``drafter_norm``: the Eagle3 drafter boundary + (opts bit 16: normalize the unrounded FP32 sum, as the drafter's stock FlashInfer fused_add_rmsnorm does); 8-token + bucket only, x only, no MXFP8 copy. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture`` = (out [>= M, H] view, fc_norm weight + [H], eps): the boundary also writes the Eagle3 capture of the updated residual, bit for bit + m3_spec_tail.capture_norm (CAPTURE_NORM_IN_BOUNDARY; not with ``drafter_norm``). + + A None makes the caller take another collective, so it must be the same on every TP rank. It is: it depends only + on the shapes, dtypes and layouts of the operands (the same code builds them on every rank) and on whether the + workspace's flags exist, which ``get_workspace`` creates collectively, before any capture, for cluster 8.""" + workspace = _bucket(workspace, x.shape[0]) + if workspace is None: + return None + if ( + x.dim() != 2 + or x.shape[0] > MAX_TOKENS + or x.dtype != torch.bfloat16 + or x.shape[1] != workspace.hidden + or not x.is_contiguous() + or not residual.is_contiguous() + or (x2 is not None and (not x2.is_contiguous() or x2.shape != x.shape)) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + if drafter_norm and ( + workspace.max_tokens != 8 or x2 is not None or quant or prepushed not in (0, 2) + ): + return None + outs = _allreduce_norm_op( + x, + x2, + residual, + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and x.shape[1] % 128 == 0), + int(prepushed), + (_ar_opts(x.shape[0], workspace.max_tokens) & (~0 if carveout else ~8)) + | (16 if drafter_norm else 0), + *_capture_args(capture, drafter_norm), + own_op=own_op(), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +def m3_mc_finalize_allreduce_norm( + fc2, + perm, + scales, + shared, + residual, + norm_weight, + eps, + workspace, + use_pdl=False, + cluster=8, + quant=False, + capture=None, +): + """MoE finalize (sum_k scale * fc2[perm] + shared) + all-reduce + residual + RMSNorm in one kernel. + Returns (normed, updated_residual) or None when unsupported. With ``quant`` the normed row is also published, + MXFP8-quantized, to m3_prequant.SLOT for the next MXFP8 GEMM. ``capture``: as ``m3_mc_allreduce_norm``.""" + m = residual.shape[0] + workspace = _bucket(workspace, m) + if workspace is None: + return None + if ( + residual.dim() != 2 + or m > MAX_TOKENS + or residual.shape[1] != workspace.hidden + or fc2.dtype != torch.bfloat16 + or not fc2.is_contiguous() + or fc2.dim() != 2 + or perm.dtype != torch.int32 + or perm.dim() != 2 + or perm.shape[0] != m + or not perm.is_contiguous() + or scales.dtype not in (torch.float32, torch.bfloat16) + or not scales.is_contiguous() + or (shared is not None and (shared.shape != residual.shape or not shared.is_contiguous())) + ): + return None + flags = workspace.flags_for(cluster) + if flags is None: + return None + outs = _finalize_allreduce_norm_op( + fc2, + perm, + scales, + shared, + residual.contiguous(), + norm_weight, + float(eps), + workspace.uc, + workspace.mc.data_ptr(), + flags, + workspace.rank, + workspace.world_size, + cluster, + use_pdl, + workspace.max_tokens, + bool(quant and residual.shape[1] % 128 == 0), + _ar_opts(m, workspace.max_tokens, finalize=True), + *_capture_args(capture, False), + own_op=own_op(), + ) + if len(outs) == 4: + from .m3_prequant import SLOT + + SLOT.publish(outs[0], outs[2], outs[3]) + return outs[0], outs[1] + + +@functools.lru_cache(maxsize=None) +def _own_op_loaded() -> bool: + torch.ops.trtllm.minimax_m3_sm107_tp8_mc_allreduce_norm + return True + + +# The boundary all-reduces run this target's op (trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm) rather than the +# image's: the same exchange and arithmetic built for the two-tray TP8 group. +OWN_OP = True +OWN_OP_PATH = "m3_mc_ar_own_op" + + +def own_op() -> bool: + """Select this target's all-reduce when OWN_OP is enabled, requiring its compiled op.""" + if not OWN_OP: + return False + _own_op_loaded() + m3_report.taken(OWN_OP_PATH, None, "the boundary all-reduces run this target's op") + return True + + +@functools.lru_cache(maxsize=None) +def capture_supported() -> bool: + """Whether the multicast all-reduce op the boundaries call (own_op) has the capture epilogue (cap_out / cap_w / + cap_eps). The same library is loaded on every rank, so the ranks agree.""" + name = ( + "minimax_m3_sm107_tp8_mc_allreduce_norm" + if own_op() + else "minimax_m3_sm107_tp2_mc_allreduce_norm" + ) + op = getattr(torch.ops.trtllm, name) + return any(a.name == "cap_out" for a in op.default._schema.arguments) + + +def _capture_args(capture, drafter_norm: bool) -> tuple: + """(cap_out, cap_w, cap_eps) for the op, from ``capture`` = (out, fc_norm weight, eps) or None.""" + if capture is None: + return (None, None, 0.0) + assert not drafter_norm, "the capture epilogue is not built for the drafter boundary" + out, weight, eps = capture + return (out, weight, float(eps)) + + +_WORKSPACE: dict = {} +_AGREE_COMMS: dict = {} + + +def _tp_comm(mapping): + """A communicator over mapping's TP group, created once (collective over that group only).""" + key = tuple(mapping.tp_group) + comm = _AGREE_COMMS.get(key) + if comm is None: + from tensorrt_llm._utils import mpi_comm, mpi_disabled + + if mpi_disabled(): + comm = mapping.tp_group_pg + else: + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + _AGREE_COMMS[key] = comm + return comm + + +def tp_agree(mapping, ok: bool) -> bool: + """Whether ``ok`` holds on every rank of mapping's TP group (one blocking host collective). + + For a fast path whose availability is decided from rank-local state (a kernel build that can fail on one rank): + every rank must take the same path, so that they issue the same collectives with the same exchange protocol and + the same kernels. Call it once per decision, at build time, at the same point on every rank, outside CUDA-graph + capture; never per step. With one rank, or attention DP (whose ranks run different batches), it returns ``ok``. + """ + if mapping.tp_size <= 1 or mapping.enable_attention_dp: + return ok + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("tp_agree must run outside CUDA-graph capture") + from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded + + return _mnnvl_workspace_all_succeeded(_tp_comm(mapping), bool(ok)) + + +def tp_agree_built(mapping, built, what: str): + """``built`` when it is not None on every TP rank, else None on every rank (see ``tp_agree``).""" + if tp_agree(mapping, built is not None) or built is None: + return built + from tensorrt_llm.logger import logger + + logger.warning(f"{what}: built here but not on another TP rank; every rank falls back") + return None + + +def get_workspace(mapping, hidden: int): + """Process-wide workspace (created collectively on first eager use); None while capturing before creation. + + The first eager call for a hidden size (the warmup forward, before any CUDA-graph capture) creates one + McAllReduceWorkspace per bucket in BUCKETS on every TP rank (MPI collective over the TP group) plus each + bucket's 8-CTA-cluster call counters; later calls, captured or not, return the same McBuckets. The workspaces + live until process exit. + """ + ws = _WORKSPACE.get(hidden) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + return None + wss = [McAllReduceWorkspace.create(mapping, hidden, b) for b in BUCKETS] + for w in wss: + w.flags_for(8) + ws = McBuckets(wss) + _WORKSPACE[hidden] = ws + return ws diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/__init__.py new file mode 100644 index 000000000000..e361dc614231 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/__init__.py @@ -0,0 +1,16 @@ +# 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. +"""MiniMax-M3 MoE layer of the TP8 target for low-batch decode as a CuTe DSL kernel: ``m3_moe`` for 1..8 tokens +(op.py) at the TP8 routed and shared intermediate of 384 per rank.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py new file mode 100644 index 000000000000..d1730239463c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py @@ -0,0 +1,3841 @@ +# 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. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP rank: ``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 # capacity: fixed: the checkpoint's hidden size +# routed intermediate per rank: 768 at TP4, 384 at TP8 (FC2 K = 1.5 256-K tiles: the last tile's A box is half out +# of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 MMAs / scale atoms are issued; no padding, no extra +# bytes) +I_TP = _cfg("i_tp", 768) +TWO_I = 2 * I_TP +assert I_TP % 128 == 0 and 128 <= I_TP <= 768, I_TP +E_LOCAL = 128 +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 # capacity: fixed: the CTA size +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# EARLY_AR (scheduling only): when the finalize pushes into the multicast all-reduce, EVERY CTA releases the +# dependents right after the grid dependency, so the all-reduce (which reads only the pushed tagged words, without a +# grid wait) and, through its own early release, the next layer's QKV GEMV become resident on each SM this kernel +# frees, instead of after the last CTA's teardown. The arithmetic is unchanged. +EARLY_AR = bool(_cfg("early_ar", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +# the shared intermediate per rank: 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate positions: its +# gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP4: 2 x (32 + 32), the order +# of m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. +SH_I = _cfg("sh_i", 768) +SH_ROWS = SH_I // 24 +# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) +SH_SPLIT = 128 // (2 * SH_ROWS) +SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N +SH_KB = H // 32 # capacity: fixed: a kernel tile constant (gate_up K blocks) +SH_DKB = SH_I // 32 # down K blocks +SH_DKT = SH_I // 128 # down k-tiles +SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the +# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction +# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) +) # (in-chain: a win at M <= 2, not 4..8) +if SHARED: + assert SH_I in (384, 768), SH_I + assert SH_PRETILE or SH_SPLIT == 2, ( + "the un-pretiled shared stream is the 32-row (TP4) layout only" + ) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # capacity: fixed: pipeline stages (6 stages per unit) +NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) +# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) +KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 + + +def _nkb2(t2: int) -> int: + """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" + return NUM_KB if t2 < KT2 - 1 else KB2_LAST + + +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy +# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few +# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. +HB_REPL = _cfg("hb_repl", 1) +assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) +HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 +PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 # capacity: fixed: a warp role index +X_WARP = 5 # capacity: fixed: a warp role index +S_WARP = 6 # capacity: fixed: a warp role index +MMA_WARP = 7 # capacity: fixed: a warp role index +# E0_M4: at M = 4 the MMA warp routes +# token 3 (NRW0 = 4), so the early group-0 stages go to FC2 B builder warp 8 instead (idle until the routing is done; +# RW_EXTRA = 0 at M = 4). Issue order only: exact. +E0_M4 = ( + bool(_cfg("e0_m4", 1)) and NRW0 == 4 and RW_EXTRA == 0 and THREADS // 32 > 8 +) # M == 4 only: at M > 4 warp 8 routes (RW_EXTRA) +E0_WARP = ( + 8 if E0_M4 else MMA_WARP +) # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and (NRW0 <= 3 or E0_M4) and bool(_cfg("e0_off", 1)) +# B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing +# warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the +# two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. +if B1_AFTER_TOPK and XSTAGE and EARLY0 and not E0_OFFLOAD: + raise ValueError( + "m3_moe: b1_after_topk=1 with xstage, early0 and no early-stage offload deadlocks (warp 4 waits on x_bar[1] " + "before TOPK2_BAR, the B1 build waits on TOPK2_BAR before x_bar[1]); set b1_after_topk=0 or early0=0" + ) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 # capacity: fixed: the SM's tensor-memory columns +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +assert M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 +_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 +_SH_FIX = ( + SH_GT * SH_NB * 128 + + SH_DKT * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + _SH_XCH_BYTES + + 2 * SH_ROWS * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * SH_NB * 128 +_SB_XC = _SB_BH + SH_DKT * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + _SH_XCH_BYTES +_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = SH_NB +SH_SFA = SH_ACC_DN + 16 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +SH_TMEM = 256 # capacity: fixed: the shared CTAs' TMEM columns +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + # (a partial last k-tile has no atoms past I_TP) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SH_NB, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(SH_DKT) + kt = gd - jt * cutlass.Int32(SH_DKT) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if cutlass.const_expr(SH_SPLIT == 2): + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + else: + # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are + # accumulator columns 8 qd + t + av = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_GU + qd * cutlass.Int32(8), + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(1): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) + * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd == cutlass.Int32(0): + # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + for s_ in cutlass.range_constexpr(1, SH_SPLIT): + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + lane * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(SH_ROWS * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(SH_DKT): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if cutlass.const_expr(EARLY_AR): + if ar_on: + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + for rr in cutlass.range_constexpr(HB_REPL): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4) + + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + hb_off = cutlass.Int32(0) + if cutlass.const_expr(HB_REPL > 1): + # this CTA's copy of the records + hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview( + (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) + ), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(_nkb2(t2)): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if cutlass.const_expr(HB_REPL == 1): + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + else: + if lane < HB_REPL: # copy `lane` + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4) + + lane * cutlass.Int32(HB_STRIDE), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py new file mode 100644 index 000000000000..70deb357f4a8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py @@ -0,0 +1,693 @@ +# 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. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP8 rank (routed and shared intermediate 384 per rank), +over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size +INTER = 384 # capacity: fixed: the routed experts' width per tp8 rank +NUM_EXPERTS = 128 +SHARED_INTER = 384 # capacity: fixed: the shared expert's width per tp8 rank (the routed one) +SHARED_CTAS = 24 # capacity: fixed: an SM split, not a size (m3_moe's shared-expert CTAs (beside num_ctas routed CTAs)) +TOP_K = 4 # capacity: fixed: the checkpoint's top-k +U1 = 2 * INTER // 32 # FC1 units per expert +# The router GEMV takes at most 2 experts per CTA (router weights held in registers). +MIN_CTAS = 64 # capacity: fixed: an SM split, not a size + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # capacity: fixed: pipeline stages (FC1 stages (items) per unit) +_ITEM_BYTES = _UNIT_BYTES // _S1 # capacity: fixed: a ring item's bytes +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters (7 and 8 tokens). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 1, 1) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit). 6 at 8 tokens measured 0.4-0.5 us faster than the +# default 7 (same MMAs, bit-identical outputs). +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 6) +# ``hb_repl``: copies of the tagged FC1 -> FC2 records (applied when > 1); CTA bx polls copy bx % hb_repl, which +# spreads the polls of the few hot record lines over L2. 7 and 8 tokens hand off through the B images instead. +HB_REPL_BY_M: Tuple[int, ...] = (32, 32, 8, 8, 1, 1, 1, 1) +# ``lg_repl``: replicas of the router logit exchange at up to LG_REPL_MAX_TOKENS tokens (fewer stores before the +# exchange; bit-identical outputs). Above that the kernel's default applies. +LG_REPL = 32 +LG_REPL_MAX_TOKENS = 4 # capacity: fixed: measured crossover +# Byte-balanced FC1 ranges below one unit per CTA (:func:`_fine_ranges`) for groups too small to give every CTA a +# whole unit; False gives such groups whole units on the CTAs with the least FC2 work. +FINE_PLAN = True +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 + = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row + R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then + down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per + layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + P = I_ // n_cta + split = 128 // (2 * P) + gt = H_ // 128 // split + dkt = I_ // 128 + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + octant = r // P + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(2 * dkt): + jt, kt = gd // dkt, gd % dkt + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, gt + gd, :16384] = img.reshape(-1) + out[s, gt + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def _fine_ranges(weights, total: int): + """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item + counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary + strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit + runs at least to that unit's end. Returns the item counts.""" + n = len(weights) + items = [0] * n + s = 0 + acc = 0.0 + for c in range(n): + acc += weights[c] + e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) + if s % _S1 and e < total: # started inside a unit: run to (at least) its end + e = max(e, min(total, (s // _S1 + 1) * _S1)) + items[c] = e - s + s = e + assert s == total + return items + + +def make_phase_plan( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = FINE_PLAN, +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [max(_S1, (limit - f) // _ITEM_BYTES) for f in fc2] + + lo, hi = 0, max(fc2) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 + lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] + sw = sum(wts) + items = _fine_ranges([w * total / sw for w in wts], total) + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + inter: int = INTER, + early_release: bool = True, + ): + """``inter``: the routed intermediate per rank; this engine is built for INTER (384) only. + ``early_release``: on a call that pushes into the boundary all-reduce (``arp``), every CTA releases the PDL + dependents right after its grid dependency wait (the kernel's EARLY_AR); False releases them after the push.""" + if inter != INTER: + raise ValueError( + f"m3_moe (TP8): routed intermediate {inter} per rank is not supported (only {INTER})" + ) + self.inter = inter + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, num_tokens, M3_MOE_LAG, 0.0 if num_tokens == 1 else 1.0, True + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "i_tp": INTER, + } + if HB_REPL_BY_M[num_tokens - 1] > 1: + self.config["hb_repl"] = HB_REPL_BY_M[num_tokens - 1] + if num_tokens <= LG_REPL_MAX_TOKENS: + self.config["lg_repl"] = LG_REPL + if shared: + self.config["sh_i"] = SHARED_INTER + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if not early_release: + self.config["early_ar"] = 0 + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + # E0@m4: at 4 tokens the early group-0 FC1 stages are issued by FC2 B builder warp 8 instead of the routing + # warp 4 (issue order only, bitwise; TP4 -0.19 / TP8 -0.21 us/call). + if num_tokens == 4: + from tensorrt_llm.logger import logger as _e0log + + _e0log.info( + "MiniMax-M3 MoE engine E0M4 (%s): early group-0 stages on warp 8, tokens=4" % "TP8" + ) + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full( + (mod.G_CAP * mod.U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw + ) + # the FC1 -> FC2 records; with HB_REPL > 1 copies, HB_STRIDE int32 words apart + hb_words = mod.G_CAP * num_tokens * mod.U1 * 4 + if mod.HB_REPL > 1: + hb_words = mod.HB_REPL * mod.HB_STRIDE + self.hbuf32 = torch.zeros(hb_words, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [768, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 384], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + key = (_ENGINE_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py new file mode 100644 index 000000000000..48a7827f59b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Hand-off of an MXFP8-quantized activation from the kernel that produced it to the next MXFP8 GEMM. + +A producer that already holds a BF16 row in registers (the multicast all-reduce + RMSNorm, the split-K attention +merge) also writes its MXFP8 form (E4M3 rows + UE8M0 scales in the 128x4-swizzled layout of +``trtllm::mxfp8_quantize(x, True)``) and publishes it here keyed by the BF16 tensor. The MXFP8 linear whose input +is that tensor takes it and skips its own quantize kernel. Every take clears the slot, so a hand-off is only ever +consumed by the first MXFP8 linear that runs after the producer. + +Module state shared across files (one instance each per process): + +* ``SLOT`` (PrequantSlot): published by ``m3_mc_allreduce`` (``quant=True``), the sparse-decode merge and the fused + decode attention (MXFP8 output); cleared by the sparse decode before it runs; taken by the MXFP8 linear dispatch + and by the one-launch QKV + attention + o_proj path (which also peeks ``pending``). +* ``AR_PUSH`` (ArPushSlot): armed by the fast forward around the attention call; the o_proj GEMV takes the target + and marks ``pushed``; the fast forward checks ``take_pushed`` before the attention boundary all-reduce. +* ``MOE_PUSH`` (MoePushSlot): armed by the fast forward before the MoE; the routed-expert engine takes ``engine`` + and the shared expert's down projection takes ``shared``, each marks its pushed output; the fast forward + checks ``take_pushed`` before the MoE boundary all-reduce. +""" + + +class PrequantSlot: + def __init__(self): + self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) + + def publish(self, x, xq, xsf): + self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) + + def clear(self): + self.pending = None + + def take(self, x): + """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" + pending, self.pending = self.pending, None + if ( + pending is None + or x.dim() != 2 + or x.data_ptr() != pending[0] + or x.shape[0] != pending[1] + or x.shape[1] != pending[2] + ): + return None + return pending[3], pending[4] + + +SLOT = PrequantSlot() + + +class ArPushSlot: + """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. + + The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output + it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" + + def __init__(self): + self.target = None + self.pushed = None # data_ptr of the pushed output + + def arm(self, target): + self.target = target + self.pushed = None + + def take_target(self): + target, self.target = self.target, None + return target + + def take_pushed(self, x): + pushed, self.pushed = self.pushed, None + return pushed is not None and x.data_ptr() == pushed + + +AR_PUSH = ArPushSlot() + + +class MoePushSlot: + """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, + ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv + ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" + + def __init__(self): + self.engine = None + self.shared = None + self.engine_pushed = None # data_ptr of the pushed routed output + self.shared_pushed = None # data_ptr of the pushed shared output + + def arm(self, engine, shared): + self.engine, self.shared = engine, shared + self.engine_pushed = self.shared_pushed = None + + def disarm(self): + self.engine = self.shared = None + + def take_engine(self): + engine, self.engine = self.engine, None + return engine + + def take_shared(self): + shared, self.shared = self.shared, None + return shared + + def take_pushed(self, shared, routed): + """(routed pushed, shared pushed) for these outputs; clears the marks.""" + e = routed is not None and self.engine_pushed == routed.data_ptr() + s = shared is not None and self.shared_pushed == shared.data_ptr() + self.engine_pushed = self.shared_pushed = None + return e, s + + +MOE_PUSH = MoePushSlot() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_report.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_report.py new file mode 100644 index 000000000000..7f864bd240c6 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_report.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fast-path reports a kernel module makes from inside a call: ``taken`` (an engagement line), ``missed`` (a covered +shape that cannot run: raises under ``modeling_v2='require'``), ``uncovered`` and ``capture_missed``. + +The policy (logging once, raising under require) is the target's: ``modeling.py`` installs its functions here at import, +so a kernel reports without importing the target module. Until then every report is a no-op. +""" + +from __future__ import annotations + +from typing import Callable, Optional + + +def _noop(*args, **kwargs) -> None: + return None + + +taken: Callable[..., None] = _noop +missed: Callable[..., None] = _noop +uncovered: Callable[..., None] = _noop +capture_missed: Callable[..., None] = _noop + + +def install( + *, + on_taken: Callable[[str, Optional[int], str], None], + on_missed: Callable[[str, Optional[int], str], None], + on_uncovered: Callable[[str, str], None], + on_capture_missed: Callable[[str, Optional[int], str], None], +) -> None: + """Route this module's reports to the target's policy.""" + global taken, missed, uncovered, capture_missed + taken, missed, uncovered, capture_missed = on_taken, on_missed, on_uncovered, on_capture_missed diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py new file mode 100644 index 000000000000..30855eb7f05d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Low-M fp32 router projection: ``logits[M, E] = float(x[M, K]) @ W[E, K]^T`` in one kernel. + +Replaces the bf16->fp32 cast, the TF32 GEMM and its split-K reduce on decode +steps. One CTA per ``ROWS`` experts: the fp32 weight rows are bulk-copied into +shared memory before ``griddepcontrol.wait`` (overlapping the producer kernel +under PDL), the bf16 activations after it. Products and sums are exact fp32 FMA. + +The kernels are ``torch.ops.trtllm.minimax_m3_router_gemv`` / ``_router_gemv_fp4`` / ``_router_tc``, reached +through the ``ops/moe`` entries. +""" + +from typing import Optional + +import torch + +from ..ops.moe.m3_router_gemv import m3_router_gemv as _router_gemv_op +from ..ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 as _router_gemv_fp4_op +from ..ops.moe.m3_router_tc import m3_router_tc as _router_tc_op + +# Opt-in dynamic shared memory left after the kernel's static reduction buffer (SM100: 227 KiB). +_MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit + + +def m3_router_gemv( + x: torch.Tensor, weight: torch.Tensor, rows: int = 1, use_pdl: bool = False +) -> Optional[torch.Tensor]: + """fp32 ``x @ weight.T`` for bf16 ``x`` with at most 16 rows; None when unsupported (only ``rows=1`` is built).""" + if ( + rows != 1 + or not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or weight.shape[0] % rows != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or rows * x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return _router_gemv_op(x, weight, rows, use_pdl) + + +def m3_router_gemv_fp4( + x: torch.Tensor, weight: torch.Tensor, sf_scale: torch.Tensor, use_pdl: bool = False +) -> Optional[tuple]: + """Router logits plus the NVFP4 activation the routed experts consume (bit-identical to + ``fp4_quantize(x, sf_scale, 16, False, False)``): (logits [M, E] fp32, x_fp4 [M, K / 2] uint8, + x_sf [M * K / 16] uint8 linear). None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 16 + or weight.dtype != torch.float32 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or x.data_ptr() % 16 != 0 + or (x.stride(0) * 2) % 16 != 0 + or sf_scale is None + or sf_scale.dtype != torch.float32 + or sf_scale.numel() != 1 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or x.shape[1] * 4 + x.shape[0] * x.shape[1] * 2 > _MAX_DYNAMIC_SMEM + ): + return None + return tuple(_router_gemv_fp4_op(x, weight, sf_scale.reshape(1), use_pdl)) + + +def m3_router_tc( + x: torch.Tensor, + weight: torch.Tensor, + sf_scale: Optional[torch.Tensor] = None, + use_pdl: bool = False, +) -> Optional[tuple]: + """Router logits (TF32 tensor cores, like the native TF32 GEMM) for 1..48 tokens with 128 experts, plus the + NVFP4 activation when ``sf_scale`` is given. Returns (logits,) or (logits, x_fp4, x_sf); None when unsupported.""" + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.dim() != 2 + or x.stride(1) != 1 + or not 1 <= x.shape[0] <= 48 + or weight.dtype != torch.float32 + or weight.shape[0] != 128 + or not weight.is_contiguous() + or weight.shape[1] != x.shape[1] + or x.shape[1] % 1024 != 0 + or (x.stride(0) * 2) % 16 != 0 + or x.data_ptr() % 16 != 0 + or torch.cuda.get_device_capability(x.device)[0] < 10 + or 16 * (x.shape[1] // 8 + 4) * 4 + x.shape[0] * (x.shape[1] // 8 + 8) * 2 + > _MAX_DYNAMIC_SMEM + or (sf_scale is not None and (sf_scale.dtype != torch.float32 or sf_scale.numel() != 1)) + ): + return None + res = _router_tc_op(x, weight, None if sf_scale is None else sf_scale.reshape(1), use_pdl) + return tuple(res) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_sampled_accept.py new file mode 100644 index 000000000000..c3c2edd5d078 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_sampled_accept.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). + +Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +active must stay on the generic path; the caller decides. + +Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +(the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens heavier +than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the nucleus +with probability >= top_p, so most rows decide on their second exchange. + +Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based Philox +keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one 4-word +record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and reads +the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the same +order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing and +two call halves keep a rank one call ahead from overwriting a slot still being read. +""" + +from typing import Optional + +import torch + +from . import m3_report + +# The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/sampledAccept.cu) and called as +# torch.ops.trtllm.minimax_m3_sampled_accept. +MAX_ROWS = 48 # capacity: fixed: the kernel's built rows (decode steps) + + +def _op_args() -> frozenset: + """The optional arguments the loaded trtllm::minimax_m3_sm107tp8_sampled_accept takes (the tp2 op is used as it + is, read-only: an older image's op has neither prewait_params nor early_trigger).""" + global _OP_ARGS + if _OP_ARGS is None: + schema = torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept.default._schema + _OP_ARGS = frozenset(a.name for a in schema.arguments) + return _OP_ARGS + + +_OP_ARGS = None + + +def one_round_k() -> int: + """The SA-1R variant (ONE_ROUND_K) when the loaded op takes it, else 0 (two rounds).""" + return int(ONE_ROUND_K) if "one_round_k" in _op_args() else 0 + + +# The vocab shard limit of trtllm::minimax_m3_sm107tp8_sampled_accept (kMaxVs, kept from the TP2 kernel it copies): +# 200064 / 2 = 100032 columns. +MAX_VS = 100352 # capacity: fixed: the vocabulary shard bound + + +class SampledAcceptWorkspace: + """Exchange slots behind one multicast mapping (two call halves), the (finished rows, call) counters, the per-row + token scratch, an extra-attempt counter and the SA-1R round-0 decided counter. Created collectively, outside + CUDA-graph capture.""" + + def __init__(self, uc, mc_ptr, rank, world, handle=None, comm=None): + self.uc, self.mc_ptr, self.rank, self.world = uc, mc_ptr, rank, world + self.handle, self.comm = handle, comm + dev = uc.device + self.counters = torch.zeros(2, dtype=torch.int32, device=dev) + self.tokens = torch.zeros(MAX_ROWS, dtype=torch.int32, device=dev) + self.extra_attempts = torch.zeros(1, dtype=torch.int32, device=dev) + self.decided = torch.zeros(1, dtype=torch.int32, device=dev) + + @classmethod + def create(cls, mapping, device) -> "SampledAcceptWorkspace": + # SA-1R's records carry the round-0 list words: the size depends on the variant + words = int( + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words( + max(mapping.tp_size, 1), one_round_k() + ) + if one_round_k() + else torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words( + max(mapping.tp_size, 1) + ) + ) + if mapping.tp_size == 1: + uc = torch.zeros(words, dtype=torch.int64, device=device) + return cls(uc, uc.data_ptr(), 0, 1) + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "SampledAcceptWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 8, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("SampledAcceptWorkspace.create failed on another rank") + return cls(uc, mc.data_ptr(), mapping.tp_rank, mapping.tp_size, handle, comm) + + +_WS: Optional[SampledAcceptWorkspace] = None + + +def workspace(mapping, device) -> Optional[SampledAcceptWorkspace]: + global _WS + if _WS is None: + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("spec_tail", None, "the sampled-accept workspace") + return None + _WS = SampledAcceptWorkspace.create(mapping, device) + return _WS + + +def sampled_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + min_ps: Optional[torch.Tensor], + seeds: torch.Tensor, + offsets: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, + ws: Optional[SampledAcceptWorkspace] = None, + use_pdl: bool = True, +): + """Sampled verification (temperature + top-p) from this rank's bf16 logits shard [nreq * (draft_len + 1), V / tp]: + outputs as shard_accept (accepted [nreq, D+1], num_accepted [nreq], gather_ids [nreq]); None + when unavailable. Per-row sampling parameters and Philox (seed, offset) are the spec metadata's, one per row. + top_ks / min_ps are not applied: rows with top-k or min-p active must be routed to the generic path by the + caller.""" + del top_ks, min_ps + nreq, draft_len = draft_tokens.shape + rows = nreq * (draft_len + 1) + if logits.dim() != 2 or logits.shape[0] != rows or rows > MAX_ROWS or logits.shape[1] > MAX_VS: + return None + if logits.dtype != torch.bfloat16 or logits.stride(1) != 1: + return None + ws = ws if ws is not None else workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept( + logits, + int(vocab_offset), + temperatures[:rows].float().contiguous(), + top_ps[:rows].float().contiguous(), + seeds[:rows].long().contiguous(), + offsets[:rows].long().contiguous(), + ws.uc, + ws.mc_ptr, + ws.counters, + ws.tokens, + ws.rank, + ws.world, + draft_tokens.int().contiguous(), + batch_idx.int().contiguous(), + accepted, + num_accepted, + gather_ids, + int(force_total), + next_new, + force_vec, + ws.extra_attempts, + bool(use_pdl), + **{ + name: bool(on) + for name, on in ( + ("prewait_params", PREWAIT_PARAMS), + ("early_trigger", EARLY_TRIGGER), + ) + if name in _op_args() + }, + **({"one_round_k": one_round_k(), "decided": ws.decided} if one_round_k() else {}), + **({"ilp_walk": True} if _ilp_walk(rows) else {}), + ) + if one_round_k(): + m3_report.taken("m3_sampled_accept_1r", rows, f"one_round_k {one_round_k()}") + return accepted, num_accepted, gather_ids + + +# Off at TP8: the sampled verify loads each row's sampling parameters after its PDL dependency wait, and releases its +# dependents at the acceptance epilogue (both are timing only; the same values and tokens either way). +PREWAIT_PARAMS = False +EARLY_TRIGGER = False +# SA-1R: each CTA's round-0 record also carries its largest logits (82: each thread's top 2 pooled, the CTA's top 8 +# plus the bound of everything it does not list), so a row whose first candidate is provably in the nucleus decides in +# round 0. The same tokens, accepted counts and Philox draws as the two-round kernel, bit for bit; 0 = two rounds. +ONE_ROUND_K = 82 # capacity: fixed: a record format, not a size +# SA-PW: the owner walk takes its tokens 8 at a time (the same draw, tokens and Philox draws as the one-token walk, +# bit for bit). Off. +ILP_WALK = False +ILP_WALK_PATH = "m3_sampled_accept_ilp_walk" + + +def _ilp_walk(rows: int) -> bool: + """ILP_WALK when the loaded op takes ilp_walk, logged once per row count as taken.""" + if not ILP_WALK or "ilp_walk" not in _op_args(): + return False + m3_report.taken( + ILP_WALK_PATH, + rows, + "the sampled verify's owner walk takes 8 tokens at a time", + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py new file mode 100644 index 000000000000..10bd966adf98 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py @@ -0,0 +1,367 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). + +Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long``, +one 1024-thread CTA per row), from a CTA of 384 threads: + +* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and + mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's + valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned + key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped + block and masks the result, so a thread's loads issue back to back. +* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of + the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. +* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by + that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate + against the whole row instead. +* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. + +The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait +(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are +compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and +loads them before the executor's warmup, so a decode step never compiles one. +""" + +from typing import Dict, Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.utils import ( + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +# The 384-thread select reads its row's valid-block count before its grid wait instead of after it, which takes one +# dependent global load off the post-wait path; the selection is the same either way. Exact: the count +# (MsaMetadata.msa_n_valid_blocks) is written only before the step's first layer, behind a plain launch. +SELECT_N_VALID_PREWAIT = True + +SELECT_THREADS = 384 +SELECT_WARPS = SELECT_THREADS // 32 +# 1M tokens + draft tokens at 128 tokens per block. +MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row +MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging +TOPK = 16 # capacity: fixed: MSA's top-k +NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) +# The key of -inf (the reference's 0xff800000): never emitted. +NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 +# (keys per thread, rows up to this many valid blocks) +ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) + +# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare +# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group +# minima [16], per-warp candidate counts [SELECT_WARPS]. +_CAND = 0 +_PICK_KEY = _CAND + 2 * MAX_CAND +_PICK_BLOCK = _PICK_KEY + TOPK +_WARP_MIN = _PICK_BLOCK + TOPK +_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 +_WARP_CNT = _GROUP_MIN + 16 +SMEM_WORDS = _WARP_CNT + SELECT_WARPS +SMEM_BYTES = 4 * SMEM_WORDS + + +def _key(score): + """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" + bits = cutlass.Float32(score).bitcast(cutlass.Int32) + return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) + + +def _imin(a, b): + return cutlass.select_(b < a, b, a) + + +def _imax(a, b): + return cutlass.select_(b > a, b, a) + + +def _forced_key(scores, head, block, query, init_blocks, local_start): + raw = scores[head, block, query] + s = cutlass.select_( + block >= local_start, + cutlass.Float32(LOCAL_SCORE), + cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), + ) + return _key(s) + + +@cute.jit +def _select_row( + row, + scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides + n_valid: cute.Tensor, # int32 [>= total_q] + out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] + init_blocks, + local_blocks, + num_kv_heads, + smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory + bar_id: cutlass.Constexpr, + loaded_trigger: cutlass.Constexpr, + raw_prewait, + n_valid_prewait: cutlass.Constexpr, +): + """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier + ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the + row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by + the caller before its grid wait; otherwise it is read here.""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % cutlass.Int32(32) + warp = tid // cutlass.Int32(32) + query = row // num_kv_heads + head = row % num_kv_heads + if cutlass.const_expr(n_valid_prewait): + raw = raw_prewait + else: + raw = n_valid[query] + valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) + local_start = _imax(raw - local_blocks, cutlass.Int32(0)) + cand = cute.make_tensor( + cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) + ) + + if tid < TOPK: + smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) + smem[_PICK_BLOCK + tid] = cutlass.Int32(0) + + if valid <= TOPK: + # Every valid block is picked (thread t holds block t). + if tid < valid: + smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) + smem[_PICK_BLOCK + tid] = tid + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): + items = ITEM_CLASSES[ci][0] + lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] + hi = ITEM_CLASSES[ci][1] + if (valid > lo) & (valid <= hi): + key = cute.make_rmem_tensor(items, cutlass.Int32) + local_min = cutlass.Int32(NO_KEY) + for s in cutlass.range_constexpr(items): + # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. + block = tid + cutlass.Int32(s * SELECT_THREADS) + k = _forced_key( + scores, head, _imin(block, valid - 1), query, init_blocks, local_start + ) + key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) + local_min = _imin(local_min, key[s]) + local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) + if lane < 16: + smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + if tid < 16: + g = cutlass.Int32(NO_KEY) + for w in cutlass.range_constexpr(SELECT_WARPS): + g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) + smem[_GROUP_MIN + tid] = g + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + bound = cutlass.Int32( + -2147483648 + ) # the smallest key: signed keys, unlike the reference's unsigned ones + for g in cutlass.range_constexpr(16): + bound = _imax(bound, smem[_GROUP_MIN + g]) + + # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(items): + cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) + incl = cnt + for d in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync( + incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) + ) + incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) + if lane == 31: + smem[_WARP_CNT + warp] = incl + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + before = cutlass.Int32(0) + total = cutlass.Int32(0) + for w in cutlass.range_constexpr(SELECT_WARPS): + c = smem[_WARP_CNT + w] + before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) + total += c + slot = before + incl - cnt + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + if slot < MAX_CAND: + cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( + tid + cutlass.Int32(s * SELECT_THREADS) + ) + slot += 1 + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if total <= MAX_CAND: + for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): + mine = cand[i] + rank = cutlass.Int32(0) + for j in cutlass.range(total, unroll=8): + rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) + smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) + else: + # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + block = tid + cutlass.Int32(s * SELECT_THREADS) + rank = cutlass.Int32(0) + for j in cutlass.range(valid, unroll=1): + kj = _forced_key(scores, head, j, query, init_blocks, local_start) + less = (kj < key[s]) | ((kj == key[s]) & (j < block)) + rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = key[s] + smem[_PICK_BLOCK + rank] = block + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if warp == 0: + # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. + v = cutlass.Int32(NO_KEY) + if lane < TOPK: + pk = smem[_PICK_KEY + lane] + if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): + v = smem[_PICK_BLOCK + lane] + for size_log in cutlass.range_constexpr(1, 6): + for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): + other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) + ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) + lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) + v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) + if lane < TOPK: + out[query, head, lane] = cutlass.select_( + v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v + ) + + +class _SelectLongDslKernel: + """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count + before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because + ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's + first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it + (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" + + def __init__( + self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool + ): + self.use_pdl = use_pdl + self.early_trigger = early_trigger + self.loaded_trigger = (not early_trigger) and trigger_phase == 1 + self.n_valid_prewait = n_valid_prewait + + @cute.kernel + def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): + row, _, _ = cute.arch.block_idx() + smem = SmemAllocator() + words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) + raw_prewait = cutlass.Int32(0) + if cutlass.const_expr(self.n_valid_prewait): + raw_prewait = n_valid[row // num_kv_heads] + griddepcontrol_wait() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + _select_row( + row, + scores, + n_valid, + out, + init_blocks, + local_blocks, + num_kv_heads, + words, + bar_id=1, + loaded_trigger=self.loaded_trigger, + raw_prewait=raw_prewait, + n_valid_prewait=self.n_valid_prewait, + ) + + @cute.jit + def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): + num_kv_heads = scores.shape[0] + self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( + grid=(scores.shape[2] * num_kv_heads, 1, 1), + block=(SELECT_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) + + +#: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). +kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} + + +def kernel_key( + use_pdl: bool, early_trigger: bool, trigger_phase: int +) -> Tuple[bool, bool, int, bool]: + """The compile key of a launch (the trigger phase does not matter with the early trigger), with this target's + ``SELECT_N_VALID_PREWAIT``.""" + return ( + bool(use_pdl), + bool(early_trigger), + 0 if early_trigger else int(trigger_phase), + bool(SELECT_N_VALID_PREWAIT), + ) + + +def _compile(key: Tuple[bool, bool, int, bool]): + if key not in kernel_cache: + heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() + scores = make_fake_tensor( + cutlass.Float32, + (heads, blocks, rows), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) + out = make_fake_tensor( + cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) + ) + kernel_cache[key] = cute.compile( + _SelectLongDslKernel(*key), + scores, + n_valid, + out, + cutlass.Int32(0), + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + return kernel_cache[key] + + +def select_long_dsl( + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool, + early_trigger: bool, + trigger_phase: int, +) -> None: + """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" + _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( + scores, n_valid_blocks, out, init_blocks, local_blocks + ) + + +def prebuild(keys, device) -> list: + """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" + scores = torch.zeros( + 1, 2048, 4, dtype=torch.float32, device=device + ) # capacity: fixed: a dummy prebuild row, not a size + n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) + out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) + built = [] + for key in keys: + key = kernel_key(*key) + _compile(key)(scores, n_valid, out, 0, 1) + built.append(key) + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py new file mode 100644 index 000000000000..8caaee389d0a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fused greedy Eagle3 verification for the MiniMax-M3 target head (vocab-sharded TP head, verify rows <= 48). + +Replaces the all-gather of the target logits, the argmax and the strict-acceptance ops of a greedy verify step with +one kernel over this rank's BF16 logits shard (the target lm_head's own local GEMM output): + +* Every logit becomes a 64-bit key ``ordered(value) << 32 | ~global_index`` (ties pick the lowest index, as + ``torch.argmax``), max-reduced per row: warp shuffles, then ``atomicMax``. +* The last CTA to finish pushes the rank's per-row key into every rank's exchange slot through a multicast mapping + (Lamport: 0 = empty, the key is never 0), polls its own slots for all ranks, takes the global argmax, re-arms the + slots and runs the strict acceptance. Call n uses slot half n % 2. + +The argmax is taken over the BF16 logits without arithmetic, so the tokens are those of the in-tree path's argmax +over the gathered logits. The kernel is compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/ +shardAccept.cu) and called as ``torch.ops.trtllm.minimax_m3_shard_accept``. +""" + +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from . import m3_report + +#: Verify rows per call, and exchange rows per call half of the workspace. +SHARD_ROWS = 48 # capacity: fixed: the kernel's built verify rows + + +@dataclass(eq=False) +class AcceptWorkspace: + """Per-rank exchange slots ([2 halves][SHARD_ROWS rows][world] uint64) behind one multicast mapping, the local + key accumulators and the (finished CTAs, call) counters.""" + + uc: torch.Tensor + mc_ptr: int + keys: torch.Tensor + counters: torch.Tensor + rank: int + world_size: int + handle: Any = None + comm: Any = None + + +_WS: Optional[AcceptWorkspace] = None +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + if not _mnnvl_workspace_all_succeeded(comm, True): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = mc.data_ptr(), mapping.tp_rank + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None + return AcceptWorkspace( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) + + +def _workspace(mapping, device) -> Optional[AcceptWorkspace]: + """The exchange workspace, created collectively outside CUDA-graph capture on first use.""" + global _WS + if _WS is not None: + return _WS + if torch.cuda.is_current_stream_capturing(): + m3_report.capture_missed("m3_shard_accept", None, "the shard-accept workspace") + return None + _WS = create_workspace(mapping, device, SHARD_ROWS) + return _WS + + +def shard_accept( + logits: torch.Tensor, + vocab_offset: int, + mapping, + draft_tokens: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor] = None, + force_vec: Optional[torch.Tensor] = None, +): + """Greedy verification from this rank's BF16 logits shard. + + Args: + logits: [nreq * (draft_len + 1) <= SHARD_ROWS, V / tp] BF16, unit column stride: this rank's columns + (global ids ``vocab_offset ..``) of the verify rows, request-major. + draft_tokens: [nreq, draft_len] int32, the drafts being verified. + batch_idx: [nreq] int32, each request's row in the verify layout. + force_total / force_vec: forced acceptance (see ``m3_shard_accept``). + next_new: optional int32 [nreq, >= 1]; column 0 receives the last accepted token. + + Returns: + (accepted_tokens [nreq, draft_len + 1] int32, num_accepted [nreq] int32, first-pass gather ids [nreq] + int64), or None when the shape is not served or the workspace is missing during a capture. Collective over + the TP group. + """ + nreq, draft_len = draft_tokens.shape + if ( + logits.dim() != 2 + or logits.shape[0] != nreq * (draft_len + 1) + or logits.shape[0] > SHARD_ROWS + or logits.dtype != torch.bfloat16 + or logits.stride(1) != 1 + ): + return None + ws = _workspace(mapping, logits.device) + if ws is None: + return None + dev = logits.device + accepted = torch.empty((nreq, draft_len + 1), dtype=torch.int32, device=dev) + num_accepted = torch.empty(nreq, dtype=torch.int32, device=dev) + gather_ids = torch.empty(nreq, dtype=torch.int64, device=dev) + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + ws.keys, + ws.counters, + ws.uc, + ws.mc_ptr, + ws.rank, + ws.world_size, + draft_tokens.contiguous(), + batch_idx, + accepted, + num_accepted, + gather_ids, + force_total, + next_new, + _num_sms(), + force_vec, + ) + return accepted, num_accepted, gather_ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_spec_tail.py new file mode 100644 index 000000000000..996371370fde --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_spec_tail.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Small fused kernels of the MiniMax-M3 Eagle3 step (GB300, decode). + +``chunk_norm`` / ``capture_norm`` apply the drafter's per-layer fc_norm RMSNorms to the captured target hidden states +(the in-tree ``apply_eagle3_fc`` norms, same fp32 arithmetic); ``draft_input`` is a fused draft pass's input (the +embedding lookup, the drafter layer's input_layernorm / hidden_norm and their concat, same arithmetic); +``step_gather`` / ``slot_scatter`` move the step's integer outputs between the spec sampler's store and the engine's +input buffers. +""" + +from typing import Optional, Tuple + +import torch + +# The kernels are compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu) and called as +# torch.ops.trtllm.minimax_m3_{draft_input, chunk_norm, capture_norm, step_gather, slot_scatter}. + + +def draft_input( + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + eps_embed: float, + w_hidden: torch.Tensor, + eps_hidden: float, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(cat([rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden], -1) [M, 2H], hidden[rows] [M, H]) for + int32 ids [M] and int64 rows [M] (rows None: hidden[:M]); ids outside the table embed to zero.""" + m, h = ids.shape[0], table.shape[1] + out = torch.empty((m, 2 * h), dtype=table.dtype, device=table.device) + residual = torch.empty((m, h), dtype=table.dtype, device=table.device) + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, use_pdl + ) + return out, residual + + +def chunk_norm( + src: torch.Tensor, weights: torch.Tensor, eps: float, use_pdl: bool = True +) -> torch.Tensor: + """Per-chunk RMSNorm of src [M, >= C*H] (row stride may exceed C*H) with weights [C, H] -> [M, C*H].""" + c, h = weights.shape + out = torch.empty((src.shape[0], c * h), dtype=src.dtype, device=src.device) + torch.ops.trtllm.minimax_m3_chunk_norm(src, weights, eps, out, use_pdl) + return out + + +# The overlap scheduler's per-step gathers and the spec sampler's slot stores as one kernel each (step_gather / +# slot_scatter) instead of ~10 and 4 small torch ops; called by the target's SpecStepFold. + + +def step_gather( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 +) -> bool: + """See step_gather_kernel; False when the inputs are outside what the kernel covers (caller falls back).""" + ts = (new_tok, next_draft, lens, slots, ppi, ids, draft, pos_off, kv_off) + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or next_draft.dim() != 2 + or new_tok.dim() < 2 + ): + return False + if (new_tok.dim() == 3 and new_tok.shape[2] != 1) or not ( + lens.is_contiguous() + and ids.is_contiguous() + and draft.is_contiguous() + and pos_off.is_contiguous() + and kv_off.is_contiguous() + ): + return False + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, next_draft, lens, slots, ppi, P, T, D, ids, ids0, draft, d0, pos_off, kv_off, b0 + ) + return True + + +def slot_scatter( + slots, src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft +) -> bool: + """See slot_scatter_kernel (2-D int32 views, zero padding); False to fall back.""" + ts = (src_new, dst_new, src_next_new, dst_next_new, src_lens, dst_lens, src_draft, dst_draft) + if any(t.dtype != torch.int32 or not t.is_cuda for t in ts) or slots.dtype != torch.long: + return False + two = (src_new, dst_new, src_next_new, dst_next_new, src_draft, dst_draft) + if any(t.dim() != 2 for t in two) or src_lens.dim() != 1 or dst_lens.dim() != 1: + return False + torch.ops.trtllm.minimax_m3_slot_scatter( + slots, + src_new, + dst_new, + src_next_new, + dst_next_new, + src_lens, + dst_lens, + src_draft, + dst_draft, + ) + return True + + +def capture_norm( + h: torch.Tensor, + residual: Optional[torch.Tensor], + weight: torch.Tensor, + eps: float, + dst: torch.Tensor, +) -> None: + """dst [M, H] (row stride may exceed H) = rmsnorm(h + residual) * weight with chunk_norm's rounding.""" + torch.ops.trtllm.minimax_m3_capture_norm(h, residual, weight, eps, dst) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py new file mode 100644 index 000000000000..f5a1b229c777 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The target lm_head shard of a verify step of at most 8 rows, streamed from a lossless 12-bit coding of its BF16 +weight (75% of the bytes; ``ops.gemm.m3_target_head`` has the format). + +``prepare_head`` codes this rank's shard once, before the executor's warmup, and keeps it only when it decodes to the +BF16 weight bit for bit (``round_trip_ok``); a shard that does not keeps the in-tree BF16 GEMM. ``logits`` runs the +kernel: BF16 x BF16 MMAs with fp32 accumulation over the full K and one BF16 rounding per logit, the in-tree head's +arithmetic in another summation order. The kernel is compiled into TensorRT-LLM +(cpp/tensorrt_llm/kernels/minimaxM3/targetHead.cu) and called through the catalog entry ``ops.gemm.m3_target_head``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +import torch + +from ..ops.gemm.m3_target_head import CTA_WARPS, MAX_ROWS, encode_chunked, round_trip_ok +from ..ops.gemm.m3_target_head import m3_target_head as _target_head_op + + +@dataclass(eq=False) +class TargetHead: + """One rank's target lm_head shard in the 12-bit coding, and the kernel's split-tile scratch. + + ``coded`` is False when the coding did not reproduce the weight bit for bit (the in-tree GEMM then runs).""" + + weight: torch.Tensor + coded: bool + w: Optional[torch.Tensor] = None + base: int = 0 + exc_off: Optional[torch.Tensor] = None + exc_pos: Optional[torch.Tensor] = None + exc_w: Optional[torch.Tensor] = None + part: Optional[torch.Tensor] = None + counter: Optional[torch.Tensor] = None + grid: int = 0 + + +_NUM_SMS: Optional[int] = None + + +def _num_sms() -> int: + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return _NUM_SMS + + +def head_supported(lm_head) -> bool: + """A bias-free 2-D BF16 head of whole 16-row tiles whose K the coding covers (a multiple of 64).""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + return ( + getattr(lm_head, "bias", None) is None + and weight.shape[0] % 16 == 0 + and weight.shape[1] % 64 == 0 + ) + + +def prepare_head(lm_head, grid: Optional[int] = None) -> TargetHead: + """Code ``lm_head``'s local shard and check it decodes to the BF16 weight bit for bit (fails closed: a shard that + does not is returned with ``coded=False`` and nothing kept). Built once, eagerly, before the executor's warmup + (``M3Eagle3Worker.prebuild_target_head``).""" + weight = lm_head.weight + with torch.no_grad(): + w = weight.detach() + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + if not round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return TargetHead(weight=weight, coded=False) + if grid is None: + grid = _num_sms() + n = weight.shape[0] + return TargetHead( + weight=weight, + coded=True, + w=packed, + base=base, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + part=torch.zeros(grid * CTA_WARPS * 2 * 4 * 32, dtype=torch.float32, device=weight.device), + counter=torch.zeros((n + 15) // 16, dtype=torch.int32, device=weight.device), + grid=grid, + ) + + +def logits(head: TargetHead, x: torch.Tensor, use_pdl: bool = True) -> torch.Tensor: + """This rank's BF16 logits shard [rows, n] of the rows x [rows <= MAX_ROWS, K] (BF16) from the coded shard.""" + assert head.coded and 1 <= x.shape[0] <= MAX_ROWS, (head.coded, x.shape) + if x.stride(1) != 1 or x.stride(0) % 8 or x.data_ptr() % 16: + x = x.contiguous() + n = head.weight.shape[0] + out = torch.empty((x.shape[0], n), dtype=torch.bfloat16, device=x.device) + return _target_head_op( + head.w, + x, + out, + head.part, + head.counter, + head.exc_off, + head.exc_pos, + head.exc_w, + head.base, + n, + head.grid, + use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py new file mode 100644 index 000000000000..db0077117fe3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The token map of ``thop.trtllm_gen_generation_preprocess`` for one request with several tokens. + +For multi-token generation the op's QKV preprocessing places every row by the workspace token map (``tokensInfo``, +(request, token) per row), but builds that map only when ``isBuildDecoderInfoKernelNeeded()`` holds, which at one +request with the default RoPE it does not. The op then applies whatever map the last multi-request call left in the +workspace: a row can land on another request's length and pages, and its RoPE position and K/V slot are wrong. This +module wraps the op so that, for one request and several tokens, the map the build would write, (0, t) for each row t, +is written into the op's own map slice of the workspace (the offset from the op's layout function) just before the +call. Every other call passes through unchanged. The rows' positions still come from the request's generation length, +exactly as with a freshly built map. + +The write is a device copy from a map prebuilt at install (no allocation, no host sync), so it is captured into the +decode graphs with the call it precedes. +""" + +from __future__ import annotations + +from typing import Callable, Dict, Optional, Tuple + +import torch + +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.logger import logger + +#: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). +_ARG_NAMES = ( + "qkv_input", + "workspace", + "sequence_lengths", + "spec_decoding_generation_lengths", + "spec_decoding_position_offsets", + "kv_cache_block_offsets", + "host_kv_cache_pool_pointers", + "host_kv_cache_pool_mapping", + "kv_scale_orig_quant", + "kv_scale_quant_orig", + "attention_output_orig_quant", + "rotary_inv_freq", + "rotary_cos_sin", + "mrope_position_deltas", + "layer_idx", + "seq_offset", + "num_heads", + "num_kv_heads", + "head_size", + "tokens_per_block", + "kv_cache_quant_mode", + "max_attention_window_size", + "cyclic_attention_window_size", + "num_tokens", + "batch_beam", + "input_seq_length", + "max_past_kv_length", + "rotary_embedding_dim", + "rotary_embedding_base", + "rotary_embedding_scale_type", + "rotary_embedding_scale", + "rotary_embedding_max_positions", + "position_embedding_type", + "bmm1_scale", + "bmm2_scale", + "fp8_context_fmha", + "predicted_tokens_per_seq", + "attention_chunk_size", + "multi_processor_count", + "total_num_blocks", + "kv_factor", + "need_build_kv_cache_metadata", + "cross_attention", + "skip_fmha_workspace", +) +_INT2_BYTES = 8 # capacity: fixed: an element size + +_original: Optional[Callable] = None +_max_rows = 0 +_maps: Dict[torch.device, torch.Tensor] = {} +_offsets: Dict[Tuple, int] = {} + + +def _no_taken(rows: int) -> None: + pass + + +def _no_missed(rows: int, reason: str) -> None: + pass + + +_on_taken: Callable[[int], None] = _no_taken +_on_missed: Callable[[int, str], None] = _no_missed + + +def _indexed(device: torch.device) -> torch.device: + """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: + """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" + device = _indexed(device) + m = _maps.get(device) + if m is None or m.shape[0] < rows: + reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" + _on_missed(rows, reason) + raise RuntimeError(f"m3_tokmap_guard: {reason}") + return m[:rows] + + +def _map_offset(args: dict, rows: int) -> int: + x = args["qkv_input"] + skip = bool(args.get("skip_fmha_workspace", False)) + key = ( + x.dtype, + rows, + int(args["num_heads"]), + int(args["head_size"]), + int(args["rotary_embedding_dim"]), + int(args["num_kv_heads"]), + skip, + ) + off = _offsets.get(key) + if off is None: + layout = thop.get_trtllm_gen_generation_workspace_layout( + x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip + ) + off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) + if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: + raise RuntimeError( + f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " + f"(offset {off}, {size} bytes): the layout changed under this guard" + ) + _offsets[key] = off + return off + + +def _guarded(*args, **kwargs): + a = dict(zip(_ARG_NAMES, args)) + a.update(kwargs) + rows = int(a["num_tokens"]) + if ( + a.get("spec_decoding_generation_lengths") is not None + and int(a["predicted_tokens_per_seq"]) > 1 + and int(a["batch_beam"]) == 1 + and rows > 1 + and not bool(a.get("cross_attention", False)) + ): + ws = a["workspace"] + src = _one_request_map(ws.device, rows) + off = _map_offset(a, rows) + # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. + dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] + dst.view(torch.int32).view(rows, 2).copy_(src) + _on_taken(rows) + return _original(*args, **kwargs) + + +def install( + device: torch.device, + max_rows: int, + on_taken: Callable[[int], None], + on_missed: Callable[[int, str], None] = _no_missed, +) -> bool: + """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to + ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this + call installed the wrapper.""" + global _original, _max_rows, _on_taken, _on_missed + device = _indexed(device) + if device not in _maps or _maps[device].shape[0] < max_rows: + t = torch.arange(max_rows, dtype=torch.int32, device=device) + _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() + _max_rows = max(_max_rows, max_rows) + _on_taken = on_taken + _on_missed = on_missed + if _original is not None: + return False + _original = thop.trtllm_gen_generation_preprocess + thop.trtllm_gen_generation_preprocess = _guarded + logger.info( + "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " + f"(up to {max_rows} rows)" + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/__init__.py new file mode 100644 index 000000000000..77d7b84d9bfe --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Mixture-of-experts entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py new file mode 100644 index 000000000000..113878a991f9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py @@ -0,0 +1,341 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3's MoE layer for 1..8 decode tokens in one CuTe DSL launch (``m3_moe``, the sm_107 TP2 build: routed +and shared expert width 1536 per rank): the router, the routed NVFP4 experts and the shared MXFP8 expert, on +caller-owned state (an ``M3MoeState`` and one ``M3MoeLayer`` per MoE layer).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +from ..m3_moe.op import ( + HIDDEN, + INTER, + SHARED_INTER, + TOP_K, + M3MoeLayer, + M3MoeState, + ar_push_desc, + is_supported, +) + +__all__ = [ + "M3MoeLayer", + "M3MoeState", + "ar_push_desc", + "is_supported", + "m3_moe_engine", + "moe_reference", + "moe_route", + "mxfp8_dequant", + "mxfp8_quant_dequant", + "unswizzle_sf", +] + +# The model's constants the kernels are built with: routed_scaling_factor, the shared expert's SwiGLU-OAI +# (swiglu_alpha, swiglu_limit, the (up + 1) offset). +ROUTED_SCALING = 2.0 +SHARED_ALPHA = 1.702 +SHARED_LIMIT = 7.0 + +_LOG2E = 1.4426950216293334961 +_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) +_F32_MAX = 3.40282346638528859812e38 + + +def _f32(x: float) -> torch.Tensor: + return torch.tensor(x, dtype=torch.float32) + + +def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: + """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" + lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) + lo = lut[(packed & 0xF).long()] + hi = lut[(packed >> 4).long()] + return torch.stack((lo, hi), dim=-1).flatten(-2) + + +def _e2m1_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" + a = x.abs() + q = torch.full_like(a, 6.0) + for hi, value, inclusive in ( + (5.0, 4.0, True), + (3.5, 3.0, False), + (2.5, 2.0, True), + (1.75, 1.5, False), + (1.25, 1.0, True), + (0.75, 0.5, False), + (0.25, 0.0, True), + ): + q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) + return torch.where(x < 0, -q, q) + + +def _e4m3_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" + return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + + +def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: + """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" + bits = x.contiguous().view(torch.int32) + code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + return code.clamp(0, 254) + + +def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: + """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" + flat = sf.reshape(-1).view(torch.uint8) + r = torch.arange(rows, device=flat.device).view(-1, 1) + b = torch.arange(blocks, device=flat.device).view(1, -1) + padded_blocks = (blocks + 3) // 4 * 4 + idx = ( + (r // 128) * (128 * padded_blocks) + + (b // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (b % 4) + ) + return flat[idx] + + +def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" + rows, k = q.shape + codes = unswizzle_sf(sf, rows, k // 32).float() + return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( + 32, dim=1 + ) + + +def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: + """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to + fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" + rows, k = x.shape + blocks = x.float().view(rows, k // 32, 32) + code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) + scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) + return (_e4m3_round(blocks / scale) * scale).view(rows, k) + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16).float() + + +def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" + return (a.double() @ b.double().t()).float() + + +def _row_perm(rows: int, gated: bool) -> torch.Tensor: + """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle + of epilogue tile 128).""" + probe = torch.empty(rows, 1) + perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) + if gated: + perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] + return perm + + +def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: + """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row + order (the global scale is not applied: it is the per-expert alpha).""" + rows, half = w.shape + codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() + vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) + return vals[perm_inv] + + +def moe_route(x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor): + """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" + logits = _gemm(x, router_weight) + sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() + keys = sig + bias + ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] + sk = keys.gather(1, ids) - bias[ids] + total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) + w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) + return ids, _bf16(w) + + +def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: + """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" + rows, k = x.shape + blocks = x.view(rows, k // 16, 16) + vmax = blocks.abs().amax(dim=-1, keepdim=True) + sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) + # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element + # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). + osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) + return (_e2m1_round(blocks * osc) * sff).view(rows, k) + + +def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: + """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 + input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with + act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" + a, c, alpha, beta, lim = (esc[i] for i in range(5)) + gc = torch.where(gate > lim, lim, gate) + gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) + uc = torch.minimum(torch.maximum(up, -lim), lim) + c1 = a * (alpha * _f32(_LOG2E).to(up.device)) + sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) + h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) + blocks = h.view(-1, 16) + amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) + sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) + osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) + return (_e2m1_round((blocks * osc) * c) * sff).view(-1) + + +def _shared_expert(x: torch.Tensor, shared_mxfp8) -> torch.Tensor: + """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the + model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" + gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 + gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) + gate, up = gu[:, :SHARED_INTER], gu[:, SHARED_INTER:] + a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) + ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 + sg = 1.0 / ( + torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 + ) + h = _bf16(ub * (sg * a)) + return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) + + +def moe_reference(hidden: torch.Tensor, layer: M3MoeLayer, shared: bool) -> torch.Tensor: + """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" + x = hidden.float() + m = x.shape[0] + ids, weights = moe_route(x, layer.router_weight, layer.bias) + xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) + inv31 = torch.argsort(_row_perm(2 * INTER, gated=True)).to(x.device) + inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) + y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) + for e in torch.unique(ids).tolist(): + w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) + w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) + esc = layer.escal[:, e] + for t, k in (ids == e).nonzero().tolist(): + acc1 = _gemm(xr[t : t + 1], w31)[0] + # The buffer holds [w3 (up); w1 (gate)]. + h = _fc1_epilogue(acc1[:INTER], acc1[INTER:], esc) + y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) + acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) + for k in range(TOP_K): + acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() + routed = acc.float() + if not shared: + return routed.to(torch.bfloat16) + return (_shared_expert(x, layer.shared_mxfp8) + _bf16(routed)).to(torch.bfloat16) + + +_ENGAGED: set = set() + + +class M3MoeEngine(OpWrapper): + """One MiniMax-M3 MoE layer for M (1..8) decode tokens: ``layer.state.run(layer, hidden, out, arp)``. + + ``hidden`` bf16 [M, 6144] contiguous, M = the state's token count. ``layer`` (from ``M3MoeState.layer``) carries the + TRTLLM-Gen NVFP4 expert buffers read in place, the per-expert scalars, the router weight / bias, the activation's + NVFP4 global scale and, on a ``shared`` state, the shared expert's MXFP8 weights; its state carries the compiled + kernel and the workspace every layer on it shares. Returns bf16 [M, 6144]: bf16(shared + routed) on a ``shared`` + state, else the routed partial. ``out``: write there instead of a new tensor. ``arp`` (``ar_push_desc``): also + push the rows into a multicast all-reduce exchange. + + The routed part is TRTLLM-Gen's NVFP4 MoE on the fp32 router logits: sigmoid (tanh form) + correction bias top-4 + of 128 (ties to the lower id), weights sigmoid / sum * 2.0 rounded to bf16; the activation quantized to NVFP4 with + the layer's global scale; per expert the raw FC1 sums, SwiGLU-OAI (alpha / beta / clamp), the NVFP4 requant with + fc31_scale_c, FC2, bf16(fc2_alpha x sum); the fma chain over the top-4 in order. The shared expert is the target's + MXFP8 SwiGLU-OAI MLP. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # One cell per token count: m_max is a trace-time constant, so each is its own build, and the target builds all + # eight (decode batches of 1..8 requests, spec-decode verify rows up to 8). The whys name what the build differs in. + CELLS: tuple[Cell, ...] = ( + Cell( + why="1 token: FC1 -> FC2 records replicated 32x, logits 32x, last phase without trailing FC2", + spec=dict(tokens=1), + ), + Cell( + why="2 tokens: records 32x, logits 32x, the one build without AR_PRE", + spec=dict(tokens=2), + ), + Cell(why="3 tokens: records 8x, logits 32x, 4 FC2-B builder warps", spec=dict(tokens=3)), + Cell( + why="4 tokens: one request verifying 3 drafts (the TP2 b1 point): 6-stage ring, B2IMG hand-off, " + "the 4 two-block FC2 CTAs stream no FC1", + spec=dict(tokens=4), + ), + Cell( + why="5 tokens: single records, a routing warp takes two tokens, logits 32x", + spec=dict(tokens=5), + ), + Cell( + why="6 tokens: single records, routing warps 8.. take tokens 4..5", spec=dict(tokens=6) + ), + Cell(why="7 tokens: the B2IMG global FC1 -> FC2 hand-off", spec=dict(tokens=7)), + Cell( + why="8 tokens: two requests verifying 3 drafts (the TP2 b2 point): 5-stage ring, x read from global " + "(no XSTAGE), no early group 0, the 4 two-block FC2 CTAs stream no FC1", + spec=dict(tokens=8), + ), + ) + + note = """ + The kernel's elementwise math uses the PTX approximations (rcp.approx, ex2.approx, tanh) where `reference` uses + torch's correctly rounded ones, and its GEMMs sum in tensor-core order where `reference` sums in fp64: a + quantization boundary (NVFP4 activation / intermediate code, routing near-tie) can fall on the other side, so the + two agree within the band the test derives, not bit for bit. The routed part is bit for bit the stock TRTLLM-Gen + NVFP4 MoE (fp4_quantize + fp4_block_scale_moe_runner) on the same routing; the test checks that too. + + Not covered: the all-reduce push (`arp`), which needs a multicast exchange (the decode forward's boundary all-reduce + exercises it); the non-shared state (`shared=False`), which this target never builds. Every layer on one state + must run in one stream order: the workspace is shared and tagged per call, never reset. + + No `is_valid`: the compiled call checks every buffer's dtype, shape and alignment against its signature and + raises on a mismatch (a misaligned `hidden`, an `out` that is not bf16 [M, 6144], a malformed `arp`). + """ + + def __call__( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + state = layer.state + if (state.device, state.m) not in _ENGAGED: + # The engagement line (once per device and token count, on every rank: the first call, eager or captured). + _ENGAGED.add((state.device, state.m)) + logger.info( + f"[M3 sm_107 tp2] MoE engine m3_moe engaged: {state.m} tokens, {state.num_ctas} routed + " + f"{state.mod.NSH} shared CTAs, {state.mod.STAGES}-stage ring, on {state.device}" + ) + return state.run(layer, hidden, out, arp=arp) + + def reference( + self, + hidden: torch.Tensor, + layer: M3MoeLayer, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Plain torch over the layer's buffers (see the class docstring); `out` and `arp` are ignored.""" + return moe_reference(hidden, layer, shared=layer.state.shared) + + +m3_moe_engine = M3MoeEngine() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py new file mode 100644 index 000000000000..636dfcd0b614 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV of the decode MoE: ``logits[M, E] = float(x) @ W^T`` at one token.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +_E4M3_MAX = 448.0 +_E2M1_MAX = 6.0 +#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. +_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) +_TIE_UP = (False, True, False, True, False, True, False) +#: The kernel builds each block's output scale through two ``rcp.approx.ftz.f32`` reciprocals (~2^-23 relative): +#: a value this close to an e2m1 midpoint may round to either neighbour, everything further out must match. +TIE_WINDOW = 2.0**-20 + + +def nvfp4_quantize( + x: torch.Tensor, sf_scale: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], + the pre-rounding scaled values [M, K] that ``near_tie`` reads). + + Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. + """ + gs = sf_scale.float().reshape(()) + m, k = x.shape + blocks = x.float().view(m, k // 16, 16) + vecmax = blocks.abs().amax(-1) + sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) + out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) + scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) + mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) + tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) + a = scaled.abs().unsqueeze(-1) + code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) + negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) + code = code | (negative.to(torch.uint8) << 3) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, sf.view(torch.uint8).reshape(-1), scaled + + +def near_tie(scaled: torch.Tensor) -> torch.Tensor: + """Elements whose e2m1 rounding the kernel's approximate reciprocals decide (see ``TIE_WINDOW``).""" + mid = torch.tensor(_E2M1_MIDPOINTS, device=scaled.device, dtype=torch.float32) + a = scaled.abs().unsqueeze(-1) + return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) + + +def fp32_logits(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + """``float(x) @ w^T`` as a true fp32 product (TF32 off).""" + with true_fp32_matmul(): + return x.float() @ w.t() + + +class M3RouterGemv(OpWrapper): + """fp32 router logits ``float(x[M, K]) @ w[E, K]^T`` with exact fp32 FMA, one CTA per expert. + + ``x`` is BF16 with 1..16 rows, unit column stride and 16-byte rows; ``w`` is the fp32 router weight; K a + multiple of 1024. ``rows`` experts per CTA; only 1 is built. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target takes the GEMV at exactly one token (2..48 go to the tensor-core router): 128 experts, hidden 6144. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token, 128 experts x hidden 6144", + spec=dict(m=1, e=128, k=6144), + ), + ) + + note = """ + The op accepts up to 16 rows; the target sends one, so wider calls are not cells. The logits differ from a true + fp32 product only in accumulation order (per-thread FMA chains over K / 1024 strided column groups, a warp tree + and an 8-warp sum). + """ + + def __call__(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + return torch.ops.trtllm.minimax_m3_router_gemv(x, w, rows, pdl) + + def reference(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: + """True fp32 product; ``rows`` and ``pdl`` do not enter the arithmetic.""" + return fp32_logits(x, w) + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> None: + # The activation rows are bulk-copied from x's base address; the op checks the row stride but not the base, + # and a misaligned bulk copy faults the context rather than raising. + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + + +m3_router_gemv = M3RouterGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py new file mode 100644 index 000000000000..0306210adaab --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""fp32 router GEMV plus the routed experts' NVFP4 activation, one token.""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_router_gemv import fp32_logits, nvfp4_quantize + + +class M3RouterGemvFp4(OpWrapper): + """``[logits fp32 [M, E], x_fp4 [M, K / 2] uint8, x_sf [M * K / 16] uint8]`` from one kernel. + + The logits are ``m3_router_gemv``'s; the NVFP4 activation is ``fp4_quantize(x, sf_scale, 16, False, False)`` + (16-element blocks, E4M3 scales in the linear layout, per-tensor ``sf_scale``) for the routed experts' FC1. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The fused-router decode step at one token: 128 experts, hidden 6144, the experts' input scale. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one decode token: logits for 128 experts and the NVFP4 activation of hidden 6144", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 8.0), + ), + Cell( + why="a global scale that puts block scales near the E4M3 top, where the scale cast saturates", + spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 2.0), + ), + ) + + note = """ + The e2m1 codes are exact except within `TIE_WINDOW` of a rounding midpoint, where the kernel's approximate + reciprocals decide the direction. Up to 16 rows are accepted; the target sends one. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_gemv_fp4(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """True fp32 logits and the NVFP4 recipe; ``pdl`` does not enter the arithmetic.""" + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return fp32_logits(x, w), packed, sf + + def is_valid(self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + # A zero or non-finite global scale makes every block scale 0 or NaN and the data codes garbage; the op + # reads the scalar on the device and cannot refuse it. + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_gemv_fp4 = M3RouterGemvFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py new file mode 100644 index 000000000000..f776afce11da --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TF32 tensor-core router of the decode MoE for 2..48 tokens, optionally with the NVFP4 activation.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from .m3_router_gemv import nvfp4_quantize + + +def tf32_rna(w: torch.Tensor) -> torch.Tensor: + """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" + bits = w.contiguous().view(torch.int32) + rounded = (bits + 0x1000) & ~0x1FFF + finite = (bits & 0x7F800000) != 0x7F800000 + return torch.where(finite, rounded, bits).view(torch.float32) + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # The target routes 2..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. + CELLS: tuple[Cell, ...] = ( + Cell( + why="two tokens, logits only (the smallest call the target makes here)", + spec=dict(m=2, k=6144, fp4=False), + ), + Cell( + why="8 tokens with the NVFP4 activation: one 8-token MMA tile, every expert tile quantizes one row", + spec=dict(m=8, k=6144, fp4=True), + ), + Cell( + why="13 tokens: a partial second token tile", + spec=dict(m=13, k=6144, fp4=True), + ), + Cell( + why="the 48-token maximum, six token tiles, logits only", + spec=dict(m=48, k=6144, fp4=False), + ), + Cell( + why="the 48-token maximum with the NVFP4 activation", + spec=dict(m=48, k=6144, fp4=True), + ), + # Every other row count the router takes (modeling's ROUTER_TC), with and without the NVFP4 activation. + *( + Cell( + why=f"{m} rows: a router row count between the ones above", + spec=dict(m=m, k=6144, fp4=m % 2 == 1), + ) + for m in range(2, 49) + if m not in (2, 8, 13, 48) + ), + ) + + note = """ + The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both + sum the same 16 slice partials in the same order. Only the one the test device picks is driven. + """ + + def __call__( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> list[torch.Tensor]: + return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) + + def reference( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> tuple[torch.Tensor, ...]: + """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" + with true_fp32_matmul(): + logits = x.float() @ tf32_rna(w).t() + if sf_scale is None: + return (logits,) + packed, sf, _ = nvfp4_quantize(x, sf_scale) + return logits, packed, sf + + def is_valid( + self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool + ) -> None: + assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" + if sf_scale is not None: + value = float(sf_scale.reshape(())) + assert value > 0.0 and value != float("inf"), ( + f"sf_scale must be positive and finite, got {value}" + ) + + +m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py new file mode 100644 index 000000000000..8c2a67356d3f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One kernel for MSA's device-side KV-length correction (MsaAttentionMetadata.on_update_kv_lens). + +The torch formulation is ~23 one-CTA kernels (min, cumsum, gathers, div / clamp / remainder, the block-table lookup, +the valid-block counts) and runs at the head of every captured decode step, twice. This kernel computes the same +int32 results in one launch: + + kv_true[b] = min(kv_lens[b], staged[b]) -> seq_lens_out[b], cu_kv[b + 1] = prefix sum (cu_kv[0] = 0) + pos[t] = kv_true[r] - qo[r] + intra[t] (r = q_batch_row[t]) + out_loc[t] = table[r, min(max(pos, 0) // page, ncols - 1)] * page + max(pos, 0) % page + n_valid[t] = (max(pos + 1, 0) + page - 1) // page + +``msa_kv_update`` is the target's ``MsaKvUpdate`` (``tensorrt_llm._torch.pyexecutor.runtime_hooks``): the MSA +metadata calls it with its own buffers and falls back to the torch formulation when it returns False. The kernel is +``torch.ops.trtllm.minimax_m3_msa_kv_update``, called through this target's ``msa_kv_update`` entry. +""" + +from typing import Optional + +import torch + +from ..ops.attention.msa_kv_update import msa_kv_update as _kv_update_op + + +def msa_kv_update( + kv_lens: torch.Tensor, + kv_lens_staged: torch.Tensor, + seq_lens: torch.Tensor, + q_batch_row: torch.Tensor, + q_intra: torch.Tensor, + block_table: torch.Tensor, + tokens_per_block: int, + batch: int, + total_q: int, + msa_seq_lens: torch.Tensor, + cu_kv_lens: Optional[torch.Tensor], + out_cache_loc: torch.Tensor, + n_valid_blocks: Optional[torch.Tensor], +) -> bool: + """The fused correction (see module doc); False, with nothing written, when an operand is outside what the + kernel covers. Arguments in ``MsaKvUpdate`` order; ``seq_lens`` are the per-request query lengths (``qo``).""" + ts = [ + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + msa_seq_lens, + out_cache_loc, + ] + ts += [t for t in (cu_kv_lens, n_valid_blocks) if t is not None] + if ( + any(t.dtype != torch.int32 or not t.is_cuda for t in ts) + or block_table.dim() != 2 + or block_table.stride(1) != 1 + ): + return False + if any(not t.is_contiguous() for t in ts if t is not block_table): + return False + if ( + kv_lens.numel() < batch + or kv_lens_staged.numel() < batch + or seq_lens.numel() < batch + or msa_seq_lens.numel() < batch + or q_batch_row.numel() < total_q + or q_intra.numel() < total_q + or out_cache_loc.numel() < total_q + or (cu_kv_lens is not None and cu_kv_lens.numel() < batch + 1) + or (n_valid_blocks is not None and n_valid_blocks.numel() < total_q) + ): + return False + _kv_update_op( + kv_lens, + kv_lens_staged, + seq_lens, + q_batch_row, + q_intra, + block_table, + tokens_per_block, + batch, + total_q, + msa_seq_lens, + cu_kv_lens, + out_cache_loc, + n_valid_blocks, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/__init__.py new file mode 100644 index 000000000000..e6c3d806a4a3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) as a CuTe DSL kernel, configured for the TP2 +projection shapes on SM107 (see op.py).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..3ba6fb0d7c8c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1472 @@ +# 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. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +``prod`` (#117, the sparse QKV + index projection at split 8, at most 8 tokens): instead of storing the packed row, the +owners stage their bf16 rows, the CTAs exchange them by st.async so that rank t holds token t's 128-dim row of this +tile's packed head, and rank t's warp 4 runs that head's producer (``m3_layer.phase_producer.producer_head``: RMSNorm, +RoPE, FP8 q / index-q / index-K and NVFP4 K / V stores) on it, bit for bit what the separate producer kernel does with +the stored row. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +from ..m3_layer.phase_producer import ProducerArgs, bf16_bits_to_f32, producer_head + +CTA_M = 128 # capacity: fixed: a kernel tile constant (weight rows per CTA = the MMA's M) +MMA_N = 8 # token columns +CTA_K = 128 # capacity: fixed: a kernel tile constant (one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row) +MMA_INST_K = 32 +# 4 MMAs per k-tile, one per scale byte of the atom. +K_BLOCKS = CTA_K // MMA_INST_K # capacity: fixed: a kernel tile constant +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _ld_shared_v4_b32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.b32 of four int32 words at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.i32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", "=r,=r,=r,=r,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [cutlass.Int32(_llvm.extractvalue(_T.i32(), out, [i], loc=loc, ip=ip)) for i in range(4)] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +# The swigluoai activation of MiniMax-M3's dense MLP, gate * sigmoid(1.702 * gate) * (up + 1) with gate clamped to <= 7 +# and up to [-7, 7], in exactly the instruction sequence the Triton silu_and_mul kernel compiles to (modules/swiglu.py +# on SM107): the gate clamp on the packed bf16 pair, the up clamp as min.xorsign.abs, two rounded multiplies into +# ex2.approx, div.full for the reciprocal, two rounded multiplies, round to bf16. No multiply feeds an add, so the +# explicit .rn forms compile to what the Triton kernel's contractable forms do. +SWIGLU_OAI_ALPHA = 1.702 +SWIGLU_OAI_LIMIT = 7.0 + + +@dsl_user_op +def _swiglu_oai_bf16x2(gate_word, up_word, *, loc=None, ip=None): + """The activation of one packed bf16 pair of gate and up values: its two bf16 results as one int32 word (the low + half from the low halves).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [cutlass.Int32(w).ir_value(loc=loc, ip=ip) for w in (gate_word, up_word)], + "{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, $1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, $2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 $0, r1, r0;\n\t}", "=r,r,r", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[ + bool + ], # x_bf16: x_words is [M, 2K] gate | up; the activation is swigluoai(gate, up) + # prod: the producer's addresses and strides (ProducerArgs fields; `packed` is unused), then num_pages and eps. + pr_q: cutlass.Int64, + pr_iq: cutlass.Int64, + pr_ik: cutlass.Int64, + pr_kvd: cutlass.Int64, + pr_kvs: cutlass.Int64, + pr_loc: cutlass.Int64, + pr_quant: cutlass.Int64, + pr_pos: cutlass.Int64, + pr_rope: cutlass.Int64, + pr_qw: cutlass.Int64, + pr_kw: cutlass.Int64, + pr_iqw: cutlass.Int64, + pr_ikw: cutlass.Int64, + pr_ips: cutlass.Int64, + pr_its: cutlass.Int64, + pr_dps: cutlass.Int64, + pr_dls: cutlass.Int64, + pr_dts: cutlass.Int64, + pr_sps: cutlass.Int64, + pr_sls: cutlass.Int64, + pr_pages: cutlass.Int32, + pr_eps: cutlass.Float32, + prod: cutlass.Constexpr[bool], +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # prod: this CTA's bf16 owner rows by token ([8][rows_per_rank]), and the packed-head row of token `rank` gathered + # from the cluster (128 bf16 as 64 words, completed on pr_full by 8 x rows_per_rank * 2 bytes of st.async). + pr_stage = None + pr_row = None + pr_full = None + if cutlass.const_expr(prod): + pr_stage = cutlass.Array( + cutlass.Uint16, 8 * rows_per_rank, space=cutlass.AddressSpace.smem, alignment=16 + ) + pr_row = cutlass.Array(cutlass.Int32, 64, space=cutlass.AddressSpace.smem, alignment=16) + pr_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if cutlass.const_expr(prod): + # Completed by the 8 ranks' st.async (this one's included) only on a rank that receives a token row. + prims.mbarrier_init(pr_full, 1) + if rank < num_tokens: + prims.mbarrier_arrive_expect_tx(pr_full, split * rows_per_rank * 2) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in if x_swiglu else k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + if cutlass.const_expr(x_swiglu): + # The up half of the same columns, K values after the gate in the row. + qz_wvecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(k_in // 2 + qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + if cutlass.const_expr(x_swiglu): + qz_word = _swiglu_oai_bf16x2( + qz_word, cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + ) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The flags are rotated by the previous all-reduce after its dependent release; this load is ordered by + # the wait above, which waits for ALL in-flight upstream grids, not only the immediate predecessor (PTX + # griddepcontrol; MoE test_pdl_transitive_wait.py). Keep the wait before this load. + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.const_expr(prod): + # The bf16 the packed row would hold (RN, as the y store below), staged by token. + if cutlass.Int32(4 * g + j) < num_tokens: + pr_stage.store( + cutlass.Uint16( + (_bf16_bits(sums[j]) >> cutlass.Int32(16)) + & cutlass.Int32(0xFFFF) + ), + idx=cutlass.Int32((4 * g + j) * rows_per_rank) + local, + ) + else: + if cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The flags are rotated by the previous all-reduce after its dependent release; this load is ordered + # by the wait above, which waits for ALL in-flight upstream grids, not only the immediate + # predecessor (PTX griddepcontrol; MoE test_pdl_transitive_wait.py). Keep the wait before this load. + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(prod): + # Transpose: thread e < 2 T sends half e % 2 (16 bytes = 8 dims) of token e / 2's staged rows to cluster + # rank e / 2, at dims [rank * rows_per_rank + 8 (e % 2), ...) of that rank's packed-head row. + pr_e = tx - cutlass.Int32(128) + if pr_e < num_tokens * cutlass.Int32(2): + pr_t = pr_e // cutlass.Int32(2) + pr_h = pr_e % cutlass.Int32(2) + pr_w = _ld_shared_v4_b32( + pr_stage.subview( + pr_t * cutlass.Int32(rows_per_rank) + pr_h * cutlass.Int32(8) + ).data_ptr() + ) + _st_async_v4( + _mapa_u32( + pr_row.subview( + rank * cutlass.Int32(rows_per_rank // 2) + pr_h * cutlass.Int32(4) + ).data_ptr(), + pr_t, + ), + pr_w[0], + pr_w[1], + pr_w[2], + pr_w[3], + _mapa_u32(pr_full.data_ptr(), pr_t), + ) + if warp_id == 4: + if rank < num_tokens: + # The producer reads step inputs (positions, slots) and writes the caches: ordered after the + # upstream grids by this thread's own grid wait (already satisfied: the activation landed). + prims.griddepcontrol(prims.GridDepAction.WAIT) + while not _test_wait_cluster(pr_full.data_ptr(), 0): + pass + pr_lo = pr_row.load(idx=lane * cutlass.Int32(2)) + pr_hi = pr_row.load(idx=lane * cutlass.Int32(2) + cutlass.Int32(1)) + mask = cutlass.Int32(0xFFFF) + producer_head( + rank, + n_tile, + lane, + bf16_bits_to_f32(pr_lo & mask), + bf16_bits_to_f32((pr_lo >> cutlass.Int32(16)) & mask), + bf16_bits_to_f32(pr_hi & mask), + bf16_bits_to_f32((pr_hi >> cutlass.Int32(16)) & mask), + ProducerArgs( + packed=cutlass.Int64(0), + q_out=pr_q, + index_q_out=pr_iq, + index_k=pr_ik, + kv_data=pr_kvd, + kv_scale=pr_kvs, + out_cache_loc=pr_loc, + kv_quant_scale=pr_quant, + position_ids=pr_pos, + rotary_cos_sin=pr_rope, + q_weight=pr_qw, + k_weight=pr_kw, + index_q_weight=pr_iqw, + index_k_weight=pr_ikw, + eps=pr_eps, + num_pages=pr_pages, + index_page_stride=pr_ips, + index_token_stride=pr_its, + data_page_stride=pr_dps, + data_plane_stride=pr_dls, + data_token_stride=pr_dts, + scale_page_stride=pr_sps, + scale_plane_stride=pr_sls, + ), + ) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[bool], + # prod: see m3_mx_gemv_kernel (zeros when off). + pr_q: cutlass.Int64, + pr_iq: cutlass.Int64, + pr_ik: cutlass.Int64, + pr_kvd: cutlass.Int64, + pr_kvs: cutlass.Int64, + pr_loc: cutlass.Int64, + pr_quant: cutlass.Int64, + pr_pos: cutlass.Int64, + pr_rope: cutlass.Int64, + pr_qw: cutlass.Int64, + pr_kw: cutlass.Int64, + pr_iqw: cutlass.Int64, + pr_ikw: cutlass.Int64, + pr_ips: cutlass.Int64, + pr_its: cutlass.Int64, + pr_dps: cutlass.Int64, + pr_dls: cutlass.Int64, + pr_dts: cutlass.Int64, + pr_sps: cutlass.Int64, + pr_sls: cutlass.Int64, + pr_pages: cutlass.Int32, + pr_eps: cutlass.Float32, + prod: cutlass.Constexpr[bool], +) -> None: + # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. + assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( + "x_swiglu needs x_bf16 and excludes din" + ) + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, x_swiglu, + pr_q, pr_iq, pr_ik, pr_kvd, pr_kvs, pr_loc, pr_quant, pr_pos, pr_rope, pr_qw, pr_kw, pr_iqw, pr_ikw, pr_ips, + pr_its, pr_dps, pr_dls, pr_dts, pr_sps, pr_sls, pr_pages, pr_eps, prod, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py new file mode 100644 index 000000000000..efe4879f9fda --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py @@ -0,0 +1,645 @@ +# 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. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) at the TP2 per-rank projection shapes (SM107). + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 # capacity: fixed: the device's shared-memory limit + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring[, contig, nacc, a_tmem, st_async]) of the TP2 per-rank projections. The optional fields default +# to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. SM107 graph us at M 1 / 4 / 8 over weights rotated beyond L2, against the +# stock MXFP8 quantize + FlashInfer CuTe DSL GEMM: QKV 4.11 / 4.12 / 4.10 vs 8.6-8.8, o_proj 3.90 / 3.96 / 3.70 vs +# 6.9-7.2, shared gate_up 3.11 / 3.09 / 3.10 vs 8.1-8.4, shared down 2.22 vs 4.7-4.9, dense gate_up 8.15 / 8.56 / 8.55 +# vs 10.2-10.5, dense down 5.89 / 5.72 / 5.46 vs 8.7-9.0. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + # The sm_107 tp8 target's per-rank projections (8 ranks over two trays), (split, ring) from the SM107 sweep 778596 + # (locked, MXFP8 in, back-to-back PDL, rows 1 / 4 / 8; the GB300 tp8 table has none of these shapes on SM107). + # sparse-layer QKV + index projection: 2.37 / 2.43 / 2.42 us + (1536, 6144): (8, 8), + # dense-layer QKV: 2.34 / 2.35 / 2.36 us + (1280, 6144): (8, 8), + # o_proj: 2.05 / 2.07 / 2.07 us + (6144, 1024): (2, 12), + # shared-expert gate_up: 2.16 / 2.17 / 2.17 us + (768, 6144): (8, 8), + # shared-expert down: 1.49 / 1.50 / 1.52 us + (6144, 384): (1, 8), + # dense-layer gate_up: 3.09 / 3.12 / 3.08 us + (3072, 6144): (4, 12), + # dense-layer down: 2.29 / 2.30 / 2.29 us + (6144, 1536): (2, 8), +} +# Per-token-count configs (N, K, M); none at TP2. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = {} +# Above 8 tokens (wider than the TP2 decode step); none at TP2. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = {} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +# (rule 1b, deterministic) The decode projections at up to STRIDED_MAX_ROWS rows split K over strided k-tiles per +# split rank (contig 0: rank r owns k-tiles r, r + split, ...) instead of contiguous ranges: another fixed fp32 order. +# SM107 TP8 4 rows (783579), us/call contiguous -> strided: sparse QKV 2.405 -> 2.371, dense QKV 2.323 -> 2.291, +# o_proj 2.061 -> 2.019, dense gate_up 3.047 -> 3.006, dense down 2.313 -> 2.287; deterministic (3 runs bitwise) and +# max error vs the fp64 product equal to the contiguous order's. Kill: False. +STRIDED_K_TILES = True +STRIDED_MAX_ROWS = 8 # capacity: fixed: the strided-input kernels' built rows +_STRIDED = frozenset({(1536, 6144), (1280, 6144), (6144, 1024), (3072, 6144), (6144, 1536)}) + +# #117: the producer epilogue serves the sparse-layer QKV + index projection, 12 packed heads of 128 (8 q, k, v, +# index-q, index-k), one head per split-8 cluster. +PROD_N = 12 * 128 # capacity: fixed: the producer's packed heads +# The kernel's producer arguments on a call without the producer epilogue (20 addresses / strides, pages, eps). +PROD_OFF = (0,) * 21 + (0.0,) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + if ( + STRIDED_K_TILES + and m is not None + and m <= STRIDED_MAX_ROWS + and (n, k) in _STRIDED + and len(cfg) == 2 + ): + cfg = cfg + (0,) + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # 17..32 rows at K < 1024 decline (an unsplit K below 1024 is slower than the FlashInfer GEMM there); no TP2 + # projection has K < 1024. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], + x_swiglu: bool = False, + prod: bool = False, +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + if prod and not ( + n == PROD_N + and split == 8 + and m <= split + and st_async + and not x_bf16 + and not x_swiglu + and ar is None + ): + # The producer epilogue: the sparse QKV + index projection (one 128-row packed head per cluster of 8), at most + # one token per cluster rank, the st.async exchange, MXFP8 input, no all-reduce push. + return None + if nacc != 1: + # More than one TMEM accumulator is broken on SM107 (783556 / 783557 / 783579: nondeterministic or garbage + # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + prod, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile( + key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple, prod_args: tuple +): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + prod, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + x_swiglu, + *prod_args, + prod, + ) + return fn + + +def _views( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight, + weight_scale, + y, + x_bf16: bool, + x_swiglu: bool = False, +) -> tuple: + # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. + k = x.shape[1] // 2 if x_swiglu else x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, + x_swiglu: bool = False, + prod_args: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + if x_swiglu: + k //= 2 + n = weight.shape[0] + if weight.shape[1] != k: + return None + prod = prod_args is not None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu, prod) + if resolved is None: + return None + key, runtime_ar = resolved + if not prod: + prod_args = PROD_OFF + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + # prod: no packed row is stored (the producer epilogue consumes it); y is the placeholder. + y = torch.empty(0 if prod else m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra, prod_args) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra, *prod_args) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, + x_swiglu: bool = False, + prod: bool = False, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu, prod) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros( + m, + 2 * k if x_swiglu else k, + dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, + device=dev, + ) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev), PROD_OFF) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_swiglu( + gate_up: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``m3_mx_gemv_bf16(swiglu(gate_up), ...)`` in one launch: ``gate_up`` bf16 [M <= 8, 2K] (contiguous, gate then + up) is activated in the kernel's quantize prologue with the swigluoai arithmetic of the Triton silu_and_mul kernel + (alpha 1.702, limit 7, the (up + 1) offset; bit for bit), then quantized as the bf16 path does.""" + if ( + gate_up.dtype != torch.bfloat16 + or gate_up.dim() != 2 + or not gate_up.is_contiguous() + or gate_up.data_ptr() % 16 != 0 + or gate_up.shape[1] % 2 != 0 + ): + return None + return _run(gate_up, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar, True) + + +def m3_mx_gemv_producer( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + q: torch.Tensor, + index_q: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + trigger_early: bool = True, + use_pdl: Optional[bool] = None, +) -> bool: + """#117: ``m3_producer(m3_mx_gemv(xq, xsf, weight, weight_scale), ...)`` in one launch, bit for bit: the QKV + + index GEMV of the sparse layer (``weight`` [12 x 128, K]) whose owner epilogue hands each token's packed head to + the producer phase (``m3_layer.phase_producer.producer_head``). Writes ``q`` FP8 [T, 8, 128] and ``index_q`` FP8 + [T, 1, 128] and the token's cache slot (NVFP4 K / V + scales, FP8 index-K) exactly as the producer does; the + caches and tensors are those ``minimax_m3_producer_pw`` takes (kv_data [pages, 2, 1, 128, 64] u8, kv_scale [pages, + 2, 1, 128, 8] u8, index_k [pages, 1, 128, 128] e4m3). Returns False (nothing launched) when the call is declined. + """ + m = xq.shape[0] + if ( + m > 8 + or weight.shape[0] != PROD_N + or q.shape != (m, 8, 128) + or index_q.shape != (m, 1, 128) + or not (q.is_contiguous() and index_q.is_contiguous()) + or kv_data.dim() != 5 + or kv_scale.dim() != 5 + or index_k.dim() != 4 + ): + return False + tensors = ( + q, + index_q, + index_k, + kv_data, + kv_scale, + out_cache_loc, + kv_quant_scale, + position_ids, + rotary_cos_sin, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + ) + prod_args = ( + *(t.data_ptr() for t in tensors), + index_k.stride(0), + index_k.stride(2), + kv_data.stride(0), + kv_data.stride(1), + kv_data.stride(3), + kv_scale.stride(0), + kv_scale.stride(1), + kv_data.shape[0], + float(eps), + ) + out = _run( + xq, xsf, weight, weight_scale, False, trigger_early, None, use_pdl, None, False, prod_args + ) + return out is not None + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/__init__.py new file mode 100644 index 000000000000..acec8fdbb26f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Normalization entries this target alone calls.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py new file mode 100644 index 000000000000..11b92b47fc50 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode-step embedding gather and the first layer's input RMSNorm in one CuTe DSL launch, into caller buffers.""" + +from __future__ import annotations + +from typing import Tuple + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm + +__all__ = ["DecodeEmbedding", "m3_embedding_firstnorm"] + +VOCAB = 200064 # MiniMax-M3's vocabulary; the decode table is the full (replicated) embedding + + +class M3EmbeddingFirstnorm(OpWrapper): + """Write ``table[ids]`` into ``raw`` (bf16 [M, 6144], a zero row for an id outside ``[0, V)``) and + ``raw * rsqrt(mean(raw^2) + eps) * weight`` into ``out``; returns ``(out, raw)`` -- the first layer's normed input + and its residual. ``ids`` int32 or int64 [M], M = 1..8; ``table`` bf16 [V, 6144] contiguous; ``weight`` bf16 + [6144]. The first call per (M, id dtype, table shape, PDL) compiles (``DecodeEmbedding`` does it at build).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target builds every (rows 1..8, id dtype) owner; the kernel is one CTA per row, so the rows differ only in + # grid size: the ends of the range and both id dtypes. + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, int32 ids", spec=dict(rows=1, id_dtype=torch.int32)), + Cell(why="eight requests, int32 ids", spec=dict(rows=8, id_dtype=torch.int32)), + Cell(why="one request, int64 ids", spec=dict(rows=1, id_dtype=torch.int64)), + Cell( + why="six rows (one request verifying a draft), int64 ids", + spec=dict(rows=6, id_dtype=torch.int64), + ), + # Every other row count the decode dispatch takes (modeling's EMBED_FIRSTNORM), both id dtypes. + *( + Cell( + why=f"{rows} rows: a decode row count between the ones above", + spec=dict(rows=rows, id_dtype=torch.int32 if rows % 2 else torch.int64), + ) + for rows in (2, 3, 4, 5, 7) + ), + ) + + note = """ + The gathered rows are exact. The norm is flashinfer's RMSNorm arithmetic in its reduction order with a fast + rsqrt, where `reference` sums in fp64 and takes torch's rsqrt: within the band the test derives, not bit for bit. + SM107 against the stock embedding + RMSNorm module (flashinfer's kernel, 768 threads): the gathered rows are + bitwise equal; the normed rows differ in one bf16 ulp on ~5e-6 of the elements (~1.6% of rows, 1-2 elements each: + the fp32 sum-of-squares tree differs, so the rsqrt input can differ in its last bits). + An id outside [0, V) gives a zero row and a zero norm (torch's embedding would raise). No accepted input is known + to compute wrongly: the host side rejects any other dtype, shape or a strided buffer. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return m3_embed_norm(ids, table, weight, eps, raw, out, use_pdl) + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return (x * rstd * weight.double()).to(torch.bfloat16), rows + + +m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/__init__.py new file mode 100644 index 000000000000..d25c41dc9557 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/__init__.py @@ -0,0 +1,4 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Speculative-decoding entries this target alone calls: the Eagle3 verify ops, the capture / fc_norm RMSNorms and +the step-input staging.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py new file mode 100644 index 000000000000..be5aa46fe774 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_chunk_norm import rmsnorm_rows + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = ( + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ) + + note = """ + fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional + one-ulp move against torch's statistic. + """ + + def __call__( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) + return dst + + def reference( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> torch.Tensor: + """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" + x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h + return rmsnorm_rows(x, w, eps) + + def is_valid( + self, + h: torch.Tensor, + r: Optional[torch.Tensor], + w: torch.Tensor, + eps: float, + dst: torch.Tensor, + ) -> None: + # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read + # past its elements without an error. + assert w.is_contiguous(), "w must be dense [H]" + + +m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py new file mode 100644 index 000000000000..2f71f1b381b1 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states +(``m3_spec_tail.chunk_norm``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS: tuple[Cell, ...] = ( + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="two requests' 8 verify rows (b2) from the capture buffer, dense [8, 3 x 6144]", + spec=dict(m=8, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ) + + note = """ + fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then + against torch's statistic; the cells' band states it. + """ + + def __call__( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) + return dst + + def reference( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ): + """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" + c, h = w.shape + x = src[:, : c * h].reshape(src.shape[0], c, h) + return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) + + def is_valid( + self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True + ) -> None: + c, h = w.shape + assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + + +m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py new file mode 100644 index 000000000000..4fe14d538a5f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat +(``m3_spec_tail.draft_input``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). +MAX_HIDDEN = 8192 + + +def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" + xf = x.float() + r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) + return (xf * r * weight.float()).to(torch.bfloat16) + + +class M3DraftInput(OpWrapper): + """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = + hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of + a vocab-local embedding). Returns ``(out, residual)``.""" + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling's DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the + # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests + # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", + spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), + ) + for n in range(1, 9) + ), + *( + Cell( + why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", + spec=dict(m=n, h=6144, gather=True, invalid_id=False), + ) + for n in range(1, 9) + ), + Cell( + why="third pass of 3 requests: one row each, hidden read in place", + spec=dict(m=3, h=6144, gather=False, invalid_id=False), + ), + Cell( + why="one row with an out-of-table id: the embedding half is zero", + spec=dict(m=1, h=6144, gather=False, invalid_id=True), + ), + ) + + note = """ + The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's + summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, + also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the + drafter's embedding is the whole table, so its ids are always in range. + """ + + def __call__( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + torch.ops.trtllm.minimax_m3_draft_input( + ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl + ) + return out, residual + + def reference( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / + ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" + valid = (ids >= 0) & (ids < table.shape[0]) + emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) + emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) + hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] + ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) + return ref, (hid.clone() if residual is not None else None) + + def is_valid( + self, + ids: torch.Tensor, + table: torch.Tensor, + hidden: torch.Tensor, + rows: Optional[torch.Tensor], + w_embed: torch.Tensor, + w_hidden: torch.Tensor, + eps_embed: float, + eps_hidden: float, + out: torch.Tensor, + residual: Optional[torch.Tensor], + pdl: bool = True, + ) -> None: + h = table.shape[1] + # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. + assert table.is_contiguous() and hidden.is_contiguous(), ( + "table / hidden must be dense [*, H]" + ) + if rows is not None: + assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" + else: + assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" + # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. + assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" + assert residual is None or ( + residual.is_contiguous() and residual.shape == (ids.shape[0], h) + ), "residual must be contiguous [M, H]" + + +m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py new file mode 100644 index 000000000000..b2487a213990 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter MLP's plain SwiGLU as one C++ launch (modeling's ``DRAFT_SILU_CPP``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3DraftSilu(OpWrapper): + """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: + trtllm::minimax_m3_sm107_tp8_silu_and_mul (minimaxM3Sm107Tp8/siluAndMul.cu), bit for bit the Triton + ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP + runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs + the Triton op).""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP8 drafter MLP shard: intermediate 18432 / 8 = 2304, gate | up 4608 wide, at the drafter's decode rows. + CELLS: tuple[Cell, ...] = ( + Cell(why="the TP8 drafter MLP at 1 row", spec=dict(rows=1, d=2304)), + Cell(why="the TP8 drafter MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=2304)), + Cell(why="the TP8 drafter MLP at 8 rows", spec=dict(rows=8, d=2304)), + ) + + note = """ + Elementwise, one rounding: bit for bit the Triton kernel (the same PTX instruction sequence), checked over every + bf16 gate and every bf16 up pattern. + """ + + def __call__(self, x: torch.Tensor) -> Optional[torch.Tensor]: + if ( + x.dtype != torch.bfloat16 + or x.dim() != 2 + or not x.is_cuda + or not x.is_contiguous() + or x.shape[1] % 16 != 0 + or x.data_ptr() % 16 != 0 + or x.shape[0] > 65535 + ): + return None + out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) + torch.ops.trtllm.minimax_m3_sm107_tp8_silu_and_mul(x, out) + return out + + def reference(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul(x) + + def is_valid(self, x: torch.Tensor) -> None: + assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" + + +m3_draft_silu = M3DraftSilu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py new file mode 100644 index 000000000000..e63389ffc6eb --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One pass of the fused greedy Eagle3 draft loop: the drafter's vocab-sharded BF16 head (streamed from its lossless +12-bit coding, ``ops.gemm.m3_target_head``, or as BF16 fragments), the cross-rank argmax and the loop's per-pass +bookkeeping in one kernel (``m3_draft_step.draft_step``).""" + +from __future__ import annotations + +import functools +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( + Arch, + Cell, + OpWrapper, + true_fp32_matmul, +) + +from ..gemm.m3_target_head import decode as decode_coded +from .m3_shard_accept import exchange_tokens, local_argmax, make_keys + +#: Requests (head rows) per pass, and exchange rows per call half of the workspace. +MAX_TOKENS = 8 +#: 16-row weight tiles one CTA can own (the kernel's shared-memory accumulator). +MAX_TILES_PER_CTA = 32 +KERNEL_WARPS = 16 # warps per CTA (draftStep.h kWarps) + + +@functools.lru_cache(maxsize=64) +def max_tile_pieces(n_tiles: int, num_ctas: int, chunks_per_tile: int) -> int: + """The most warps that stream one weight tile under the kernel's whole-tile split: CTA c owns tiles + [c * n / grid, (c + 1) * n / grid), and its (tile, chunk) range is cut into KERNEL_WARPS equal contiguous + segments. The shared-memory atomics add a tile's pieces in arrival order: deterministic for at most 2 pieces (two + addends from 0 commute), not for more (fixed_order then sums them in warp order).""" + worst = 1 + for c in range(num_ctas): + lo, hi = c * n_tiles // num_ctas, (c + 1) * n_tiles // num_ctas + span = (hi - lo) * chunks_per_tile + bounds = [w * span // KERNEL_WARPS for w in range(KERNEL_WARPS + 1)] + for t in range(hi - lo): + a, b = t * chunks_per_tile, (t + 1) * chunks_per_tile + worst = max( + worst, sum(1 for w in range(KERNEL_WARPS) if bounds[w] < b and bounds[w + 1] > a) + ) + return worst + + +#: Draft length of the fused loop, and the drafter's TP8 head shard (200064-id vocabulary) and hidden size. +DRAFT_LEN = 3 +DRAFT_SHARD = 25008 +HIDDEN = 6144 + + +def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the + 12-bit coded stream with its exception lists decoded.""" + if exc_off is None: + return from_fragments(wf) + n = (exc_off.numel() - 1) * 16 + return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + + +def shard_logits( + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() + + +class M3DraftStep(OpWrapper): + """One fused greedy draft pass over ``rows <= 8`` requests (one head row each). + + Head input ``x[x_rows[b]]`` (pass 0: the last accepted verify row of each request) or ``x[b]``. The head is this + rank's BF16 lm_head shard in fragment order (``m3_draft_step.bf16_fragments``, global ids ``vocab_offset ..``), + or with ``exc_off`` / ``exc_pos`` / ``exc_w`` / ``exp_base`` its lossless 12-bit coding (``wf`` = the uint8 stream + of ``m3_target_head.encode_chunked``, which must round-trip to the BF16 shard); each logit is the fp32 dot product + rounded to BF16 and the row's token is the global argmax over all ranks' + shards (ties, -0 / +0 included, to the lowest global id), exchanged through the multicast slots + ``[2][8][world]``. The last CTA then writes ``draft_out[b, step]`` and ``ids_next[b]`` (the token), ``pos_next[b] + = pos_src[x_rows[b] or b] + 1``, and advances the draft KV lengths as the in-tree draft loop does between passes + (pass 0 with ``num_acc``: ``kv_lens -= draft_len - num_acc`` and ``seq_lens = 1``; later passes: ``kv_lens += + 1``). ``out_keys`` (optional) receives each row's winning key. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every pass of the fused draft loop (modeling's DRAFT_STEP): draft_len 3 at the TP8 GEN worker's 1..2 requests, + # the drafter's full-vocabulary BF16 head (25008 rows per rank of 8, hidden 6144; 1563 tiles, 8 per CTA on 212 + # SMs). Pass 0 reads the verify rows' hidden (4 per request) through x_rows and rolls the KV back; the later passes + # read one row per request. BF16 fragments only: on sm_107 the 12-bit coded head is slower than the BF16 stream + # (modeling's DRAFT_HEAD_CODED is False). + CELLS: tuple[Cell, ...] = ( + *( + Cell( + why=f"pass 0 at {n} request(s): head rows gathered from {4 * n} verify rows, KV rollback, seq_lens = 1", + spec=dict( + nreq=n, + draft_len=DRAFT_LEN, + step=0, + rank=7 * (n - 1), + world=8, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ) + for n in (1, 2) + ), + Cell( + why="pass 1 at two requests: one row each, kv_lens + 1, rank 6's ids", + spec=dict( + nreq=2, + draft_len=DRAFT_LEN, + step=1, + rank=6, + world=8, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ), + Cell( + why="last pass (step 2) at one request, rank 0", + spec=dict( + nreq=1, + draft_len=DRAFT_LEN, + step=2, + rank=0, + world=8, + vs=DRAFT_SHARD, + k=HIDDEN, + coded=False, + ), + ), + Cell( + why="pass 2 at 2 requests with exact cross-rank ties (the lower global id wins)", + spec=dict( + nreq=2, + draft_len=DRAFT_LEN, + step=2, + rank=3, + world=8, + vs=DRAFT_SHARD, + k=HIDDEN, + ties=True, + coded=False, + ), + ), + ) + + note = """ + The head is the in-tree BF16 head's arithmetic (BF16 weights and activations, fp32 accumulation, a BF16 logit; + the coded head's weight is the BF16 shard bit for bit; its exceptions decode to +0 in the MMAs and their exact + products are added at the tile's flush) + and the argmax is the in-tree one over BF16 logits, so a token differs from the in-tree path's only where the fp32 + sum, accumulated in another order, rounds to the other side of a BF16 boundary and that decides the row's maximum. + The cells' rows are decidable and the peers' keys sit 0.25 away (or tie exactly); the edge tests (forced BF16 + ties, -0 / +0, top-2 one fp32 step apart that round to one BF16 value, adjacent BF16 values) use integer data + whose fp32 sums are exact in any order and hold the kernel to the in-tree path bit for bit. The cross-rank + exchange is certified on one device with the peers' keys written into the slots (see m3_shard_accept). + """ + + def __call__( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> None: + # fold: (the 14 fold-tail tensors in buildParams' order, rows, tokens per step) -- the step's forward-tail store + # and next-step staging, run by the pass's last CTA after its tokens (modeling's DRAFT_STEP_FOLD_TAIL). + tensors, fold_rows, fold_t = ( + (list(fold[0]), int(fold[1]), int(fold[2])) if fold is not None else (None, 0, 0) + ) + torch.ops.trtllm.minimax_m3_sm107tp8_draft_step( + x, + x_rows, + rows, + wf, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft_out, + step, + ids_next, + pos_src, + pos_next, + kv_lens, + num_acc, + seq_lens, + draft_len, + num_ctas, + pdl, + out_keys, + exc_off, + exc_pos, + exc_w, + exp_base, + tile_lo, + tile_hi, + total_ctas, + None, + None, + tensors, + fold_rows, + fold_t, + fixed_order, + ) + + def reference( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> dict[str, torch.Tensor]: + """Every buffer the pass writes, computed on copies: "logits" (this rank's BF16 logits as fp32), "local_keys" + (this rank's key per row), "keys" (the winning key), "tokens", "draft_out", "ids_next", "pos_next", "kv_lens" + and "seq_lens". + + Reads ``ws_uc``, ``counters`` and ``kv_lens`` before the call. ``keys``, ``ws_mc_ptr``, ``num_ctas``, ``pdl``, + ``out_keys`` and ``fixed_order`` do not change the result; ``fold`` adds the fold tail's writes (m3_fold_tail), + not these. + """ + h = x[x_rows.long()] if x_rows is not None else x[:rows] + logits = shard_logits(h, wf, exc_off, exc_pos, exc_w, exp_base) + best, idx = local_argmax(logits, vocab_offset) + local_keys = make_keys(best, idx) + tokens = exchange_tokens(local_keys, ws_uc, counters, rank, world, MAX_TOKENS) + half = int(counters[1].item()) & 1 + slots = ws_uc.view(torch.int64)[ + half * MAX_TOKENS * world : (half * MAX_TOKENS + rows) * world + ] + slots = slots.view(rows, world).clone() + slots[:, rank] = local_keys + sign = torch.iinfo(torch.int64).min + win = (slots ^ sign).max(dim=1).values ^ sign + d_out = draft_out.clone() + d_out[:rows, step] = tokens + src = x_rows.long() if x_rows is not None else torch.arange(rows, device=x.device) + lens = kv_lens.clone() + sl = seq_lens.clone() if seq_lens is not None else None + if step == 0 and num_acc is not None: + lens[:rows] -= draft_len - num_acc[:rows] + sl[:rows] = 1 + elif step > 0: + lens[:rows] += 1 + out = { + "logits": logits, + "local_keys": local_keys, + "keys": win, + "tokens": tokens, + "draft_out": d_out, + "ids_next": tokens.clone(), + "pos_next": pos_src[src] + 1, + "kv_lens": lens, + } + if sl is not None: + out["seq_lens"] = sl + return out + + def is_valid( + self, + x: torch.Tensor, + x_rows: Optional[torch.Tensor], + rows: int, + wf: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft_out: torch.Tensor, + step: int, + ids_next: torch.Tensor, + pos_src: torch.Tensor, + pos_next: torch.Tensor, + kv_lens: torch.Tensor, + num_acc: Optional[torch.Tensor], + seq_lens: Optional[torch.Tensor], + draft_len: int, + num_ctas: int, + pdl: bool = True, + out_keys: Optional[torch.Tensor] = None, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, + tile_lo: int = 0, + tile_hi: int = 0, + total_ctas: int = 0, + fold: Optional[tuple] = None, + fixed_order: bool = False, + ) -> None: + assert not bool(keys[:rows].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + # The kernel reads the head row x + x_rows[b] * K and the position pos_src[x_rows[b]] with no bound. + if x_rows is not None: + assert bool(((x_rows[:rows] >= 0) & (x_rows[:rows] < x.shape[0])).all()), ( + "x_rows must index x" + ) + assert bool((x_rows[:rows] < pos_src.shape[0]).all()), ( + "pos_src must cover the head rows" + ) + else: + assert pos_src.shape[0] >= rows, "pos_src must cover the head rows" + # Each CTA accumulates at most MAX_TILES_PER_CTA tiles in shared memory. + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + assert -(-tiles // num_ctas) <= MAX_TILES_PER_CTA, "too many weight tiles per CTA" + # D-FIX: the atomics' arrival-order sum is deterministic only while no tile takes more than 2 warp pieces. + assert ( + fixed_order + or max_tile_pieces(tiles, num_ctas, x.shape[1] // (16 if exc_off is None else 64)) <= 2 + ), ( + "a tile takes more than 2 warp pieces on this grid: the atomics' sum is order-dependent (fixed_order)" + ) + if exc_off is not None: + # The exceptions are read per tile in list order with no bound against the coded stream. + assert int(exc_off[-1].item()) == exc_pos.numel(), ( + "exc_off must end at the exception count" + ) + assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" + # A NaN logit maps to a key above every finite one and wins the row silently. + h = x[x_rows.long()] if x_rows is not None else x[:rows] + assert not bool(torch.isnan(h).any()), "NaN head inputs win the argmax" + if fold is not None: + assert len(fold) == 3 and len(fold[0]) == 14, ( + "fold is (the 14 fold-tail tensors, rows, tokens per step)" + ) + if num_acc is not None: + na = num_acc[:rows] + assert bool(((na >= 1) & (na <= draft_len + 1)).all()), ( + "num_acc counts must lie in [1, draft_len + 1]" + ) + + +m3_draft_step = M3DraftStep() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py new file mode 100644 index 000000000000..a192e89c5c76 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from ..m3_force_vec import force_vec_kernel, force_vec_req_kernel + + +class M3ForceVec(OpWrapper): + """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` + (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 + ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + force_vec_kernel[(1,)]( + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, + pool, + counter, + out, + batch_size, + base_total, + max_total, + frac, + counter_stride, + slot_stride, + pool_mask, + block, + ): + """The counts (a new tensor; ``counter`` and ``out`` are not written).""" + c = int(counter.item()) + 1 + slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) + idx = (c * counter_stride + slots * slot_stride) & pool_mask + r = pool[idx] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec = M3ForceVec() + + +class M3ForceVecReq(OpWrapper): + """For each of the ``batch_size`` requests ``b`` write + ``min(base_total + (pool[(keys[b] * key_stride) & pool_mask] < frac), max_total)`` into ``out`` (int32): + one draw per request at its own key (its KV length at the verify), no shared counter. ``block`` is the launch + width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). + + Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) + runs it. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), + Cell(why="two requests, fraction 0.3", spec=dict(batch=2, block=16, frac=0.3)), + ) + + note = """ + Integer arithmetic and one fp32 comparison: `reference` must match exactly. + """ + + def __call__( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + force_vec_req_kernel[(1,)]( + pool, + keys, + out, + batch_size, + base_total, + max_total, + frac, + key_stride, + pool_mask, + BLOCK=block, + ) + return out + + def reference( + self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block + ): + """The counts (a new tensor; ``out`` is not written).""" + k = keys[:batch_size].to(torch.int64) + r = pool[(k * key_stride) & pool_mask] + return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) + + +m3_force_vec_req = M3ForceVecReq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py new file mode 100644 index 000000000000..d86ee7e23625 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Sampled (temperature + top-p) Eagle3 verification over the TP-sharded vocabulary in one kernel +(``m3_sampled_accept.sampled_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +from .m3_shard_accept import strict_accept + +#: Verify rows per call. +MAX_ROWS = 48 +#: The fused sampling kernel's 1 / T guard. +TEMP_EPSILON = 1e-6 + + +def nucleus_probs( + logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +class M3SampledAccept(OpWrapper): + """Sampled verification of ``rows = nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard. + + Each row draws one token from its top-p nucleus of softmax(logits / T) over the full (all-rank) vocabulary by + distributed rejection sampling (Philox keyed by the row's seed / offset, the row, the attempt and the CTA), then + the strict acceptance runs against ``draft`` as in ``m3_shard_accept``: ``acc_tokens``, ``num_acc``, + ``gather_ids``, optional ``next_new[:, 0]``. ``tokens`` receives the sampled token per row; ``extra_attempts`` + accumulates exchange rounds beyond the second. The exchange goes through the multicast workspace (two call + halves of ``minimax_m3_sampled_accept_ws_words(world)`` words) with the (finished rows, call) ``counters``. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3 (T = 4 rows per request) at the TP8 GEN worker's batch: 1 / 2 requests (4 / 8 rows, 8 CTAs per + # row), and the 48-row maximum (4 CTAs per row, 12.5 KB of staged logits per CTA); the 25008-column target shard + # (vocab 200064 / 8). Single rank (see note). + CELLS: tuple[Cell, ...] = ( + Cell( + why="b1 (4 rows), top_p 0.8 over a 12-token head: the draws follow the nucleus distribution", + spec=dict( + nreq=1, vs=25008, vocab_offset=0, check="distribution", top_p=0.8, temperature=0.7 + ), + ), + Cell( + why="b2 (8 rows): top_p ~ 0 keeps only the argmax, token exact", + spec=dict( + nreq=2, + vs=25008, + vocab_offset=7 * 25008, + check="argmax", + top_p=1e-7, + temperature=1.0, + ), + ), + Cell( + why="b2 (8 rows) with top_p = 1 (no cut: the first draw stands), next_new and forced counts", + spec=dict( + nreq=2, + vs=25008, + vocab_offset=7 * 25008, + check="distribution", + top_p=1.0, + temperature=1.0, + force="vec", + ), + ), + Cell( + why="AgentX sampling (T 1.0, top_p 0.95) at b1 with forced counts: the draws follow the nucleus", + spec=dict( + nreq=1, + vs=25008, + vocab_offset=0, + check="distribution", + top_p=0.95, + temperature=1.0, + force="total", + ), + ), + Cell( + why="48 rows (4 CTAs per row): the second launch geometry, nucleus distribution at top_p 0.5", + spec=dict( + nreq=12, + vs=25008, + vocab_offset=0, + check="distribution", + top_p=0.5, + temperature=1.3, + force="total", + ), + ), + ) + + note = """ + What is sampled is gated statistically, the rest exactly. Every drawn token must lie in the reference nucleus; + the empirical distribution over many calls (distinct Philox offsets) must sit within a stated total-variation + band of the reference's; the acceptance outputs must equal the strict acceptance of the drawn tokens exactly. + The Philox stream itself is not reproduced in torch: the kernel's draw uses the hardware exp / log + approximations (__expf, __logf) and a CTA-order inverse CDF, which torch does not reproduce bit for bit. + + Certified at world 1 only. The decision is replicated from every rank's per-round records, which another rank + computes from its own shard on the fly, so peers cannot be stood in for on one device; the cross-rank path + needs the multi-device job and is not covered here. Rows with top-k or min-p must not reach this op (it applies + neither); is_valid states it. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept( + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, + ) -> None: + rows = logits.shape[0] + # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the + # wrong rows. + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + # Per-row parameters are read at the row index, unchecked. + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), ( + "per-row parameters must cover every row" + ) + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + # The finished-row counter elects the CTA that runs the acceptance. + assert int(counters[0].item()) == 0, ( + "counters[0] (finished rows) must be zero between calls" + ) + + +m3_sampled_accept = M3SampledAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept_ws_words.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept_ws_words.py new file mode 100644 index 000000000000..3c3974cca0d8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept_ws_words.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Size of the sampled-verify exchange workspace (``m3_sampled_accept.SampledAcceptWorkspace.create``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Exchange rounds per row, verify rows per call, CTAs per row (at most), uint64 words per record. +MAX_ATTEMPTS = 16 +MAX_ROWS = 48 +MAX_SPLITS = 8 +RECORD_WORDS = 4 +#: SA-1R (one_round_k 9 / 82): 9 listed bf16 values ride each record, 3 per extra word. +LIST_WORDS = {0: 0, 9: 3, 82: 3} + + +class M3SampledAcceptWsWords(OpWrapper): + """uint64 words of the multicast workspace ``minimax_m3_sm107tp8_sampled_accept`` exchanges through for + ``world`` ranks: two call halves of ``[attempts][rows][world * splits][record words]``.""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="the tp8 target's workspace", spec=dict(world=8)), + Cell(why="a single rank (local workspace, no multicast)", spec=dict(world=1)), + Cell( + why="the tp8 target's workspace with SA-1R pooled (82)", + spec=dict(world=8, one_round_k=82), + ), + Cell( + why="the tp8 target's workspace with SA-1R top-9 (9)", spec=dict(world=8, one_round_k=9) + ), + ) + + note = """ + A host computation (no device work); it is an op so the workspace size and the kernel's indexing come from one + definition. + """ + + def __call__(self, world: int, one_round_k: int = 0) -> int: + return torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words(world, one_round_k) + + def reference(self, world: int, one_round_k: int = 0) -> int: + """Two halves of attempts x rows x (world x splits) records x record words.""" + return ( + 2 + * MAX_ATTEMPTS + * MAX_ROWS + * world + * MAX_SPLITS + * (RECORD_WORDS + LIST_WORDS[one_round_k]) + ) + + +m3_sampled_accept_ws_words = M3SampledAcceptWsWords() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py new file mode 100644 index 000000000000..6df70b3f467b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Greedy Eagle3 verification from this rank's logits shard: per-row argmax, the cross-rank (value, index) exchange +and the strict acceptance in one kernel (``m3_shard_accept.shard_accept``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Verify rows per call, and exchange rows per call half of its workspace. +SHARD_ROWS = 48 + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out + + +class M3ShardAccept(OpWrapper): + """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` + (global ids ``vocab_offset ..``). + + Each row's token is its global argmax over all ranks' shards: keys ``ordered(value) << 32 | ~index`` (ties take + the lowest index, as ``torch.argmax``) are max-reduced over the shard, pushed into every rank's exchange slot + through the multicast slots ``[2][48][world]`` (Lamport: 0 = empty) and max-reduced over the ranks by the last + CTA, which then runs the strict acceptance: ``acc_tokens [nreq, D + 1]``, ``num_acc [nreq]`` (``force_vec`` / + ``force_total`` override the counts), ``gather_ids [nreq]`` (the first draft pass's row per request) and, when + given, ``next_new[:, 0]`` (the last accepted token). + """ + + ARCHS = frozenset({Arch.SM_107}) + + # Every greedy verify step of the target (modeling's SHARD_ACCEPT): draft_len 3 at the TP8 GEN worker's 1..2 + # requests (4..8 rows), plus the 48-row maximum; the target lm_head's BF16 shard logits (25008 columns per rank + # at TP8, 8 exchange slots per row). The logits come from the head GEMM, so the op only reduces. + CELLS: tuple[Cell, ...] = ( + Cell( + why="b1: 4 verify rows (1 request x T = 4), rank 0 of 8, next_new staged", + spec=dict(nreq=1, draft_len=3, vs=25008, world=8, rank=0, force="none", strided=False), + ), + Cell( + why="b2: 8 verify rows, rank 7, global ids offset by 7 x 25008", + spec=dict(nreq=2, draft_len=3, vs=25008, world=8, rank=7, force="none", strided=False), + ), + Cell( + why="8 rows (2 requests) read from a row-strided view of a wider logits buffer, forced counts", + spec=dict(nreq=2, draft_len=3, vs=25008, world=8, rank=3, force="vec", strided=True), + ), + Cell( + why="the 48-row maximum (12 requests at draft 3), force_total, rows with exact cross-rank ties", + spec=dict( + nreq=12, + draft_len=3, + vs=25008, + world=8, + rank=5, + force="total", + strided=False, + ties=True, + ), + ), + ) + + note = """ + The op reduces BF16 logits exactly, so the reference is bit-exact, including cross-rank ties (lowest global + index wins). The cross-rank exchange is certified on one device: the other ranks' keys are written into this + rank's slots before the call, which is what their multicast pushes would have put there; the multicast mapping + itself (ws_mc_ptr = this rank's slots, as at world 1) is not exercised across devices here. + """ + + def __call__( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + torch.ops.trtllm.minimax_m3_shard_accept( + logits, + vocab_offset, + keys, + counters, + ws_uc, + ws_mc_ptr, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + num_ctas, + force_vec, + ) + + def reference( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Lowest-index maximum of each row's BF16 logits, the largest key among it and the peers' keys in + ``ws_uc``, then the strict acceptance. Reads ``ws_uc`` / ``counters`` before the call; the scratch and + output buffers and ``num_ctas`` do not change the result.""" + best, idx = local_argmax(logits.float(), vocab_offset) + tokens = exchange_tokens(make_keys(best, idx), ws_uc, counters, rank, world, SHARD_ROWS) + out = strict_accept(tokens, draft, batch_idx, force_total, next_new, force_vec) + out["tokens"] = tokens + return out + + def is_valid( + self, + logits: torch.Tensor, + vocab_offset: int, + keys: torch.Tensor, + counters: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + num_ctas: int, + force_vec: Optional[torch.Tensor] = None, + ) -> None: + assert not bool(keys[: logits.shape[0]].any()), "keys must start zeroed" + assert int(counters[0].item()) == 0, ( + "counters[0] (finished CTAs) must be zero between calls" + ) + assert ws_uc.numel() >= 2 * SHARD_ROWS * world, "ws_uc must hold [2][48][world] slots" + # A NaN logit maps to a key above every finite one and wins the row silently. + assert not bool(torch.isnan(logits).any()), "NaN logits win the argmax" + nreq = draft.shape[0] + assert bool(((batch_idx[:nreq] >= 0) & (batch_idx[:nreq] < nreq)).all()), ( + "batch_idx must index the nreq requests" + ) + if force_vec is not None: + fv = force_vec[:nreq] + assert bool(((fv >= 1) & (fv <= draft.shape[1] + 1)).all()), ( + "force_vec counts must lie in [1, D + 1]" + ) + + +m3_shard_accept = M3ShardAccept() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py new file mode 100644 index 000000000000..ff07c54de1f0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py @@ -0,0 +1,72 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> +``m3_spec_tail.slot_scatter``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +def padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: + """src rows zero-padded (or cut) to ``width`` columns.""" + out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) + w = min(width, src.shape[1]) + out[:, :w] = src[:, :w] + return out + + +class M3SlotScatter(OpWrapper): + """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, + ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The + destinations are the store's transposed [slots, width] views.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of + # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) + # is zero-padded. + CELLS: tuple[Cell, ...] = ( + Cell( + why="8 rows into shuffled slots, widths equal to the store's", + spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), + ), + Cell( + why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", + spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), + ), + ) + + note = """ + Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. + """ + + def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): + torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) + + def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): + """The four destinations after the call, on copies.""" + s = slots.long() + out = {} + for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): + r = dst.clone() + r[s] = padded_rows(src[: s.numel()], dst.shape[1]) + out[name] = r + lr = dl.clone() + lr[s] = sl[: s.numel()] + out["dl"] = lr + return out + + def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): + assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( + "slots must index the store" + ) + # Two rows into one slot race (which one lands is unordered). + assert slots.unique().numel() == slots.numel(), "slots must be distinct" + assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + + +m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py new file mode 100644 index 000000000000..56c5f54b1161 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's speculative step correction as one launch (``RuntimeHooks.spec_offsets_op``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm.logger import logger + +# The engagement line is logged once per process (module state: __call__ reads no instance state). +_ENGAGED: list = [] + + +class M3SpecOffsets(OpWrapper): + """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp8_spec_offsets (minimaxM3Sm107Tp8/specOffsets.cu), the + same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other + layouts (the caller runs the ATen adds).""" + + ARCHS = frozenset({Arch.SM_107}) + + CELLS: tuple[Cell, ...] = ( + Cell(why="1 request, 4 tokens (draft length 3)", spec=dict(tokens=4, seqs=1)), + Cell(why="8 requests, 32 tokens", spec=dict(tokens=32, seqs=8)), + Cell(why="wrap-around values", spec=dict(tokens=8, seqs=2, wrap=True)), + ) + + note = "Integer adds: the reference (the two in-place ATen adds) is exact." + + def __call__(self, pos, pos_off, kv, kv_off) -> bool: + ts = (pos, pos_off, kv, kv_off) + if ( + any( + t.dtype != torch.int32 or t.dim() != 1 or t.stride(0) != 1 or not t.is_cuda + for t in ts + ) + or pos.numel() != pos_off.numel() + or kv.numel() != kv_off.numel() + ): + return False + torch.ops.trtllm.minimax_m3_sm107_tp8_spec_offsets(pos, pos_off, kv, kv_off) + if not _ENGAGED: + _ENGAGED.append(True) + logger.info("[M3 host] spec step correction: one launch taken (spec_offsets_op)") + return True + + def reference(self, pos, pos_off, kv, kv_off) -> None: + pos += pos_off + kv += kv_off + + def is_valid(self, pos, pos_off, kv, kv_off) -> None: + assert pos.numel() == pos_off.numel() and kv.numel() == kv_off.numel(), ( + "equal lengths per pair" + ) + + +m3_spec_offsets = M3SpecOffsets() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py new file mode 100644 index 000000000000..a0da8b337574 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch +(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" + +from __future__ import annotations + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS: tuple[Cell, ...] = ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + Cell( + why="an unchanged batch of 3 previous requests (c3) in shuffled slots", + spec=dict(p=3, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="an unchanged batch of 4 previous requests (c4) in shuffled slots", + spec=dict(p=4, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + ) + + note = """ + Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid + states them. + """ + + def __call__( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + torch.ops.trtllm.minimax_m3_step_gather( + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ) + + def reference( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" + nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) + s = slots[:P].long() + ids_r, draft_r = ids.clone(), draft.clone() + ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) + draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) + pos_r = torch.zeros_like(pos_off) + pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] + kv_r = torch.zeros_like(kv_off) + kv_r[b0 : b0 + P] = lens[s] - T + return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} + + def is_valid( + self, + new_tok, + next_draft, + lens, + slots, + ppi, + P, + T, + D, + ids, + ids0, + draft, + d0, + pos_off, + kv_off, + b0, + ): + # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or + # writes another buffer. + s = slots[:P] + assert bool( + ( + (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) + ).all() + ), "slots must index the store" + assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( + "ppi must index lens" + ) + assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" + assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" + assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( + "offset buffers too short" + ) + + +m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py new file mode 100644 index 000000000000..a6c626ff2d2f --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode graph's step-input prologue in one launch (this target's ``RuntimeHooks.pack_prologue_op``).""" + +from __future__ import annotations + +from typing import Optional + +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + + +class M3StepPrologue(OpWrapper): + """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch + (trtllm::minimax_m3_sm107_tp8_step_prologue, minimaxM3Sm107Tp8/stepPrologue.cu). + + ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one + ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the + window, a device copy's a device address). The launch writes up to ``max_desc`` copies of each kind, every host + copy before any device copy, then zeroes the header; with ``epoch`` it increments ``epoch[0]`` once. The default + prologue does the same with two parallel-scatter launches (``step_inputs._graph_scatter_kernel``) and a header + fill.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The TP8 decode step's packs: a few dozen host copies of 4 to a few hundred bytes (positions, KV lengths, slots, + # the spec sampler's uploads), a handful of device copies (the graph runner's static inputs), some reading what a + # host copy of the same window wrote. + CELLS: tuple[Cell, ...] = ( + Cell(why="an empty window (a replay with no new upload)", spec=dict(h2d=0, d2d=0)), + Cell(why="host copies only, odd sizes and alignments", spec=dict(h2d=40, d2d=0)), + Cell( + why="host copies feeding device copies (a device copy reads a host copy's destination)", + spec=dict(h2d=24, d2d=8, chained=True), + ), + Cell( + why="every descriptor slot (128 + 128 of the pack's 256)", spec=dict(h2d=128, d2d=128) + ), + ) + + note = """ + Byte copies: the reference (the default prologue's Triton launches and fill on the same window) is exact. + Destinations are raw addresses; is_valid cannot check them. + """ + + def __call__(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + torch.ops.trtllm.minimax_m3_sm107_tp8_step_prologue(window, epoch, max_desc) + + def reference(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + """The default prologue on the same arguments: it writes the same destinations.""" + from tensorrt_llm._torch.pyexecutor.step_inputs import _graph_scatter_kernel + + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + bump = epoch is not None and not device_src + _graph_scatter_kernel[(max_desc,)]( + window, + stage_i64, + epoch if bump else stage_i64, + DEVICE_SRC=device_src, + BLOCK=256, + BUMP_EPOCH=bump, + num_warps=1, + ) + window[:16].zero_() + + def is_valid(self, window: torch.Tensor, epoch: Optional[torch.Tensor], max_desc: int) -> None: + assert window.dtype == torch.uint8 and window.is_contiguous(), "window: contiguous uint8" + assert window.numel() >= 16 + 24 * max_desc, "window: header and max_desc descriptors" + counts = window[:16].view(torch.int64).tolist() + assert counts[0] + counts[1] <= max_desc, "at most max_desc copies in all" + + +m3_step_prologue = M3StepPrologue() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py new file mode 100644 index 000000000000..e3694274c124 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py @@ -0,0 +1,1754 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Triton block-sparse GQA decode attention for MiniMax-M3. + +Flash-decoding over the blocks the indexer selected: one CTA per (query token, +top-k chunk, KV head) accumulates a partial output plus its log-sum-exp, and a +second kernel merges the chunks by LSE weight. That beats running the +context-schedule FMHA kernel at decode, where a single query token leaves most +of a 128-row Q tile idle. + +GB300/SM103 tuning targets GQA group 16, top-k 16 and four-token Eagle3 +verification. Split-K, query pairing and launch tables require four KV heads; +the linear-softmax policy also covers one and two KV heads. Other hardware +and shapes retain the default policies, including supported SM100 NVFP4. + +Vendored from the vLLM source linked in the file header (v0.26.1rc0-77-g6f91edf96). +Differences from upstream: + +* K and V are separate HND paged views ([num_pages, num_kv_heads, page_size, + head_dim]) with independent strides rather than one cache fused along the + last dim, because that is how the M3 KV pool is laid out here. +* Only the scalar KV-scale mode is kept; M3 stores unscaled E4M3 K/V. +* A packed NVFP4 cache is read in place, dequantized in-register from its E2M1 + nibbles and E4M3 block scales. +* Partial-output and LSE scratch come from the persistent buffer arena so their + addresses survive CUDA graph replay. +* A single top-k chunk writes the final output directly, avoiding scratch and + the otherwise redundant LSE merge. +* The merge kernel zeroes rows whose chunks are all empty instead of letting + them produce NaN, since CUDA-graph padding rows flow on into the rest of the + network here. + +This is the MiniMax-M3 NVFP4 / sm_107 / tp8 target's copy. On top of the +in-tree decode it routes long-context TP8 steps (8 query heads, one KV head) to the one-launch +``m3_attn_lc8`` kernel, short-sequence verify steps of at least +``MQ_MIN_TOKENS`` rows to the multi-query producer, releases PDL dependents +early, and has the merge emit the MXFP8 o_proj hand-off (``m3_prequant.SLOT``). +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + check_decode_span_shape, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers + +from . import m3_attn_lc8, m3_prequant, m3_report + +# One sparse block is exactly one KV page. +SPARSE_BLOCK_SIZE = 128 # capacity: fixed: the MSA block size of the checkpoint + +_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +_TL_DTYPES = { + torch.bfloat16: tl.bfloat16, + torch.float16: tl.float16, + torch.float32: tl.float32, +} + +# Total CTAs the split-K partitioning aims for before it stops splitting. +_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size + +# NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale +# byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. +NVFP4_SF_VEC_SIZE = 16 + + +# Fewest query rows for which the merge writes the MXFP8 o_proj hand-off: every decode step, so the MXFP8 o_proj +# never runs its own quantize kernel after attention. +MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound +# The sparse decode (LC8, the Triton merge) also writes the MXFP8 o_proj hand-off (bit for bit +# torch.ops.trtllm.mxfp8_quantize of its BF16 rows) and publishes it to m3_prequant.SLOT. +MERGE_MXFP8_HANDOFF = True +# The one-launch long-context decode (m3_attn_lc8) at the TP8 heads: its shapes (modeling's SPARSE_DECODE_LC8 keys: +# "prewait" = up to m3_attn_lc8.PREWAIT_MAX_ROWS query rows, "plain" above) and its report path. Short-sequence +# identity steps run it too, taking the identity selection from the device lengths (LC8_IDENTITY, with its merge-once +# form). +LC8_KEYS = frozenset({"plain", "prewait"}) +LC8_PATH = "m3_sparse_decode_lc8" +LC8_IDENTITY = True +# The identity-selection producer's shapes (modeling's SPARSE_DECODE_PRODUCER keys) and its report path. +PRODUCER_KEYS = frozenset({"multi_query", "per_row"}) +PRODUCER_PATH = "sparse_decode_producer" + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and swizzled UE8M0 + scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips its own + quantization. One 128-row buffer per width serves every layer: each o_proj consumes it before the next layer's + merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + +MERGE_MXFP8 = MergeMxfp8() + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + if IDENTITY: + # Slot `it` of the identity selection, bounded by this row's own kv extent above. + blk = it + else: + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, + # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + if IDENTITY: + blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) + blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) + else: + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for cand in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + if IDENTITY: + blk_0 = tl.where(valid_0, idx_0, sentinel) + blk_1 = tl.where(valid_1, idx_1, sentinel) + else: + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the + # pool, reads nothing. + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. Because this trigger precedes this grid's own + # wait, it says nothing about what upstream has finished: a dependent may read only step-invariant data + # (weights) before its griddepcontrol.wait, never this step's attention output, q, top-k table or lengths. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +# NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. +_SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) +# Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) + + +# Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query +# tokens; below it the per-token producer is as fast. +MQ_MIN_TOKENS = 17 # capacity: fixed: the multi-query producer's crossover +# Warps per multi-query producer program. +MQ_NUM_WARPS = 8 + + +def _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + *, + sm_scale, + output, + decode_query_len, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + identity=False, + identity_step=False, +) -> None: + """minimax_m3_sparse_attn_decode through m3_attn_lc8: same table, output and MXFP8 hand-off contract. ``identity`` + takes the short-sequence identity selection from the device lengths (the table is not read); ``identity_step`` + marks a step with no select (the kernel reads q after its grid wait), also when the staged table is read.""" + use_pdl = get_env_enable_pdl() + total_q, num_heads, head_dim = q.shape + mxfp8 = None + m3_prequant.SLOT.clear() + if MERGE_MXFP8_HANDOFF and MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS: + mxfp8 = MERGE_MXFP8.buffers(num_heads * head_dim, output.device) + # The early trigger lets the o_proj launch and stream its weights while attention runs. + m3_attn_lc8.m3_lc8_attn( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + linear_softmax=_sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=1, + gqa_group_size=num_heads, + max_topk=int(topk_idx.shape[-1]), + decode_query_len=decode_query_len, + ), + mxfp8=mxfp8, + use_pdl=use_pdl, + identity=identity, + identity_step=identity_step, + prefetch=m3_attn_lc8.PREFETCH.take(), + next_kv=m3_attn_lc8.NEXT_KV.take(), + ) + if m3_attn_lc8.ONE_DIV_MERGE and m3_attn_lc8.MERGE_ONCE: + m3_report.taken(ONE_DIV_MERGE_PATH, total_q) + if mxfp8 is not None: + m3_prequant.SLOT.publish( + output.view(total_q, num_heads * head_dim), mxfp8[0][:total_q], mxfp8[1] + ) + + +# Engagement path of the identity-selection producers reading their blocks from the kv lengths (IDENTITY). +DEVICE_IDENTITY_PATH = "sparse_decode_device_identity" +# Engagement path of the LC's one-division split-K merge (MSA_ONE_DIV_MERGE). +ONE_DIV_MERGE_PATH = "sparse_decode_lc_one_div_merge" + + +@torch.no_grad() +def minimax_m3_sparse_attn_decode( + q: torch.Tensor, # [total_q, num_heads, head_dim] + k_paged: torch.Tensor, # [num_pages, num_kv_heads, page_size, head_dim] + v_paged: torch.Tensor, # same layout as k_paged + topk_idx: torch.Tensor, # [num_kv_heads, total_q, topk] + block_table: torch.Tensor, # [num_reqs, max_blocks] + seq_lens: torch.Tensor, # [num_reqs] int32 + *, + sm_scale: float, + output: torch.Tensor, # [total_q, num_heads, head_dim] + decode_query_len: int, + kv_scale: Optional[torch.Tensor] = None, + k_block_scale: Optional[torch.Tensor] = None, + v_block_scale: Optional[torch.Tensor] = None, + k_global_scale: Optional[torch.Tensor] = None, + v_global_scale: Optional[torch.Tensor] = None, + num_topk_chunks: Optional[int] = None, + identity_blocks: bool = False, +) -> None: + """Block-sparse GQA decode attention, written into output in place. + + q may be FP8, as a fused QK-norm/RoPE producer emits it. The kernel widens + it to output.dtype in-register rather than making the caller materialize a + widened copy, which would cost a standalone kernel on every sparse layer. + The attention math is unchanged either way, since E4M3 widens exactly. + + kv_scale is an optional scalar dequantization factor for an FP8 cache; + MiniMax-M3 stores unscaled E4M3, so it is normally None. + + Passing the four NVFP4 scale tensors selects the packed-NVFP4 cache instead: + k_paged/v_paged then hold two E2M1 elements per byte, the block scales are + paged E4M3 bytes in the layouts msa_scatter writes, and the two per-tensor + dequant scales are single-element FP32 device tensors. + + num_topk_chunks overrides the split-K factor and exists for tests: the + merged result must not depend on it. + """ + total_q, num_heads, head_dim = q.shape + num_kv_heads = int(k_paged.shape[1]) + check_decode_span_shape( + "MiniMax-M3 Triton sparse decode", total_q, int(seq_lens.shape[0]), decode_query_len + ) + if int(k_paged.shape[2]) != SPARSE_BLOCK_SIZE: + raise ValueError( + f"MiniMax-M3 sparse decode requires page_size={SPARSE_BLOCK_SIZE}; " + f"got {int(k_paged.shape[2])}." + ) + if num_heads % num_kv_heads: + # The kernel covers heads pid_kh * gqa_group_size + [0, group), so a + # head count that is not a whole number of groups leaves the tail heads + # unwritten and the merge kernel reads uninitialized partials. + raise ValueError( + f"MiniMax-M3 sparse decode requires num_heads ({num_heads}) to be a " + f"multiple of num_kv_heads ({num_kv_heads})." + ) + nvfp4_args = (k_block_scale, v_block_scale, k_global_scale, v_global_scale) + kv_nvfp4 = any(arg is not None for arg in nvfp4_args) + scale_cols = head_dim // NVFP4_SF_VEC_SIZE + if kv_nvfp4: + capability = torch.cuda.get_device_capability(q.device) + if capability not in _SM100F_CAPABILITIES: + raise NotImplementedError( + "MiniMax-M3 NVFP4 sparse decode requires SM100/SM103/SM107; " + f"got SM{capability[0] * 10 + capability[1]}." + ) + _check_nvfp4_inputs(k_paged, v_paged, nvfp4_args, head_dim=head_dim, scale_cols=scale_cols) + # Two E2M1 elements per byte; Triton must see them unsigned so that the + # high-nibble shift does not sign-extend. + k_paged, v_paged = k_paged.view(torch.uint8), v_paged.view(torch.uint8) + k_block_scale, v_block_scale = ( + k_block_scale.view(torch.uint8), + v_block_scale.view(torch.uint8), + ) + max_topk = int(topk_idx.shape[-1]) + gqa_group_size = num_heads // num_kv_heads + use_scale = k_paged.dtype in _FP8_DTYPES and kv_scale is not None + # Triton needs a real pointer even for the unused arguments. + scale_arg = kv_scale if use_scale else output + k_sc_arg, v_sc_arg = (k_block_scale, v_block_scale) if kv_nvfp4 else (output, output) + k_gs_arg, v_gs_arg = (k_global_scale, v_global_scale) if kv_nvfp4 else (output, output) + + widen_q = q.dtype in _FP8_DTYPES + if widen_q and output.dtype not in _TL_DTYPES: + raise ValueError( + f"MiniMax-M3 sparse decode cannot widen FP8 q into {output.dtype}; " + f"supported compute dtypes are {sorted(d.__str__() for d in _TL_DTYPES)}." + ) + # Packed NVFP4 dequantizes to FP16. The production producer supplies FP8 + # q, which widens exactly to FP16, so keep both dot operands there instead + # of converting every K/V tile to BF16. Accumulation remains FP32. + if widen_q: + compute_dtype = tl.float16 if kv_nvfp4 else _TL_DTYPES[output.dtype] + else: + compute_dtype = tl.float32 + + # Long-context TP8 steps: attention, split-K merge and MXFP8 hand-off in one launch. A selected table on an NVFP4 + # cache with FP8 q at the TP8 heads (8 query heads, one KV head) is the kernel's (LC8_KEYS: "prewait" up to + # m3_attn_lc8.PREWAIT_MAX_ROWS query rows); operands it declines are a miss. + lc_key = "prewait" if total_q <= m3_attn_lc8.PREWAIT_MAX_ROWS else "plain" + lc_covered = ( + kv_nvfp4 + and widen_q + and (not identity_blocks or (LC8_IDENTITY and m3_attn_lc8.MERGE_ONCE)) + and num_topk_chunks is None + and q.dim() == 3 + and int(q.shape[1]) == m3_attn_lc8.NUM_HEADS + and num_kv_heads == 1 + and lc_key in LC8_KEYS + ) + if lc_covered and not m3_attn_lc8.lc8_applies( + q, k_paged, v_paged, topk_idx, k_block_scale, v_block_scale, output + ): + m3_report.missed( + LC8_PATH, + total_q, + "the step's q / KV pages / scales / top-k table / output are not the layout the kernel is built for", + ) + lc_covered = False + if lc_covered: + _longattn_lc_decode( + q, + k_paged, + v_paged, + topk_idx, + block_table, + seq_lens, + sm_scale=sm_scale, + output=output, + decode_query_len=decode_query_len, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, + k_global_scale=k_global_scale, + v_global_scale=v_global_scale, + # The identity selection from the device lengths (the staged table is not read). + identity=identity_blocks, + # Every identity step (no select ran): q after the grid wait. + identity_step=identity_blocks, + ) + m3_report.taken(LC8_PATH, total_q) + return + + query_group_size = ( + _sm100f_nvfp4_query_group_size( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else 1 + ) + if num_topk_chunks is None: + if query_group_size == 2: + measured_chunks = 8 + else: + measured_chunks = ( + _sm100f_nvfp4_num_topk_chunks( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else None + ) + num_topk_chunks = ( + measured_chunks + if measured_chunks is not None + else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) + ) + elif num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + + # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one + # program per (request, block) covers all of them and dequantizes each block once. + use_mq = ( + identity_blocks + and total_q >= MQ_MIN_TOKENS + and kv_nvfp4 + and num_kv_heads == 1 + and decode_query_len > 1 + and head_dim == 128 + and gqa_group_size * decode_query_len <= 128 + ) + if identity_blocks and kv_nvfp4 and widen_q: + # The identity-selection route is modeling's SPARSE_DECODE_PRODUCER ("per_row"; "multi_query" needs one KV + # head, so it never runs at the TP2 heads); a variant the range does not certify is a miss (raises under + # require). + producer_key = "multi_query" if use_mq else "per_row" + if producer_key not in PRODUCER_KEYS: + m3_report.missed( + PRODUCER_PATH, + total_q, + f"the identity-selection producer variant {producer_key!r} is not certified at {num_kv_heads} KV heads", + ) + if use_mq: + num_topk_chunks = max_topk + direct_output = num_topk_chunks == 1 + if direct_output: + # With one chunk, acc_o is already normalized and final. Point the + # producer at the caller's output and compile away the LSE store. + # Keeping a real dummy pointer preserves Triton's pointer contract. + o_arg = output + lse_arg = output + o_strides = (0, output.stride(0), output.stride(1), output.stride(2)) + lse_strides = (0, 0, 0) + else: + # Persistent arena rather than torch.empty, so the partials keep one + # address across CUDA graph replays. FP32 keeps the split-K factor from + # perturbing the merged result. + reserve = torch.cuda.is_current_stream_capturing() + o_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads, head_dim], + torch.float32, + buffer_name="m3_sparse_decode_o_partial", + reserve_buffer=reserve, + ) + lse_arg = get_memory_buffers().get_buffer( + [num_topk_chunks, total_q, num_heads], + torch.float32, + buffer_name="m3_sparse_decode_lse_partial", + reserve_buffer=reserve, + ) + o_strides = tuple(o_arg.stride()) + lse_strides = tuple(lse_arg.stride()) + + use_pdl = get_env_enable_pdl() + pdl_launch = {"launch_pdl": True} if use_pdl else {} + use_linear_softmax = kv_nvfp4 and _sm100f_nvfp4_use_linear_softmax( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if query_group_size == 2: + launch_options = {"num_warps": 4, "num_stages": 1} + else: + launch_options = ( + _sm100f_nvfp4_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + + # The pair path merges the two sorted selector rows on device. Equal blocks + # share one K/V dequantization; divergent blocks update only their owning + # query, so the optimization does not rely on runtime selection equality. + if total_q % query_group_size or decode_query_len % query_group_size: + raise ValueError( + f"query group {query_group_size} must divide total_q={total_q} and " + f"decode_query_len={decode_query_len}." + ) + + producer_kernel = ( + _gqa_sparse_decode_pair_kernel if query_group_size == 2 else _gqa_sparse_decode_kernel + ) + device_identity = identity_blocks + if device_identity and not use_mq: + m3_report.taken(DEVICE_IDENTITY_PATH, total_q) + single_kernel_args = ( + { + # No split-K arrival counters: the separate merge kernel combines the chunks. + "merge": (output, None), + "tail": (output.stride(0), output.stride(1), output.stride(2), num_kv_heads), + # An early release pays off when the dependent is the consumer of the + # final output (it can prefetch), or a merge grid that itself releases + # its dependent at entry: the o_proj then launches during this kernel and + # streams its weights while attention runs. Above 8 query rows the + # concurrent weight stream slows the larger attention grid more than the + # prefetch saves (GB300 TP4: -1.1% at 48 rows). + "options": {"FUSED_MERGE": False, "EARLY_TRIGGER": direct_output or total_q <= 8}, + } + if query_group_size == 1 + else {"merge": (), "tail": (), "options": {}} + ) + pair_kernel_options = ( + {"QUERY_GROUP_SIZE": query_group_size, "MAX_TOPK": max_topk} + if query_group_size == 2 + else {} + ) + if use_mq: + from .triton_sparse_decode_mq import launch_mq_producer + + launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + sm_scale=sm_scale, + decode_query_len=decode_query_len, + use_pdl=use_pdl, + early_trigger=False, + num_warps=MQ_NUM_WARPS, + ) + else: + producer_kernel[ + (total_q // query_group_size) * num_topk_chunks, + num_kv_heads, + ]( + q, + k_paged, + v_paged, + scale_arg, + k_sc_arg, + v_sc_arg, + k_gs_arg, + v_gs_arg, + topk_idx, + o_arg, + lse_arg, + *single_kernel_args["merge"], + block_table, + seq_lens, + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + q.stride(0), + q.stride(1), + q.stride(2), + k_paged.stride(0), + k_paged.stride(1), + k_paged.stride(2), + k_paged.stride(3), + v_paged.stride(0), + v_paged.stride(1), + v_paged.stride(2), + v_paged.stride(3), + k_sc_arg.stride(0), + k_sc_arg.stride(1), + v_sc_arg.stride(0), + v_sc_arg.stride(1), + topk_idx.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + *o_strides, + *lse_strides, + block_table.stride(0), + block_table.shape[1], + k_paged.shape[0], + *single_kernel_args["tail"], + BLOCK_SIZE_K=SPARSE_BLOCK_SIZE, + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_SCALE=use_scale, + WIDEN_Q=widen_q, + COMPUTE_DTYPE=compute_dtype, + USE_PDL=use_pdl, + KV_NVFP4=kv_nvfp4, + USE_LINEAR_SOFTMAX=use_linear_softmax, + SCALE_COLS=scale_cols, + DIRECT_OUTPUT=direct_output, + # Identity selection from the lengths this kernel reads, not the staged table. + IDENTITY=device_identity, + **single_kernel_args["options"], + **pair_kernel_options, + **pdl_launch, + **launch_options, + ) + m3_prequant.SLOT.clear() + if not direct_output: + width = num_heads * head_dim + mxfp8 = ( + MERGE_MXFP8.buffers(width, output.device) + if MERGE_MXFP8_HANDOFF + and MERGE_MXFP8_MIN_TOKENS <= total_q <= MergeMxfp8.MAX_ROWS + and head_dim == 128 + and output.dim() == 3 + and output.is_contiguous() + and output.dtype == torch.bfloat16 + else None + ) + mxfp8_args = ( + { + "q8_ptr": mxfp8[0], + "sf_ptr": mxfp8[1], + "num_sf_k_tiles": width // 128, + "MXFP8_OUT": True, + } + if mxfp8 is not None + else {} + ) + merge_launch_options = ( + _sm100f_nvfp4_merge_launch_options( + total_q=total_q, + num_kv_heads=num_kv_heads, + gqa_group_size=gqa_group_size, + max_topk=max_topk, + decode_query_len=decode_query_len, + ) + if kv_nvfp4 + else {} + ) + _merge_topk_attn_out_kernel[(total_q, num_heads)]( + o_arg, + lse_arg, + output, + head_dim, + *o_strides, + *lse_strides, + output.stride(0), + output.stride(1), + output.stride(2), + NUM_TOPK_CHUNKS=num_topk_chunks, + USE_PDL=use_pdl, + # Release the o_proj at the merge's entry so it streams its weights during the merge. + EARLY_TRIGGER=True, + **mxfp8_args, + **pdl_launch, + **merge_launch_options, + ) + if mxfp8 is not None: + m3_prequant.SLOT.publish(output.view(total_q, width), mxfp8[0][:total_q], mxfp8[1]) + + +__all__ = [ + "MERGE_MXFP8", + "MERGE_MXFP8_MIN_TOKENS", + "MergeMxfp8", + "NVFP4_SF_VEC_SIZE", + "SPARSE_BLOCK_SIZE", + "_sm100f_nvfp4_merge_launch_options", + "_sm100f_nvfp4_num_topk_chunks", + "_sm100f_nvfp4_query_group_size", + "_sm100f_nvfp4_use_linear_softmax", + "minimax_m3_sparse_attn_decode", + "resolve_num_topk_chunks", +] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py new file mode 100644 index 000000000000..048d31b82d68 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multi-query block-sparse decode producer for MiniMax-M3 short-sequence verify steps. + +In short-sequence mode every request attends at most ``topk`` KV blocks and block selection is the identity, so +all ``decode_query_len`` query tokens of a request read the same blocks. The per-token producer +(``_gqa_sparse_decode_kernel``) then dequantizes every block once per query token. This producer runs one program +per (request, block) for all of the request's query tokens: rows are (token, head) pairs, each block is +dequantized once, and the causal limit is a per-row mask. + +It writes the same per-block partials (output scaled by 1/l and the V global scale, log2 LSE) as the per-token +producer with one block per chunk, so ``_merge_topk_attn_out_kernel`` merges them unchanged. +""" + +import triton +import triton.language as tl + +from .triton_sparse_decode import _dequant_nvfp4_rows + + +@triton.jit( + do_not_specialize=["num_reqs"], do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"] +) +def _gqa_sparse_decode_mq_kernel( + q_ptr, # [num_reqs * QLEN, num_heads, head_dim], FP8 or FP16/BF16 + k_ptr, # packed NVFP4 K: [num_pages, 1, page_size, head_dim / 2] uint8 + v_ptr, # packed NVFP4 V: same layout + k_sc_ptr, # K block scales [num_pages, 1, page_size, head_dim / 16], token-linear + v_sc_ptr, # V block scales, 4x4 token-quad order + k_gs_ptr, + v_gs_ptr, + t_ptr, # topk table [1, total_q, topk]; the request's last token's row is read + o_ptr, # partials [topk, total_q, num_heads, head_dim] fp32 + lse_ptr, # partials [topk, total_q, num_heads] fp32 (log2) + block_table_ptr, + seq_lens, + num_reqs, + sm_scale, + stride_qn, + stride_qh, + stride_k_blk, + stride_k_pos, + stride_v_blk, + stride_v_pos, + stride_ksc_blk, + stride_vsc_blk, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + QLEN: tl.constexpr, + GROUP: tl.constexpr, # query heads per KV head (16) + BLOCK_R: tl.constexpr, # next power of two of QLEN * GROUP + HEAD_DIM: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + SCALE_COLS: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr, +): + pid_r = tl.program_id(0) + pid_c = tl.program_id(1) # top-k position == block index under identity selection + sm_scale_log2e = sm_scale * 1.4426950409 * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + + off_r = tl.arange(0, BLOCK_R) + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, HEAD_DIM) + off_dp = tl.arange(0, HEAD_DIM // 2) + off_s = tl.arange(0, SCALE_COLS) + row_t = off_r // GROUP + row_h = off_r % GROUP + r_mask = off_r < QLEN * GROUP + tok = pid_r * QLEN + row_t + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + pid_r) + # Token t of the request sits at seq_len - QLEN + t and attends positions [0, kv_len_t). + kv_len = tl.maximum(seq_len - QLEN + row_t + 1, 0) + blk = tl.load(t_ptr + (pid_r * QLEN + QLEN - 1) * stride_tn + pid_c * stride_tk) + blk_ok = blk >= 0 + blk = tl.where(blk_ok, blk, 0) + page = tl.load(block_table_ptr + pid_r * stride_bt_b + blk, mask=blk_ok, other=0).to(tl.int64) + # Rows past the last token's position are never read: their cache bytes may be stale (0 * NaN in the PV dot). + n_valid = tl.where( + blk_ok, tl.minimum(tl.maximum(seq_len - blk * BLOCK_SIZE_K, 0), BLOCK_SIZE_K), 0 + ) + n_mask = off_n < n_valid + + q = tl.load( + q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :], + mask=r_mask[:, None], + other=0.0, + ).to(tl.float16) + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + page * stride_k_blk + off_n[:, None] * stride_k_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + page * stride_ksc_blk + off_n[:, None] * SCALE_COLS + off_s[None, :], + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ) + valid = ( + (blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None]) & n_mask[None, :] & r_mask[:, None] + ) + qk = tl.dot(q, tl.trans(k_deq).to(tl.float16)) * sm_scale_log2e + qk = tl.where(valid, qk, float("-inf")) + m = tl.max(qk, axis=1) + m_safe = tl.where(m == float("-inf"), 0.0, m) + p = tl.exp2(qk - m_safe[:, None]) + p_sum = tl.sum(p, axis=1) + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + page * stride_v_blk + off_n[:, None] * stride_v_pos + off_dp[None, :], + mask=n_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=n_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=HEAD_DIM, + SCALE_COLS=SCALE_COLS, + SF_VEC=HEAD_DIM // SCALE_COLS, + ).to(tl.float16) + acc = tl.dot(p.to(tl.float16), v) + # An empty partial stores zero with LSE -inf, so the merge gives it weight 0. + scale = tl.where(p_sum > 0.0, 1.0 / p_sum, 0.0) * v_global_scale + lse = tl.where(p_sum > 0.0, m + tl.log2(p_sum), float("-inf")) + acc = acc * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + tok[:, None] * stride_o_b + + row_h[:, None] * stride_o_h + + off_d[None, :], + acc, + mask=r_mask[:, None], + ) + tl.store(lse_ptr + pid_c * stride_l_c + tok * stride_l_b + row_h * stride_l_h, lse, mask=r_mask) + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +def launch_mq_producer( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + block_table, + seq_lens, + o_arg, + lse_arg, + *, + sm_scale, + decode_query_len, + use_pdl, + early_trigger, + num_warps=8, + num_stages=1, +): + """Launch the multi-query producer; arguments as minimax_m3_sparse_attn_decode after its NVFP4 views.""" + total_q, num_heads, head_dim = q.shape + num_reqs = total_q // decode_query_len + max_topk = int(topk_idx.shape[-1]) + rows = decode_query_len * num_heads + pdl_launch = {"launch_pdl": True} if use_pdl else {} + _gqa_sparse_decode_mq_kernel[(num_reqs, max_topk)]( + q, + k_paged, + v_paged, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk_idx, + o_arg, + lse_arg, + block_table, + seq_lens, + num_reqs, + sm_scale, + q.stride(0), + q.stride(1), + k_paged.stride(0), + k_paged.stride(2), + v_paged.stride(0), + v_paged.stride(2), + k_block_scale.stride(0), + v_block_scale.stride(0), + topk_idx.stride(1), + topk_idx.stride(2), + o_arg.stride(0), + o_arg.stride(1), + o_arg.stride(2), + lse_arg.stride(0), + lse_arg.stride(1), + lse_arg.stride(2), + block_table.stride(0), + QLEN=decode_query_len, + GROUP=num_heads, + BLOCK_R=triton.next_power_of_2(rows), + HEAD_DIM=head_dim, + BLOCK_SIZE_K=int(k_paged.shape[2]), + SCALE_COLS=head_dim // 16, + USE_PDL=use_pdl, + EARLY_TRIGGER=early_trigger, + num_warps=num_warps, + num_stages=num_stages, + **pdl_launch, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/weights.py new file mode 100644 index 000000000000..871d42635e24 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/weights.py @@ -0,0 +1,884 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Weight tables and loaders: MiniMax-M3 NVFP4 / sm_107 / tp8. + +Two tables, one per checkpoint the target reads: + +* ``MODEL_WEIGHTS`` (``MiniMaxM3Tp8Weights``): the vision-language checkpoint's language model, + ``language_model.*``. Every rank is handed the whole checkpoint, so every tensor-parallel split is this + table's: the Q heads, the K/V/index-Q heads (4 each over 8 ranks: rank r holds head r // 2, as the built-in + loader duplicates them), the dense MLPs and the shared expert's + intermediate, the o_proj / down-projection input columns, every routed expert's intermediate (all 128 + experts live on every rank), and the vocabulary of the shell's ``lm_head``. The index-K head is replicated. +* ``DRAFT_WEIGHTS`` (``MiniMaxM3Tp8DraftWeights``): the one-layer BF16 Eagle3 drafter's own checkpoint + (``speculative_model``), loaded through the shell's ``load_draft_weights`` when speculation is on. + +Storage is HF ``[out, in]`` row-major. Five families need a transform, each a layout a kernel demands: + +* **MXFP8 block scales** (every attention projection, the dense MLPs and the shared expert): the checkpoint + stores UE8M0 bytes ``[N, K / 32]``; the GEMMs read the 128x4-swizzled order + (``trtllm::block_scale_interleave``). A fused projection concatenates its shards' scales first and swizzles + once. +* **The five-way sparse-layer projection** ``[Q | K | V | index-Q | index-K]`` and the three-way dense one + ``[Q | K | V]``: the per-rank shards concatenated along N. +* **Gated MLPs** (dense and shared expert): ``[gate ; up]`` rows. +* **NVFP4 routed experts** (layers 3..59): ``[up (w3) ; gate (w1)]`` rows, then the trtllm-gen row shuffle of + the weight and of its E4M3 scales (gated-activation reorder + 128-row epilogue tile shuffle, + ``trtllmgen_maybe_get_cached_*_permute_indices``) and the 128x4 swizzle of the scales. +* **Gemma boundary norms**: ``input_layernorm``, ``post_attention_layernorm`` and ``model.norm`` are stored as + Gemma norms, ``(1 + w) * x``; the fused all-reduce + RMSNorm kernels apply ``w * x``, so ``1`` is added + (in BF16) at load. The per-head Q/K/index norms stay Gemma: their kernels apply the offset themselves. + +The routed experts' per-tensor scalars (input scales and second-level weight scales, every expert) load raw; +``modeling`` derives the kernels' global scales, alphas and activation constants from them once the weights +are real. + +Loading contract: ``load(model, weights)`` consumes the engine-provided checkpoint dict, fills every declared +parameter exactly once and asserts coverage in both directions: every declared parameter written, and every +checkpoint key either consumed or in ``expected_unconsumed`` (the vision tower and its projector, which this +text-only target does not build, and the language model's multi-token-prediction head). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, Callable + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.moe.fused_moe.quantization import ( + trtllmgen_maybe_get_cached_w2_permute_indices, + trtllmgen_maybe_get_cached_w3_w1_permute_indices, +) + +#: ``layers=ALL_LAYERS`` is the common case: a weight every decoder layer carries. +ALL_LAYERS = object() + +_LM = "language_model." +#: Checkpoint key families this target does not build: the vision tower and its two projectors (text only), and the +#: language model's multi-token-prediction head (speculation, when on, is the separate Eagle3 drafter). +_UNBUILT_PREFIXES = ( + "vision_tower.", + "multi_modal_projector.", + "patch_merge_mlp.", + _LM + "model.mtp.", +) + +MX_BLOCK = 32 # MXFP8: one UE8M0 scale per 32 elements along K +NV_BLOCK = 16 # NVFP4: one E4M3 scale per 16 elements along K +EPILOGUE_TILE_M = ( + 128 # the trtllm-gen MoE runner's row tile: fixes its weight and scale row shuffles +) + + +@dataclass(frozen=True) +class W: + """One weight role: what it is, where its bytes come from, what happens on the way. + + ``shape`` takes the dimension bundle and returns the declared shape. ``src`` is a callable + ``(d, i) -> [(key or key tuple, index into the param | None), ...]`` with keys relative to the checkpoint + root; a key tuple hands the transform those tensors together. ``layers`` is ``ALL_LAYERS``, ``None`` (not + per-layer) or ``d -> layer indices``. ``transform(src, d)`` runs on the destination's device before the copy. + """ + + name: str + shape: Callable[[Any], tuple[int, ...]] + src: Callable + dtype: torch.dtype | None = None + transform: Callable | None = None + layers: object | None = ALL_LAYERS + + def layer_indices(self, d: SimpleNamespace) -> list[int] | None: + if self.layers is None: + return None + if self.layers is ALL_LAYERS: + return list(range(d.num_layers)) + return list(self.layers(d)) + + +# ------------------------------------------------------------------------------------------------ shards + + +def _col( + t: torch.Tensor, d: SimpleNamespace, tp: int | None = None, rank: int | None = None +) -> torch.Tensor: + """This rank's rows (column-parallel output split), even split, sliced on the host and moved to the + destination device (only this rank's bytes are read).""" + tp = d.tp if tp is None else tp + rank = d.rank if rank is None else rank + n = t.shape[0] // tp + return _dev(t[rank * n : (rank + 1) * n], d) + + +def _heads(t: torch.Tensor, d: SimpleNamespace, head_dim: int) -> torch.Tensor: + """This rank's rows of a projection split by heads of ``head_dim`` rows: an even split when there are at least as + many heads as ranks, else the head of rank // (ranks per head), duplicated on those ranks.""" + n = t.shape[0] // head_dim + if n >= d.tp: + return _col(t, d) + h = d.rank // (d.tp // n) + return _dev(t[h * head_dim : (h + 1) * head_dim], d) + + +def _row(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """This rank's input columns (row-parallel input split), even split, on the destination device.""" + n = t.shape[1] // d.tp + return _dev(t[:, d.rank * n : (d.rank + 1) * n], d) + + +def _dev(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """``t`` on the destination device (``d.dev``, set per fill).""" + return t.to(d.dev, non_blocking=True) + + +def _swizzle(scale: torch.Tensor) -> torch.Tensor: + """128x4 swizzle of a row-major ``[N, K / 32]`` UE8M0 scale matrix, flat.""" + return torch.ops.trtllm.block_scale_interleave(scale.contiguous()) + + +# --------------------------------------------------------------------------------------------- transforms + + +def _gemma_fold(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """A Gemma boundary norm as the plain norm the fused all-reduce epilogue applies: ``w + 1`` in BF16.""" + return t + 1.0 + + +def _f32(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return t.to(torch.float32) + + +def _qkv_sparse_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[Q | K | V | index-Q | index-K]``: Q by attention heads, K/V/index-Q by their heads, index-K replicated.""" + q, k, v, iq, ik = t + hd, xd = d.head_dim, d.index_dim + return torch.cat( + [_col(q, d), _heads(k, d, hd), _heads(v, d, hd), _heads(iq, d, xd), _dev(ik, d)], dim=0 + ).contiguous() + + +def _qkv_sparse_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v, iq, ik = t + hd, xd = d.head_dim, d.index_dim + return _swizzle( + torch.cat( + [_col(q, d), _heads(k, d, hd), _heads(v, d, hd), _heads(iq, d, xd), _dev(ik, d)], dim=0 + ) + ) + + +def _qkv_dense_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0).contiguous() + + +def _qkv_dense_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return _swizzle(torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0)) + + +def _index_qk_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """``[index-Q | index-K]`` in BF16: this rank's index-Q heads and the replicated index-K, the E4M3 values cast + (the scale tensors, the last two sources, are consumed only).""" + iq, ik = t[0], t[1] + return torch.cat([_heads(iq, d, d.index_dim), _dev(ik, d)], dim=0).to(d.dtype).contiguous() + + +def _gate_up_w(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return torch.cat([_col(gate, d), _col(up, d)], dim=0).contiguous() + + +def _gate_up_s(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + gate, up = t + return _swizzle(torch.cat([_col(gate, d), _col(up, d)], dim=0)) + + +def _row_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _row(t, d).contiguous() + + +def _row_s(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _swizzle(_row(t, d)) + + +def _col_w(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + return _col(t, d).contiguous() + + +#: Permute indices of the trtllm-gen expert relayouts, keyed by (shape, family, scale group): one expert's are +#: every expert's, so each is computed once per load. +_PERMUTE_CACHE: dict = {} + + +def _expert_w31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 operand: ``[up ; gate]`` packed-FP4 rows of this rank, then the gated-activation row + reorder and the epilogue-tile shuffle.""" + up, gate = t + w = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + w, _PERMUTE_CACHE, EPILOGUE_TILE_M, is_gated_act_gemm=True + ) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s31(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + """One expert's FC1 E4M3 scales: ``[up ; gate]`` rows of this rank, the same row reorder at scale + granularity, then the 128x4 swizzle.""" + up, gate = t + s = torch.cat([_col(up, d), _col(gate, d)], dim=0).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w3_w1_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK, is_gated_act_gemm=True + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled.view(torch.uint8).reshape(shape)) + .view(torch.float8_e4m3fn) + .reshape(shape) + ) + + +def _expert_w2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 operand: this rank's input columns, then the epilogue-tile row shuffle.""" + w = _row(t, d).contiguous() + perm = trtllmgen_maybe_get_cached_w2_permute_indices(w, _PERMUTE_CACHE, EPILOGUE_TILE_M) + return torch.ops.trtllm.shuffle_matrix(w, perm.to(w.device)).view(w.dtype) + + +def _expert_s2(t: torch.Tensor, d: SimpleNamespace) -> torch.Tensor: + """One expert's FC2 E4M3 scales: this rank's columns, the row shuffle at scale granularity, the swizzle.""" + s = _row(t, d).contiguous() + shape = s.shape + sf = s.view(torch.uint8) + perm = trtllmgen_maybe_get_cached_w2_permute_indices( + sf, _PERMUTE_CACHE, EPILOGUE_TILE_M, num_elts_per_sf=NV_BLOCK + ) + shuffled = torch.ops.trtllm.shuffle_matrix(sf, perm) + return ( + torch.ops.trtllm.block_scale_interleave(shuffled).view(torch.float8_e4m3fn).reshape(shape) + ) + + +# ------------------------------------------------------------------------------------------------ sources + + +def _one(template: str): + return lambda d, i, t=template: [(_LM + t.format(i=i), None)] + + +def _many(*templates: str): + return lambda d, i, ts=templates: [(tuple(_LM + t.format(i=i) for t in ts), None)] + + +def _per_expert(*suffixes: str): + """One source per routed expert, written at its expert index (all experts are local at tp8).""" + + def build(d, i, sx=suffixes): + out = [] + for e in range(d.num_experts): + keys = tuple(_LM + f"model.layers.{i}.block_sparse_moe.experts.{e}.{s}" for s in sx) + out.append((keys if len(keys) > 1 else keys[0], (e,))) + return out + + return build + + +def _dense_layers(d): + return range(d.first_sparse) + + +def _sparse_layers(d): + return range(d.first_sparse, d.num_layers) + + +def _three_way_qkv_layers(d): + """Layers whose QKV is the three-way [Q | K | V]: the dense layers, and every layer without the five-way packed + projection (``fuse_qkv_index_projection=False``).""" + return range(d.first_sparse) if d.fused else range(d.num_layers) + + +def _five_way_qkv_layers(d): + return range(d.first_sparse, d.num_layers) if d.fused else range(0) + + +def _unfused_sparse_layers(d): + return range(0) if d.fused else range(d.first_sparse, d.num_layers) + + +_ATTN = "model.layers.{i}.self_attn" +_MOE = "model.layers.{i}.block_sparse_moe" + + +class MiniMaxM3Tp8Weights: + """MiniMax-M3 NVFP4 / sm_107 / tp8: the language model's table and the passes over it. + + ``declare`` allocates every weight it names, ``manifest`` maps checkpoint keys onto them, ``load`` fills + them; ``_targets`` is the one place that decides which parameters exist. + """ + + #: Free the per-expert relayout scratch after each expert stack. + RELEASE_AFTER = "_e_s2" + + WEIGHTS: tuple[W, ...] = ( + # --- the two boundary norms of every layer (Gemma, folded) --- + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.input_layernorm.weight"), + transform=_gemma_fold, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_one("model.layers.{i}.post_attention_layernorm.weight"), + transform=_gemma_fold, + ), + # --- attention: per-head Q/K norms (Gemma, unfolded), the packed QKV and the output projection --- + W("q_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".q_norm.weight")), + W("k_norm", shape=lambda d: (d.head_dim,), src=_one(_ATTN + ".k_norm.weight")), + W( + "qkv_w", + shape=lambda d: (d.qkv_dense_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in "qkv")), + transform=_qkv_dense_w, + layers=_three_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_dense_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many(*(f"{_ATTN}.{p}_proj.weight_scale_inv" for p in "qkv")), + transform=_qkv_dense_s, + layers=_three_way_qkv_layers, + ), + W( + "qkv_w", + shape=lambda d: (d.qkv_sparse_rows, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many(*(f"{_ATTN}.{p}_proj.weight" for p in ("q", "k", "v", "index_q", "index_k"))), + transform=_qkv_sparse_w, + layers=_five_way_qkv_layers, + ), + W( + "qkv_s", + shape=lambda d: (d.qkv_sparse_rows * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + *( + f"{_ATTN}.{p}_proj.weight_scale_inv" + for p in ("q", "k", "v", "index_q", "index_k") + ) + ), + transform=_qkv_sparse_s, + layers=_five_way_qkv_layers, + ), + # The separate BF16 index projection [index-Q | index-K] of the unfused sparse layers. The built-in model + # builds it as an unquantized Linear and fills it through the fused gate/up loader, which casts the + # checkpoint's E4M3 values to BF16 and does not apply their block scales: so does this table (the scale + # tensors are consumed, not applied). + W( + "idx_w", + shape=lambda d: (d.index_heads * d.index_dim + d.index_dim, d.hidden), + src=_many( + _ATTN + ".index_q_proj.weight", + _ATTN + ".index_k_proj.weight", + _ATTN + ".index_q_proj.weight_scale_inv", + _ATTN + ".index_k_proj.weight_scale_inv", + ), + transform=_index_qk_w, + layers=_unfused_sparse_layers, + ), + W( + "idx_q_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_q_norm.weight"), + layers=_sparse_layers, + ), + W( + "idx_k_norm", + shape=lambda d: (d.index_dim,), + src=_one(_ATTN + ".index_k_norm.weight"), + layers=_sparse_layers, + ), + W( + "o_w", + shape=lambda d: (d.hidden, d.q_width), + dtype=torch.float8_e4m3fn, + src=_one(_ATTN + ".o_proj.weight"), + transform=_row_w, + ), + W( + "o_s", + shape=lambda d: (d.hidden * (d.q_width // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_ATTN + ".o_proj.weight_scale_inv"), + transform=_row_s, + ), + # --- dense SwiGLU-OAI MLP (layers 0..2) --- + W( + "mlp_gu_w", + shape=lambda d: (2 * d.dense_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight", "model.layers.{i}.mlp.up_proj.weight" + ), + transform=_gate_up_w, + layers=_dense_layers, + ), + W( + "mlp_gu_s", + shape=lambda d: (2 * d.dense_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + "model.layers.{i}.mlp.gate_proj.weight_scale_inv", + "model.layers.{i}.mlp.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_dense_layers, + ), + W( + "mlp_dn_w", + shape=lambda d: (d.hidden, d.dense_inter), + dtype=torch.float8_e4m3fn, + src=_one("model.layers.{i}.mlp.down_proj.weight"), + transform=_row_w, + layers=_dense_layers, + ), + W( + "mlp_dn_s", + shape=lambda d: (d.hidden * (d.dense_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one("model.layers.{i}.mlp.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_dense_layers, + ), + # --- MoE (layers 3..59): the fp32 router and its bias, the shared expert --- + W( + "router", + shape=lambda d: (d.num_experts, d.hidden), + dtype=torch.float32, + src=_one(_MOE + ".gate.weight"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "router_bias", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_one(_MOE + ".e_score_correction_bias"), + transform=_f32, + layers=_sparse_layers, + ), + W( + "sh_gu_w", + shape=lambda d: (2 * d.shared_inter, d.hidden), + dtype=torch.float8_e4m3fn, + src=_many( + _MOE + ".shared_experts.gate_proj.weight", _MOE + ".shared_experts.up_proj.weight" + ), + transform=_gate_up_w, + layers=_sparse_layers, + ), + W( + "sh_gu_s", + shape=lambda d: (2 * d.shared_inter * (d.hidden // MX_BLOCK),), + dtype=torch.uint8, + src=_many( + _MOE + ".shared_experts.gate_proj.weight_scale_inv", + _MOE + ".shared_experts.up_proj.weight_scale_inv", + ), + transform=_gate_up_s, + layers=_sparse_layers, + ), + W( + "sh_dn_w", + shape=lambda d: (d.hidden, d.shared_inter), + dtype=torch.float8_e4m3fn, + src=_one(_MOE + ".shared_experts.down_proj.weight"), + transform=_row_w, + layers=_sparse_layers, + ), + W( + "sh_dn_s", + shape=lambda d: (d.hidden * (d.shared_inter // MX_BLOCK),), + dtype=torch.uint8, + src=_one(_MOE + ".shared_experts.down_proj.weight_scale_inv"), + transform=_row_s, + layers=_sparse_layers, + ), + # --- the routed experts, NVFP4 (trtllm-gen layouts), all 128 on every rank --- + W( + "e_w31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // 2), + dtype=torch.uint8, + src=_per_expert("w3.weight", "w1.weight"), + transform=_expert_w31, + layers=_sparse_layers, + ), + W( + "e_w2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // 2), + dtype=torch.uint8, + src=_per_expert("w2.weight"), + transform=_expert_w2, + layers=_sparse_layers, + ), + W( + "e_s31", + shape=lambda d: (d.num_experts, 2 * d.moe_inter, d.hidden // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w3.weight_scale", "w1.weight_scale"), + transform=_expert_s31, + layers=_sparse_layers, + ), + W( + "e_s2", + shape=lambda d: (d.num_experts, d.hidden, d.moe_inter // NV_BLOCK), + dtype=torch.float8_e4m3fn, + src=_per_expert("w2.weight_scale"), + transform=_expert_s2, + layers=_sparse_layers, + ), + # The per-tensor scalars of every expert, raw; modeling derives the kernels' scalars from them. + *( + W( + f"e_{name}", + shape=lambda d: (d.num_experts,), + dtype=torch.float32, + src=_per_expert(suffix), + layers=_sparse_layers, + ) + for name, suffix in ( + ("isc1", "w1.input_scale"), + ("isc3", "w3.input_scale"), + ("isc2", "w2.input_scale"), + ("ws2_1", "w1.weight_scale_2"), + ("ws2_3", "w3.weight_scale_2"), + ("ws2_2", "w2.weight_scale_2"), + ) + ), + # --- not per-layer --- + W( + "embed", + shape=lambda d: (d.vocab, d.hidden), + src=_one("model.embed_tokens.weight"), + layers=None, + ), + W( + "final_norm", + shape=lambda d: (d.hidden,), + src=_one("model.norm.weight"), + transform=_gemma_fold, + layers=None, + ), + ) + + def dims(self, core) -> SimpleNamespace: + cfg = core.model_config.pretrained_config + mapping = core.model_config.mapping + sparse = cfg.sparse_attention_config + freq = [int(f) for f in sparse["sparse_attention_freq"]] + tp = mapping.tp_size + head_dim = int(cfg.head_dim) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + index_heads = max(int(sparse["sparse_num_index_heads"]) // tp, 1) + index_dim = int(sparse["sparse_index_dim"]) + return SimpleNamespace( + fused=bool( + getattr( + core.model_config.sparse_attention_config, "fuse_qkv_index_projection", False + ) + ), + num_layers=cfg.num_hidden_layers, + first_sparse=freq.index(1), + hidden=cfg.hidden_size, + head_dim=head_dim, + heads=heads, + kv_heads=kv_heads, + index_heads=index_heads, + index_dim=index_dim, + q_width=heads * head_dim, + qkv_dense_rows=(heads + 2 * kv_heads) * head_dim, + qkv_sparse_rows=(heads + 2 * kv_heads) * head_dim + index_heads * index_dim + index_dim, + dense_inter=int(cfg.dense_intermediate_size) // tp, + shared_inter=int(cfg.shared_intermediate_size) * int(cfg.n_shared_experts) // tp, + moe_inter=int(cfg.intermediate_size) // tp, + num_experts=int(cfg.num_local_experts), + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + layers = entry.layer_indices(d) + if layers is None: + return [(entry.name, None)] + return [(f"l{i}_{entry.name}", i) for i in layers] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + """Allocate every weight the table declares (meta under the engine's meta init).""" + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, core) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + """param key -> [(checkpoint key or key tuple, index into the param | None, transform | None)].""" + d = self.dims(core) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, core, weights) -> set: + """The vision tower and its projectors, and the multi-token-prediction head: in the checkpoint, never built + by this target.""" + return {k for k in weights.keys() if k.startswith(_UNBUILT_PREFIXES)} + + def load(self, model, weights) -> None: + """Fill every declared weight and the shell's ``lm_head`` shard; assert coverage both ways.""" + core = model.model + d = self.dims(core) + manifest = self.manifest(core) + assert set(manifest) == set(core.w.keys()), ( + "manifest/parameter drift", + set(manifest) ^ set(core.w.keys()), + ) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(core, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(core.w[param_key], ckpt, index, transform, weights, d, consumed) + if param_key.endswith(self.RELEASE_AFTER): + torch.cuda.empty_cache() + # The shell owns the vocab-parallel head: this rank's rows of the full table. + _fill(model.lm_head.weight, _LM + "lm_head.weight", None, _col_w, weights, d, consumed) + shell = { + n for n, _ in model.named_parameters() if not n.startswith(("model.w.", "draft_model.")) + } + assert shell == {"lm_head.weight"}, f"unfed shell parameters: {sorted(shell)}" + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + "checkpoint coverage: leftover keys are not exactly the unbuilt families (vision, MTP head); " + f"unexpected {sorted(leftover - expected)[:8]}, missed {sorted(expected - leftover)[:8]}" + ) + + +def _materialize(entry) -> torch.Tensor: + """``[:]`` on a lazy safetensors slice; a 0-dim tensor (the per-tensor scalars) rejects that index.""" + return entry if getattr(entry, "ndim", 1) == 0 else entry[:] + + +def _fill(param, ckpt, index, transform, weights, d, consumed) -> None: + keys = ckpt if isinstance(ckpt, tuple) else (ckpt,) + for key in keys: + assert key in weights, f"checkpoint key missing: {key}" + dst = param.data if index is None else param.data[index] + d.dev = dst.device + # Host tensors: a transform slices this rank's part on the host and moves only that to the device. + src_all = tuple(_materialize(weights[key]) for key in keys) + if transform is not None: + src = transform(src_all if isinstance(ckpt, tuple) else src_all[0], d) + else: + assert len(src_all) == 1, "a multi-key row needs a transform" + src = src_all[0] + if src.dim() == 0 and dst.dim() == 0: + pass + elif src.numel() == 1 and dst.numel() == 1: + src = src.reshape(dst.shape) + src = src.to(dst.device, non_blocking=True) + assert tuple(dst.shape) == tuple(src.shape), (ckpt, tuple(dst.shape), tuple(src.shape)) + assert src.dtype == dst.dtype, (ckpt, src.dtype, dst.dtype) + dst.copy_(src, non_blocking=True) + consumed.update(keys) + _release(weights, keys) + + +def _release(weights, keys) -> None: + """Drop consumed checkpoint tensors from the engine's dict as the load goes (bounded host memory).""" + if hasattr(weights, "mark_consumed_keys"): + weights.mark_consumed_keys(keys) + else: + for key in keys: + weights.pop(key, None) + + +MODEL_WEIGHTS = MiniMaxM3Tp8Weights() + + +# -------------------------------------------------------------------------------------------- the drafter + + +def _draft_one(template: str): + return lambda d, i, t=template: [(t, None)] + + +def _draft_many(*templates: str): + return lambda d, i, ts=templates: [(tuple(ts), None)] + + +def _draft_fc_norms(d): + return range(d.capture_layers) + + +def _draft_qkv(t: tuple[torch.Tensor, ...], d: SimpleNamespace) -> torch.Tensor: + q, k, v = t + hd = d.head_dim + return torch.cat([_col(q, d), _heads(k, d, hd), _heads(v, d, hd)], dim=0).contiguous() + + +class MiniMaxM3Tp8DraftWeights: + """The one-layer BF16 Eagle3 drafter (``speculative_model``), tp8. + + The fc over the three captured target layers and its per-layer fc_norms, the layer's three norms, its QKV + (input ``[embeds | hidden]``, so twice the hidden width; Q by heads, K/V by KV heads, a KV head on several + ranks when there are fewer than ranks) and o_proj, its gated MLP, the final norm and the full-vocabulary head + (vocab-parallel). The checkpoint's ``embed_tokens`` is the target's table (the drafter reads the target's); it is + a predicted non-load. + """ + + WEIGHTS: tuple[W, ...] = ( + W( + "fc", + shape=lambda d: (d.hidden, d.capture_layers * d.hidden), + src=_draft_one("fc.weight"), + layers=None, + ), + W( + "fc_norm", + shape=lambda d: (d.hidden,), + src=lambda d, i: [(f"fc_norm.{i}.weight", None)], + layers=_draft_fc_norms, + ), + W( + "in_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.input_layernorm.weight"), + layers=None, + ), + W( + "hidden_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.hidden_norm.weight"), + layers=None, + ), + W( + "post_norm", + shape=lambda d: (d.hidden,), + src=_draft_one("layers.0.post_attention_layernorm.weight"), + layers=None, + ), + W( + "qkv", + shape=lambda d: (d.qkv_rows, 2 * d.hidden), + src=_draft_many(*(f"layers.0.self_attn.{p}_proj.weight" for p in "qkv")), + transform=_draft_qkv, + layers=None, + ), + W( + "o", + shape=lambda d: (d.hidden, d.q_width), + src=_draft_one("layers.0.self_attn.o_proj.weight"), + transform=_row_w, + layers=None, + ), + W( + "gu", + shape=lambda d: (2 * d.inter, d.hidden), + src=_draft_many("layers.0.mlp.gate_proj.weight", "layers.0.mlp.up_proj.weight"), + transform=_gate_up_w, + layers=None, + ), + W( + "dn", + shape=lambda d: (d.hidden, d.inter), + src=_draft_one("layers.0.mlp.down_proj.weight"), + transform=_row_w, + layers=None, + ), + W("final_norm", shape=lambda d: (d.hidden,), src=_draft_one("norm.weight"), layers=None), + ) + + def dims(self, draft) -> SimpleNamespace: + cfg = draft.draft_config.pretrained_config + mapping = draft.draft_config.mapping + tp = mapping.tp_size + head_dim = int(getattr(cfg, "head_dim", cfg.hidden_size // cfg.num_attention_heads)) + heads = cfg.num_attention_heads // tp + kv_heads = max(cfg.num_key_value_heads // tp, 1) + return SimpleNamespace( + hidden=cfg.hidden_size, + capture_layers=draft.capture_layers, + head_dim=head_dim, + q_width=heads * head_dim, + qkv_rows=(heads + 2 * kv_heads) * head_dim, + inter=cfg.intermediate_size // tp, + vocab=cfg.vocab_size, + tp=tp, + rank=mapping.tp_rank, + dtype=cfg.torch_dtype, + ) + + def _targets(self, entry: W, d: SimpleNamespace) -> list[tuple[str, int | None]]: + if entry.layers is None: + return [(entry.name, None)] + return [(f"{entry.name}{i}", i) for i in entry.layers(d)] + + def declare(self, d: SimpleNamespace) -> nn.ParameterDict: + w = nn.ParameterDict() + for entry in self.WEIGHTS: + dtype = entry.dtype or d.dtype + for key, _ in self._targets(entry, d): + w[key] = nn.Parameter( + torch.empty(*entry.shape(d), dtype=dtype), requires_grad=False + ) + return w + + def manifest(self, draft) -> dict[str, list[tuple[Any, Any, Callable | None]]]: + d = self.dims(draft) + rows: dict[str, list[tuple[Any, Any, Callable | None]]] = {} + for entry in self.WEIGHTS: + for key, i in self._targets(entry, d): + rows[key] = [(ckpt, index, entry.transform) for ckpt, index in entry.src(d, i)] + return rows + + def expected_unconsumed(self, draft, weights) -> set: + """The drafter checkpoint's own embedding: the drafter reads the target's table.""" + return {"embed_tokens.weight"} & set(weights.keys()) + + def load(self, draft, weights) -> None: + d = self.dims(draft) + manifest = self.manifest(draft) + assert set(manifest) == set(draft.w.keys()), set(manifest) ^ set(draft.w.keys()) + all_keys = set(weights.keys()) + expected = self.expected_unconsumed(draft, weights) + consumed: set = set() + for param_key, sources in manifest.items(): + for ckpt, index, transform in sources: + _fill(draft.w[param_key], ckpt, index, transform, weights, d, consumed) + _fill(draft.lm_head.weight, "lm_head.weight", None, _col_w, weights, d, consumed) + torch.cuda.synchronize() + leftover = all_keys - consumed + assert leftover == expected, ( + f"drafter checkpoint coverage: unexpected {sorted(leftover - expected)[:8]}, " + f"missed {sorted(expected - leftover)[:8]}" + ) + + +DRAFT_WEIGHTS = MiniMaxM3Tp8DraftWeights() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py index 9b60e91f3165..36649748133a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/routing.py @@ -21,6 +21,7 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING, Any, Optional, Tuple from tensorrt_llm.llmapi.llm_args import EagleDecodingConfig @@ -32,10 +33,17 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.llmapi.llm_args import TorchLlmArgs -# The one GPU architecture these targets are written for: the decode kernels -# (the tcgen05 GEMVs, the fused MoE engines, the attention megakernels and the -# multicast all-reduce) are built and measured on GB300 only. +# The GPU architectures these targets are written for: the tp4 / tp8 decode +# kernels (the tcgen05 GEMVs, the fused MoE engines, the attention megakernels +# and the multicast all-reduce) are built and measured on GB300 (sm103); the +# tp2 and tp4 Rubin (sm107) targets are ports, switched on one certified path +# at a time. _SM = (10, 3) +_SM_RUBIN = (10, 7) + +# The routing knobs of the former separate Rubin tp2 wide-decode target. Its 9..192-row decode cells are now part of the +# one Rubin tp2 target, so there is nothing left to select: a deployment that still sets either knob is stale. +_REMOVED_TP2_KNOBS = ("TRTLLM_M3_TP2_WIDE", "TRTLLM_M3_TP2_WIDE_MIN_BATCH") # The checkpoint's config shape, read from ``text_config``; see the note in the # gpt_oss routing module for what shape-sniffing does and does not pin. @@ -109,9 +117,9 @@ def _parallel(m) -> Optional[str]: Attention is head-split and the routed experts are tensor-parallel with every expert on every rank: the expert split decides which expert weights each rank loads and into which shapes (per-rank expert width 768 at tp4, - 384 at tp8), so it selects a target rather than a runtime branch. + 384 at tp8, 1536 at tp2), so it selects a target rather than a runtime branch. """ - for tp in (4, 8): + for tp in (2, 4, 8): if ( m.world_size == tp and m.tp_size == tp @@ -200,6 +208,57 @@ def route( return ModelingV2MiniMaxM3Nvfp4Sm103Tp8 + # The same checkpoint on 2 Rubin GPUs (sm107), tensor-parallel attention and + # experts, one-model Eagle3 or no speculative decoding, at any batch size (decode rows 1..192). + if ( + compute_capability == _SM_RUBIN + and shape == _SHAPE + and layout_miss is None + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp2" + and speculation + and kv_cache + ): + stale = [knob for knob in _REMOVED_TP2_KNOBS if knob in os.environ] + if stale: + raise ValueError( + f"{', '.join(stale)}: removed; the MiniMax-M3 sm107 tp2 target serves every batch size (its " + "wide-decode cells are part of it), so unset these variables" + ) + from .minimax_m3_nvfp4__sm_107__tp2.modeling import ModelingV2MiniMaxM3Nvfp4Sm107Tp2 + + return ModelingV2MiniMaxM3Nvfp4Sm107Tp2 + + # The same checkpoint on 4 Rubin GPUs (sm107), tensor-parallel attention and + # experts, one-model Eagle3 or no speculative decoding. + if ( + compute_capability == _SM_RUBIN + and shape == _SHAPE + and layout_miss is None + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp4" + and speculation + and kv_cache + ): + from .minimax_m3_nvfp4__sm_107__tp4.modeling import ModelingV2MiniMaxM3Nvfp4Sm107Tp4 + + return ModelingV2MiniMaxM3Nvfp4Sm107Tp4 + + # The same checkpoint on 8 Rubin GPUs over two trays (sm107), tensor-parallel + # attention and experts, one-model Eagle3 or no speculative decoding. + if ( + compute_capability == _SM_RUBIN + and shape == _SHAPE + and layout_miss is None + and quant_algo == QuantAlgo.MIXED_PRECISION + and parallel == "tp8" + and speculation + and kv_cache + ): + from .minimax_m3_nvfp4__sm_107__tp8.modeling import ModelingV2MiniMaxM3Nvfp4Sm107Tp8 + + return ModelingV2MiniMaxM3Nvfp4Sm107Tp8 + raise NoTarget( f"compute capability {compute_capability}, text_config shape {shape}" f"{f' ({layout_miss})' if layout_miss else ''}, " @@ -209,7 +268,7 @@ def route( f"attention_dp={mapping.enable_attention_dp}, speculative decoding " f"{spec_config.spec_dec_mode.name if spec_config is not None else 'off'}" f", KV cache {kv_cache_quant_algo or 'bf16'}; the " - "minimax_m3 targets are the MiniMax-M3 NVFP4 checkpoint on sm103 (GB300), 4 or 8 ranks of " - "tensor-parallel attention and experts, with one-model Eagle3 or no speculative decoding, and an FP8 or " - "NVFP4 KV cache" + "minimax_m3 targets are the MiniMax-M3 NVFP4 checkpoint on sm103 (GB300) with 4 or 8 ranks, or on sm107 " + "(Rubin) with 2, 4 or 8 ranks, of tensor-parallel attention and experts, with one-model Eagle3 or no " + "speculative decoding, and an FP8 or NVFP4 KV cache" ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index e05523c35224..12164dcefaaf 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -28,6 +28,7 @@ from __future__ import annotations +import os from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple import torch @@ -156,6 +157,11 @@ def set_index_v(self, layer_idx: int, out_cache_loc: torch.Tensor, idx_v: torch. buf.index_copy_(0, out_cache_loc.to(torch.long), idx_v.to(buf.dtype)) +# SM107: price M3's real per-token footprint in the pre-allocation estimate (index-K on every sparse layer, FP8 dense +# and shared Eagle layers) instead of the generic all-NVFP4 K/V sum. Kill switch: False keeps the generic estimate. +_SM107_EXACT_KV_SIZE_ESTIMATE = True + + def shared_draft_layer_count(spec_config, layer_mask) -> int: """How many one-model draft layers the base manager appends to this one. @@ -270,6 +276,11 @@ def extend_attention_op_pools_for_shared_draft_layers( ) +# Zero every KV pool view (NVFP4 + FP8 K/V, NVFP4 block scales, paged index keys) once at allocation, so no reader +# ever sees uninitialized bytes. Kill switch (diagnostics only): TRTLLM_M3_KV_POOL_ZERO_INIT=0. +_KV_POOL_ZERO_INIT = os.environ.get("TRTLLM_M3_KV_POOL_ZERO_INIT", "1") == "1" + + def _require_trtllm_gen_extra_page_route(tokens_per_block: int) -> None: """Fail at setup, not at the first prefill, when the P128 opt-in has no FMHA (SM107). @@ -328,6 +339,7 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2): # Decode-step shortcuts a model opts into (RuntimeHooks.msa); see configure_decode_shortcuts. skip_unchanged_block_offsets = False reuse_pool_views = False + block_ids_from_views = False _main_kv_mapper_kind = MapperKind.HND _main_kv_view_layout = "NHD" @@ -472,18 +484,19 @@ def __init__( self._block_scale_buffers[layer_idx, kv_layout] = self.get_block_scale_buffers( layer_idx, kv_layout ) - # The KV pools are written only at live slots, and readers that consume a whole block or page (the CSR - # sparse prefill's V dequant, the full-context decode on a last page's tail) would turn uninitialized - # bytes into 0 * NaN = NaN. Zero every layer's K/V bytes (NVFP4 sparse, FP8 dense, the drafter's in its - # own manager), the NVFP4 block scales and the paged index keys once, before any request runs: every - # byte a kernel writes is unchanged. - for view in ( - *self._kv_buffers.values(), - *self._block_scale_buffers.values(), - *self._index_k_buffers.values(), - ): - if view is not None: - view.zero_() + if _KV_POOL_ZERO_INIT: + # The KV pools are written only at live slots, and readers that consume a whole block or page (the CSR + # sparse prefill's V dequant, the full-context decode on a last page's tail) would turn uninitialized + # bytes into 0 * NaN = NaN. Zero every layer's K/V bytes (NVFP4 sparse, FP8 dense, the drafter's in its + # own manager), the NVFP4 block scales and the paged index keys once, before any request runs: every + # byte a kernel writes is unchanged. + for view in ( + *self._kv_buffers.values(), + *self._block_scale_buffers.values(), + *self._index_k_buffers.values(), + ): + if view is not None: + view.zero_() def _validate_speculative_config(self, spec_config: DecodingBaseConfig | None) -> None: # Dynamic-tree relocation moves packed K/V bytes without their NVFP4 @@ -542,6 +555,72 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): * self.head_dim_per_layer[int(local_layer_idx)] ) + @classmethod + def get_cache_size_per_token(cls, model_config, mapping, num_layers=None, **kwargs): + """Pre-allocation bytes per token, including what the generic estimate leaves out. + + The base classmethod sums K/V for the target layers only, all at the cache dtype. This manager also holds + one index-K head per sparse layer (replicated on every TP rank), keeps the dense target layers and the + appended one-model Eagle layers in FP8 when the cache is NVFP4 (see :meth:`_build_cache_config`), and + allocates those shared draft layers itself. On SM107 the difference is added to the per-token term. + """ + base = super().get_cache_size_per_token(model_config, mapping, num_layers, **kwargs) + if ( + not _SM107_EXACT_KV_SIZE_ESTIMATE + or get_sm_version() != 107 + or kwargs.get("is_draft", False) + ): + return base + extra = cls._extra_bytes_per_token(model_config, mapping, num_layers, kwargs) + if isinstance(base, tuple): + return (base[0] + extra,) + tuple(base[1:]) + return base + extra + + @classmethod + def _extra_bytes_per_token(cls, model_config, mapping, num_layers, kwargs) -> int: + config = model_config.pretrained_config + config = getattr(config, "text_config", None) or config + num_target_layers = int(num_layers if num_layers is not None else config.num_hidden_layers) + local_layers = ( + list(mapping.pp_layers(num_target_layers)) + if mapping.pp_size > 1 + else list(range(num_target_layers)) + ) + sparse_layers = [layer for layer in local_layers if layer >= 3] + dense_layers = [layer for layer in local_layers if layer < 3] + tp = 1 if mapping.enable_attention_dp else mapping.tp_size + kv_heads = max(1, int(config.num_key_value_heads) // tp) + head_dim = int( + getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + ) + + sparse_config = getattr(model_config, "sparse_attention_config", None) + index_dim = int(getattr(sparse_config, "sparse_index_dim", None) or 128) + index_elem = 1 if str(getattr(sparse_config, "indexer_kv_dtype", "bf16")) == "fp8" else 2 + extra = len(sparse_layers) * index_dim * index_elem + + kv_elems = 2 * kv_heads * head_dim # K and V + quant_config = model_config.quant_config + nvfp4 = quant_config is not None and quant_config.quant_mode.has_fp4_kv_cache() + if nvfp4: + # The base prices every layer as NVFP4 data + block scales; dense layers are FP8. + extra += len(dense_layers) * (kv_elems - (kv_elems // 2 + kv_elems // 16)) + draft_layer_bytes = kv_elems + else: + draft_layer_bytes = None + spec_config = kwargs.get("spec_config") + if spec_config is not None and not kwargs.get("use_separate_draft_kv_cache", False): + num_draft = shared_draft_layer_count(spec_config, None) + if draft_layer_bytes is None: + elem = ( + 1 + if quant_config is not None and quant_config.quant_mode.has_fp8_kv_cache() + else 2 + ) + draft_layer_bytes = kv_elems * elem + extra += num_draft * draft_layer_bytes + return extra + def is_nvfp4_layer(self, layer_idx: int) -> bool: """Whether ``layer_idx`` stores packed NVFP4 data and block scales.""" return self.dtype == DataType.NVFP4 and int(layer_idx) in self.sparse_layer_ids @@ -955,7 +1034,10 @@ def get_buffers( ) def configure_decode_shortcuts( - self, skip_unchanged_block_offsets: bool, reuse_pool_views: bool + self, + skip_unchanged_block_offsets: bool, + reuse_pool_views: bool, + block_ids_from_views: bool = False, ) -> None: """Turn on the decode-step shortcuts a model opted into. @@ -963,9 +1045,14 @@ def configure_decode_shortcuts( whose inputs equal the previous copy's into the same destination. ``reuse_pool_views``: hand out one pool view per geometry instead of re-wrapping the pool memory on every attention call. + ``block_ids_from_views``: fill :meth:`get_block_ids_per_seq` rows from the + page-index views directly (see there). """ self.skip_unchanged_block_offsets = skip_unchanged_block_offsets self.reuse_pool_views = reuse_pool_views + self.block_ids_from_views = block_ids_from_views + if block_ids_from_views: + logger.info("[M3 host] block-id tables filled from page-index views") def _pool_view( self, @@ -1192,7 +1279,22 @@ def get_block_ids_per_seq(self, request_ids): The rows are written through a numpy view of a single zero-filled, pinned result, so the attention metadata builders ship it to the device in one asynchronous copy. + + With ``block_ids_from_views`` each row is slice-copied from the request's + page-index view (an int32 array) rather than built as a Python list of the + whole padded row first; the table is the same. """ + if self.block_ids_from_views: + views = [self.kv_cache_map[req_id].get_base_page_indices(0) for req_id in request_ids] + max_blocks = max((len(view) for view in views), default=0) + padded_tensor = torch.zeros( + (len(views), max_blocks), dtype=torch.int32, pin_memory=prefer_pinned() + ) + rows = padded_tensor.numpy() + for row, view in zip(rows, views): + row[: len(view)] = view + rows[rows == BAD_PAGE_INDEX] = 0 + return padded_tensor block_ids_per_seq = self.get_batch_cache_indices(request_ids) batch = len(block_ids_per_seq) max_blocks = max((len(block_ids) for block_ids in block_ids_per_seq), default=0) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py index b3cbb2716240..2f40c6b090f2 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/paged_cache.py @@ -8,6 +8,42 @@ import torch +# E4M3 casts saturate to +-448, the bits the CUDA FP8 writers produce (__NV_SATFINITE / cvt.rn.satfinite); NaN stays +# NaN. Used by the E4M3 cache writes below and by the targets' transient FP8 query casts. Where torch's own E4M3 cast +# already saturates (c10's float8_e4m3fn conversion, shared by its CPU and CUDA kernels) the cast alone is used: an +# explicit clamp would add a kernel per cast for the same bytes. A torch whose cast does not saturate gets the clamp. +# Kill: False (the plain cast, whatever it does). +FP8_SATURATE = True +_E4M3_MAX = 448.0 + + +def _cast_saturates() -> bool: + """Whether this torch's E4M3 cast maps every out-of-range value, infinities included, to +-448.""" + # An explicit device: the first import can run under a default-device context (meta while a model is built). + probe = torch.tensor( + [449.0, 463.0, 464.0, 1e4, float("inf"), -464.0, -float("inf")], device="cpu" + ) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0], device="cpu") + return torch.equal(probe.to(torch.float8_e4m3fn).float(), want) + + +# Probed once at import on the CPU, so no CUDA work runs here (nor later inside a graph capture). +_CAST_SATURATES = _cast_saturates() + + +def fp8_e4m3_saturate(values: torch.Tensor) -> torch.Tensor: + """``values`` as float8_e4m3fn, saturated to +-448 (bitwise the plain cast for every finite in-range value). + Values already in float8_e4m3fn are returned as they are (they are finite by construction or already NaN).""" + if values.dtype == torch.float8_e4m3fn: + return values + if FP8_SATURATE and not _CAST_SATURATES: + values = values.clamp(-_E4M3_MAX, _E4M3_MAX) + return values.to(torch.float8_e4m3fn) + + +def _to_cache_dtype(values: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return fp8_e4m3_saturate(values) if dtype == torch.float8_e4m3fn else values.to(dtype) + def write_kv_slots( cache: torch.Tensor, @@ -60,11 +96,11 @@ def write_kv_slots( if layout == "HND": # Advanced indices on dims 0 and 2 broadcast to [num_tokens] and # move front, giving a [num_tokens, num_heads, channel] target. - cache[page, :, within, :] = values.to(cache.dtype) + cache[page, :, within, :] = _to_cache_dtype(values, cache.dtype) else: - cache[page, within] = values.to(cache.dtype) + cache[page, within] = _to_cache_dtype(values, cache.dtype) else: - cache.index_copy_(0, out_cache_loc.to(torch.long), values.to(cache.dtype)) + cache.index_copy_(0, out_cache_loc.to(torch.long), _to_cache_dtype(values, cache.dtype)) __all__ = ["write_kv_slots"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py index 2ba83440c671..fb6e33b3612c 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/kernels/trtllm_gen_dense_decode.py @@ -25,6 +25,9 @@ import torch +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers from .msa_utils import check_decode_span_shape @@ -276,7 +279,7 @@ def minimax_m3_trtllm_gen_dense_decode( # q/k/v, so an FP8 pool needs FP8 Q. M3 stores unscaled E4M3, so this is a # plain cast and a no-op when the fused producer already emitted FP8. if kv_pool.dtype == torch.float8_e4m3fn and q.dtype != torch.float8_e4m3fn: - q = q.to(torch.float8_e4m3fn) + q = fp8_e4m3_saturate(q) reserve = torch.cuda.is_current_stream_capturing() workspace_size = _workspace(q.dtype, num_heads, int(q.shape[2]), int(kv_pool.shape[1])) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py index 7e28b67fef30..73e492fd2649 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py @@ -39,10 +39,11 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs -from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.h2d_cache import copy_h2d_if_changed, copy_h2d_pool_shared from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy from tensorrt_llm._utils import maybe_pin_memory from tensorrt_llm.bindings import DataType +from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig from .common import ( @@ -111,9 +112,23 @@ def _host_prefix_sum(lens_cpu: torch.Tensor, dtype: torch.dtype) -> torch.Tensor return maybe_pin_memory(torch.cat([sums.new_zeros(1), sums]).to(dtype)) +# Kill switches for A/B only. +# SUBPAGE_IN_GRAPH: the sub-page block table is expanded from msa_block_table on the device inside the forward, once, +# by the first dense layer that reads it (msa_subpage_rows), so the block table can join the step-input pack and the +# step boundary issues no expansion; False expands it in prepare() after an unpacked block-table upload, as before. +SUBPAGE_IN_GRAPH = True +# SKIP_UNCHANGED_BLOCK_TABLE: msa_block_table takes the H2D skip cache. Its only writer is prepare() through that cache +# (the device and every captured graph only read it), so the shadow describes the device rows; False always uploads it. +SKIP_UNCHANGED_BLOCK_TABLE = True +# TAIL_FILL_IN_FORWARD: prepare() records where msa_out_cache_loc's live slots end and the forward invalidates the rest +# (on_update_kv_lens, which opens every forward, so a captured forward fills its own tail at each replay) instead of an +# eager fill between the decode graphs; False fills it in prepare() as before. +TAIL_FILL_IN_FORWARD = True + + def _fill_unused_slots(buf: torch.Tensor, start: int) -> None: - """buf[start:] = -1 after a write of the live slots buf[:start] (see h2d_cache.fill_tail_if_changed).""" - h2d_cache.fill_tail_if_changed(buf, start, -1) + """buf[start:] = -1 after a write of the live slots buf[:start] (see h2d_cache.fill_tail_pool_shared).""" + h2d_cache.fill_tail_pool_shared(buf, start, -1) class MsaDecodeSpan(NamedTuple): @@ -196,6 +211,11 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata): # factor, or 0 where the pool has no single one; see msa_subpage_rows. msa_subpage_block_table: Optional[torch.Tensor] = None _msa_subpages_per_slot: int = 0 + # SUBPAGE_IN_GRAPH: the rows prepare() staged for the forward to expand (0: none), and whether this forward did. + _msa_subpage_rows: int = 0 + _msa_subpage_expanded: bool = False + # TAIL_FILL_IN_FORWARD: the first msa_out_cache_loc row past this step's live slots (None: prepare() filled it). + _msa_tail_start: Optional[int] = None # Inputs for on_update_kv_lens: each query token's request row and offset # within the request, and the kv_lens prepare() staged (the upper bound the # correction is clamped to). The write slots are re-derived from @@ -377,6 +397,15 @@ def msa_subpage_rows(self, row_first: int, row_last: int) -> Tuple[Optional[torc table = self.msa_subpage_block_table if table is None: return None, 0 + rows = self._msa_subpage_rows + if rows > 0 and not self._msa_subpage_expanded: + # Once per forward (on_update_kv_lens, which opens every forward, resets it), from the block table the + # pack landed, over the rows prepare() staged; a captured forward records it once, so every replay + # expands its own rows. + write_subpage_block_table( + self.msa_block_table[:rows], self._msa_subpages_per_slot, table[:rows] + ) + self._msa_subpage_expanded = True return table[row_first:row_last], self._msa_subpages_per_slot def _msa_main_kv_is_fp8(self) -> bool: @@ -713,9 +742,9 @@ def _stage_msa_short_table(self, n_valid: torch.Tensor, params) -> None: ar = torch.arange(topk, dtype=torch.int32) table = torch.where(ar[None, :] < n_valid.to(torch.int32)[:, None], ar[None, :], -1) table = table.unsqueeze(0).expand(num_kv_heads, total_q, topk).contiguous() - copy_h2d_if_changed( - self.msa_short_table[: table.numel()], maybe_pin_memory(table.flatten()), pack=True - ) + # Uploaded every step, never skipped: the buffer is shared through the CUDA graph buffer pool, so the skip + # cache's host shadow is not proof of its device contents, and a stale table selects too few blocks. + step_inputs.copy_(self.msa_short_table[: table.numel()], maybe_pin_memory(table.flatten())) self._msa_short_ok = True def prepare(self) -> None: @@ -889,6 +918,18 @@ def _msa_kv_lens_may_change(self) -> bool: or (runtime_features is not None and runtime_features.has_speculative_draft_tokens) ) + def _fill_tail_in_forward(self) -> None: + """TAIL_FILL_IN_FORWARD: -1 over msa_out_cache_loc past the live slots prepare() staged (see the prepare() + comment on padded rows), before anything in the forward reads it.""" + start = self._msa_tail_start + buf = self.msa_out_cache_loc + if start is not None and buf is not None and start < buf.shape[0]: + logger.info_once( + "[M3 host] MSA out_cache_loc tail filled inside the forward", + key="msa_tail_fill_in_forward", + ) + buf[start:].fill_(-1) + def on_update_kv_lens(self) -> None: """Re-derive lengths, slots and valid-block counts from the corrected kv_lens_cuda. @@ -908,6 +949,9 @@ def on_update_kv_lens(self) -> None: extend_ctx promotes speculative generation rows into the context prefix. Staged host bounds remain valid upper bounds as lengths shrink. """ + # A new forward: the first dense layer that reads the sub-page table expands it again (msa_subpage_rows). + self._msa_subpage_expanded = False + self._fill_tail_in_forward() super().on_update_kv_lens() if not self._msa_fields_ready or not self._msa_kv_lens_dynamic: return @@ -1134,6 +1178,8 @@ def _build_msa_fields(self) -> None: decoding the inputs for on_update_kv_lens are staged as well. """ self._msa_fields_ready = False + # Set again below once this step's live slots are staged; until then the forward fills nothing. + self._msa_tail_start = None if not self._msa_buffers_ready: return request_ids = self.request_ids @@ -1211,7 +1257,11 @@ def _build_msa_fields(self) -> None: # would make its shape dynamic, and a negative slot is what makes those # rows cache-write no-ops instead. Invalidate the unwritten tail so they # cannot reuse the previous step's live slots, which address real pages. - _fill_unused_slots(self.msa_out_cache_loc, total_new_tokens) + if TAIL_FILL_IN_FORWARD: + self._msa_tail_start = total_new_tokens + else: + self._msa_tail_start = None + _fill_unused_slots(self.msa_out_cache_loc, total_new_tokens) if kv_indices is not None: step_inputs.copy_(self.msa_kv_indices[: int(kv_indices.shape[0])], kv_indices) @@ -1219,12 +1269,22 @@ def _build_msa_fields(self) -> None: # from the same host block ids the flat page table was built from. # Columns past a request's page count are left stale rather than # cleared: every consumer bounds its walk by seq_lens. - block_table_changed = copy_h2d_if_changed( - self.msa_block_table[:batch_size, :block_table_cols], - maybe_pin_memory(block_ids_cpu.to(torch.int32)), - # write_subpage_block_table below reads it on the device before the replay. - pack=self.msa_subpage_block_table is None, - ) + # With SUBPAGE_IN_GRAPH nothing reads the table on the device before the replay, so it can join the pack; + # the prepare-time expansion below reads it first. + block_table_host = maybe_pin_memory(block_ids_cpu.to(torch.int32)) + pack_block_table = self.msa_subpage_block_table is None or SUBPAGE_IN_GRAPH + if SKIP_UNCHANGED_BLOCK_TABLE: + block_table_changed = copy_h2d_if_changed( + self.msa_block_table[:batch_size, :block_table_cols], + block_table_host, + pack=pack_block_table, + ) + else: + block_table_changed = copy_h2d_pool_shared( + self.msa_block_table[:batch_size, :block_table_cols], + block_table_host, + pack=pack_block_table, + ) # Packed only when the device scan below (h2d_cache off) does not read it before the replay. if h2d_cache.enabled(): step_inputs.copy_(self.msa_seq_lens_cuda[:batch_size], kv_lens_cpu) @@ -1236,8 +1296,8 @@ def _build_msa_fields(self) -> None: if h2d_cache.enabled(): # Both prefix sums from the host lengths (the values the device scans would give), staged with # their leading zero: copies instead of zero / scan kernels between the decode graphs. - copy_h2d_if_changed(qo_lens_cuda, qo_lens_cpu, pack=True) - copy_h2d_if_changed( + copy_h2d_pool_shared(qo_lens_cuda, qo_lens_cpu, pack=True) + copy_h2d_pool_shared( self.msa_cu_q_lens[: batch_size + 1], _host_prefix_sum(qo_lens_cpu, self.msa_cu_q_lens.dtype), pack=True, @@ -1295,7 +1355,14 @@ def _build_msa_fields(self) -> None: # Sub-page expansion for the trtllm-gen dense layers, staged once here # instead of once per layer, outside capture into a graph-stable # buffer as with the slot table above. - if self.msa_subpage_block_table is not None and block_table_changed: + # With SUBPAGE_IN_GRAPH the forward expands this step's rows instead (msa_subpage_rows). + self._msa_subpage_expanded = False + self._msa_subpage_rows = batch_size if SUBPAGE_IN_GRAPH else 0 + if ( + not SUBPAGE_IN_GRAPH + and self.msa_subpage_block_table is not None + and block_table_changed + ): write_subpage_block_table( self.msa_block_table[:batch_size], self._msa_subpages_per_slot, @@ -1319,10 +1386,10 @@ def _build_msa_fields(self) -> None: torch.arange(total_new_tokens, dtype=torch.int64) - torch.repeat_interleave(starts, qo_long) ).to(torch.int32) - copy_h2d_if_changed( + copy_h2d_pool_shared( self.msa_q_batch_row[:total_new_tokens], maybe_pin_memory(batch_row_cpu), pack=True ) - copy_h2d_if_changed( + copy_h2d_pool_shared( self.msa_q_intra[:total_new_tokens], maybe_pin_memory(intra_cpu), pack=True ) # The staged lens are the upper bound on_update_kv_lens clamps to: the diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 4e95b945e8ae..6d7c540d9119 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -48,7 +48,7 @@ from ...kimi_k3_cache_policy import get_kimi_k3_fp8_kv_layer_ids, is_kimi_k3 from ...pyexecutor import step_inputs from ...pyexecutor.config_utils import is_mla -from ...pyexecutor.h2d_cache import copy_h2d_if_changed +from ...pyexecutor.h2d_cache import copy_h2d_pool_shared from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) from ..kernels.sage_quant import SageQuantized, sage_quantize @@ -1023,8 +1023,9 @@ def prepare(self) -> None: device='cpu', ) self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens) - copy_h2d_if_changed(self.prompt_lens_cuda[:self.num_seqs], - self.prompt_lens_cpu[:self.num_seqs]) + copy_h2d_pool_shared(self.prompt_lens_cuda[:self.num_seqs], + self.prompt_lens_cpu[:self.num_seqs], + pack=True) # number of tokens in the kv cache for each sequence in the batch cached_token_lens = torch.tensor( @@ -1228,7 +1229,7 @@ def prepare_encoder_decoder_from_precomputed_lengths( f"({self.kv_cache_manager.max_seq_len}).") num_seqs = self.num_seqs - copy_h2d_if_changed(self.prompt_lens_cuda[:num_seqs], prompt_lens) + copy_h2d_pool_shared(self.prompt_lens_cuda[:num_seqs], prompt_lens) self.kv_lens_cuda[:num_seqs].copy_(kv_lens, non_blocking=True) self.host_total_kv_lens[0] = context_kv_tokens self.host_total_kv_lens[1] = generation_kv_tokens diff --git a/tensorrt_llm/_torch/custom_ops/mxfp8_cute_pinned_tactics.py b/tensorrt_llm/_torch/custom_ops/mxfp8_cute_pinned_tactics.py new file mode 100644 index 000000000000..1511adaed71a --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/mxfp8_cute_pinned_tactics.py @@ -0,0 +1,160 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Pinned FlashInfer CuTe DSL MXFP8 GEMM tactics for SM107 (MiniMax-M3 TP2 projections), above 32 rows. + +TRT-LLM runs FlashInfer's CuTe DSL MXFP8 GEMM with FlashInfer's config sweep skipped, so above 32 rows (the +split-K low-M kernel's limit) every call takes the default tile (128, 128), cluster (1, 1): ceil(N / 128) CTAs, a +quarter of the SMs streaming the weight. This table holds FlashInfer's own tuned tactic per (N, K, M bucket) on +SM107. Every entry is non-split-K (split_k 1, each output keeps its full-K accumulation order), so its output equals +the default tactic's bit for bit (measured on all entries), at 18..40% less time from 33 to 16384 rows. A fixed +table rather than live tuning keeps the choice identical across runs and nodes. Used where TRT-LLM calls the +FlashInfer runner (decode-graph capture, autotune); eager prefill takes other GEMMs. Off unless armed: the executor +calls arm() for a model whose RuntimeHooks set mxfp8_pinned_tactics. Kill: TRTLLM_MXFP8_CUTE_PIN=0. +""" + +import os +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger + +ENABLED = os.environ.get("TRTLLM_MXFP8_CUTE_PIN", "1") == "1" +MIN_ROWS = 33 # 1..32 rows keep FlashInfer's split-K low-M path (its tuned tactics there are split-K: not bitwise) +MAX_BUCKET = 16384 + +Tactic = Tuple[Tuple[int, int], Tuple[int, int], bool, bool, int] +# (N, K, M bucket = next power of two >= M, 64..16384) +# -> (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch, split_k) +TACTICS: Dict[Tuple[int, int, int], Tactic] = { + (3072, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (3072, 6144, 128): ((256, 64), (2, 1), True, False, 1), + (3072, 6144, 256): ((256, 64), (2, 1), True, False, 1), + (3072, 6144, 512): ((256, 64), (2, 1), False, False, 1), + (3072, 6144, 1024): ((256, 128), (2, 1), False, False, 1), + (3072, 6144, 2048): ((256, 256), (2, 1), False, False, 1), + (3072, 6144, 4096): ((256, 192), (4, 2), False, False, 1), + (3072, 6144, 8192): ((256, 192), (2, 1), False, False, 1), + (3072, 6144, 16384): ((256, 192), (2, 1), True, False, 1), + (4608, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (4608, 6144, 128): ((256, 64), (2, 1), True, False, 1), + (4608, 6144, 256): ((256, 64), (2, 4), False, True, 1), + (4608, 6144, 512): ((256, 128), (2, 2), False, False, 1), + (4608, 6144, 1024): ((256, 192), (2, 1), False, False, 1), + (4608, 6144, 2048): ((256, 192), (2, 1), False, False, 1), + (4608, 6144, 4096): ((256, 256), (2, 1), False, False, 1), + (4608, 6144, 8192): ((256, 192), (2, 4), True, False, 1), + (4608, 6144, 16384): ((256, 256), (2, 1), True, False, 1), + (4992, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (4992, 6144, 128): ((256, 64), (2, 1), True, False, 1), + (4992, 6144, 256): ((256, 64), (2, 4), False, True, 1), + (4992, 6144, 512): ((256, 128), (4, 2), False, True, 1), + (4992, 6144, 1024): ((256, 192), (2, 1), False, False, 1), + (4992, 6144, 2048): ((256, 192), (2, 1), False, False, 1), + (4992, 6144, 4096): ((256, 192), (2, 1), False, False, 1), + (4992, 6144, 8192): ((256, 192), (2, 1), False, False, 1), + (4992, 6144, 16384): ((256, 192), (2, 1), True, False, 1), + (6144, 1536, 64): ((256, 64), (2, 1), True, False, 1), + (6144, 1536, 128): ((256, 64), (2, 1), True, False, 1), + (6144, 1536, 256): ((256, 64), (2, 1), True, False, 1), + (6144, 1536, 512): ((256, 128), (2, 1), False, False, 1), + (6144, 1536, 1024): ((256, 256), (2, 1), False, False, 1), + (6144, 1536, 2048): ((256, 192), (2, 4), False, False, 1), + (6144, 1536, 4096): ((256, 192), (2, 1), False, False, 1), + (6144, 1536, 8192): ((256, 192), (2, 1), False, False, 1), + (6144, 1536, 16384): ((256, 192), (2, 1), True, False, 1), + (6144, 4096, 64): ((256, 64), (2, 1), True, False, 1), + (6144, 4096, 128): ((256, 64), (4, 1), True, False, 1), + (6144, 4096, 256): ((256, 64), (2, 1), True, False, 1), + (6144, 4096, 512): ((256, 128), (2, 1), False, False, 1), + (6144, 4096, 1024): ((256, 256), (2, 1), False, False, 1), + (6144, 4096, 2048): ((256, 192), (2, 4), False, False, 1), + (6144, 4096, 4096): ((256, 192), (2, 1), False, False, 1), + (6144, 4096, 8192): ((256, 192), (2, 1), True, False, 1), + (6144, 4096, 16384): ((256, 192), (2, 1), True, False, 1), + (6144, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (6144, 6144, 128): ((256, 64), (2, 1), True, False, 1), + (6144, 6144, 256): ((256, 64), (2, 1), True, False, 1), + (6144, 6144, 512): ((256, 128), (2, 1), True, False, 1), + (6144, 6144, 1024): ((256, 256), (2, 1), False, False, 1), + (6144, 6144, 2048): ((256, 192), (2, 4), False, False, 1), + (6144, 6144, 4096): ((256, 192), (2, 1), False, False, 1), + (6144, 6144, 8192): ((256, 192), (2, 1), False, False, 1), + (6144, 6144, 16384): ((256, 256), (2, 1), True, False, 1), + (12288, 6144, 64): ((256, 64), (2, 1), True, False, 1), + (12288, 6144, 128): ((128, 64), (1, 2), True, False, 1), + (12288, 6144, 256): ((256, 128), (2, 1), False, False, 1), + (12288, 6144, 512): ((256, 256), (2, 1), False, False, 1), + (12288, 6144, 1024): ((256, 192), (2, 2), False, False, 1), + (12288, 6144, 2048): ((256, 192), (2, 1), False, False, 1), + (12288, 6144, 4096): ((256, 192), (2, 1), False, False, 1), + (12288, 6144, 8192): ((256, 256), (2, 1), False, False, 1), + (12288, 6144, 16384): ((256, 256), (2, 1), True, False, 1), +} + +_STATE: dict = {} +# Process-wide once armed: another model later or concurrently in the same process, at a table shape on SM107, +# also takes the pinned tactic. Its outputs are unchanged (each entry is bitwise the default tactic), only faster. +# The LLM API runs one executor per worker process. +_ARMED = [False] + + +def arm() -> None: + """Engage the table for this process (a model opted in through RuntimeHooks.mxfp8_pinned_tactics).""" + _ARMED[0] = True + + +def _bucket(m: int) -> int: + return max(64, 1 << (m - 1).bit_length()) + + +def tactic_for(m: int, n: int, k: int) -> Optional[Tactic]: + """The pinned tactic of an [m, k] x [n, k]^T call, or None (FlashInfer's own path).""" + if not ENABLED or not _ARMED[0] or m < MIN_ROWS or m > MAX_BUCKET or get_sm_version() != 107: + return None + return TACTICS.get((n, k, _bucket(m))) + + +def run( + act: torch.Tensor, + act_scale: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> Optional[torch.Tensor]: + """The MXFP8 GEMM act @ weight^T on FlashInfer's CuTe DSL runner at the pinned tactic, or None when no entry covers + the call (the caller runs FlashInfer's default path). Bitwise equal to that default path where an entry exists.""" + m, k = act.shape + n = weight.shape[0] + tactic = tactic_for(m, n, k) + if tactic is None: + return None + try: + from flashinfer.gemm import gemm_base as _gb + + key = (out_dtype, act.device) + st = _STATE.get(key) + if st is None: + major, minor = _gb.get_compute_capability(act.device) + st = _STATE[key] = ( + _gb._cute_dsl_gemm_mxfp8_runner(major, minor, True, out_dtype), + _gb._get_cache_buf("mm_mxfp8_workspace", _gb.DEFAULT_WORKSPACE_SIZE, act.device), + ) + logger.info( + "[MXFP8 GEMM] pinned CuTe DSL tactics engaged above 32 rows (TRTLLM_MXFP8_CUTE_PIN)" + ) + runner, workspace = st + out = torch.empty((m, n), dtype=out_dtype, device=act.device) + runner( + inputs=[act, weight.t(), act_scale, weight_scale, out_dtype, out, workspace], + tactic=tactic, + ) + return out + except Exception as e: # noqa: BLE001 - a FlashInfer API change falls back to the default path, once loudly + if "failed" not in _STATE: + _STATE["failed"] = True + logger.warning( + f"[MXFP8 GEMM] pinned CuTe DSL tactic unavailable ({type(e).__name__}: {e}); default path" + ) + return None diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index b248495311f2..672c9b9f20e8 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -41,6 +41,7 @@ from ..cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl +from . import mxfp8_cute_pinned_tactics as _mxfp8_cute_pinned from .fast_custom_op import fast_custom_op if IS_FLASHINFER_AVAILABLE: @@ -811,6 +812,13 @@ def forward( backend = "cutlass" context = nullcontext() else: + # Above 32 rows on SM107, for a model that opts in: the pinned tuned tactic (bitwise the + # heuristic's output, faster); see mxfp8_cute_pinned_tactics. Otherwise, or with no entry: + # FlashInfer's heuristic config. + pinned = _mxfp8_cute_pinned.run(act, act_scale, weight, + weight_scale, self.output_dtype) + if pinned is not None: + return pinned # Skip FlashInfer's CuTeDSL config sweep; use its heuristic config. backend = "cute-dsl" context = _flashinfer_autotune(tune_mode=False, @@ -2697,6 +2705,16 @@ def _( return gate.new_empty(gate.shape, dtype=o_dtype) +# The AllReduce tuner profiles NCCL_SYMMETRIC by pre-allocating two NCCL window buffers of [2 x max_num_tokens, hidden] +# (plus registration) inside the memory-profiling warm-up, so they shrink the KV pool even where NCCL_SYMMETRIC never +# wins (agg TP2: only NCCL / ONESHOT are picked; ~1.24 GiB per rank of KV, GEMM-AR 803388). +# TRTLLM_AR_TUNE_NCCL_SYMMETRIC=0 leaves it out of the candidates (no windows) and makes plain NCCL the cache-miss +# fallback; the default keeps stock behavior. +_AR_TUNE_NCCL_SYMMETRIC = os.environ.get("TRTLLM_AR_TUNE_NCCL_SYMMETRIC", + "1") == "1" +_AR_TUNE_NCCL_SYMMETRIC_LOGGED = False + + class AllReduceRunner(TunableRunner): _prealloc_lock: ClassVar[threading.Lock] = threading.Lock() _prealloc_done: ClassVar[set] = set() @@ -2745,7 +2763,7 @@ def unique_id( self.input_dtype, self.op, self.input_uses_nccl_window, - self.exclude_nccl_symmetric, + self.exclude_nccl_symmetric or not _AR_TUNE_NCCL_SYMMETRIC, ) @classmethod @@ -2807,10 +2825,16 @@ def get_valid_tactics( profile: OptimizationProfile, **kwargs, ) -> List[int]: - valid_strategies = [ - AllReduceStrategy.NCCL_SYMMETRIC.value, - AllReduceStrategy.NCCL.value, - ] + global _AR_TUNE_NCCL_SYMMETRIC_LOGGED + valid_strategies = [AllReduceStrategy.NCCL.value] + if _AR_TUNE_NCCL_SYMMETRIC: + valid_strategies.insert(0, AllReduceStrategy.NCCL_SYMMETRIC.value) + elif not _AR_TUNE_NCCL_SYMMETRIC_LOGGED: + _AR_TUNE_NCCL_SYMMETRIC_LOGGED = True + logger.info( + "[tunable_allreduce] AR tuner: NCCL_SYMMETRIC left out (no NCCL window buffers " + "pre-allocated); cache-miss fallback is NCCL (TRTLLM_AR_TUNE_NCCL_SYMMETRIC=0)" + ) if self.exclude_nccl_symmetric: valid_strategies = [AllReduceStrategy.NCCL.value] # Fallback in allreduceOp is set to NCCL_SYMMETRIC as default @@ -2870,6 +2894,9 @@ def _register_native_tactics(self, inputs: List[torch.Tensor]) -> None: if cache_entry is None: continue _, tactic, _ = cache_entry + if ((self.exclude_nccl_symmetric or not _AR_TUNE_NCCL_SYMMETRIC) + and tactic == AllReduceStrategy.NCCL_SYMMETRIC.value): + tactic = AllReduceStrategy.NCCL.value torch.ops.trtllm.register_allreduce_tactic(input, residual, norm_weight, scale, bias, workspace, self.group, @@ -2900,11 +2927,12 @@ def forward( return input if tactic == -1: # tactic == -1 means the autotuner cache missed for this shape. - # Fall back to NCCL_SYMMETRIC; asymmetric ncclMemAlloc failures are handled + # Honor the environment policy; asymmetric ncclMemAlloc failures are handled # by a cross-rank barrier in NCCLWindowAllocator which falls back # to plain NCCL. - tactic = AllReduceStrategy.NCCL_SYMMETRIC.value - if (self.exclude_nccl_symmetric + tactic = (AllReduceStrategy.NCCL_SYMMETRIC.value if + _AR_TUNE_NCCL_SYMMETRIC else AllReduceStrategy.NCCL.value) + if ((self.exclude_nccl_symmetric or not _AR_TUNE_NCCL_SYMMETRIC) and tactic == AllReduceStrategy.NCCL_SYMMETRIC.value): tactic = AllReduceStrategy.NCCL.value diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py new file mode 100644 index 000000000000..d6a67225eb01 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CuTe DSL decode kernels of the MiniMax-M3 NVFP4 / sm_103 / tp4 modeling_v2 target.""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py new file mode 100644 index 000000000000..4c6fc92cd49c --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 decode embedding gather + first-layer RMSNorm as one CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py new file mode 100644 index 000000000000..acf27a95316b --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/m3_embedding_firstnorm_kernel.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. + +128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original +K3 wrapper this target does not promise FlashInfer's bit-identical reduction. +""" + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +from cutlass.experimental import primitives as prims + +NORM_THREADS = 128 +NORM_VEC = 8 +NORM_WARPS = 4 + + +@cute.kernel +def m3_embed_norm_kernel( + ids: cute.Tensor, # int32 or int64 [N] + table: cute.Tensor, # bf16 [V, H] + weight: cute.Tensor, # bf16 [H] + raw: cute.Tensor, # bf16 [N, H], out: the rows + out: cute.Tensor, # bf16 [N, H], out: the normed rows + vocab: cutlass.Int64, + eps: cutlass.Float32, + tv_layout: cute.Layout, + tiler_mn: cute.Shape, + hidden: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(prims.GridDepAction.WAIT) + + smem = cutlass.utils.SmemAllocator() + reduction_buffer = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 + ) + + tok = cutlass.Int64(ids[bidx]) + valid = (tok >= cutlass.Int64(0)) & (tok < vocab) + row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) + + gX = cute.local_tile(table, tiler_mn, (row, 0)) + gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) + gY = cute.local_tile(out, tiler_mn, (bidx, 0)) + w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) + gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) + + copy_atom_load = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 + ) + copy_atom_store = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 + ) + thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) + thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) + + tXgX = thr_copy_X.partition_S(gX) + tXrX = cute.make_fragment_like(tXgX) + tWgW = thr_copy_W.partition_S(gW) + tWrW = cute.make_fragment_like(tWgW) + tXrW = thr_copy_X.retile(tWrW) + tXgR = thr_copy_O.partition_D(gR) + tXgO = thr_copy_O.partition_D(gY) + tXrO = cute.make_fragment_like(tXgO) + + # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. + tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) + if valid: + cute.copy(copy_atom_load, tXgX, tXrX) + cute.copy(copy_atom_load, tWgW, tWrW) + cute.copy(copy_atom_store, tXrX, tXgR) + + # flashinfer's RMSNormKernel arithmetic, in its order (see the header). + x = tXrX.load().to(cutlass.Float32) + x_sq = x * x + sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) + for i in cutlass.range_constexpr(5): + sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + if lane == 0: + reduction_buffer[0, warp] = sum_sq + cute.arch.barrier() + total = cutlass.Float32(0.0) + if lane < NORM_WARPS: + total = reduction_buffer[0, lane] + for i in cutlass.range_constexpr(5): + total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) + mean_sq = total / cutlass.Float32(hidden) + rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) + w = tXrW.load().to(cutlass.Float32) + y = x * rstd * (w + cutlass.Float32(0.0)) + tXrO.store(y.to(out.element_type)) + cute.copy(copy_atom_store, tXrO, tXgO) + + +@cute.jit +def m3_embed_norm( + ids: cute.Tensor, + table: cute.Tensor, + weight: cute.Tensor, + raw: cute.Tensor, + out: cute.Tensor, + vocab: cutlass.Int64, + eps: cutlass.Float32, + n_tokens: cutlass.Constexpr[int], + hidden: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, +) -> None: + # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. + blocks = hidden // (NORM_VEC * NORM_THREADS) + tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), + stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip + m3_embed_norm_kernel( + ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden + ).launch( + grid=[n_tokens, 1, 1], + block=[NORM_THREADS, 1, 1], + smem=NORM_WARPS * 4, + stream=stream, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py new file mode 100644 index 000000000000..0acfdecb196e --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host side of the MiniMax-M3 decode embedding + first-layer RMSNorm (m3_embedding_firstnorm_kernel.py): H 6144, +1..8 rows, into caller-owned buffers. + +Compiled on the first call per (rows, id dtype, table shape, PDL), which must happen outside CUDA-graph capture +(:class:`DecodeEmbedding` compiles at construction). +""" + +import threading +from typing import Dict, Tuple + +import cuda.bindings.driver as cuda_driver +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit + +HIDDEN = 6144 +MAX_ROWS = 8 + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + + +def arg(tensor): + return from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _check(ids, table, weight, raw, out) -> None: + rows = ids.shape[0] if ids.dim() == 1 else 0 + if not ( + 1 <= rows <= MAX_ROWS + and ids.dtype in (torch.int32, torch.int64) + and ids.is_contiguous() + and table.dtype == torch.bfloat16 + and table.dim() == 2 + and table.shape[1] == HIDDEN + and table.is_contiguous() + and weight.dtype == torch.bfloat16 + and tuple(weight.shape) == (HIDDEN,) + and all( + t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() + for t in (raw, out) + ) + ): + raise ValueError( + "m3_embed_norm: expects ids int32/int64 [1..8], a contiguous bf16 table [V, 6144], bf16 weight [6144] " + "and contiguous bf16 raw / out [rows, 6144]" + ) + + +def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = _compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with _lock: + fn = _compiled.get(key) + if fn is None: + stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) + fn = _compiled[key] = cute.compile( + _m3_embed_norm_jit, + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + _check(ids, table, weight, raw, out) + fn = _compile(ids, table, weight, eps, raw, out, use_pdl) + stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + arg(ids), + arg(table), + arg(weight), + arg(raw), + arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return out, raw + + +class DecodeEmbedding: + """Caller-owned buffers of one (rows, id dtype) decode embedding: ``raw`` (the gathered rows, the first residual) + and ``output`` (their RMSNorm), over ``table`` and the first layer's input ``norm``. Built (and compiled) before + CUDA-graph capture.""" + + def __init__(self, table, norm, rows, id_dtype, use_pdl=True): + assert not torch.cuda.is_current_stream_capturing() + assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN + assert ( + norm.weight.dtype == torch.bfloat16 + and norm.weight.shape == (HIDDEN,) + and not norm.use_gemma + ) + assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.norm, self.rows, self.id_dtype, self.use_pdl = ( + table, + norm, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(norm.variance_epsilon) + self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _compile(ids, table, norm.weight, self.eps, self.raw, self.output, use_pdl) + + def __call__(self, ids): + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return m3_embed_norm( + ids, self.table, self.norm.weight, self.eps, self.raw, self.output, self.use_pdl + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py new file mode 100644 index 000000000000..0aa75c0bf82c --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/__init__.py @@ -0,0 +1,16 @@ +# 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. +"""MiniMax-M3 MoE layers for low-batch decode as CuTe DSL kernels: ``m3_moe`` for 1..8 tokens (op.py) and ``m3_moe16`` +for 9..16 tokens (op16.py).""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py new file mode 100644 index 000000000000..e317e5e3087a --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe16_kernel.py @@ -0,0 +1,3871 @@ +# 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. +"""MiniMax-M3 routed + shared experts for 9..16 decode tokens on one TP rank: ``m3_moe16`` (v2). + +Same math as the 1..8-token engine (m3_moe_kernel.py), operation for operation for the routed part (bitwise equal to +TRTLLM-Gen's fused MoE on the fp32 router-GEMV logits); a different schedule for 9..16 rows: + +- Front: the hidden states land in smem by bulk copies in K chunks (the idle ring); all 16 warps run the router + GEMV (m3_router_gemv's per-thread order) and the NVFP4 quantization of B1 from smem; the logits cross CTAs as + tagged words; top-k one warp per token. +- Ring: FC1 / FC2 weight stages by TMA; their scale words go by 4-byte cp.async straight into the instruction atom + layout (SFA_CP; no raw scale box, no splicer warp), completing on the stage's full barrier. +- FC1 -> FC2 hand-off (B2IMG): the FC1 epilogue writes each unit's FP4 values and scale byte into a per-group global + image laid out as the FC2 B tile + its scale atoms; one counter per group (release add per finished unit, every + counter gains exactly 192 per call); warp 5 loads a completed image with two bulk copies. +- Shared expert on NSH more CTAs (MXFP8 tcgen05), as the 1..8 engine. +Configuration via M3_CONFIG (m_max, num_ctas, stages, nb2_cap, lag, shared, sfa_cp, b2buf, ...). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +H = 6144 +I_TP = 768 +TWO_I = 2 * I_TP +E_LOCAL = 128 +TOP_K = 4 +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 12) +assert 9 <= M_MAX <= 16, M_MAX +NCTA = _cfg("num_ctas", 128) +FUSED_ROUTER = True +NROUTE = E_LOCAL +NGEMV = (NROUTE + NCTA - 1) // NCTA +assert NGEMV == 1, "m3_moe16: one router expert per CTA (num_ctas >= 128)" +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +LAG = _cfg("lag", 3) +G_CAP = TOP_K * M_MAX +NG_CAP = M_MAX +THREADS = 512 +NWARPS = THREADS // 32 +GEMV_THREADS = 256 +SHARED = bool(_cfg("shared", 0)) +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +SFA_CP = bool( + _cfg("sfa_cp", 1) +) # A scales by 4-byte cp.async into the atom layout (else raw box + splicer warp) +B2IMG = True +PERM = bool(_cfg("perm", 1)) +PERM_MIN_NG = max( + 3, _cfg("perm_min_ng", 5) +) # (measured: helps at random routing G >= 23, hurts at spec G 13-14) +# a partial last slot group becomes group 1 (the plan must match: op16) +DBG_DELAY = bool( + _cfg("dbg_delay", 0) +) # (stress builds) pseudo-random producer / epilogue / release stalls +DBG_NOSTORE = bool(_cfg("dbg_nostore", 0)) # (timing experiment only: wrong results) +EPI_FENCE = bool( + _cfg("epi_fence", 0) +) # (rel_warp 0) writer lanes' fence.acq_rel.gpu before the release add +EPI_PFENCE = bool( + _cfg("epi_pfence", 1) +) # producer-side fence.proxy.async.global before the release add +REL_WARP = bool( + _cfg("rel_warp", 1) +) # warp 8 does the per-group release adds (off the epilogue's path) +NSH = 24 if SHARED else 0 +SH_I = 768 +SH_KB = H // 32 +SH_DKB = SH_I // 32 +SH_SPLIT = 2 +SH_GT = SH_KB // 4 // SH_SPLIT +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * (SH_I // 128) +SHX_BAR = 6 +SH_PRETILE = False +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = _cfg("n1", 48 if _cfg("m_max", 12) <= 12 else 64) # B rows QS q + t +N1_LD = 1 << (N1 - 1).bit_length() # tcgen05.ld .num (a power of two) covering N1 columns +QS = N1 // 4 +assert QS >= M_MAX +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 +K1_TILES = H // 256 +S1 = K1_TILES // C1 +NUM_BYTES_A = MMA_M * KT_BYTES +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 +KT2 = I_TP // 256 +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = _cfg("n2", (4 * M_MAX + 15) // 16 * 16) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +IMG_B2 = KT2 * NUM_BYTES_B2 # one group's B2 tile image (global and smem) +IMG_SF = KT2 * NUM_KB * 512 # its scale atoms +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) # every group counter gains 192 per call +B1IMG_OFF = None # set below (after N1): the activation image after the NG_CAP group images +TK_WARPS = [0, 1, 2, 3, 6, 8, 9, 10, 11, 12, 13, 14, 5, 7, 15, 4][ + : _cfg("m_max", 12) +] # top-k warp of token t (4 last: early0) +B1_WARP = 15 if _cfg("m_max", 12) <= 14 else -1 # (-1: warp 0 after the routing) +NB2_CAP = _cfg("nb2_cap", 2) +B1IMG_OFF = M_MAX * IMG_BYTES +B1IMG_BYTES = (S1 * N1 * 128 + M_MAX * (H // 16) + 15) // 16 * 16 +IMG_TOTAL = B1IMG_OFF + B1IMG_BYTES +UNITS_ROW = TWO_I // R1 +PARTIAL_EMPTY = -1 +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +TOPK_BAR = 4 +ALL_BAR = 5 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 16) +LG_REPL = _cfg("lg_repl", 64) +LG_STRIDE = M_MAX * E_LOCAL + 32 +W_WARP = 4 +X_WARP = 5 # B2 image loader; slots + plan +S_WARP = 6 # (SFA_CP = 0) scale splicer +MMA_WARP = 7 +EARLY0 = bool(_cfg("early0", 1)) +EVICT_FIRST = 0x12F0000000000000 +EVICT_LAST = 0x14F0000000000000 +_LOG2E = 1.4426950408889634 +NB2BUF = _cfg("b2buf", 2) +SPB = 1 +# x staging: K chunks of XCK elements, NXB smem buffers (rotating if fewer than the chunks) +XCK = _cfg("xck", 2048) +NXC = H // XCK +XCB = M_MAX * XCK * 2 # bytes of one chunk (all tokens) +STAGE_A = NUM_BYTES_A +STAGE_SF = NUM_BYTES_SFA if SFA_CP else NUM_BYTES_SFA_RAW + NUM_BYTES_SFA +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + NB2BUF * IMG_B2 + + IMG_SF + + NB2_CAP * M_MAX * TOP_K * R2 * 2 + + SPB * M_MAX * C1 * R1 * 4 + + (G_CAP + 1) * (NG_CAP + 1) * 4 + + 6 * 1024 +) +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // (STAGE_A + STAGE_SF) +ACC1_COL = 0 +ACC2_COL = 2 * N1 +ACC2_STRIDE = 1 << (N2 - 1).bit_length() +SFB1_COL = ACC2_COL + 2 * ACC2_STRIDE +SFB1_WSTEP = 1 +SFB2_COL = ( + SFB1_COL + NUM_SFB1_ATOMS * 4 +) # one TMEM slot (loaded after the previous group's FC2 MMAs) +SFA1_COL = SFB2_COL + KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +# smem: [ring A | ring SF | B2 bufs | B2 SF staging | B1]; x chunks alias [0, _OFF_B1) during the front +_OFF_SFA = STAGES * NUM_BYTES_A +_OFF_B2 = (_OFF_SFA + STAGES * STAGE_SF + 1023) // 1024 * 1024 +_OFF_SFB2 = _OFF_B2 + NB2BUF * IMG_B2 +_OFF_B1 = (_OFF_SFB2 + IMG_SF + 1023) // 1024 * 1024 +BIG = _OFF_B1 + S1 * NUM_BYTES_B1 +NXB = min(NXC, _OFF_B1 // XCB) +assert NXB >= 2, (NXB, _OFF_B1, XCB) +# shared CTA layout in BIG: one pass over all M tokens; B rows SHN q + t (gate_up, 2 K slices), t (down) +SHN = 16 +SH_NPASS = 1 +SH_HINT = EVICT_FIRST +SH_PF_POST = bool( + _cfg("sh_pf_post", 1) +) # L2 prefetch of the shared weights after the grid dependency (else before) +_SB_A = 0 +_SH_FIX = ( + SH_GT * 2 * SHN * 128 + + 6 * SHN * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + 64 * SHN * 4 + + 2 * 32 * SHN * 4 + + 2048 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * 2 * SHN * 128 +_SB_XC = _SB_BH + 6 * SHN * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + 64 * SHN * 4 +_SB_END = _SB_GU + 2 * 32 * SHN * 4 +SH_ACC_GU = 0 +SH_ACC_DN = 2 * SHN +SH_SFA = SH_ACC_DN + 2 * SHN +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +SH_TMEM = 256 +if SHARED: + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) +TD_STRIDE = 64 +PFD = _cfg("pfd", 0) # L2 prefetch distance in ring stages (0: off) +PF_WARP = 12 +TICKET = bool(_cfg("ticket", 0)) # CTA roles by entry order +# priority tokens: token 0 and PT1 (M / 2: the second request's first row at spec c2) get their logits published first +NPT = _cfg("npt", 0) +PT1 = M_MAX // 2 +PF_PRIO = bool( + _cfg("pf_prio", 0) +) # L2 prefetch of the priority tokens' experts (this CTA's 1/NCTA of each) +PF_W1 = TWO_I * (H // 2) // NCTA +PF_S1 = TWO_I * (H // 16) // NCTA +PF_W2 = H * (I_TP // 2) // NCTA +PF_S2 = H * (I_TP // 16) // NCTA +assert PF_W1 % 16 == 0 and PF_S1 % 16 == 0 and PF_W2 % 16 == 0 and PF_S2 % 16 == 0 + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + # B rows n >= 32 (N1 = 64) go to the atom's next 32-bit column (row n: lane n % 32, column n // 32) + for w in cutlass.range_constexpr((N1 + 31) // 32): + n = lane + cutlass.Int32(32 * w) + q = n // cutlass.Int32(QS) + tl = n - q * cutlass.Int32(QS) + sf_ok = (n < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr( + base_row + warp * cutlass.Int32(32), + base_col + SFB1_COL + a * 4 + w * SFB1_WSTEP, + ), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + gpos, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = gpos % cutlass.Int32(SH_NST) # ring position (stage g of this pass) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=SH_HINT, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue( + lane, + sid, + cutlass.Int32(g), + s_big, + sh_full, + tma_gu, + tma_dn, + gusf, + dnsf, + cutlass.Int32(g), + ) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + elif cutlass.const_expr(not SH_PF_POST): + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(SH_PF_POST): + if warp == 0: # after the grid dependency: DRAM is idle during the routed front + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + for ps in cutlass.range_constexpr(SH_NPASS): + TT0 = 0 # one pass: B rows SHN q + t (gate_up), t (down) + MS = M_MAX + GB = ps * SH_STAGES # ring positions before this pass + PAR = ps & 1 # parity of the per-pass barriers (B ready, accumulators) + if warp == 0: + for g in cutlass.range(SH_NST if ps == 0 else 0, SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) # ring position over the passes + _wait( + sh_empty.subview(gg % cutlass.Int32(SH_NST)), + ((gg // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf, gg) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=2 * SHN, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=SHN, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + gg = g + cutlass.Int32(GB) + st = gg % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (gg // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32((2 * SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32((SHN * 128) >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(SHN), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + if cutlass.const_expr(ps > 0): + _wait( + acc_dn, PAR ^ 1 + ) # the previous pass's MMAs are done: its B tiles and TMEM scales are free + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((MS * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(MS * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=(t + cutlass.Int32(TT0)) * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(SHN) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + + it * cutlass.Int32(2 * SHN * 128) + + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of + # k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(SHN - 1) + ok = (n < cutlass.Int32(2 * SHN)) & (t < cutlass.Int32(MS)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n // cutlass.Int32(SHN)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t + # at column 8 q + t + _wait(acc_gu, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=2 * SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(MS): + bigf.store( + cutlass.Float32(av[SHN + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) + * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(MS): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(SHN) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((32 * MS + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(32 * MS): + p = pq // cutlass.Int32(MS) + tt = pq - p * cutlass.Int32(MS) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(SHN) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(32) + p) * cutlass.Int32(SHN) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview( + (tt + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + + sid * cutlass.Int32(32) + + p + ) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((MS * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(MS * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = (t + cutlass.Int32(TT0)) * cutlass.Int32(SH_I) + kb * cutlass.Int32( + 32 + ) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + + kt * cutlass.Int32(SHN * 128) + + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (t & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + ( + ((c0 + cutlass.Int32(1)) ^ (t & cutlass.Int32(7))) + << cutlass.Int32(2) + ), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(MS) + for kt in cutlass.range_constexpr(6): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, PAR) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * SHN, + cutlass.Float32, + ), + num=SHN, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(MS): + _st_relaxed_u32( + dnx.subview(cutlass.Int32((tt + TT0) * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +# ============================================================================= +# v2 helpers: global release / acquire, proxy fences. +# ============================================================================= +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _atom_add_u64(addr, val, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr), cutlass.Int64(val)], + "atom.relaxed.gpu.global.add.u64 $0, [$1], $2;", + "=l,l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.proxy.async.global;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fence_acq_rel_gpu(*, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + "fence.acq_rel.gpu;", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _dbg_delay(salt): + """(DBG_DELAY stress builds) a pseudo-random ~0..4 us stall: protocol robustness under skew.""" + if cutlass.const_expr(DBG_DELAY): + r = cutlass.Int32(_gtime() >> cutlass.Int64(6)) * cutlass.Int32(1103515245) + salt + if ((r >> cutlass.Int32(12)) & cutlass.Int32(7)) == cutlass.Int32(0): + for _ in cutlass.range((r >> cutlass.Int32(16)) & cutlass.Int32(31), unroll=1): + _nanosleep(128) + + +@cute.jit +def _waitt(bar: cutlass.Array, parity, acc): + """_wait that (TIMING) adds the waited ns to acc.""" + t0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + if cutlass.const_expr(TIMING): + acc = acc + (_gtime() - t0) + return acc + + +@cute.jit +def _tstore(tdbg: cutlass.Array, bx, i: cutlass.Constexpr, v): + if cutlass.const_expr(TIMING): + tdbg.store(cutlass.Int64(v), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@cute.jit +def _issue_fc1( + lane, + stage, + e, + c, + j, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a1, + tma_s1, + s1w: cutlass.Array, +): + """FC1 stage j of unit c of expert e into ring stage `stage` (the caller waited its empty barrier): 4 TMA boxes + (rows 32 c .., k-tiles 4 j + q) and the stage's scale atoms (SFA_CP: word [r][q] of atom kb <- raw atom + (k-tile 4 j + q, kb) word [r][c % 4], by this warp's 32 lanes).""" + row = c * cutlass.Int32(R1) + kx = j * cutlass.Int32(C1 * KT_BYTES) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview(stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(q * NUM_BYTES_BOX1)), + tma_a1, + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32(SF1_KA) + + j * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview(stage * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1, + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + atom0 = (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) * cutlass.Int32( + SF1_KA + ) + j * cutlass.Int32(16) + src0 = atom0 * cutlass.Int32(128) + lane * cutlass.Int32(4) + (c & cutlass.Int32(3)) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _issue_fc2( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sA: cutlass.Array, + sSF: cutlass.Array, + ab_full: cutlass.Array, + tma_a2, + tma_s2, + s2w: cutlass.Array, +): + """FC2 k-tile t2 of output block blk for the group's n_here slots: box jj = rows 32 blk .. of slot jj's w2 at + k-tile t2; scale word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) word [r][blk % 4].""" + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32(NUM_BYTES_BOX1 if SFA_CP else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2, + (cutlass.Int32(t2 * KT_BYTES), blk * cutlass.Int32(R2), e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSF.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2, + ( + cutlass.Int32(0), + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSF.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + # one group per ring stage (racecheck: the issuer waits it before the slot's reuse) + prims.cp_async_commit_group() + cute.arch.sync_warp() + + +@cute.jit +def _quant_words(w, sc, rsc, rsix): + """NVFP4 of one 16-value block given as 8 bf16-pair words: (lo word, hi word, scale e4m3 bits).""" + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.jit +def _quant_block(xs32: cutlass.Array, widx, sc, rsc, rsix): + """NVFP4 of one 16-value block (8 bf16 words at widx): (lo word, hi word, scale e4m3 bits), fp4_quantize's ops.""" + w = [] + for h in cutlass.range_constexpr(2): + v4 = xs32.load(idx=widx + cutlass.Int32(4 * h), vector_size=4, alignment=16) + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), cutlass.Float32(0.0) + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + return lo, hi, sfb + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_gu_desc: cutlass.GridConstant[cuda.TensorMap], + tma_dn_desc: cutlass.GridConstant[cuda.TensorMap], + bias: cutlass.Array, # f32 [128] + escal: cutlass.Array, # f32 [6, 128] + plan: cutlass.Array, # int32 [4 M + 1, NCTA, M + 1] + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + img32: cutlass.Array, # int32 [NG_CAP, IMG_BYTES / 4]: per-group FC2 B image (tile + scale atoms) + cnt: cutlass.Array, # int32 [NG_CAP * 8]: per-group unit counters (index 8 g) + s1w: cutlass.Array, # int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # int32 words of w2_weight_scale + w1b: cutlass.Array, # int8 bytes of w3_w1_weight (L2 prefetch) + w2b: cutlass.Array, # int8 bytes of w2_weight (L2 prefetch) + epochs: cutlass.Array, # int32 [NCTA + NSH] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8] + tdbg: cutlass.Array, # int64 [NCTA + NSH, 64] + xb32: cutlass.Array, # int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # f32 router weight [128, 6144] + sfsc: cutlass.Array, # f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # int64 [LG_REPL, LG_STRIDE] + guw: cutlass.Array, + gusf: cutlass.Array, + dnw: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if cutlass.const_expr(TICKET): + # role by entry order: the CTAs that become resident last (behind the predecessor's CTAs) take the shared + # expert roles (most slack), the first NCTA the routed ones (the router GEMV / logits exchange) + s_tk = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + if tidx == 0: + old = _atom_add_u64( + cnt.subview(cutlass.Int32(NG_CAP * 8 + 8)).data_ptr().toint(), cutlass.Int64(1) + ) + s_tk.store(cutlass.Int32(old % cutlass.Int64(NCTA + NSH)), idx=0) + prims.barrier_cta_sync(0) + bx = s_tk.load(idx=0) + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + # ring scale words: SFA_CP: the instruction atoms [STAGES][4 x 512 B]; else raw boxes [STAGES][8 KB] + spliced + sSF = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(STAGES * STAGE_SF // 4,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA + (0 if SFA_CP else STAGES * NUM_BYTES_SFA_RAW)), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=512, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(NB2BUF * IMG_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), shape=(IMG_SF // 4,), dtype=cutlass.Int32, alignment=1024 + ) + ys = cutlass.Array( + cutlass.BFloat16, + NB2_CAP * M_MAX * TOP_K * R2, + space=cutlass.AddressSpace.smem, + alignment=16, + ) # bf16 values + s_part = cutlass.Array( + cutlass.Float32, SPB * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array( + cutlass.Float32, 9 * M_MAX + 16, space=cutlass.AddressSpace.smem, alignment=16 + ) + # the activation's scale codes [M, 384] until they are in TMEM: aliases ys (first written by the FC2 epilogue) + s_xsf = cutlass.Array( + ys.data_ptr(0), shape=(M_MAX * (H // 64),), dtype=cutlass.Int32, alignment=16 + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + img_bar = cutlass.Array(cutlass.Int64, NB2BUF, space=cutlass.AddressSpace.smem, alignment=8) + sfb2_free = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_full = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + x_empty = cutlass.Array(cutlass.Int64, NXB, space=cutlass.AddressSpace.smem, alignment=8) + route_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + ys_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + grp_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_load = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + topk_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tok0_bar = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + px_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + s_prog = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array(cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + s_gn = cutlass.Array( + cutlass.Int32, NG_CAP, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots per group + s_emap = cutlass.Array(cutlass.Int32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array(cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32) + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + if tidx < NB2BUF: + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), 1) + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < NXB: + prims.mbarrier_init(x_full.subview(tidx), 1) + prims.mbarrier_init(x_empty.subview(tidx), NWARPS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + prims.mbarrier_init(sfb2_free.subview(0), 1) + prims.mbarrier_init(route_bar.subview(0), 32) + prims.mbarrier_init(ys_done.subview(0), EPI_THREADS) + prims.mbarrier_init(b1_load.subview(0), 1) + prims.mbarrier_init(topk_bar.subview(0), M_MAX * TOP_K) + prims.mbarrier_init(tok0_bar.subview(0), TOP_K) + prims.mbarrier_init(px_full.subview(0), 1) + if tidx < 4: + prims.mbarrier_init(grp_done.subview(tidx), EPI_THREADS) + if tidx == 0: + s_prog.store(cutlass.Int32(0), idx=0) + if warp == X_WARP: # slot tables: empty (the slot warp fills them after the top-k) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + for z in cutlass.range_constexpr(E_LOCAL // 32): + s_emap.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + if cutlass.const_expr(not SFA_CP): + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + # router weights (expert bx) in registers: thread tg = tid % 256 owns elements 4 tg + 1024 j + tg = tidx % cutlass.Int32(GEMV_THREADS) + half = tidx // cutlass.Int32(GEMV_THREADS) + rwv = [ + rw.load( + idx=bx * cutlass.Int32(H) + tg * cutlass.Int32(4) + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + # ---- x chunks into the idle ring (warp 6): chunk c = K [XCK c, XCK (c + 1)) of every token + if warp == S_WARP: + if cutlass.const_expr(NPT > 0): + # the priority tokens' rows first (their logits are published ahead of the others) + if lane == 0: + prims.mbarrier_arrive_expect_tx(px_full.subview(0), NPT * H * 2) + cute.arch.sync_warp() + if lane < NPT: + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(lane * cutlass.Int32(H * 2)), + xb32.subview( + cutlass.select_( + lane == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1) + ) + * cutlass.Int32(H // 2) + ), + px_full.subview(0), + H * 2, + ) + for c in cutlass.range_constexpr(NXB): + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(c), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(c * XCB) + lane * cutlass.Int32(XCK * 2)), + xb32.subview(lane * cutlass.Int32(H // 2) + cutlass.Int32(c * XCK // 2)), + x_full.subview(c), + XCK * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + img8f = cutlass.Array( + img32.data_ptr(0), shape=(IMG_TOTAL,), dtype=cutlass.Int8, alignment=16 + ) + if cutlass.const_expr(NPT > 0): + # ---- priority tokens (0 and PT1): their logits first (m3_router_gemv's per-thread / warp order) + _wait(px_full.subview(0), 0) + ph = cutlass.select_(half < cutlass.Int32(NPT), half, cutlass.Int32(0)) + pacc = cutlass.Float32(0.0) + px32 = cutlass.Array( + sB1.data_ptr(0), shape=(NPT * H // 2,), dtype=cutlass.Int32, alignment=16 + ) + for j in cutlass.range_constexpr(H // 1024): + pxw = px32.load( + idx=ph * cutlass.Int32(H // 2) + cutlass.Int32(512 * j) + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + px0 = cutlass.Int32(pxw[0]) + px1 = cutlass.Int32(pxw[1]) + pacc = _fma_rn( + (px0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][0]), + pacc, + ) + pacc = _fma_rn( + (px0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][1]), + pacc, + ) + pacc = _fma_rn( + (px1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][2]), + pacc, + ) + pacc = _fma_rn( + (px1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[j][3]), + pacc, + ) + for off in (16, 8, 4, 2, 1): + pacc = _add_rn(pacc, cute.arch.shuffle_sync_bfly(pacc, off)) + if lane == 0: + s_red.store(pacc, idx=cutlass.Int32(9 * M_MAX) + warp) # [half][8 warps] + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if tidx < cutlass.Int32(NPT * LG_REPL): + hh = tidx // cutlass.Int32(LG_REPL) + rr = tidx - hh * cutlass.Int32(LG_REPL) + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn( + sacc, s_red.load(idx=cutlass.Int32(9 * M_MAX + w) + hh * cutlass.Int32(8)) + ) + tp = cutlass.select_(hh == cutlass.Int32(0), cutlass.Int32(0), cutlass.Int32(PT1)) + _st_relaxed_u64( + lgbuf.subview(rr * cutlass.Int32(LG_STRIDE) + tp * cutlass.Int32(E_LOCAL) + bx) + .data_ptr() + .toint(), + ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)), + ) + if tidx == 0: + _stamp(tdbg, bx, 43) + # ---- this CTA's share of the NVFP4 activation (QB blocks, from global) for the global B1 image: loads first + QB = (M_MAX * (H // 16) + NCTA - 1) // NCTA + qb = bx * cutlass.Int32(QB) + tidx + q_ok = (tidx < cutlass.Int32(QB)) & (qb < cutlass.Int32(M_MAX * (H // 16))) + qb = cutlass.select_(q_ok, qb, cutlass.Int32(0)) + qtok = qb // cutlass.Int32(H // 16) + qkb = qb - qtok * cutlass.Int32(H // 16) + xq_v = [ + xb32.load( + idx=qtok * cutlass.Int32(H // 2) + qkb * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + for h in range(2) + ] + # ---- router GEMV from the staged chunks, all 512 threads (rolled: the front runs once per call) + MH = (M_MAX + 1) // 2 # tokens per GEMV half + acc_r = cute.make_rmem_tensor((MH,), cutlass.Float32) + for i in cutlass.range_constexpr(MH): + acc_r[i] = cutlass.Float32(0.0) + big32 = cutlass.Array( + s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16 + ) + JPC = XCK // 1024 # router j-steps per chunk + for xc in cutlass.range(NXC, unroll=1): + buf = xc % cutlass.Int32(NXB) + _wait(x_full.subview(buf), (xc // cutlass.Int32(NXB)) & cutlass.Int32(1)) + if (tidx == 0) & (xc == cutlass.Int32(0)): + _stamp(tdbg, bx, 41) + xw0 = buf * cutlass.Int32(XCB // 4) # chunk word base in s_big + for jj in cutlass.range_constexpr(JPC): + jg = xc * cutlass.Int32(JPC) + cutlass.Int32(jj) + wsel = [cutlass.Float32(rwv[0][z]) for z in range(4)] + for jx in cutlass.range_constexpr(1, H // 1024): + for z in cutlass.range_constexpr(4): + wsel[z] = cutlass.select_( + jg == cutlass.Int32(jx), cutlass.Float32(rwv[jx][z]), wsel[z] + ) + for i in cutlass.range_constexpr(MH): + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + xw = big32.load( + idx=xw0 + + t * cutlass.Int32(XCK // 2) + + cutlass.Int32(512 * jj) + + tg * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + acc_g = acc_r[i] + acc_g = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[0], acc_g + ) + acc_g = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[1], acc_g + ) + acc_g = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), wsel[2], acc_g + ) + acc_g = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), wsel[3], acc_g + ) + acc_r[i] = acc_g + if cutlass.const_expr(NXB < NXC): + if xc + cutlass.Int32(NXB) < cutlass.Int32(NXC): + # every thread's reads of this chunk are done (CTA barrier, then a generic -> async proxy fence by + # the refilling warp: the bulk copy overwrites what the generic proxy read); racecheck-clean + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == S_WARP: + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_full.subview(buf), XCB) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + buf * cutlass.Int32(XCB) + lane * cutlass.Int32(XCK * 2) + ), + xb32.subview( + lane * cutlass.Int32(H // 2) + + (xc + cutlass.Int32(NXB)) * cutlass.Int32(XCK // 2) + ), + x_full.subview(buf), + XCK * 2, + ) + if tidx == 0: + _stamp(tdbg, bx, 42) + for i in cutlass.range_constexpr(MH): + v = acc_r[i] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + t = half * cutlass.Int32(MH) + cutlass.Int32(i) + if (lane == 0) & (t < cutlass.Int32(M_MAX)): + s_red.store(v, idx=(warp % cutlass.Int32(8)) * cutlass.Int32(M_MAX) + t) + # this CTA's quantization blocks into the global B1 image (B1 tile bytes, then the [M, 384] scale codes) + if q_ok: + qw = [] + for h in cutlass.range_constexpr(2): + for z in cutlass.range_constexpr(4): + qw.append(cutlass.Int32(xq_v[h][z])) + qlo, qhi, qsf = _quant_words(qw, sc, rsc, rsix) + img32.store( + (qlo, qhi), idx=cutlass.Int32(B1IMG_OFF // 4) + _b1_addr(qtok, qkb), alignment=8 + ) + img8f.store(cutlass.Int8(qsf), idx=cutlass.Int32(B1IMG_OFF + S1 * NUM_BYTES_B1) + qb) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ALL_BAR, number_of_threads=THREADS) + if warp == cutlass.Int32(0): + _dbg_delay(bx * cutlass.Int32(3)) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX + t))) + lw = ltag | (cutlass.Int64(sacc.bitcast(cutlass.Int32)) & cutlass.Int64(0xFFFFFFFF)) + _st_relaxed_u64( + lgbuf.subview(tidx * cutlass.Int32(LG_STRIDE) + cutlass.Int32(t * E_LOCAL) + bx) + .data_ptr() + .toint(), + lw, + ) + if tidx == 0: + _stamp(tdbg, bx, 21) + if tidx == cutlass.Int32( + THREADS - 32 + ): # B1 share released after the logits (its gpu fence stalls warp 15 only) + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() # once, after the join (ALL_BAR) of the writer threads + _red_release_add( + cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint(), cutlass.Int32(1) + ) + # ---- the B1 image (every CTA's share) into smem (warp 13; the top-k warps are 0 .. M - 1) + if cutlass.const_expr(B1_WARP >= 0): + if warp == cutlass.Int32(B1_WARP): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + + # ---- top-k: token t on warp t (keys sigmoid + bias, ties to the lower id), TRTLLM-Gen MiniMax2 ops + tk_t = cutlass.Int32(-1) # this warp's top-k token (TK_WARPS[t] = warp) + for i_tk in cutlass.range_constexpr(M_MAX): + tk_t = cutlass.select_(warp == cutlass.Int32(TK_WARPS[i_tk]), cutlass.Int32(i_tk), tk_t) + if tk_t >= cutlass.Int32(0): + t = tk_t + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 22) + lgv = [cutlass.Int32(res[z]).bitcast(cutlass.Float32) for z in range(4)] + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hk = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hk[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + if lane < TOP_K: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + if (lane == 0) & (t == cutlass.Int32(0)): + _stamp(tdbg, bx, 20) + if cutlass.const_expr(PF_PRIO): + if (t == cutlass.Int32(0)) | (t == cutlass.Int32(PT1)): + if lane < TOP_K: + # this CTA's 1/NCTA of each of the token's experts into L2 (DRAM is idle until the routing is + # complete; every byte is streamed later anyway) + _l2_prefetch( + w1b.subview( + e_l * cutlass.Int32(TWO_I * H // 2) + bx * cutlass.Int32(PF_W1) + ) + .data_ptr() + .toint(), + PF_W1, + ) + _l2_prefetch( + s1w.subview( + (e_l * cutlass.Int32(TWO_I * H // 16) + bx * cutlass.Int32(PF_S1)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S1, + ) + _l2_prefetch( + w2b.subview( + e_l * cutlass.Int32(H * I_TP // 2) + bx * cutlass.Int32(PF_W2) + ) + .data_ptr() + .toint(), + PF_W2, + ) + _l2_prefetch( + s2w.subview( + (e_l * cutlass.Int32(H * I_TP // 16) + bx * cutlass.Int32(PF_S2)) + // cutlass.Int32(4) + ) + .data_ptr() + .toint(), + PF_S2, + ) + cute.arch.sync_warp() + if lane < TOP_K: # every writer lane arrives (its own release of s_ids / s_w) + prims.mbarrier_arrive(topk_bar.subview(0)) # this token's ids / weights are in smem + if t == cutlass.Int32(0): + prims.mbarrier_arrive( + tok0_bar.subview(0) + ) # token 0's experts: the early group-0 stages + cute.arch.sync_warp() + + if cutlass.const_expr(EARLY0): + if warp == W_WARP: + # slots 0..3 = token 0's four distinct experts; group 0's stage ranges are the same for every G + _wait(tok0_bar.subview(0), 0) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(STAGES): + _issue_fc1( + lane, + n0, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + n0 = n0 + cutlass.Int32(1) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if warp == X_WARP: + # slots: entry i < 4 M is (token i / 4, top-k i % 4); distinct experts in entry order + NCH = (TOP_K * M_MAX + 31) // 32 + _wait(topk_bar.subview(0), 0) + gcnt = cutlass.Int32(0) + rec = [] # per chunk: (first, rank, slot rank of the entry, expert, entry valid) + for ch in cutlass.range_constexpr(NCH): + ei = lane + cutlass.Int32(32 * ch) + valid = ei < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, ei, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + known = cutlass.Boolean(False) + prev = cutlass.Int32(-1) + key = my_e + if cutlass.const_expr(ch > 0): + prev = cutlass.select_( + valid, + s_emap.load(idx=cutlass.select_(valid, my_e, cutlass.Int32(0))), + cutlass.Int32(-1), + ) + known = prev >= cutlass.Int32(0) + key = cutlass.select_(known, cutlass.Int32(-1000) - lane, my_e) + cute.arch.sync_warp() # every lane's s_emap read before any lane's write + same = _match_any(key) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (known == cutlass.Boolean(False)) & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = gcnt + cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + slot_l = cutlass.select_(known, prev, slot_l) + if first: + s_emap.store(rank, idx=my_e) + if valid: + if bx == 0: + route_dbg.store( + my_e, + idx=(ei // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + ei % cutlass.Int32(TOP_K), + ) + rec.append((first, rank, slot_l, my_e, valid)) + cute.arch.sync_warp() + gcnt = gcnt + cute.arch.popc(bal) + # slot of rank r: with a partial last group (tail < 4) and >= 3 groups, that group becomes group 1 (its + # FC2 then is not the last hand-off behind a short FC1 phase); the full groups 1.. shift up by one + ngr = (gcnt + cutlass.Int32(3)) >> cutlass.Int32(2) + tail = gcnt - (ngr - cutlass.Int32(1)) * cutlass.Int32(4) + do_perm = ( + cutlass.Boolean(PERM) + & (tail < cutlass.Int32(4)) + & (ngr >= cutlass.Int32(PERM_MIN_NG)) + ) + for z in cutlass.range_constexpr((G_CAP + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP): + s_el.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + for ch in cutlass.range_constexpr(NCH): + first, rank, slot_r, my_e, valid = rec[ch] + ei = lane + cutlass.Int32(32 * ch) + for which in cutlass.range_constexpr(2): + r = rank if which == 0 else slot_r + sl_m = cutlass.select_( + r < cutlass.Int32(4), + r, + cutlass.select_( + r >= (ngr - cutlass.Int32(1)) * cutlass.Int32(4), + r - (ngr - cutlass.Int32(2)) * cutlass.Int32(4), + r + cutlass.Int32(4), + ), + ) + sl_m = cutlass.select_(do_perm, sl_m, r) + if cutlass.const_expr(which == 0): + if first: + s_el.store(my_e, idx=sl_m) + else: + if valid: + s_kpos.store( + ei % cutlass.Int32(TOP_K), + idx=sl_m * cutlass.Int32(M_MAX) + ei // cutlass.Int32(TOP_K), + ) + # slots per group + if lane < cutlass.Int32(NG_CAP): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + nh = cutlass.select_(do_perm & (lane == cutlass.Int32(1)), tail, nh) + nh = cutlass.select_( + do_perm & (lane == ngr - cutlass.Int32(1)), cutlass.Int32(4), nh + ) + s_gn.store(nh, idx=lane) + cute.arch.sync_warp() + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + prims.mbarrier_arrive( + route_bar.subview(0) + ) # every lane: its slots / kpos / plan-row writes are in smem + cute.arch.sync_warp() + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = s_gn.load(idx=lane) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add(cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad) + else: + _wait(route_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(B1_WARP < 0): + if warp == cutlass.Int32(0): + if lane == 0: + b1addr = cnt.subview(cutlass.Int32(NG_CAP * 8)).data_ptr().toint() + b1tgt = (ep + cutlass.Int32(1)) * cutlass.Int32(NCTA) + cv = _ld_acquire_u32(b1addr) + while (cv - b1tgt) < cutlass.Int32(0): + _nanosleep(64) + cv = _ld_acquire_u32(b1addr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx( + b1_load.subview(0), S1 * NUM_BYTES_B1 + M_MAX * (H // 16) + ) + prims.cp_async_bulk_shared_cluster_global( + sB1.subview(0), + img32.subview(cutlass.Int32(B1IMG_OFF // 4)), + b1_load.subview(0), + S1 * NUM_BYTES_B1, + ) + prims.cp_async_bulk_shared_cluster_global( + s_xsf.subview(0), + img32.subview(cutlass.Int32((B1IMG_OFF + S1 * NUM_BYTES_B1) // 4)), + b1_load.subview(0), + M_MAX * (H // 16), + ) + _stamp(tdbg, bx, 23) + cute.arch.sync_warp() + if warp < 4: # B1 scales to TMEM (the MMA needs them only with the first ring stage) + _wait(b1_load.subview(0), 0) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + n_iss = cutlass.Int32( + 0 + ) # stages issued (early ones included): the L2-prefetch warp's throttle + w_emp = cutlass.Int64(0) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _dbg_delay(bx) + _issue_fc1( + lane, + stage, + e, + c, + j, + sA, + sSF, + ab_full, + tma_a1_desc.get_ptr(), + tma_s1_desc.get_ptr(), + s1w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + w_emp = _waitt(ab_empty.subview(stage), phase, w_emp) + prims.cp_async_wait_group(STAGES - 1) + _issue_fc2( + lane, + stage, + g, + blk, + t2, + n_here, + s_el, + sA, + sSF, + ab_full, + tma_a2_desc.get_ptr(), + tma_s2_desc.get_ptr(), + s2w, + ) + stage, phase = _adv(stage, phase) + if cutlass.const_expr(PFD > 0): + n_iss = n_iss + cutlass.Int32(1) + if lane == 0: + prims.store_ext( + n_iss.ir_value(), + s_prog.subview(0), + order="relaxed", + scope="cta", + ) + if lane == 0: + _stamp(tdbg, bx, 3) + _tstore(tdbg, bx, 30, w_emp) + + # ---- (SFA_CP = 0) scale splicer (warp 6): atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if cutlass.const_expr(not SFA_CP): + if warp == S_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt_ = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt_ = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt_ = cutlass.Int32(KT2) + for _j in cutlass.range(cnt_, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32( + 4 + ) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSF.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSF.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSF.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSF.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + # ---- FC2 B loader (warp 5): group g's image once its counter is complete + if warp == X_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for g in range(ng): + buf = g % cutlass.Int32(NB2BUF) + if g >= cutlass.Int32(NB2BUF): + _waitd( + b2_free.subview(buf), + ((g // cutlass.Int32(NB2BUF)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 29) + caddr = cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2_POLL_NS) + cv = _ld_acquire_u32(caddr) + cute.arch.sync_warp() + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 51) + _fence_proxy_async_global() + if g >= cutlass.Int32(1): + _waitd( + sfb2_free.subview(0), (g - cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 107 + ) + if lane == 0: + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(0), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), tdbg, 108 + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 52) + if g >= cutlass.Int32(1): + # the single TMEM scale slot: the previous group's FC2 MMAs are done + gp = g - cutlass.Int32(1) + _waitd( + b2_free.subview(gp % cutlass.Int32(NB2BUF)), + (gp // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), + tdbg, + 109, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 53) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for aa in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFB2_COL + aa * 4), + sd0 + cutlass.Int64(aa * (512 >> 4)), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + prims.tcgen05_commit(sfb2_free.subview(0)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + + # ---- (PFD) L2 prefetch warp (12): the same stage sequence as the producer, PFD stages ahead of it + if cutlass.const_expr(PFD > 0): + if warp == cutlass.Int32(PF_WARP): + pf = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + for j in cutlass.range(j0, j1, unroll=1): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + kx = j * cutlass.Int32(C1 * KT_BYTES) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_prefetch( + tma_a1_desc.get_ptr(), + [ + kx + cutlass.Int32(q * KT_BYTES), + c * cutlass.Int32(R1), + e, + ], + [], + ) + atom0 = ( + e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2)) + ) * cutlass.Int32(SF1_KA) + j * cutlass.Int32(16) + _l2_prefetch( + s1w.subview(atom0 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 8192, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = s_gn.load(idx=g) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + if lane == 0: + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + while prg + cutlass.Int32(PFD) <= pf: + _nanosleep(100) + prg = cutlass.Int32( + prims.load_ext( + s_prog.subview(0), + dtype=cutlass.Int32, + order="relaxed", + scope="cta", + ) + ) + if pf >= prg: + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load( + idx=g * cutlass.Int32(4) + cutlass.Int32(jj) + ) + prims.cp_async_bulk_tensor_prefetch( + tma_a2_desc.get_ptr(), + [ + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ], + [], + ) + a2 = ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) * cutlass.Int32(SF2_KA) + cutlass.Int32(4 * t2) + _l2_prefetch( + s2w.subview(a2 * cutlass.Int32(128)) + .data_ptr() + .toint(), + 2048, + ) + cute.arch.sync_warp() + pf = pf + cutlass.Int32(1) + + # ---- (REL_WARP) group releases (warp 8): this CTA's finished units of group p, one release add per phase + if cutlass.const_expr(REL_WARP): + if warp == cutlass.Int32(8): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + grp_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + _dbg_delay(bx * cutlass.Int32(13) + p) + if (lane == 0) & (n_done > cutlass.Int32(0)): + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + _red_release_add( + cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + if (lane == 0) & (p == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 54) + cute.arch.sync_warp() + + # ---- MMA (warp 7) + if warp == MMA_WARP: + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + m_b2 = cutlass.Int64(0) + m_full = cutlass.Int64(0) + m_acc = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + m_acc = _waitt( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + cute.arch.sync_warp() + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g % cutlass.Int32(NB2BUF) + m_b2_0 = m_b2 + m_b2 = _waitt( + b2_ready.subview(buf), (g // cutlass.Int32(NB2BUF)) & cutlass.Int32(1), m_b2 + ) + if cutlass.const_expr(TIMING): + if (lane == 0) & (g < cutlass.Int32(4)): + tdbg.store( + m_b2 - m_b2_0, + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(44) + g, + ) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + m_acc = _waitt( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + m_acc, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + if cutlass.const_expr(SFA_CP): + m_full = _waitt(ab_full.subview(stage), phase, m_full) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB2_COL + cutlass.Int32(t2 * NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + cute.arch.sync_warp() + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + cute.arch.sync_warp() + k2 = k2 + cutlass.Int32(1) + if prims.elect_sync(): + prims.tcgen05_commit( + b2_free.subview(buf) + ) # this group's B2 tile and TMEM scales are free + cute.arch.sync_warp() + if lane == 0: + _tstore(tdbg, bx, 26, m_b2) + _tstore(tdbg, bx, 27, m_full) + _tstore(tdbg, bx, 28, m_acc) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + img8 = cutlass.Array( + img32.data_ptr(0), shape=(NG_CAP * IMG_BYTES,), dtype=cutlass.Int8, alignment=16 + ) + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if qi < cutlass.Int32(6 * G_CAP): + e_sl = s_el.load(idx=sl_e) + if e_sl >= cutlass.Int32(0): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32(escal.load(idx=fld * cutlass.Int32(E_LOCAL) + e_sl)), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + e_wait = cutlass.Int64(0) + e_busy = cutlass.Int64(0) + e_a = cutlass.Int64(0) + e_pw = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + unit_done = (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)) + e_wait = _waitt( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + tb0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1_LD, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=(t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + tc0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_a = e_a + (tc0 - tb0) + g_u = slot >> cutlass.Int32(2) + jrow = slot & cutlass.Int32(3) + # this warp's tokens t = warp + 4 tr, all in flight together (latency-bound otherwise) + NT = (M_MAX + 3) // 4 + xr = cute.make_rmem_tensor((NT,), cutlass.Float32) + for tr in cutlass.range_constexpr(NT): + tq = cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + xr[tr] = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=(tq * C1 + 0) * R1 + lane), + s_part.load(idx=(tq * C1 + 1) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 2) * R1 + lane), + ), + s_part.load(idx=(tq * C1 + 3) * R1 + lane), + ) + if (j0 == cutlass.Int32(0)) & (j1 != cutlass.Int32(S1)): + # head of a unit split with the next CTA: its partial sums + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + _st_relaxed_u32( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint(), + xr[tr].bitcast(cutlass.Int32), + ) + else: + if j0 != cutlass.Int32(0): + # tail: the head's partial sums (head first, as TRTLLM-Gen's K order) + for tr in cutlass.range_constexpr(NT): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + pbits = _ld_relaxed_u32(paddr) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xr[tr] = _add_rn(pbits.bitcast(cutlass.Float32), xr[tr]) + cute.arch.sync_warp() + _dbg_delay(bx * cutlass.Int32(7) + warp) + es = slot * cutlass.Int32(8) + ea = s_esc.load(idx=es) + ec = s_esc.load(idx=es + cutlass.Int32(1)) + eal = s_esc.load(idx=es + cutlass.Int32(2)) + ebe = s_esc.load(idx=es + cutlass.Int32(3)) + eli = s_esc.load(idx=es + cutlass.Int32(4)) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + n0 = jrow * cutlass.Int32(M_MAX) + t2 = cblk >> cutlass.Int32(4) + cc = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = g_u * cutlass.Int32(IMG_BYTES) + sfoff = ( + gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2 * cutlass.Int32(NUM_KB) + + ((cblk & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(512) + ) + + (cblk & cutlass.Int32(3)) + ) + for tr in cutlass.range_constexpr(NT): + # (skip) only tokens routed to this slot (warp-uniform); other B2 rows are never used + if s_kpos.load( + idx=slot * cutlass.Int32(M_MAX) + + cutlass.select_( + warp + cutlass.Int32(4 * tr) < cutlass.Int32(M_MAX), + warp + cutlass.Int32(4 * tr), + cutlass.Int32(0), + ) + ) >= cutlass.Int32(0): + xv = xr[tr] + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + qv, sfb = _swiglu_quant( + xv, gate_raw, is_up, ea, ec, eal, ebe, eli + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + t = warp + cutlass.Int32(4 * tr) + if ( + (lane == 0) + & (t < cutlass.Int32(M_MAX)) + & cutlass.Boolean(not DBG_NOSTORE) + ): + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), row n + # at byte 16 (n % 32) + 4 (n / 32), byte c % 4 + n = n0 + t + img32.store( + (lo, hi), + idx=( + gbase + + t2 * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=sfoff + + (n & cutlass.Int32(31)) * cutlass.Int32(16) + + (n >> cutlass.Int32(5)) * cutlass.Int32(4), + ) + tc1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(TIMING): + e_pw = e_pw + (tc1 - tc0) + if cutlass.const_expr(not REL_WARP): + if unit_done: + if lane == 0: + if cutlass.const_expr(EPI_PFENCE): + _fence_proxy_async_global() + if cutlass.const_expr(EPI_FENCE): + _fence_acq_rel_gpu() + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=EPI_BAR, number_of_threads=EPI_THREADS + ) # s_part reads done + if cutlass.const_expr(not REL_WARP): + if unit_done & (tidx == 0): + _red_release_add( + cnt.subview(g_u * cutlass.Int32(8)).data_ptr().toint(), + cutlass.Int32(1), + ) + if cutlass.const_expr(TIMING): + e_busy = e_busy + (_gtime() - tb0) + k1 = k1 + cutlass.Int32(1) + if cutlass.const_expr(REL_WARP): + prims.mbarrier_arrive( + grp_done.subview(p & cutlass.Int32(3)) + ) # this CTA's group-p units are out + if tidx == 0: + _stamp(tdbg, bx, 10) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + e_wait = _waitt( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + e_wait, + ) + if cutlass.const_expr(TIMING): + if ( + (tidx == 0) + & (g == ng - cutlass.Int32(1)) + & (bb == cutlass.Int32(0)) + ): + _stamp(tdbg, bx, 8) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = [] + ya = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + ACC2_COL + a2 * cutlass.Int32(ACC2_STRIDE), + cutlass.Float32, + ), + num=ACC2_STRIDE, + ) + for t in cutlass.range_constexpr(M_MAX): + v = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + v = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(ya[w * M_MAX + t]), v + ) + yv.append(v) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if warp < s_gn.load(idx=g): + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + cutlass.BFloat16( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)) + ), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + if tidx == 0: + _stamp(tdbg, bx, 6) + _tstore(tdbg, bx, 31, e_wait) + _tstore(tdbg, bx, 34, e_busy) + _tstore(tdbg, bx, 36, e_a) + _tstore(tdbg, bx, 38, e_pw) + prims.mbarrier_arrive(ys_done.subview(0)) # this thread's ys values are in smem + + # ---- finalize on all 16 warps: (block bb, token t) pairs, the fma chain over t's top-4 in order, + y_sh + NPAIR = NB2_CAP * M_MAX + NPW = (NPAIR + NWARPS - 1) // NWARPS + wyl = [] + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + ok = (pair < cutlass.Int32(NPAIR)) & (bb < nb) + ch = cutlass.select_( + ok, + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wyl.append( + _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + if cutlass.const_expr(SHARED) + else cutlass.Int32(0) + ) + _wait(ys_done.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 19) + for bt in cutlass.range_constexpr(NPW): + pair = warp + cutlass.Int32(bt * NWARPS) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if (pair < cutlass.Int32(NPAIR)) & (bb < nb): + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ), + acc_o, + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SHARED): + wy = wyl[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + bias: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + img32: cute.Tensor, + cnt: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + w1b: cute.Tensor, + w2b: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + bias, + escal, + plan, + partials, + img32, + cnt, + s1w, + s2w, + w1b, + w2b, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py new file mode 100644 index 000000000000..ac9892f31ed2 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/m3_moe_kernel.py @@ -0,0 +1,3714 @@ +# 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. +"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP rank: ``m3_moe`` (group-phased weight stream). + +One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), +FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It +writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. + +- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then + makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. +- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a + stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 + slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group + p - LAG's intermediate between CTAs. +- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. + FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction + atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). +- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data + bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). +- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the + call's tag; warp 5 copies the scales to TMEM. +- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its + top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange + (AR_PRE: call parities and descriptor loaded during the last group's FC2). + +Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", +"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. +Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before +griddepcontrol.wait; every producer-output read and every global write after it. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass._mlir.dialects import llvm as _llvm +from cutlass._mlir.extras import types as _T +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +_CFG = globals().get("M3_CONFIG") or {} + + +def _cfg(key: str, default): + return type(default)(_CFG.get(key, default)) + + +# ============================================================================= +# Problem shape (trace-time constants). +# ============================================================================= +H = 6144 +I_TP = 768 +TWO_I = 2 * I_TP +E_LOCAL = 128 +TOP_K = 4 +ROUTE_SCALE = 2.0 +SUM_EPS = 1e-20 +M_MAX = _cfg("m_max", 1) +NCTA = _cfg("num_ctas", 152) +FUSED_ROUTER = bool(_cfg("fused_router", 0)) +NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes +NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. +TIMING = bool(_cfg("timing", 0)) +HANG_DBG = bool(_cfg("hang_dbg", 0)) +XVDUMP = bool( + _cfg("xvdump", 0) +) # debug only (needs a plan without split units): FC1 sums into `partials` +YDUMP = bool( + _cfg("ydump", 0) +) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) +LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG +G_CAP = TOP_K * M_MAX # distinct experts per call +NG_CAP = M_MAX # groups of 4 slots +XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. +THREADS = 256 + 32 * (XW - 1) +GEMV_THREADS = 256 +# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more +# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). +# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales +# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns +# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; +# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before +# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). +SHARED = bool(_cfg("shared", 0)) +# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next +# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed +# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. +LATE_TRIGGER = bool(_cfg("late_trigger", 1)) +# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's +# FC2 runs (not two dependent global loads per push on the kernel's tail). +AR_PRE = bool(_cfg("ar_pre", 1)) +# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the +# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). +XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) +NSH = 24 if SHARED else 0 +SH_I = 768 +SH_KB = H // 32 # gate_up K blocks +SH_DKB = SH_I // 32 # down K blocks +SH_SPLIT = 2 # gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges): MMA rows 64 q + .. = slice q +SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice +SH_STAGES_GU = SH_GT +SH_STAGES = SH_GT + 2 * (SH_I // 128) # + 2 down tiles x 6 k-tiles +SHX_BAR = 6 +# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its 36 stages, each the 16 KB +# shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction scale +# atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. +SH_PRETILE = bool( + _cfg("sh_pretile", 1 if M_MAX <= 2 else 0) +) # (in-chain: a win at M <= 2, not at 4..8) +SH_STAGE_BYTES = 16384 + 512 +MMA_M = 128 +KT_BYTES = 128 +INC_K = 2 +NUM_KB = 4 +# FC1 +N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) +QS = N1 // 4 +R1 = 32 +C1 = 4 +U1 = TWO_I // R1 # 48 units per slot +K1_TILES = H // 256 +S1 = K1_TILES // C1 # 6 stages per unit +NUM_BYTES_A = MMA_M * KT_BYTES +NUM_BYTES_BOX1 = R1 * KT_BYTES +SF1_KA = H // 64 +SF1_RB = TWO_I // 128 +NUM_BYTES_SFA_RAW = 16 * 512 +NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 +NUM_BYTES_SFA = NUM_KB * 512 +NUM_WORDS_SFA = NUM_BYTES_SFA // 4 +NUM_BYTES_B1 = N1 * KT_BYTES +NUM_SFB1_ATOMS = S1 * NUM_KB +# FC2 +R2 = 32 +NB2 = H // R2 # 192 +KT2 = I_TP // 256 # 3 +SF2_KA = I_TP // 64 +SF2_RB = H // 128 +N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t +NUM_BYTES_B2 = N2 * KT_BYTES +NB2_CAP = _cfg("nb2_cap", 2) +UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row +H_ROW = UNITS_ROW * 16 +PARTIAL_EMPTY = -1 +EPI_THREADS = 128 +EPI_BAR = 1 +ROUTE_BAR = 2 +GEMV_BAR = 3 +B2_BAR = 5 +TD_STRIDE = 64 # tdbg columns per CTA +TOPK_BAR = 4 +B1_AFTER_TOPK = bool( + _cfg("b1_after_topk", 0) +) # (experiment) the B1 build waits for the routing warps' top-k +TOPK2_BAR = 7 +EARLY0 = bool( + _cfg("early0", 1) +) # warp 4 streams group 0 (token 0's experts) right after its own top-k +EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) +# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token +RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 +NRW0 = ( + min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 +) # warps 4.. with a token (no idle pollers in the barrier) +NRW = NRW0 + RW_EXTRA +LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) +GEMV_SMEM = bool( + _cfg("gemv_smem", 1) +) # (XSTAGE) the router GEMV reads the staged hidden states from smem +LG_STRIDE = M_MAX * E_LOCAL + 32 +POLL_NS = _cfg("poll_ns", 150) +B2_POLL_NS = _cfg("b2_poll_ns", 200) +# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight +# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled +# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one +# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains +# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. +B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) +IMG_B2 = KT2 * NUM_BYTES_B2 +IMG_SF = KT2 * NUM_KB * 512 +IMG_BYTES = IMG_B2 + IMG_SF +CNT_PER_CALL = 4 * (TWO_I // R1) +B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off +B2_DLY = _cfg( + "b2_dly", 0 +) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first +W_WARP = 4 +X_WARP = 5 +S_WARP = 6 +MMA_WARP = 7 +E0_WARP = MMA_WARP # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) +E0_BAR = 8 +E1_BAR = 9 +E0_OFFLOAD = EARLY0 and NRW0 <= 3 and bool(_cfg("e0_off", 1)) +EVICT_FIRST = 0x12F0000000000000 +_LOG2E = 1.4426950408889634 +LD_POW2 = ( + 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 +) # tcgen05.ld .num covering M columns + +_FIXED_SMEM = ( + S1 * NUM_BYTES_B1 + + 2 * KT2 * NUM_BYTES_B2 + + 2 * KT2 * NUM_KB * 512 + + NB2_CAP * M_MAX * TOP_K * R2 * 4 + + 2 * M_MAX * C1 * R1 * 4 + + (M_MAX * (H // 16) if FUSED_ROUTER else 0) + + 6 * 1024 +) +# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no +# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB +SFA_CP = bool(_cfg("sfa_cp", 0)) +_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW +_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM +ACC1_COL = 0 +ACC2_COL = 2 * N1 +SFB1_COL = ACC2_COL + 2 * N2 +SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 +SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 +STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) +TMEM_COLS = 512 +assert 1 <= M_MAX <= 8, M_MAX +assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) +assert M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( + "x staging must fit the ring" +) +# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) +_OFF_SFAR = STAGES * NUM_BYTES_A +_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE +_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA +_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 +_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 +BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 +# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales +# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] +# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange +_SB_A = 0 +_SB_RAW = None +_SH_FIX = ( + SH_GT * 16 * 128 + + 6 * 8 * 128 + + M_MAX * SH_KB + + M_MAX * SH_DKB + + 64 * 8 * 4 + + 2 * 32 * 8 * 4 + + 1024 +) +SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 +_SB_RAW = SH_NST * NUM_BYTES_A +_SB_SPL = _SB_RAW + SH_NST * 2048 +_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 +_SB_BH = _SB_BX + SH_GT * 16 * 128 +_SB_XC = _SB_BH + 6 * 8 * 128 +_SB_HC = _SB_XC + M_MAX * SH_KB +_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 +_SB_GU = _SB_XCH + 64 * 8 * 4 +_SB_END = _SB_GU + 2 * 32 * 8 * 4 +# shared CTA TMEM columns: gate_up accumulator (N 16), down accumulators (2 x N 8), A scales per stage, B scales +SH_ACC_GU = 0 +SH_ACC_DN = 16 +SH_SFA = 32 +SH_SFB_GU = SH_SFA + 4 * 8 +SH_SFB_DN = SH_SFB_GU + 4 * SH_GT +SH_TMEM = 256 +if SHARED: + assert FUSED_ROUTER, "shared expert: fused router only" + assert SH_STAGES - SH_NST <= 64, SH_STAGES + assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) + + +# ============================================================================= +# Inline-PTX helpers. +# ============================================================================= +def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): + return _llvm.inline_asm( + ret, + [a.ir_value(loc=loc, ip=ip) for a in args], + ptx, + cons, + has_side_effects=side, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +# ---- all-reduce push (the MoE boundary's multicast exchange, set A): +# arp = int64 [enable, mc, flags, rank, world, maxt] +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): + """st.global.b32 when pred != 0.""" + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), + cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _ar_push(arp, v, low, t, cl): + """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on + lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call + parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + world = arp.load(idx=4) + parity = cutlass.Int64( + _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) + ) + & cutlass.Int32(1) + ) + off = ( + parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +def _ar_parity(arp, t, cl): + """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the + next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" + return _ld_volatile_s32( + arp.load(idx=2) + + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) + ) + + +def _ar_push_pre(ard, parity32, v, low, t, cl): + """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" + hi = cute.arch.shuffle_sync_bfly(v, 8) + pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) + cs = cutlass.select_(low, cl, cutlass.Int32(0)) + mc, rk, world, maxt = ard + parity = cutlass.Int64(parity32 & cutlass.Int32(1)) + off = ( + parity * maxt * world * cutlass.Int64(H // 2) + + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) + + cutlass.Int64(cs // cutlass.Int32(2)) + ) + _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) + + +@dsl_user_op +def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): + """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" + _llvm.inline_asm( + None, + [], + f"bar.arrive {bar_id}, {threads};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_relaxed_u32(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "st.relaxed.gpu.global.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _st_relaxed_u64(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int64(val)], + "st.relaxed.gpu.global.u64 [$0], $1;", + "l,l", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_relaxed_u64(addr, *, loc=None, ip=None): + return cutlass.Int64( + _asm( + _T.i64(), + [cutlass.Int64(addr)], + "ld.relaxed.gpu.global.u64 $0, [$1];", + "=l,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _nanosleep(ns: int, *, loc=None, ip=None): + _llvm.inline_asm( + None, + [], + f"nanosleep.u32 {ns};", + "", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _gtime(*, loc=None, ip=None): + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _smid(*, loc=None, ip=None): + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): + if cutlass.const_expr(TIMING): + tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) + + +@dsl_user_op +def _rcp(a, *, loc=None, ip=None): + """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _div_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fma_rn(a, b, c, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _add_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _sub_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _mul_rn(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e4m3_bits(x, *, loc=None, ip=None): + """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3_to_f32(b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Int32(b)], + "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" + "cvt.f32.f16 $0, l;\n}", + "=f,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _e2m1x2(hi, lo, *, loc=None, ip=None): + """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _bf16_round(x, *, loc=None, ip=None): + """x rounded to bf16 (nearest even), as f32.""" + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", + "=f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _redux_max_u32(v, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(v)], + "redux.sync.max.u32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16 (nearest even) bits of x in the low 16 bits.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ue8m0_rp(x, *, loc=None, ip=None): + """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(x)], + "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", + "=r,f", + False, + loc, + ip, + ) + ) & cutlass.Int32(0xFF) + + +@dsl_user_op +def _e4m3x2(hi, lo, *, loc=None, ip=None): + """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Float32(hi), cutlass.Float32(lo)], + "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", + "=r,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _min_xorsign_abs(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "min.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _div_full(a, b, *, loc=None, ip=None): + return cutlass.Float32( + _asm( + _T.f32(), + [cutlass.Float32(a), cutlass.Float32(b)], + "div.full.f32 $0, $1, $2;", + "=f,f,f", + False, + loc, + ip, + ) + ) + + +@dsl_user_op +def _ex2_full(x, *, loc=None, ip=None): + """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" + return cutlass.Float32( + _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) + ) + + +@dsl_user_op +def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): + """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" + _llvm.inline_asm( + None, + [ + cutlass.Int64(addr).ir_value(loc=loc, ip=ip), + cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.prefetch.L2.global [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def _mx_osc(amax): + """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" + code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) + cf = cutlass.Int32( + cutlass.select_( + code == cutlass.Int32(255), + cutlass.Int32(0x7FFFFFFF), + cutlass.select_( + code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) + ), + ) + ).bitcast(cutlass.Float32) + return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) + + +def _mx_words(v, osc): + """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" + words = [] + for t in range(len(v) // 4): + lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) + hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) + words.append(lo | (hi << cutlass.Int32(16))) + return words + + +def _twiddle(score): + """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" + bits = ( + score.bitcast(cutlass.Int32) + if hasattr(score, "bitcast") + else cutlass.Float32(score).bitcast(cutlass.Int32) + ) + return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) + + +@cute.jit +def _wait(bar: cutlass.Array, parity): + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +def _tptr(base_row, col, dtype=cutlass.Int32): + return cutlass.inttoptr((base_row << 16) | col, 6, dtype) + + +@dsl_user_op +def _match_any(value, *, loc=None, ip=None): + """The mask of the warp's lanes holding the same 32-bit value.""" + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int32(value)], + "match.any.sync.b32 $0, $1, 0xffffffff;", + "=r,r", + False, + loc, + ip, + ) + ) + + +@cute.jit +def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): + """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" + if cutlass.const_expr(HANG_DBG): + n = cutlass.Int32(0) + while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( + n < cutlass.Int32(1 << 22) + ): + n = n + cutlass.Int32(1) + if n >= cutlass.Int32(1 << 22): + bx, _, _ = cute.arch.block_idx() + w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tdbg.store( + cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w + ) + else: + while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): + pass + cute.arch.sync_warp() + + +@cute.jit +def _adv(stage, phase): + """The next ring position's (stage, phase).""" + stage = stage + cutlass.Int32(1) + wrap = stage == cutlass.Int32(STAGES) + return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( + wrap, phase ^ cutlass.Int32(1), phase + ) + + +@cute.jit +def _seg(i0, i1, k): + """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: + the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" + uf = i0 // cutlass.Int32(S1) + ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) + unit = uf + k - cutlass.Int32(1) + unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) + unit = cutlass.select_(ul == uf, uf, unit) + j0 = i0 - unit * cutlass.Int32(S1) + j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) + j1 = i1 - unit * cutlass.Int32(S1) + j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) + return unit, j0, j1 + + +@cute.jit +def _nsegs(i0, i1): + n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) + return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) + + +@cute.jit +def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): + """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile + 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" + q = lane // cutlass.Int32(QS) + tl = lane - q * cutlass.Int32(QS) + sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) + sfv = [] + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + j = a // NUM_KB + kb = a % NUM_KB + wsf = src.load( + idx=cutlass.select_( + sf_ok, + tl * cutlass.Int32(H // 64) + + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) + + cutlass.Int32(kb), + cutlass.Int32(0), + ) + ) + sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) + for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): + prims.tcgen05_st( + "32x32b", + _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), + sfv[a], + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(b1_ready) + + +@cute.jit +def _b1_addr(t, kb16): + """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" + kt = kb16 >> cutlass.Int32(4) + j = kt >> cutlass.Int32(2) + n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t + cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) + return ( + j * cutlass.Int32(NUM_BYTES_B1) + + n * cutlass.Int32(KT_BYTES) + + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) + + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) + ) // cutlass.Int32(4) + + +@cute.jit +def _build_b1( + tidx, + warp, + lane, + base_row, + base_col, + x32: cutlass.Array, + xsf32: cutlass.Array, + sB1: cutlass.Array, + b1_ready: cutlass.Array, +): + """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): + b = tidx + cutlass.Int32(r * EPI_THREADS) + if b < cutlass.Int32(M_MAX * (H // 16)): + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) + sB1_32.store( + (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 + ) + _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) + + +@cute.jit +def _build_b1_fused( + tidx, + warp, + lane, + base_row, + base_col, + xb32: cutlass.Array, + sfsc: cutlass.Array, + sB1: cutlass.Array, + s_xsf: cutlass.Array, + b1_ready: cutlass.Array, +): + """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" + sB1_32 = cutlass.Array( + sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 + ) + s_xsf8 = cutlass.Array( + s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 + ) + sc = cutlass.Float32(sfsc.load(idx=0)) + rsc = _rcp(sc) + rsix = _rcp(cutlass.Float32(6.0)) + NBLK = M_MAX * (H // 16) // EPI_THREADS + QB = ( + 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) + ) # blocks whose loads are in flight together + # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) + for rg in cutlass.range(NBLK // QB, unroll=1): + wall = [] + for rr in cutlass.range_constexpr(QB): + bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + for h in cutlass.range_constexpr(2): + wall.append( + xb32.load( + idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), + vector_size=4, + alignment=16, + ) + ) + for rr in cutlass.range_constexpr(QB): + b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) + t = b // cutlass.Int32(H // 16) + kb = b - t * cutlass.Int32(H // 16) + w = [] + for h in cutlass.range_constexpr(2): + v4 = wall[rr * 2 + h] + for z in cutlass.range_constexpr(4): + w.append(cutlass.Int32(v4[z])) + xs = [] + for z in cutlass.range_constexpr(8): + xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) + xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + vmax = cute.math.abs(xs[0]) + for z in cutlass.range_constexpr(1, 16): + vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) + sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) + osc = cutlass.select_( + vmax != cutlass.Float32(0.0), + _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), + cutlass.Float32(0.0), + ) + by = [ + _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) + for z in range(8) + ] + lo = ( + by[0] + | (by[1] << cutlass.Int32(8)) + | (by[2] << cutlass.Int32(16)) + | (by[3] << cutlass.Int32(24)) + ) + hi = ( + by[4] + | (by[5] << cutlass.Int32(8)) + | (by[6] << cutlass.Int32(16)) + | (by[7] << cutlass.Int32(24)) + ) + sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) + s_xsf8.store(cutlass.Int8(sfb), idx=b) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) + + +@cute.jit +def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): + """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, + gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 + and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" + neg_max = cutlass.Float32(-3.40282346638528859812e38) + gc = cutlass.select_(gate_raw > lim, lim, gate_raw) + gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) + uc = cutlass.select_(xv > lim, lim, xv) + uc = cutlass.select_(xv < -lim, -lim, uc) + c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) + sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) + gs = _mul_rn(gc, sgm) + h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) + h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) + amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) + for off in (16, 8, 4, 2, 1): + amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) + sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) + sff = _e4m3_to_f32(sfb) + osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) + return _mul_rn(_mul_rn(h, osc), c), sfb + + +@cute.jit +def _plan_range(s_plan: cutlass.Array, p): + w = s_plan.load(idx=p) + return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) + + +@cute.jit +def _sh_issue( + lane, + sid, + g, + s_big: cutlass.Array, + sh_full: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, +): + """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both + slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its + raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile + (g - SH_GT) % 6.""" + st = g % cutlass.Int32(SH_NST) + bar = sh_full.subview(st) + if cutlass.const_expr(SH_PRETILE): + if lane == 0: + src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) + prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + gusf.subview(src + cutlass.Int32(16384)), + bar, + 512, + ) + cute.arch.sync_warp() + else: + if g < cutlass.Int32(SH_GT): + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) + cute.arch.sync_warp() + if lane < 4: + q = lane >> cutlass.Int32(1) + u = lane & cutlass.Int32(1) + row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) + kt = q * cutlass.Int32(SH_GT) + g + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), + tma_gu, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview( + cutlass.Int32(_SB_RAW) + + st * cutlass.Int32(2048) + + lane * cutlass.Int32(512) + ), + gusf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) + if lane == 0: + prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) + prims.cp_async_bulk_tensor_shared_cta_global( + s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), + tma_dn, + (kt * cutlass.Int32(128), row0), + bar, + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), + dnsf.subview( + ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) + * cutlass.Int32(512) + ), + bar, + 512, + ) + cute.arch.sync_warp() + + +@dsl_user_op +def _red_release_add(addr, val, *, loc=None, ip=None): + _asm( + None, + [cutlass.Int64(addr), cutlass.Int32(val)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + True, + loc, + ip, + ) + + +@dsl_user_op +def _ld_acquire_u32(addr, *, loc=None, ip=None): + return cutlass.Int32( + _asm( + _T.i32(), + [cutlass.Int64(addr)], + "ld.acquire.gpu.global.u32 $0, [$1];", + "=r,l", + True, + loc, + ip, + ) + ) + + +@dsl_user_op +def _fence_proxy_async_global(*, loc=None, ip=None): + _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) + + +@cute.jit +def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): + """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + + +@cute.jit +def _sf1_cp( + lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array +): + """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) + word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their + arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" + src0 = ( + (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (c & cutlass.Int32(3)) + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for q in cutlass.range_constexpr(C1): + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), + s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _sf2_cp( + lane, + stage, + g, + blk, + t2: cutlass.Constexpr, + n_here, + s_el: cutlass.Array, + sSFA: cutlass.Array, + s2w: cutlass.Array, + ab_full: cutlass.Array, +): + """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot + jj's expert, word [r][blk % 4].""" + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + src0 = ( + ( + (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(128) + + lane * cutlass.Int32(4) + + (blk & cutlass.Int32(3)) + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.cp_async_shared_global( + sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), + s2w.subview(src0 + cutlass.Int32(kb * 128)), + 4, + "ca", + ) + prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) + cute.arch.sync_warp() + + +@cute.jit +def _shared_cta( + tidx, + warp, + lane, + sid, + s_big: cutlass.Array, + sh_bar: cutlass.Array, + epochs: cutlass.Array, + tdbg: cutlass.Array, + bx, + xb32: cutlass.Array, + tma_gu, + tma_dn, + gusf: cutlass.Array, + dnsf: cutlass.Array, + hx: cutlass.Array, + dnx: cutlass.Array, + arp: cutlass.Array, + tmem_ptr_i32: cutlass.Array, + gu8: cutlass.Array, + dn8: cutlass.Array, +): + """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one + 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles + of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" + big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) + bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) + big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) + sh_full = sh_bar # [0, SH_NST): stage full + sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) + b_gu_rdy = sh_bar.subview(16) + b_dn_rdy = sh_bar.subview(17) + acc_gu = sh_bar.subview(18) + acc_dn = sh_bar.subview(19) + if warp == 0: + prims.prefetch_tensormap(tma_gu) + prims.prefetch_tensormap(tma_dn) + if lane < SH_NST: + prims.mbarrier_init(sh_full.subview(lane), 1) + prims.mbarrier_init( + sh_empty.subview(lane), 33 + ) # the MMAs' commit + each splicer lane's reads done + if lane == 0: + prims.mbarrier_init(b_gu_rdy, 1) + prims.mbarrier_init(b_dn_rdy, 1) + prims.mbarrier_init(acc_gu, 1) + prims.mbarrier_init(acc_dn, 1) + if warp == 1: + prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + if warp == 0: + for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency + _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) + if cutlass.const_expr(SH_PRETILE): + # this CTA's stages beyond the ring, 4 per lane + for r in cutlass.range_constexpr(2): + pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) + if pg < cutlass.Int32(SH_STAGES): + _l2_prefetch( + gusf.subview( + (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) + ) + .data_ptr() + .toint(), + SH_STAGE_BYTES, + ) + else: + for r in cutlass.range_constexpr(2): + prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane + _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) + if lane == 0: + _l2_prefetch( + dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I + ) + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(not LATE_TRIGGER): + if arp.load(idx=0) == cutlass.Int64(0): + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + ep = epochs.load(idx=bx) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + NQ = THREADS - 64 # activation quant threads (warps 2 ..) + if warp == 0: + for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): + _wait( + sh_empty.subview(g % cutlass.Int32(SH_NST)), + ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), + ) + _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) + elif warp == 1: + # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down + idesc_gu = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=16, + m_dim=128, + ) + idesc_dn = prims.Tcgen05MxInstrDesc.build( + a_dtype=cutlass.Float8E4M3FN, + b_dtype=cutlass.Float8E4M3FN, + scale_format=1, + n_dim=8, + m_dim=128, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + desc_a0 = prims.Tcgen05SmemDesc.build( + s_big.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bx = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BX), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_bh = prims.Tcgen05SmemDesc.build( + s_big.subview(_SB_BH), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + sq = sid & cutlass.Int32(3) + _wait(b_gu_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for g in cutlass.range(SH_STAGES, unroll=1): + st = g % cutlass.Int32(SH_NST) + if g == cutlass.Int32(SH_GT): + _wait(b_dn_rdy, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) + raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) + if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): + # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 + w4 = [ + big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) + for qq in range(4) + ] + big32.store( + ( + cutlass.Int32(w4[0]), + cutlass.Int32(w4[1]), + cutlass.Int32(w4[2]), + cutlass.Int32(w4[3]), + ), + idx=cutlass.Int32(_SB_SPL // 4) + + st * cutlass.Int32(128) + + lane * cutlass.Int32(4), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + prims.mbarrier_arrive( + sh_empty.subview(st) + ) # this lane's generic reads of the slot are done + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + if g < cutlass.Int32(SH_GT): + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bx + cutlass.Int64(g * cutlass.Int32(2048 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), + da + 2 * kb, + db + 2 * kb, + idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_GT - 1): + prims.tcgen05_commit(acc_gu) + else: + gd = g - cutlass.Int32(SH_GT) + jt = gd // cutlass.Int32(6) + kt = gd - jt * cutlass.Int32(6) + dsf = prims.Tcgen05SmemDesc.build( + s_big.subview( + cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) + + st * cutlass.Int32(512 if SH_PRETILE else 2048) + ), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + prims.tcgen05_cp( + s2t_shape, + _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), + dsf, + multicast=s2t_mc, + ) + da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) + for kb in cutlass.range_constexpr(4): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, + prims.CTAGroup.CTA_1, + cutlass.inttoptr( + tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 + ), + da + 2 * kb, + db + 2 * kb, + idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), + cutlass.inttoptr( + tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 + ), + cutlass.inttoptr( + tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 + ), + ) + prims.tcgen05_commit(sh_empty.subview(st)) + if g == cutlass.Int32(SH_STAGES - 1): + prims.tcgen05_commit(acc_dn) + cute.arch.sync_warp() + if warp >= 2: + # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = + # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes + qt = tidx - cutlass.Int32(64) + for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): + q_ = qt + cutlass.Int32(r * NQ) + if q_ < cutlass.Int32(M_MAX * SH_KB): + t = q_ // cutlass.Int32(SH_KB) + b = q_ - t * cutlass.Int32(SH_KB) + wv = [ + xb32.load( + idx=t * cutlass.Int32(H // 2) + + b * cutlass.Int32(16) + + cutlass.Int32(4 * z), + vector_size=4, + alignment=16, + ) + for z in range(4) + ] + v = [] + for z in cutlass.range_constexpr(4): + for zz in cutlass.range_constexpr(4): + w = cutlass.Int32(wv[z][zz]) + v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) + v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(v[z])) + code, osc = _mx_osc(amax) + words = _mx_words(v, osc) + kt = b >> cutlass.Int32(2) + sl = kt // cutlass.Int32(SH_GT) + it = kt - sl * cutlass.Int32(SH_GT) + n = sl * cutlass.Int32(8) + t + c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BX) + it * cutlass.Int32(2048) + n * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) + if warp >= 4: + if warp < 8: + qd = warp - cutlass.Int32(4) + # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i + n = lane + t = n & cutlass.Int32(7) + ok = (n < cutlass.Int32(16)) & (t < cutlass.Int32(M_MAX)) + for i in cutlass.range_constexpr(SH_GT): + kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) + wsf = big32.load( + idx=( + cutlass.Int32(_SB_XC) + + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) + + kt * cutlass.Int32(4) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), + cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + et = tidx - cutlass.Int32(128) + if et == 0: + prims.mbarrier_arrive(b_gu_rdy) + _stamp(tdbg, bx, 54) + # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t + _wait(acc_gu, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + av = prims.tcgen05_ld( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), + num=16, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + if qd >= cutlass.Int32(2): + for tt in cutlass.range_constexpr(8): + bigf.store( + cutlass.Float32(av[8 + tt]), + idx=cutlass.Int32(_SB_XCH // 4) + + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if qd < cutlass.Int32(2): + # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 + for tt in cutlass.range_constexpr(8): + tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) + tot = _add_rn( + tot, + bigf.load( + idx=cutlass.Int32(_SB_XCH // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt) + ), + ) + bigf.store( + _bf16_round(tot), + idx=cutlass.Int32(_SB_GU // 4) + + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) + + cutlass.Int32(tt), + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag + for r in cutlass.range_constexpr((32 * M_MAX + 127) // 128): + pq = et + cutlass.Int32(r * 128) + if pq < cutlass.Int32(32 * M_MAX): + p = pq // cutlass.Int32(M_MAX) + tt = pq - p * cutlass.Int32(M_MAX) + ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) + up = bigf.load( + idx=cutlass.Int32(_SB_GU // 4) + + (cutlass.Int32(32) + p) * cutlass.Int32(8) + + tt + ) + q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) + q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) + q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) + q_e2 = _ex2_full( + _mul_rn( + _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), + cutlass.Float32(1.4426950216293334961), + ) + ) + q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) + q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) + _st_relaxed_u32( + hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(32) + p) + .data_ptr() + .toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), + ) + if et == 0: + _stamp(tdbg, bx, 50) + # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes + for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): + q_ = et + cutlass.Int32(rr * 128) + if q_ < cutlass.Int32(M_MAX * SH_DKB): + t = q_ // cutlass.Int32(SH_DKB) + kb = q_ - t * cutlass.Int32(SH_DKB) + base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) + hv = [cutlass.Int32(0) for _ in range(32)] + pend = cutlass.Int32(1) + while pend != cutlass.Int32(0): + for q4 in cutlass.range_constexpr(8): + v4 = prims.load_ext( + hx.subview(base + cutlass.Int32(4 * q4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for z in cutlass.range_constexpr(4): + hv[4 * q4 + z] = cutlass.Int32(v4[z]) + okw = cutlass.Boolean(True) + for z in cutlass.range_constexpr(32): + okw = okw & ( + ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 + ) + pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) + vals = [ + (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) + ] + amax = cutlass.Float32(0.0) + for z in cutlass.range_constexpr(32): + amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) + code, osc = _mx_osc(amax) + words = _mx_words(vals, osc) + kt = kb >> cutlass.Int32(2) + c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) + rw = ( + cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) + ) >> cutlass.Int32(2) + big32.store( + (words[0], words[1], words[2], words[3]), + idx=rw + ((c0 ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big32.store( + (words[4], words[5], words[6], words[7]), + idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), + alignment=16, + ) + big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + okh = lane < cutlass.Int32(M_MAX) + for kt in cutlass.range_constexpr(6): + wsf = big32.load( + idx=( + cutlass.Int32(_SB_HC) + + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) + + cutlass.Int32(4 * kt) + ) + >> cutlass.Int32(2) + ) + prims.tcgen05_st( + "32x32b", + _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), + cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), + ) + prims.tcgen05_wait(prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) + if et == 0: + prims.mbarrier_arrive(b_dn_rdy) + _stamp(tdbg, bx, 51) + # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l + _wait(acc_dn, 0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for jt in cutlass.range_constexpr(2): + dv = prims.tcgen05_ld( + "32x32b", + _tptr( + base_row + qd * cutlass.Int32(32), + base_col + SH_ACC_DN + jt * 8, + cutlass.Float32, + ), + num=8, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + rowo = ( + sid * cutlass.Int32(256) + + cutlass.Int32(128 * jt) + + qd * cutlass.Int32(32) + + lane + ) + for tt in cutlass.range_constexpr(M_MAX): + _st_relaxed_u32( + dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), + (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), + ) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if et == 0: + _stamp(tdbg, bx, 52) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if warp == 1: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) + if tidx == 0: + epochs.store(ep + cutlass.Int32(1), idx=bx) + + +@cute.kernel +def m3_moe_kernel( + tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], + tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], + tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines + tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines + tma_gu_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows + tma_dn_desc: cutlass.GridConstant[ + cuda.TensorMap + ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows + logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) + bias: cutlass.Array, # f32 [128] + x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) + xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) + escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha + plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 + partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums + hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records + epochs: cutlass.Array, # int32 [NCTA] + out: cutlass.Array, # bf16 [M, 6144] + route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) + tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) + xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] + rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] + sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale + lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits + guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] + gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] + dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales + hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written + dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written + arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off + s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale + s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale + img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) + b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) +): + tidx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = tidx % 32 + if tidx == 0: + _stamp(tdbg, bx, 0) + if cutlass.const_expr(TIMING): + if tidx == 0: + tdbg.store( + cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) + ) + + s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) + sA = cutlass.Array( + s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 + ) + sSFAraw = cutlass.Array( + s_big.data_ptr(_OFF_SFAR), + shape=(max(1, _RAW_STAGE // 4) * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sSFA = cutlass.Array( + s_big.data_ptr(_OFF_SFA), + shape=(NUM_WORDS_SFA * STAGES,), + dtype=cutlass.Int32, + alignment=1024, + ) + sB1 = cutlass.Array( + s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 + ) + sB2 = cutlass.Array( + s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 + ) + sSFB2 = cutlass.Array( + s_big.data_ptr(_OFF_SFB2), + shape=(2 * KT2 * NUM_KB * 128,), + dtype=cutlass.Int32, + alignment=1024, + ) + ys = cutlass.Array( + cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_part = cutlass.Array( + cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) + s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) + s_xsf = cutlass.Array( + cutlass.Int32, + M_MAX * (H // 64) if FUSED_ROUTER else 4, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) + acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) + # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice + # while warp 5 still works on an earlier group, and its parity wait would then miss them) + fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) + b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + x_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # x landed, x consumed + img_bar = cutlass.Array( + cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 + ) # (B2IMG) image landed + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + s_el = cutlass.Array( + cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 + ) # slots, then G + s_kpos = cutlass.Array( + cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_w = cutlass.Array( + cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_ids = cutlass.Array( + cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 + ) + s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) + # this CTA's plan rows for every G (static data: loaded before the grid dependency) + s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) + s_plan_all = cutlass.Array( + cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 + ) + s_esc = cutlass.Array( + cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 + ) # slot s at 8 s + + is_routed = cutlass.Boolean(True) + if cutlass.const_expr(SHARED): + sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) + is_routed = bx < cutlass.Int32(NCTA) + if bx >= cutlass.Int32(NCTA): + _shared_cta( + tidx, + warp, + lane, + bx - cutlass.Int32(NCTA), + s_big, + sh_full, + epochs, + tdbg, + bx, + xb32, + tma_gu_desc.get_ptr(), + tma_dn_desc.get_ptr(), + gusf, + dnsf, + hx, + dnx, + arp, + tmem_ptr_i32, + guw, + dnw, + ) + if is_routed: + # ---- before the grid dependency + if warp == 0: + if tidx < STAGES: + prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) + prims.mbarrier_init(ab_empty.subview(tidx), 1) + prims.mbarrier_init(sfa_ready.subview(tidx), 1) + if tidx < 2: + prims.mbarrier_init(acc_full.subview(tidx), 1) + prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(acc2_full.subview(tidx), 1) + prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) + prims.mbarrier_init(b2_ready.subview(tidx), 1) + prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) + if cutlass.const_expr(B2IMG): + prims.mbarrier_init(img_bar.subview(tidx), 1) + if tidx < 4: + prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) + if tidx == 0: + prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) + if cutlass.const_expr(XSTAGE): + prims.mbarrier_init(x_bar.subview(0), 1) + prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) + if warp == MMA_WARP: + prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) + prims.tcgen05_relinquish_alloc_permit() + if warp == W_WARP: + prims.prefetch_tensormap(tma_a1_desc.get_ptr()) + prims.prefetch_tensormap(tma_a2_desc.get_ptr()) + prims.prefetch_tensormap(tma_s1_desc.get_ptr()) + prims.prefetch_tensormap(tma_s2_desc.get_ptr()) + for i in cutlass.range_constexpr(E_LOCAL // 32): + s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) + for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): + q = lane + cutlass.Int32(32 * i) + if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): + gq = q // cutlass.Int32(NG_CAP + 1) + s_plan_all.store( + plan.load( + idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) + + q + - gq * cutlass.Int32(NG_CAP + 1) + ), + idx=q, + ) + rwv = [] + if cutlass.const_expr(FUSED_ROUTER): + tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) + rwv.append( + [ + rw.load( + idx=rbx * cutlass.Int32(H) + + tg * cutlass.Int32(4) + + cutlass.Int32(1024 * j), + vector_size=4, + alignment=16, + ) + for j in range(H // 1024) + ] + ) + if cutlass.const_expr(TIMING): + if tidx == 0: + s_el.store(cutlass.Int32(-7), idx=G_CAP) + prims.fence_mbarrier_init() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + + cute.arch.griddepcontrol_wait() + if cutlass.const_expr(XSTAGE): + if warp == 0: + if lane == 0: + prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) + cute.arch.sync_warp() + if lane < M_MAX: + prims.cp_async_bulk_shared_cluster_global( + s_big.subview(lane * cutlass.Int32(H * 2)), + xb32.subview(lane * cutlass.Int32(H // 2)), + x_bar.subview(0), + H * 2, + ) + ar_on = arp.load(idx=0) != cutlass.Int64(0) + ar_off = arp.load(idx=0) == cutlass.Int64(0) + if cutlass.const_expr(not LATE_TRIGGER): + if ar_off: + cute.arch.griddepcontrol_launch_dependents() + if tidx == 0: + _stamp(tdbg, bx, 1) + + ep = epochs.load(idx=bx) + tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) + tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) + ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) + if cutlass.const_expr(FUSED_ROUTER): + # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly + for rg in cutlass.range_constexpr(NGEMV): + eg = bx + cutlass.Int32(rg * NCTA) + if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): + accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] + xsrc = xb32 + if cutlass.const_expr(GEMV_SMEM and XSTAGE): + _wait(x_bar.subview(0), 0) + xsrc = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + for j in cutlass.range_constexpr(H // 1024): + for t in cutlass.range_constexpr(M_MAX): + xw = xsrc.load( + idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), + vector_size=2, + alignment=8, + ) + x0 = cutlass.Int32(xw[0]) + x1 = cutlass.Int32(xw[1]) + a = accs[t] + a = _fma_rn( + (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][0]), + a, + ) + a = _fma_rn( + (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][1]), + a, + ) + a = _fma_rn( + (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][2]), + a, + ) + a = _fma_rn( + (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + cutlass.Float32(rwv[rg][j][3]), + a, + ) + accs[t] = a + for t in cutlass.range_constexpr(M_MAX): + v = accs[t] + for off in (16, 8, 4, 2, 1): + v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) + if lane == 0: + s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < M_MAX: + sacc = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(8): + sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) + s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) + if tidx < LG_REPL: + for t in cutlass.range_constexpr(M_MAX): + lw = ltag | ( + cutlass.Int64( + cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( + cutlass.Int32 + ) + ) + & cutlass.Int64(0xFFFFFFFF) + ) + _st_relaxed_u64( + lgbuf.subview( + tidx * cutlass.Int32(LG_STRIDE) + + cutlass.Int32(t * E_LOCAL) + + eg + ) + .data_ptr() + .toint(), + lw, + ) + if cutlass.const_expr(NGEMV > 1): + cute.arch.sync_warp() + cute.arch.barrier( + barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS + ) # s_red reused next + if tidx == 0: + _stamp(tdbg, bx, 21) + + # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile + if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( + (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) + ): + rw_i = cutlass.select_( + warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) + ) + # ILP: poll every token's logits first, then run the independent top-k chains back to back + lgv_all = [] + t_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = rw_i + cutlass.Int32(NRW * tr) + if cutlass.const_expr(M_MAX % NRW != 0): + t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) + t_all.append(t) + lgv = [] + if cutlass.const_expr(FUSED_ROUTER): + want = ltag >> cutlass.Int64(32) + rbase = ( + (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) + + t * cutlass.Int32(E_LOCAL) + + lane + ) + done = cutlass.Int32(0) + res = [cutlass.Int64(0) for _ in range(4)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(4): + res[z] = _ld_relaxed_u64( + lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() + ) + ok = ok & ((res[z] >> cutlass.Int64(32)) == want) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( + -1 + ): + done = cutlass.Int32(1) + else: + _nanosleep(POLL_NS) + for z in cutlass.range_constexpr(4): + lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 22) + else: + for i in cutlass.range_constexpr(4): + lgv.append( + cutlass.Float32( + logits.load( + idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane + ) + ) + ) + lgv_all.append(lgv) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 61) + res_all = [] + for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): + t = t_all[tr] + lgv = lgv_all[tr] + own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) + # top-4 (keys sigmoid + bias, ties to the lower id) + keys = [] + for i in cutlass.range_constexpr(4): + sg = _fma_rn( + cutlass.Float32(0.5), + cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), + cutlass.Float32(0.5), + ) + keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) + hi = [_twiddle(keys[i]) for i in range(4)] + taken = [cutlass.Boolean(False) for _ in range(4)] + ids_t = [] + kt_t = [] + for k in cutlass.range_constexpr(TOP_K): + bh = cutlass.Int32(0) + bl = cutlass.Int32(0) + for i in cutlass.range_constexpr(4): + lo_i = cutlass.Int32(65535 - i * 32) - lane + hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) + lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) + better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( + (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) + ) + bh = cutlass.select_(better, hv, bh) + bl = cutlass.select_(better, lv, bl) + mh = _redux_max_u32(bh) + ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) + e_k = cutlass.Int32(65535) - ml + for i in cutlass.range_constexpr(4): + taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) + ids_t.append(e_k) + kv = cutlass.Float32(0.0) + for i in cutlass.range_constexpr(4): + kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) + kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) + sig = cutlass.Float32(0.0) + e_l = cutlass.Int32(0) + for k in cutlass.range_constexpr(TOP_K): + sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) + sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) + e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) + total = sig + for off in (16, 8, 4, 2, 1): + total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) + wv = _bf16_round( + _div_rn( + _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), + _add_rn(total, cutlass.Float32(SUM_EPS)), + ) + ) + res_all.append((t, own, wv, e_l)) + if (lane == 0) & (warp == cutlass.Int32(W_WARP)): + _stamp(tdbg, bx, 62) + for tr in cutlass.range_constexpr( + (M_MAX + NRW - 1) // NRW + ): # stores after every chain (one branch-free block) + t, own, wv, e_l = res_all[tr] + if (lane < TOP_K) & own: + s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) + s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) + if bx == 0: + route_dbg.store( + wv.bitcast(cutlass.Int32), + idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, + ) + cute.arch.sync_warp() + if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(W_WARP): + _bar_arrive( + E0_BAR, 64 + ) # token 0's experts are in s_ids: warp 7 issues the early stages + if cutlass.const_expr(B1_AFTER_TOPK): + _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) + cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) + if cutlass.const_expr(E0_OFFLOAD): + if warp == cutlass.Int32(E0_WARP): + cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) + # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same + # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) + if lane == 0: + _stamp(tdbg, bx, 56) + n0 = cutlass.Int32(0) + w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) + ei0 = w0 & cutlass.Int32(0xFFFF) + ei1 = w0 >> cutlass.Int32(16) + for k in range(_nsegs(ei0, ei1)): + ug, j0, j1 = _seg(ei0, ei1, k) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_ids.load(idx=ug // cutlass.Int32(U1)) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(n0), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + n0 * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(n0), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(n0), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) + n0 = n0 + cutlass.Int32(1) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if lane == 0: + s_misc.store(n0, idx=0) + cute.arch.sync_warp() + _bar_arrive(E1_BAR, 64) + if warp == W_WARP: + if lane == 0: + _stamp(tdbg, bx, 20) + # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) + valid = lane < cutlass.Int32(TOP_K * M_MAX) + my_e = cutlass.select_( + valid, + s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), + cutlass.Int32(-1) - lane, + ) + same = _match_any(my_e) + leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) + first = valid & (lane == leader) + bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) + rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) + slot_l = cute.arch.shuffle_sync(rank, leader) + for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): + if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): + s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) + cute.arch.sync_warp() + if first: + s_el.store(my_e, idx=rank) + if valid: + s_kpos.store( + lane % cutlass.Int32(TOP_K), + idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), + ) + if bx == 0: + route_dbg.store( + my_e, + idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) + + lane % cutlass.Int32(TOP_K), + ) + gcnt = cute.arch.popc(bal) + if lane == 0: + s_el.store(gcnt, idx=G_CAP) + if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) + s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) + cute.arch.sync_warp() + if cutlass.const_expr(B2IMG): + # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) + if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): + nh = gcnt - lane * cutlass.Int32(4) + nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) + nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) + pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) + if pad > cutlass.Int32(0): + _red_release_add( + b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad + ) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 24) + cute.arch.sync_warp() + _bar_arrive(ROUTE_BAR, THREADS) + if warp < 4: + if cutlass.const_expr(B1_AFTER_TOPK): + cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) + if cutlass.const_expr(FUSED_ROUTER): + if tidx == 0: + _stamp(tdbg, bx, 57) + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(0), 0) + if tidx == 0: + _stamp(tdbg, bx, 58) + xs32 = cutlass.Array( + s_big.data_ptr(0), + shape=(M_MAX * H // 2,), + dtype=cutlass.Int32, + alignment=16, + ) + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready + ) + prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights + else: + _build_b1_fused( + tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready + ) + if tidx == 0: + _stamp(tdbg, bx, 23) + else: + _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) + if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) + if tidx == 0: + _stamp(tdbg, bx, 2) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(s_el.load(idx=G_CAP)), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), + ) + + g_n = s_el.load(idx=G_CAP) + ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) + pr2 = s_plan.load(idx=NG_CAP) + blk0 = pr2 & cutlass.Int32(0xFFFF) + nb = pr2 >> cutlass.Int32(16) + + # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG + if warp == W_WARP: + if cutlass.const_expr(XSTAGE): + _wait(x_bar.subview(1), 0) # the staged hidden states are quantized + stage = cutlass.Int32(0) + phase = cutlass.Int32(1) + skip = cutlass.Int32(0) + if cutlass.const_expr(E0_OFFLOAD): + cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) + if cutlass.const_expr(EARLY0): + skip = s_misc.load(idx=0) # group-0 stages already in the ring + stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) + phase = cutlass.select_( + skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) + ) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + e = s_el.load(idx=slot) + sf_line = ( + (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) + * cutlass.Int32(SF1_KA) + + j0 * cutlass.Int32(C1 * NUM_KB) + ) * cutlass.Int32(4) + kx = j0 * cutlass.Int32(C1 * KT_BYTES) + row = c * cutlass.Int32(R1) + for j in cutlass.range(j0, j1, unroll=1): + if skip > cutlass.Int32(0): + skip = skip - cutlass.Int32(1) + else: + _waitd(ab_empty.subview(stage), phase, tdbg, 1) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, + ) + for q in cutlass.range_constexpr(C1): + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(q * NUM_BYTES_BOX1) + ), + tma_a1_desc.get_ptr(), + (kx + cutlass.Int32(q * KT_BYTES), row, e), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + ), + tma_s1_desc.get_ptr(), + (cutlass.Int32(0), sf_line), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) + stage, phase = _adv(stage, phase) + kx = kx + cutlass.Int32(C1 * KT_BYTES) + sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + for bb in range(nb): + blk = blk0 + bb + for t2 in cutlass.range_constexpr(KT2): + _waitd(ab_empty.subview(stage), phase, tdbg, 2) + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx( + ab_full.subview(stage), + n_here + * cutlass.Int32( + NUM_BYTES_BOX1 + if SFA_CP + else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 + ), + ) + for jj in cutlass.range_constexpr(4): + if cutlass.Int32(jj) < n_here: + e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) + prims.cp_async_bulk_tensor_shared_cta_global( + sA.subview( + stage * cutlass.Int32(NUM_BYTES_A) + + cutlass.Int32(jj * NUM_BYTES_BOX1) + ), + tma_a2_desc.get_ptr(), + ( + cutlass.Int32(t2 * KT_BYTES), + blk * cutlass.Int32(R2), + e, + ), + ab_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + if cutlass.const_expr(not SFA_CP): + prims.cp_async_bulk_tensor_shared_cta_global( + sSFAraw.subview( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) + ), + tma_s2_desc.get_ptr(), + ( + cutlass.Int32(0), + ( + ( + e * cutlass.Int32(SF2_RB) + + (blk >> cutlass.Int32(2)) + ) + * cutlass.Int32(SF2_KA) + + cutlass.Int32(4 * t2) + ) + * cutlass.Int32(4), + ), + ab_full.subview(stage), + ) + if cutlass.const_expr(SFA_CP): + _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) + stage, phase = _adv(stage, phase) + if lane == 0: + _stamp(tdbg, bx, 3) + + # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] + if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + for p in range(ng + cutlass.Int32(LAG)): + i0 = cutlass.Int32(0) + i1 = cutlass.Int32(0) + for part in cutlass.range_constexpr(2): + cnt = cutlass.Int32(0) + m1 = cutlass.Int32(0) + nk = cutlass.Int32(0) + if cutlass.const_expr(part == 0): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + nk = _nsegs(i0, i1) + else: + if p >= cutlass.Int32(LAG): + nk = nb + for k in range(nk): + if cutlass.const_expr(part == 0): + ug, j0, j1 = _seg(i0, i1, k) + m1 = ( + ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + ) & cutlass.Int32(3) + cnt = j1 - j0 + else: + m1 = (blk0 + k) & cutlass.Int32(3) + cnt = cutlass.Int32(KT2) + for _j in cutlass.range(cnt, unroll=1): + _waitd(ab_full.subview(stage), phase, tdbg, 3) + raw0 = ( + stage * cutlass.Int32(NUM_WORDS_SFA_RAW) + + lane * cutlass.Int32(4) + + m1 + ) + dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) + for kb in cutlass.range_constexpr(NUM_KB): + w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) + w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) + w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) + w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) + sSFA.store( + ( + cutlass.Int32(w0), + cutlass.Int32(w1), + cutlass.Int32(w2), + cutlass.Int32(w3), + ), + idx=dst0 + cutlass.Int32(kb * 128), + alignment=16, + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + cute.arch.sync_warp() + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFA1_COL + + stage * cutlass.Int32(16) + + kb * 4, + ), + sd0 + (kb * 512 >> 4), + multicast=s2t_mc, + ) + prims.tcgen05_commit(sfa_ready.subview(stage)) + stage, phase = _adv(stage, phase) + + if cutlass.const_expr(B2IMG): + # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group + if warp == cutlass.Int32(S_WARP): + for p in range(ng): + i0, i1 = _plan_range(s_plan, p) + n_done = cutlass.Int32(0) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + n_done = n_done + cutlass.select_( + (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), + cutlass.Int32(1), + cutlass.Int32(0), + ) + _wait( + fc1_done.subview(p & cutlass.Int32(3)), + (p >> cutlass.Int32(2)) & cutlass.Int32(1), + ) + if cutlass.const_expr(B2_DLY > 0): + t0d = _gtime() + dl = cutlass.Int64( + ( + (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) + & cutlass.Int32(63) + ) + * cutlass.Int32(B2_DLY // 64) + ) + while _gtime() - t0d < dl: + pass + if (lane == 0) & (n_done > cutlass.Int32(0)): + _fence_proxy_async_global() + _red_release_add( + b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done + ) + cute.arch.sync_warp() + # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete + if warp == cutlass.Int32(X_WARP): + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + if lane == 0: + _stamp(tdbg, bx, 29) + caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() + if lane == 0: + cv = _ld_acquire_u32(caddr) + while (cv - cnt_target) < cutlass.Int32(0): + _nanosleep(B2IMG_NS) + cv = _ld_acquire_u32(caddr) + _fence_proxy_async_global() + prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) + prims.cp_async_bulk_shared_cluster_global( + sB2.subview(buf * cutlass.Int32(IMG_B2)), + img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), + img_bar.subview(buf), + IMG_B2, + ) + prims.cp_async_bulk_shared_cluster_global( + sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), + img32.subview( + g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) + ), + img_bar.subview(buf), + IMG_SF, + ) + cute.arch.sync_warp() + _waitd( + img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 + ) + if prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + cute.arch.sync_warp() + if lane == 0: + _stamp(tdbg, bx, 25) + # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on + # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then + # tcgen05.cp to TMEM. + if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): + xi = cutlass.select_( + warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) + ) + if cutlass.const_expr(G_CAP > TOP_K): + if (bx == 0) & (xi == cutlass.Int32(0)): + # rows of slots this call does not use: tag 0, never taken for a later call's + for i in range( + g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, + cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), + 32, + ): + hbuf32.store( + ( + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ), + idx=i * cutlass.Int32(4), + alignment=16, + ) + s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 + sSFB2_8 = cutlass.Array( + sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 + ) + sB2_32 = cutlass.Array( + sB2.data_ptr(0), + shape=(2 * KT2 * NUM_BYTES_B2 // 4,), + dtype=cutlass.Int32, + alignment=16, + ) + nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group + NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) + for g in range(ng): + buf = g & cutlass.Int32(1) + if g >= cutlass.Int32(2): + _waitd( + b2_free.subview(buf), + ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 101, + ) + _waitd( + fc1_done.subview(g & cutlass.Int32(3)), + (g >> cutlass.Int32(2)) & cutlass.Int32(1), + tdbg, + 4, + ) # this CTA's own share of group g is out + if cutlass.const_expr(TIMING): + if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g + ) + n_here = g_n - g * cutlass.Int32(4) + n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) + nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) + rec0 = g * cutlass.Int32( + 4 * M_MAX * UNITS_ROW + ) # rows (4 g + j) M + t are contiguous + done = cutlass.Int32(0) + rspin = cutlass.Int32(0) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 29) + npoll = cutlass.Int32(0) + got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] + while done == cutlass.Int32(0): + ok = cutlass.Boolean(True) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + v4 = prims.load_ext( + hbuf32.subview((rec0 + qi) * cutlass.Int32(4)), + dtype=cutlass.Int32, + count=4, + order="relaxed", + scope="gpu", + ) + for wq in cutlass.range_constexpr(3): + got[z][wq] = cutlass.Int32(v4[wq]) + ok = ok & ( + ( + ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) + & cutlass.Int32(-16777216) + ) + == cutlass.Int32(0) + ) + npoll = npoll + cutlass.Int32(1) + if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): + done = cutlass.Int32(1) + else: + _nanosleep(B2_POLL_NS) + if cutlass.const_expr(HANG_DBG): + rspin = rspin + cutlass.Int32(1) + if rspin > cutlass.Int32(1 << 14): + done = cutlass.Int32(1) + tdbg.store( + cutlass.Int64(200), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, + ) + for z in cutlass.range_constexpr(NZ): + qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi + if qi < nrec: + w0, w1, w2 = got[z][0], got[z][1], got[z][2] + rr = qi // cutlass.Int32( + UNITS_ROW + ) # row within the group: j M + t = B row n + c = qi - rr * cutlass.Int32(UNITS_ROW) + t2 = c >> cutlass.Int32(4) + cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) + lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) + hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( + (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) + ) + sB2_32.store( + (lo, hi), + idx=( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) + + rr * cutlass.Int32(KT_BYTES) + + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) + + (c & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 + sSFB2_8.store( + cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), + idx=( + ( + (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) + + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) + ) + * cutlass.Int32(128) + + rr * cutlass.Int32(4) + ) + * cutlass.Int32(4) + + (c & cutlass.Int32(3)), + ) + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + if cutlass.const_expr(XW > 1): + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) + cute.arch.sync_warp() + if (xi == cutlass.Int32(0)) & prims.elect_sync(): + sd0 = prims.Tcgen05SmemDesc.build( + sSFB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=128, + base_offset=0, + layout=0, + ) + for a in cutlass.range_constexpr(KT2 * NUM_KB): + prims.tcgen05_cp( + s2t_shape, + _tptr( + base_row, + base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, + ), + sd0 + + cutlass.Int64( + (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) + * cutlass.Int32(512 >> 4) + ), + multicast=s2t_mc, + ) + prims.tcgen05_commit(b2_ready.subview(buf)) + if (lane == 0) & (xi == cutlass.Int32(0)): + _stamp(tdbg, bx, 25) + if cutlass.const_expr(TIMING): + if g < cutlass.Int32(8): + tdbg.store( + _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g + ) + if cutlass.const_expr(TIMING): + tdbg.store( + cutlass.Int64(npoll), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), + ) + + # ---- MMA (warp 7) + if warp == MMA_WARP: + idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N1, + m_dim=MMA_M, + sparsity_version=0, + ) + idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + scale_format=0, + n_dim=N2, + m_dim=MMA_M, + sparsity_version=0, + ) + desc_a0 = prims.Tcgen05SmemDesc.build( + sA.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b10 = prims.Tcgen05SmemDesc.build( + sB1.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + desc_b20 = prims.Tcgen05SmemDesc.build( + sB2.subview(0), + leading_byte_offset=16, + stride_byte_offset=1024, + base_offset=0, + layout=2, + ) + tcol = (base_row << 16) | base_col + _waitd(b1_ready, 0, tdbg, 5) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + stage = cutlass.Int32(0) + phase = cutlass.Int32(0) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + acc_b2 = cutlass.Int64(0) + acc_sf = cutlass.Int64(0) + for p in range(ng + cutlass.Int32(LAG)): + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + _waitd( + acc_empty.subview(ab), + ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 102, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc_ptr = cutlass.inttoptr( + tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 + ) + for j in cutlass.range(j0, j1, unroll=1): + tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 6) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 6) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw0) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc1, + (j > j0) | cutlass.Boolean(kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc_full.subview(ab)) + k1 = k1 + cutlass.Int32(1) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + buf = g & cutlass.Int32(1) + tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + _waitd( + b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 + ) + if cutlass.const_expr(TIMING): + acc_b2 = acc_b2 + (_gtime() - tw2) + if (lane == 0) & (g == ng - cutlass.Int32(1)): + _stamp(tdbg, bx, 4) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_empty.subview(a2), + ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), + tdbg, + 104, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc2_ptr = cutlass.inttoptr( + tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 + ) + for t2 in cutlass.range_constexpr(KT2): + tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) + if cutlass.const_expr(SFA_CP): + _waitd(ab_full.subview(stage), phase, tdbg, 7) + prims.fence_proxy( + "async_shared", space=prims.SharedSpace.shared_cta + ) + else: + _waitd(sfa_ready.subview(stage), phase, tdbg, 7) + if cutlass.const_expr(TIMING): + acc_sf = acc_sf + (_gtime() - tw1) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) + db = desc_b20 + cutlass.Int64( + (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) + * cutlass.Int32(NUM_BYTES_B2 >> 4) + ) + sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) + sfb_c = ( + tcol + + SFB2_COL + + buf * cutlass.Int32(KT2 * NUM_KB * 4) + + cutlass.Int32(t2 * NUM_KB * 4) + ) + if prims.elect_sync(): + if cutlass.const_expr(SFA_CP): + _sfa_to_tmem(sSFA, stage, base_row, base_col) + for kb in cutlass.range_constexpr(NUM_KB): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF4NVF4, + prims.CTAGroup.CTA_1, + acc2_ptr, + da + INC_K * kb, + db + INC_K * kb, + idesc2, + cutlass.Boolean(t2 > 0 or kb > 0), + cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), + cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), + scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, + ) + prims.tcgen05_commit(ab_empty.subview(stage)) + stage, phase = _adv(stage, phase) + if prims.elect_sync(): + prims.tcgen05_commit(acc2_full.subview(a2)) + k2 = k2 + cutlass.Int32(1) + if cutlass.const_expr(TIMING): + if lane == 0: + tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) + tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) + + # ---- epilogue warps (0-3) + if warp < 4: + row_w = base_row + warp * 32 + for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): + qi = tidx + cutlass.Int32(r * EPI_THREADS) + sl_e = qi // cutlass.Int32(6) + if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): + fld = qi - sl_e * cutlass.Int32(6) + s_esc.store( + cutlass.Float32( + escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) + ), + idx=sl_e * cutlass.Int32(8) + fld, + ) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + k1 = cutlass.Int32(0) + k2 = cutlass.Int32(0) + SH_ADD = SHARED + wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] + ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) + for p in range(ng + cutlass.Int32(LAG)): + if cutlass.const_expr(AR_PRE): + if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: + # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads + # at 8-byte offsets: misaligned) + ard = tuple( + cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) + for i_ in (1, 3, 4, 5) + ) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_cl = cutlass.select_( + f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) + ) + pars[bt] = _ar_parity(arp, f_t, f_cl) + if cutlass.const_expr(SH_ADD): + if p == ng + cutlass.Int32(LAG - 1): + # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + f_pair = warp + cutlass.Int32(bt * 4) + f_bb = f_pair // cutlass.Int32(M_MAX) + f_t = f_pair - f_bb * cutlass.Int32(M_MAX) + f_ch = cutlass.select_( + f_bb < nb, + f_t * cutlass.Int32(H) + + (blk0 + f_bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + cutlass.Int32(0), + ) + wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) + if p < ng: + i0, i1 = _plan_range(s_plan, p) + for k in range(_nsegs(i0, i1)): + ug, j0, j1 = _seg(i0, i1, k) + ab = k1 & cutlass.Int32(1) + slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) + cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) + unit = slot * cutlass.Int32(U1) + cblk + _waitd( + acc_full.subview(ab), + (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 105, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + acc = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 + ), + num=N1, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc_empty.subview(ab)) + pb = ab * cutlass.Int32(M_MAX * C1 * R1) + for t in cutlass.range_constexpr(M_MAX): + part = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(4): + part = cutlass.select_( + warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part + ) + s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + for tr in cutlass.range_constexpr((M_MAX + 3) // 4): + t = warp + cutlass.Int32(4 * tr) + if t < cutlass.Int32(M_MAX): + xv = _add_rn( + _add_rn( + _add_rn( + s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), + s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), + ), + s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), + ) + finish = cutlass.Int32(1) + if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): + paddr = ( + partials.subview( + (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane + ) + .data_ptr() + .toint() + ) + if j0 == cutlass.Int32(0): + _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) + finish = cutlass.Int32(0) + else: + pbits = _ld_relaxed_u32(paddr) + spins = cutlass.Int32(0) + while pbits == cutlass.Int32(PARTIAL_EMPTY): + pbits = _ld_relaxed_u32(paddr) + spins = spins + cutlass.Int32(1) + _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) + xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) + cute.arch.sync_warp() + if cutlass.const_expr( + XVDUMP + ): # debug: FC1 sums into the (unused) partial slots + partials.store( + xv.bitcast(cutlass.Int32), + idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) + + lane, + ) + if finish == cutlass.Int32(1): + gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) + is_up = ( + (lane >> cutlass.Int32(3)) & cutlass.Int32(1) + ) == cutlass.Int32(0) + es = slot * cutlass.Int32(8) + qv, sfb = _swiglu_quant( + xv, + gate_raw, + is_up, + s_esc.load(idx=es), + s_esc.load(idx=es + cutlass.Int32(1)), + s_esc.load(idx=es + cutlass.Int32(2)), + s_esc.load(idx=es + cutlass.Int32(3)), + s_esc.load(idx=es + cutlass.Int32(4)), + ) + byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) + bt = [ + cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) + for i in range(8) + ] + rw0 = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | tagv + ) + rw1 = ( + bt[3] + | (bt[4] << cutlass.Int32(8)) + | (bt[5] << cutlass.Int32(16)) + | tagv + ) + rw2 = ( + bt[6] + | (bt[7] << cutlass.Int32(8)) + | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) + | tagv + ) + if cutlass.const_expr(B2IMG): + if lane == 0: + # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 + # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), + # row n at byte 16 n, byte c % 4 + n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t + t2i = cblk >> cutlass.Int32(4) + cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) + gbase = p * cutlass.Int32(IMG_BYTES) + lo = ( + bt[0] + | (bt[1] << cutlass.Int32(8)) + | (bt[2] << cutlass.Int32(16)) + | (bt[3] << cutlass.Int32(24)) + ) + hi = ( + bt[4] + | (bt[5] << cutlass.Int32(8)) + | (bt[6] << cutlass.Int32(16)) + | (bt[7] << cutlass.Int32(24)) + ) + img32.store( + (lo, hi), + idx=( + gbase + + t2i * cutlass.Int32(NUM_BYTES_B2) + + n * cutlass.Int32(KT_BYTES) + + (cci ^ (n & cutlass.Int32(7))) + * cutlass.Int32(16) + + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) + ) + // cutlass.Int32(4), + alignment=8, + ) + img8 = cutlass.Array( + img32.data_ptr(0), + shape=(NG_CAP * IMG_BYTES,), + dtype=cutlass.Int8, + alignment=16, + ) + img8.store( + cutlass.Int8(sfb & cutlass.Int32(0xFF)), + idx=gbase + + cutlass.Int32(IMG_B2) + + ( + ( + t2i * cutlass.Int32(NUM_KB) + + ( + (cblk & cutlass.Int32(15)) + >> cutlass.Int32(2) + ) + ) + * cutlass.Int32(512) + ) + + n * cutlass.Int32(16) + + (cblk & cutlass.Int32(3)), + ) + else: + if lane == 0: + hbuf32.store( + (rw0, rw1, rw2, tagv), + idx=( + (slot * cutlass.Int32(M_MAX) + t) + * cutlass.Int32(UNITS_ROW) + + cblk + ) + * cutlass.Int32(4), + alignment=16, + ) + k1 = k1 + cutlass.Int32(1) + prims.mbarrier_arrive( + fc1_done.subview(p & cutlass.Int32(3)) + ) # warp 5 may poll group p + if tidx == 0: + _stamp(tdbg, bx, 10) + if cutlass.const_expr(TIMING): + if p < cutlass.Int32(8): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, + ) + if p >= cutlass.Int32(LAG): + g = p - cutlass.Int32(LAG) + for bb in range(nb): + a2 = k2 & cutlass.Int32(1) + _waitd( + acc2_full.subview(a2), + (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), + tdbg, + 106, + ) + if cutlass.const_expr(TIMING): + if (tidx == 0) & (g == ng - cutlass.Int32(1)): + tdbg.store( + _gtime(), + idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, + ) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + yv = prims.tcgen05_ld( + "32x32b", + _tptr( + row_w, + base_col + + ACC2_COL + + a2 * cutlass.Int32(N2) + + warp * cutlass.Int32(M_MAX), + cutlass.Float32, + ), + num=LD_POW2, + ) + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + prims.mbarrier_arrive(acc2_empty.subview(a2)) + sl = g * cutlass.Int32(4) + warp + if sl < g_n: + a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) + for t in cutlass.range_constexpr(M_MAX): + kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + if kp >= cutlass.Int32(0): + ys.store( + _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), + idx=( + (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) + * cutlass.Int32(TOP_K) + + kp + ) + * cutlass.Int32(R2) + + lane, + ) + k2 = k2 + cutlass.Int32(1) + prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) + if tidx == 0: + _stamp(tdbg, bx, 6) + cute.arch.sync_warp() + cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) + if tidx == 0: + _stamp(tdbg, bx, 19) + # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order + for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): + pair = warp + cutlass.Int32(bt * 4) + bb = pair // cutlass.Int32(M_MAX) + t = pair - bb * cutlass.Int32(M_MAX) + if bb < nb: + acc_o = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(TOP_K): + acc_o = _fma_rn( + s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ), + acc_o, + ) + if cutlass.const_expr( + YDUMP + ): # debug: y of (token, k) into the partials workspace + partials.store( + cutlass.Float32( + ys.load( + idx=( + (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) + + cutlass.Int32(k) + ) + * cutlass.Int32(R2) + + lane + ) + ).bitcast(cutlass.Int32), + idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)), + ) + ch = ( + t * cutlass.Int32(H) + + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(SH_ADD): + # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) + wy = wys[bt] + while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: + wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) + acc_o = _add_rn( + (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) + ) + out.store(cutlass.BFloat16(acc_o), idx=ch) + if ar_on: + cl = ( + (blk0 + bb) * cutlass.Int32(R2) + + cutlass.Int32(4) * (lane & cutlass.Int32(7)) + + (lane >> cutlass.Int32(3)) + ) + if cutlass.const_expr(AR_PRE): + _ar_push_pre( + ard, + pars[bt], + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + else: + _ar_push( + arp, + _bf16_round(acc_o), + ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), + t, + cl, + ) + + # ---- teardown + if tidx == 0: + _stamp(tdbg, bx, 7) + epochs.store(ep + cutlass.Int32(1), idx=bx) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + cute.arch.sync_warp() + prims.barrier_cta_sync(0) + if ar_on | cutlass.Boolean(LATE_TRIGGER): + # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA + cute.arch.griddepcontrol_launch_dependents() + if warp == MMA_WARP: + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS + ) + + +@cute.jit +def m3_moe( + a1_tensor: cute.Tensor, + sfa1_tensor: cute.Tensor, + a2_tensor: cute.Tensor, + sfa2_tensor: cute.Tensor, + logits: cute.Tensor, + bias: cute.Tensor, + x32: cute.Tensor, + xsf32: cute.Tensor, + escal: cute.Tensor, + plan: cute.Tensor, + partials: cute.Tensor, + hbuf32: cute.Tensor, + epochs: cute.Tensor, + out: cute.Tensor, + route_dbg: cute.Tensor, + tdbg: cute.Tensor, + xb32: cute.Tensor, + rw: cute.Tensor, + sfsc: cute.Tensor, + lgbuf: cute.Tensor, + guw: cute.Tensor, + gusf: cute.Tensor, + dnw: cute.Tensor, + dnsf: cute.Tensor, + hx: cute.Tensor, + dnx: cute.Tensor, + arp: cute.Tensor, + s1w: cute.Tensor, + s2w: cute.Tensor, + img32: cute.Tensor, + b2cnt: cute.Tensor, + stream: cuda_driver.CUstream, +): + tma_a1_desc = cuda.create_tensor_map_tiled( + global_address=a1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H // 2, TWO_I, E_LOCAL], + global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], + box_dims=(KT_BYTES, R1, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_a2_desc = cuda.create_tensor_map_tiled( + global_address=a2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[I_TP // 2, H, E_LOCAL], + global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], + box_dims=(KT_BYTES, R2, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_s1_desc = cuda.create_tensor_map_tiled( + global_address=sfa1_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_s2_desc = cuda.create_tensor_map_tiled( + global_address=sfa2_tensor.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], + global_strides=[128 // 16], + box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_gu_desc = tma_a1_desc + tma_dn_desc = tma_a2_desc + if cutlass.const_expr(SHARED and not SH_PRETILE): + tma_gu_desc = cuda.create_tensor_map_tiled( + global_address=guw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[H, 2 * SH_I], + global_strides=[H // 16], + box_dims=(KT_BYTES, 32), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + tma_dn_desc = cuda.create_tensor_map_tiled( + global_address=dnw.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[SH_I, H], + global_strides=[SH_I // 16], + box_dims=(KT_BYTES, 128), + swizzle=cuda.TensorMapSwizzle.s128b, + ) + m3_moe_kernel( + tma_a1_desc, + tma_a2_desc, + tma_s1_desc, + tma_s2_desc, + tma_gu_desc, + tma_dn_desc, + logits, + bias, + x32, + xsf32, + escal, + plan, + partials, + hbuf32, + epochs, + out, + route_dbg, + tdbg, + xb32, + rw, + sfsc, + lgbuf, + guw, + gusf, + dnw, + dnsf, + hx, + dnx, + arp, + s1w, + s2w, + img32, + b2cnt, + ).launch( + grid=(NCTA + NSH, 1, 1), + block=(THREADS, 1, 1), + cluster=(1, 1, 1), + stream=stream, + use_pdl=True, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py new file mode 100644 index 000000000000..12f6008d439f --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py @@ -0,0 +1,602 @@ +# 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. +"""Host side of the MiniMax-M3 routed-expert kernel on one TP4 rank, over the TRTLLM-Gen NVFP4 buffers read in place. + +``m3_moe`` (m3_moe_kernel.py, :class:`M3MoeState`): 1..8 decode tokens (decode batches, spec-decode verify). The call +takes the bf16 hidden states and folds in the router GEMV (bit-exact with m3_router_gemv_fp4) and the activation's +NVFP4 quantization. With ``shared=True`` it also runs the shared expert on dedicated CTAs and returns +bf16(shared + routed); otherwise it returns the rank's routed partial, bit for bit what the stock fused MoE returns. + +Usage (build before CUDA-graph capture; every call is capture-safe afterwards):: + + state = M3MoeState( + device, num_tokens=6, shared=True + ) # workspace + compiled kernel, shared by all layers + layer = state.layer( + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf), + ) + out = layer(hidden) # bf16 [M, 6144] +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from typing import Dict, Optional, Tuple + +import torch + +HIDDEN = 6144 +INTER = 768 +NUM_EXPERTS = 128 +SHARED_INTER = 768 # the shared expert's intermediate per TP rank +SHARED_CTAS = 24 # m3_moe's shared-expert CTAs (beside num_ctas routed CTAs) +TOP_K = 4 +U1 = 2 * INTER // 32 # FC1 units per expert +MIN_CTAS = 64 # the router GEMV takes at most 2 experts per CTA (router weights held in registers) + +_UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) # one FC1 unit's weight + scale bytes +_ROW2_BYTES = INTER // 2 + INTER // 16 # one FC2 weight row's bytes (values + scales) +_S1 = 6 # FC1 stages (items) per unit +_ITEM_BYTES = _UNIT_BYTES // _S1 +_modules: Dict[tuple, object] = {} +_compiled: Dict[tuple, object] = {} +# m3_moe trace-time options by token count 1..8 (index num_tokens - 1). +# ``sfa_cp``: 18 KB ring stages, the A scales copied by cp.async into the instruction layout (every token count). +SFA_CP_BY_M: Tuple[int, ...] = (1, 1, 1, 1, 1, 1, 1, 1) +# ``b2img``: FC1 -> FC2 hand-off through per-group global B images + release counters (7 and 8 tokens). +B2IMG_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 1, 1) +# ``ar_pre``: the finalize's all-reduce push with its call parities / descriptor preloaded (all but 2 tokens). +AR_PRE_BY_M: Tuple[int, ...] = (1, 0, 1, 1, 1, 1, 1, 1) +# Ring ``stages`` (0 = the kernel's default, the most that fit). 6 at 8 tokens measured 0.4-0.5 us faster than the +# default 7 (same MMAs, bit-identical outputs). +STAGES_BY_M: Tuple[int, ...] = (0, 0, 0, 0, 0, 0, 0, 6) +# m3_moe phase lag: phase p streams FC1 of group p and FC2 of group p - M3_MOE_LAG (the kernel's default lag). +M3_MOE_LAG = 2 + + +_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") + + +def _kernel_module(config: dict, path: str): + key = (os.path.basename(path),) + tuple(sorted(config.items())) + mod = _modules.get(key) + if mod is None: + name = ( + f"{__name__}_" + + os.path.splitext(os.path.basename(path))[0] + + "_" + + "_".join(f"{k}{v}" for k, v in key[1:]) + ) + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + mod.M3_CONFIG = dict(config) + sys.modules[name] = mod + spec.loader.exec_module(mod) + _modules[key] = mod + return mod + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def _shared_views(shared_weights, who): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). Stages 0..23: gate_up k-tiles + g (slice 0) / 24 + g (slice 1) of rows [gate 32 s .., up 768 + 32 s .., gate (slice 1), up (slice 1)]; 24..35: + down tile (g - 24) / 6 (rows 256 s + 128 j ..), k-tile (g - 24) % 6. Built once per layer.""" + dev = gu8.device + H_, I_ = HIDDEN, SHARED_INTER + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty(n_cta, 36, 16384 + 512, dtype=torch.uint8, device=dev) + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = torch.where((q & 1) == 0, 32 * s + ri, I_ + 32 * s + ri) + sl = q >> 1 + for g in range(24): + kt = sl * 24 + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(12): + jt, kt = gd // 6, gd % 6 + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, 24 + gd, :16384] = img.reshape(-1) + out[s, 24 + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +# --------------------------------------------------------------------------------------------------------------------- +# m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. +# --------------------------------------------------------------------------------------------------------------------- +# Pre-tiled shared weights per layer, keyed by the source buffers (which the engines read in place): every token-count +# state of a layer streams the same copy. An entry holds its sources, so their addresses are not reused while cached. +_PRETILED: Dict[tuple, tuple] = {} + + +def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: + key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + hit = _PRETILED.get(key) + if hit is None: + hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + return hit[1] + + +def make_phase_plan( + num_ctas: int, m_max: int, lag: int = 2, tail_w: float = 1.0, early0: bool = False +): + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work. Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = TOP_K * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + blk_bytes = 32 * _ROW2_BYTES # one block, one slot + for g in range(TOP_K, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * U1 * _S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and not ( + early0 and p == 0 and m_max >= 2 + ): # early0: group 0's ranges the same for every G + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + items = [0] * num_ctas + if total >= _S1 * num_ctas: + + def counts_for(limit): + return [max(_S1, (limit - f) // _ITEM_BYTES) for f in fc2] + + lo, hi = 0, max(fc2) + total * _ITEM_BYTES + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > _S1), + key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), + ) + items[c] -= 1 + extra -= 1 + else: + chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) + for c in chosen: + items[c] = _S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +class M3MoeState: + """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run + on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with + m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated + CTAs beside the ``num_ctas`` routed ones and the call returns bf16(shared + routed). Build before CUDA-graph + capture; the first call compiles.""" + + fused_router = True + + def __init__( + self, + device: torch.device, + num_tokens: int, + num_ctas: Optional[int] = None, + shared: bool = False, + ): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3MoeState allocates and compiles: build it before CUDA-graph capture" + ) + sms_all = torch.cuda.get_device_properties(device).multi_processor_count + if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + raise ValueError( + f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " + f"(all resident at once)" + ) + if not 1 <= num_tokens <= 8: + raise ValueError(f"m3_moe takes 1..8 tokens per call, not {num_tokens}") + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or sms + if not MIN_CTAS <= self.num_ctas <= sms: + raise ValueError( + f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" + ) + self.m = num_tokens + # One token: the trailing FC2s weigh nothing in the last phase's balance. Group 0's FC1 ranges are the same + # for every distinct-expert count (early0). + plan, nb_cap = make_phase_plan( + self.num_ctas, num_tokens, M3_MOE_LAG, 0.0 if num_tokens == 1 else 1.0, True + ) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + } + if shared: + self.config["shared"] = 1 + self.shared = shared + # 18 KB ring stages (A scales by cp.async into the instruction layout) per token count, unless configured + self.config["sfa_cp"] = SFA_CP_BY_M[num_tokens - 1] + # FC1 -> FC2 hand-off through per-group global B images + release counters, per token count + self.config["b2img"] = B2IMG_BY_M[num_tokens - 1] + # the finalize's all-reduce push with its call parities / descriptor preloaded, per token count + self.config["ar_pre"] = AR_PRE_BY_M[num_tokens - 1] + if STAGES_BY_M[num_tokens - 1] > 0: + self.config["stages"] = STAGES_BY_M[num_tokens - 1] + try: + self.mod = mod = _kernel_module(self.config, _ENGINE_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe: num_tokens={num_tokens}, num_ctas={self.num_ctas} does not fit the shared " + f"memory / TMEM budget (fewer CTAs means more FC2 blocks per CTA): {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + self.partials = torch.full((mod.G_CAP * U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw) + self.hbuf32 = torch.zeros(mod.G_CAP * num_tokens * U1 * 4, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros(mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw) + self._dummy = torch.zeros(4, dtype=torch.int32, **kw) + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # B2IMG: per-group FC2 B images and their unit counters (int32; every group gains CNT_PER_CALL per call) + self.img32 = torch.zeros(max(1, mod.NG_CAP * mod.IMG_BYTES // 4), dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 8, dtype=torch.int32, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [1536, 6144], its swizzled UE8M0 + scales, down e4m3 [6144, 768], its swizzled UE8M0 scales).""" + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + if self.shared: + layer.shared_weights = _shared_views(shared_weights, "m3_moe") + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + _pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + key = (_ENGINE_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState` (or op16's state): ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py new file mode 100644 index 000000000000..446776419315 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py @@ -0,0 +1,288 @@ +# 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. +"""Host side of ``m3_moe16`` (m3_moe16_kernel.py): the MiniMax-M3 MoE layer (fused router + routed experts + shared +expert) for 9..16 decode tokens on one TP4 rank. The 1..8-token engine (op.M3MoeState) is untouched; select this +state only for num_tokens > 8. + +Output: bf16 [M, 6144] = bf16(shared + routed), the routed part bit for bit TRTLLM-Gen's fused MoE on the fp32 router +GEMV logits (m3_router_gemv's math), the shared part as the in-model m3_mx_gemv shared expert. + + state = M3Moe16State(device, num_tokens=12) # before CUDA-graph capture; one per token count + layer = state.layer(w3_w1_weight, w3_w1_weight_scale, w2_weight, w2_weight_scale, fc31_alpha, fc31_scale_c, + act_alpha, act_beta, act_clamp, fc2_alpha, e_score_correction_bias, + router_weight=router_weight, fc31_input_scale=fc31_input_scale, + shared_weights=(gate_up_w, gate_up_sf, down_w, down_sf)) + out = layer(hidden) # first call compiles: call once before capture + +Protocol state (per state object, shared by the layers that run on it in stream order): per-group FC2-B images and +unit counters, the B1 image + counter, per-CTA epochs. Every call adds exactly 192 to each group counter and NCTA to +the B1 counter, so nothing is reset between calls (graph replay safe). +""" + +from __future__ import annotations + +import os +from typing import Optional + +import torch + +from .op import ( + HIDDEN, + SHARED_CTAS, + SHARED_INTER, + U1, + M3MoeLayer, + _compiled, + _kernel_module, + _shared_views, + _view, + is_supported, +) + + +def _group_slots(g: int, p: int, perm: bool, min_ng: int = 5) -> int: + """Slots of group p when the call has g distinct experts (the kernel's slot order).""" + ng = (g + 3) // 4 + tail = g - 4 * (ng - 1) + if p >= ng: + return 0 + if perm and tail < 4 and ng >= max(3, min_ng): + return tail if p == 1 else 4 + return min(4, g - 4 * p) + + +def make_phase_plan16(num_ctas: int, m_max: int, lag: int = 2, perm: bool = True, min_ng: int = 5): + """op.make_phase_plan (early0, tail_w 1) with the kernel's group sizes (PERM: a partial last group is group 1).""" + S1 = 6 + UNIT_BYTES = 32 * (HIDDEN // 2 + HIDDEN // 16) + ITEM = UNIT_BYTES // S1 + blk_bytes = 32 * (768 // 2 + 768 // 16) + gcap = 4 * m_max + nb2 = HIDDEN // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + for g in range(4, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) + for p in range(ng): + n_p = _group_slots(g, p, perm, min_ng) + total = n_p * U1 * S1 + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * _group_slots(g, p - lag, perm, min_ng) * blk_bytes + if p == ng - 1 and not (p == 0 and m_max >= 2): + for q in range(max(0, ng - lag), ng): + fc2[c] += nbs[c] * _group_slots(g, q, perm, min_ng) * blk_bytes + items = [0] * num_ctas + if total >= S1 * num_ctas: + + def counts_for(limit): + return [max(S1, (limit - f) // ITEM) for f in fc2] + + lo, hi = 0, max(fc2) + total * ITEM + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > S1), + key=lambda c: (-(fc2[c] + items[c] * ITEM), c), + ) + items[c] -= 1 + extra -= 1 + else: + for c in sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // S1]): + items[c] = S1 + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | ((i0 + items[c]) << 16) + i0 += items[c] + assert i0 == total + return plan, max(nbs) + + +_ENGINE16_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe16_kernel.py") +M16_MIN_TOKENS, M16_MAX_TOKENS = 9, 16 +# Default routed CTAs (beside the SHARED_CTAS shared-expert CTAs). +M16_NUM_CTAS = 128 +# Phase lag: phase p streams FC1 of group p and FC2 of group p - M16_LAG (the kernel's default lag). +M16_LAG = 3 + + +class M3Moe16State: + """Workspace + compiled kernel of ``m3_moe16`` for ``num_tokens`` (9..16) tokens per call.""" + + fused_router = True + shared = True + + def __init__(self, device: torch.device, num_tokens: int, num_ctas: Optional[int] = None): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "M3Moe16State allocates and compiles: build it before CUDA-graph capture" + ) + if not M16_MIN_TOKENS <= num_tokens <= M16_MAX_TOKENS: + raise ValueError( + f"m3_moe16 takes {M16_MIN_TOKENS}..{M16_MAX_TOKENS} tokens per call, not {num_tokens}" + ) + self.device = torch.device(device) + ok, why = is_supported(self.device) + if not ok: + raise RuntimeError(f"m3_moe16: {why}") + sms = torch.cuda.get_device_properties(self.device).multi_processor_count + self.num_ctas = num_ctas or M16_NUM_CTAS + if not 128 <= self.num_ctas <= sms - SHARED_CTAS: + raise ValueError( + f"m3_moe16 runs 128 .. {sms - SHARED_CTAS} routed CTAs + {SHARED_CTAS} shared (all resident)" + ) + self.m = num_tokens + plan, nb_cap = make_phase_plan16(self.num_ctas, num_tokens, M16_LAG) + self.config = { + "m_max": num_tokens, + "num_ctas": self.num_ctas, + "nb2_cap": nb_cap, + "fused_router": 1, + "timing": 0, + "shared": 1, + } + try: + self.mod = mod = _kernel_module(self.config, _ENGINE16_PATH) + except AssertionError as err: + raise ValueError( + f"m3_moe16: num_tokens={num_tokens} does not fit the smem / TMEM budget: {err}" + ) from err + kw = dict(device=self.device) + self.plan = plan.to(self.device) + # split units' head sums: every slot empty (0xFFFFFFFF) between calls + self.partials = torch.full((mod.G_CAP * U1 * num_tokens * 32,), -1, dtype=torch.int32, **kw) + # FC2-B images per group + the B1 image; counters: group g at 8 g, B1 at 8 NG_CAP, CTA tickets at 8 NG_CAP + 8 + self.img = torch.zeros(mod.IMG_TOTAL // 4, dtype=torch.int32, **kw) + self.cnt = torch.zeros(mod.NG_CAP * 8 + 16, dtype=torch.int32, **kw) + self.epochs = torch.zeros(self.num_ctas + mod.NSH, dtype=torch.int32, **kw) + self.route_dbg = torch.zeros(num_tokens, 8, dtype=torch.int32, **kw) + self.tdbg = torch.zeros(self.num_ctas + mod.NSH, 64, dtype=torch.int64, **kw) + self.lgbuf = torch.zeros( + mod.LG_REPL * mod.LG_STRIDE, dtype=torch.int64, **kw + ) # logits: tag << 32 | f32 + self._ar_off = torch.zeros(8, dtype=torch.int64, **kw) + # shared expert: h and y_sh as (call tag << 16 | bf16), tag 0 = never written + self.hx = torch.zeros(num_tokens * SHARED_INTER, dtype=torch.int32, **kw) + self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) + self._fn = None + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> M3MoeLayer: + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + ) + layer.shared_weights = _shared_views(shared_weights, "m3_moe16") + layer.shared_mxfp8 = tuple(shared_weights) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError(f"m3_moe16: expects contiguous bf16 hidden states [{m}, {HIDDEN}]") + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + layer.bias, + layer.escal, + self.plan.view(-1), + self.partials, + self.img, + self.cnt, + layer.w3_w1_weight_scale.reshape(-1).view(torch.int32), + layer.w2_weight_scale.reshape(-1).view(torch.int32), + layer.w3_w1_weight.reshape(-1).view(torch.int8), + layer.w2_weight.reshape(-1).view(torch.int8), + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + *layer.shared_weights, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_moe16 compiles on its first call: call it once before capture" + ) + import cutlass.cute as cute + + sig = [_view(t, 16 if t.element_size() * t.numel() % 16 == 0 else 4) for t in args] + key = (_ENGINE16_PATH,) + tuple(sorted(self.config.items())) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py new file mode 100644 index 000000000000..ef3a4c5847ea --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) as a CuTe DSL kernel (see op.py).""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py new file mode 100644 index 000000000000..4247a736e8ca --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/m3_mx_gemv_kernel.py @@ -0,0 +1,1297 @@ +# 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. +"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. + +The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), +fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 +scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a +``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at +most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. + +A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... +through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale +atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in +one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 +partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the +SPLIT partials in rank order and stores bf16. + +Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i +reads TMEM lanes 32 i .. 32 i + 31). +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute +import cutlass.experimental.cuda as cuda +from cutlass import dsl_user_op +from cutlass.experimental import primitives as prims + +CTA_M = 128 # weight rows per CTA = the MMA's M +MMA_N = 8 # token columns +CTA_K = 128 # one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row +MMA_INST_K = 32 +K_BLOCKS = CTA_K // MMA_INST_K # 4 MMAs per k-tile, one per scale byte of the atom +THREADS = 256 +SF_ATOM_BYTES = 512 +SF_ATOM_U16 = SF_ATOM_BYTES // 2 +SF_COLS = 4 # TMEM columns of one scale atom +W_TILE_BYTES = CTA_M * CTA_K +X_TILE_BYTES = MMA_N * CTA_K +EVICT_FIRST = 0x12F0000000000000 +A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A + +a_dtype = cutlass.Float8E4M3FN +b_dtype = cutlass.Float8E4M3FN + + +def num_k_tiles(k_in: int) -> int: + return k_in // CTA_K + + +def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: + # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 + # MMAs do not form one serial dependency chain; the epilogue adds them. + return max(16, nacc * mma_n) + + +def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: + # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] + return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 + + +def tmem_cols( + ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N +) -> int: + need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS + if a_tmem: + need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA + cols = 32 + while cols < need: + cols *= 2 + return cols + + +@dsl_user_op +def read_globaltimer(*, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int64( + _llvm.inline_asm( + _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): + """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], + "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): + """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, + [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), + cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), + cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): + """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + out = _llvm.inline_asm( + _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + return [ + cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) + ] + + +@dsl_user_op +def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): + """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + done = cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" + "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + return done != cutlass.Int32(0) + + +@dsl_user_op +def _absmax(acc, v, *, loc=None, ip=None): + """max(acc, |v|).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, + is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rcp_approx_ftz(x, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Float32( + _llvm.inline_asm( + _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): + """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], + "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _ld_volatile_s32(addr, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", + has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _st_global_b32(addr, value, *, loc=None, ip=None): + from cutlass._mlir.dialects import llvm as _llvm + + _llvm.inline_asm( + None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], + "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) # fmt: skip + + +@dsl_user_op +def _bf16x2_word(lo, hi, *, loc=None, ip=None): + """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" + "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" + "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" + "mov.b32 $0, w;\n\t}", "=r,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _bf16_bits(x, *, loc=None, ip=None): + """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], + "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _bf16_round(x): + return _bf16_bits(x).bitcast(cutlass.Float32) + + +# The swigluoai activation of MiniMax-M3's dense MLP, gate * sigmoid(1.702 * gate) * (up + 1) with gate clamped to <= 7 +# and up to [-7, 7], in exactly the instruction sequence the Triton silu_and_mul kernel compiles to (modules/swiglu.py +# on SM107): the gate clamp on the packed bf16 pair, the up clamp as min.xorsign.abs, two rounded multiplies into +# ex2.approx, div.full for the reciprocal, two rounded multiplies, round to bf16. No multiply feeds an add, so the +# explicit .rn forms compile to what the Triton kernel's contractable forms do. +SWIGLU_OAI_ALPHA = 1.702 +SWIGLU_OAI_LIMIT = 7.0 + + +@dsl_user_op +def _swiglu_oai_bf16x2(gate_word, up_word, *, loc=None, ip=None): + """The activation of one packed bf16 pair of gate and up values: its two bf16 results as one int32 word (the low + half from the low halves).""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), + [cutlass.Int32(w).ir_value(loc=loc, ip=ip) for w in (gate_word, up_word)], + "{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" + "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, $1, glim;\n\t" + "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" + "mov.b32 {hl, hh}, $2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" + "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" + "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" + "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" + "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" + "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" + "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" + "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" + "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" + "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" + "cvt.rn.bf16x2.f32 $0, r1, r0;\n\t}", "=r,r,r", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +@dsl_user_op +def _rope_bits(c, a, s, b, *, loc=None, ip=None): + """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" + from cutlass._mlir.dialects import llvm as _llvm + from cutlass._mlir.extras import types as _T + + return cutlass.Int32( + _llvm.inline_asm( + _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], + "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" + "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", + has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, + ) + ) # fmt: skip + + +def _din_load4(table, hidden, hidden_half, word_idx): + """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform + base-address select instead of a branch.""" + addr = cutlass.select_( + hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() + ) + ptr = cutlass.inttoptr( + addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 + ) + return ptr.load(count=4, alignment=16) + + +def _stamp(timers, timed: bool, bx, slot: int): + # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). + if timed: + timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) + + +def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): + # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... + if contig: + return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) + return rank + cutlass.Int32(i * split) + + +@cute.kernel +def m3_mx_gemv_kernel( + tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 + tma_desc_wsf: cutlass.GridConstant[ + cuda.TensorMap + ], # W scale atoms [N / 128 * K / 128, 256] u16 + tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 + tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 + y: cutlass.Array, # [M * N] bf16, token-major + x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words + timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) + ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[ + int + ], # 1: push the output rows into the next all-reduce; 2: and trigger after it + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n + din_i: cutlass.Array, # din: int32 token ids [M] + din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words + din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words + din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words + din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] + din_vocab: cutlass.Int32, + din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words + din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words + din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res + rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] + rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] + rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] + rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 + rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 + rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale + rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[ + bool + ], # x_bf16: x_words is [M, 2K] gate | up; the activation is swigluoai(gate, up) +): + k_tiles = num_k_tiles(k_in) + x_tile_bytes = mma_n * CTA_K + my_tiles = k_tiles // split + rows_per_rank = CTA_M // split + tx, _, _ = cute.arch.thread_idx() + bx, _, _ = cute.arch.block_idx() + warp_id = cute.arch.warp_idx() + rank = cutlass.Int32(0) + if cutlass.const_expr(split > 1): + rank = cute.arch.block_idx_in_cluster() + tma_ptr_w = tma_desc_w.get_ptr() + tma_ptr_wsf = tma_desc_wsf.get_ptr() + tma_ptr_x = tma_desc_x.get_ptr() + tma_ptr_xsf = tma_desc_xsf.get_ptr() + n_tile = bx // cutlass.Int32(split) + m_offset = n_tile * cutlass.Int32(CTA_M) + + smem_a = cutlass.Array( + cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_b = cutlass.Array( + cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 + ) + smem_sfa = cutlass.Array( + cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + smem_sfb = cutlass.Array( + cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 + ) + w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) + act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) + tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) + # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded + # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a + # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). + mb_stride = mma_n + 4 + mailbox = cutlass.Array( + cutlass.Float32, + split * rows_per_rank * mb_stride, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. + din_rs = None + if cutlass.const_expr(din): + din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) + + if warp_id == 0: + prims.prefetch_tensormap(tma_ptr_w) + prims.prefetch_tensormap(tma_ptr_wsf) + prims.prefetch_tensormap(tma_ptr_x) + prims.prefetch_tensormap(tma_ptr_xsf) + if prims.elect_sync(): + for s in cutlass.range_constexpr(ring): + prims.mbarrier_init(w_full.subview(s), 1) + prims.mbarrier_init(mma_done.subview(s), 1) + prims.mbarrier_init(act_full, 1) + prims.mbarrier_init(acc_done, 1) + if cutlass.const_expr(split > 1 and st_async): + # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. + prims.mbarrier_init(reduce_ready, 1) + prims.mbarrier_arrive_expect_tx( + reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 + ) + elif cutlass.const_expr(split > 1): + prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) + if warp_id == 2: + prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) + prims.tcgen05_relinquish_alloc_permit() + prims.fence_mbarrier_init() + if cutlass.const_expr(split > 1): + prims.barrier_cluster_arrive_relaxed() + prims.barrier_cluster_wait() + prims.barrier_cta_sync(0) + if tx == 0: + _stamp(timers, timed, bx, 0) + + if warp_id == 0: + if prims.elect_sync(): + for i in cutlass.range_constexpr(ring, my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] + ) + if cutlass.const_expr((n_out // CTA_M) * split >= 80): + # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which + # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's + # (bx + grid / 2) weights into L2. + if bx < cutlass.Int32((n_out // CTA_M) * split // 2): + pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) + p_rank = pb % cutlass.Int32(split) + p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) + for i in cutlass.range_constexpr(my_tiles): + p_k = _k_tile(p_rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_prefetch( + tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] + ) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + if cutlass.const_expr(i >= ring): + while not cute.arch.mbarrier_try_wait( + mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 + ): + pass + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_a.subview(stage * W_TILE_BYTES), + tma_ptr_w, + (k_global * cutlass.Int32(CTA_K), m_offset), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfa.subview(stage * SF_ATOM_BYTES), + tma_ptr_wsf, + (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), + w_full.subview(stage), + l2_cache_hint=EVICT_FIRST, + ) + elif warp_id == 1: + prims.griddepcontrol(prims.GridDepAction.WAIT) + if cutlass.const_expr(trigger_early): + # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read + # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its + # wait, so the chain of releases never runs ahead of a producer still in flight. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if tx == 32: + _stamp(timers, timed, bx, 1) + if cutlass.const_expr(not x_bf16): + if prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) + for i in cutlass.range_constexpr(my_tiles): + k_global = _k_tile(rank, i, split, my_tiles, contig) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_b.subview(i * x_tile_bytes), + tma_ptr_x, + (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), + act_full, + ) + prims.cp_async_bulk_tensor_shared_cta_global( + smem_sfb.subview(i * SF_ATOM_BYTES), + tma_ptr_xsf, + (cutlass.Int32(0), k_global), + act_full, + ) + elif warp_id == 2: + tmem_raw = tmem_ptr_i32.load() + base_col = tmem_raw & 0xFFFF + base_row = tmem_raw >> 16 + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + ring * SF_COLS + idesc = prims.Tcgen05MxInstrDesc.build( + a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M + ) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + if cutlass.const_expr(a_tmem): + # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they + # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, + # its scale atoms and the MMAs, A read from TMEM. + base_row = tmem_raw >> 16 + base_col = tmem_raw & 0xFFFF + sfa_col0 = base_col + acc_cols(nacc, mma_n) + sfb_col0 = sfa_col0 + my_tiles * SF_COLS + a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) + s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + if prims.elect_sync(): + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_128X256B, + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), + 6, + cutlass.Int32, + ), + desc_a + ((MMA_INST_K >> 4) * kb), + group=prims.CTAGroup.CTA_1, + ) + prims.tcgen05_cp( + s2t_shape, + cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), + desc_sfa, multicast=s2t_multicast, + ) # fmt: skip + if cutlass.const_expr(i + ring < my_tiles): + prims.tcgen05_commit(mma_done.subview(stage)) + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + if prims.elect_sync(): + for i in cutlass.range_constexpr(my_tiles): + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + sfa_tmem = cutlass.inttoptr( + (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 + ) + sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) + prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) + for kb in cutlass.range_constexpr(K_BLOCKS): + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), + cutlass.inttoptr( + a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 + ), + desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), + i > 0 or kb > 0, + sfa_tmem, sfb_tmem, + ) # fmt: skip + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + else: + if cutlass.const_expr(x_bf16): + # Completed by a thread's arrive: a suspended try_wait wakes late on those. + while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): + pass + if tx == 64: + _stamp(timers, timed, bx, 2) + for i in cutlass.range_constexpr(my_tiles): + stage = i % ring + while not cute.arch.mbarrier_try_wait( + w_full.subview(stage).data_ptr(), (i // ring) % 2 + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) + sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) + desc_sfa = prims.Tcgen05SmemDesc.build( + smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_sfb = prims.Tcgen05SmemDesc.build( + smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, + base_offset=0, layout=0, + ) # fmt: skip + desc_a = prims.Tcgen05SmemDesc.build( + smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + desc_b = prims.Tcgen05SmemDesc.build( + smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, + base_offset=0, layout=2, + ) # fmt: skip + if prims.elect_sync(): + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, + multicast=s2t_multicast) # fmt: skip + prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, + multicast=s2t_multicast) # fmt: skip + for kb in cutlass.range_constexpr(K_BLOCKS): + inc = ((MMA_INST_K * 8 // 8) >> 4) * kb + j = i * K_BLOCKS + kb + prims.tcgen05_mma_block_scale( + prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, + cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), + desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, + cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), + ) # fmt: skip + prims.tcgen05_commit(mma_done.subview(stage)) + if tx == 64: + _stamp(timers, timed, bx, 3) + if prims.elect_sync(): + prims.tcgen05_commit(acc_done) + elif warp_id >= 4: + lane = tx % 32 + row = (warp_id % 4) * 32 + lane + if cutlass.const_expr(x_bf16): + # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as + # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), + # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, + # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. + qz_tid = tx - cutlass.Int32(128) + prims.griddepcontrol(prims.GridDepAction.WAIT) + din_half_k = k_in // 2 + din_hidden_half = cutlass.Int32(0) + if cutlass.const_expr(din): + # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to + # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 + # threads per token, 16-byte loads, all of a thread's loads in flight at once. + din_hidden_half = cutlass.select_( + rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), + cutlass.Int32(1), + cutlass.Int32(0), + ) + din_tok = qz_tid // cutlass.Int32(16) + din_part = qz_tid % cutlass.Int32(16) + din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) + din_id = din_i.load(idx=din_tok_ld) + din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) + din_hrow = din_tok_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) + din_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + din_hrow, + cutlass.select_(din_valid, din_id, cutlass.Int32(0)), + ) + # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. + din_copy = ( + ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) + & (bx == cutlass.Int32(split // 2)) + & (din_hidden_half == cutlass.Int32(1)) + & (din_tok < num_tokens) + ) + din_base = din_row * cutlass.Int32(din_half_k // 2) + din_ss = cutlass.Float32(0.0) + din_vecs_per_part = din_half_k // 8 // 16 + for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): + din_n = min(16, din_vecs_per_part - din_round) + din_loaded = [] + for din_j in cutlass.range_constexpr(din_n): + din_widx = din_base + ( + din_part + cutlass.Int32((din_round + din_j) * 16) + ) * cutlass.Int32(4) + din_loaded.append( + _din_load4(din_table, din_hidden, din_hidden_half, din_widx) + ) + if din_copy: + for din_j in cutlass.range_constexpr(din_n): + din_v = din_loaded[din_j] + din_res.store( + (din_v[0], din_v[1], din_v[2], din_v[3]), + idx=din_tok * cutlass.Int32(din_half_k // 2) + + (din_part + cutlass.Int32((din_round + din_j) * 16)) + * cutlass.Int32(4), + alignment=16, + ) + for din_j in cutlass.range_constexpr(din_n): + for din_e in cutlass.range_constexpr(4): + din_word = cutlass.Int32(din_loaded[din_j][din_e]) + din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + din_ss = din_ss + din_lo * din_lo + din_hi * din_hi + for din_o in cutlass.range_constexpr(4): + din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) + if din_part == cutlass.Int32(0): + din_eps = din_f.load(idx=din_hidden_half) + din_r = cute.math.rsqrt( + din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True + ) + din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) + din_rs.store( + cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok + ) + prims.barrier_cta_sync(3, thread_count=128) + qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, + alignment=16) # fmt: skip + qz_items = my_tiles * K_BLOCKS * mma_n + qz_rounds = (qz_items + 127) // 128 + # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of + # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. + QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 + qz_loaded = [] + for qz_r in cutlass.range_constexpr(qz_rounds): + qz_item = qz_tid + cutlass.Int32(qz_r * 128) + qz_t = qz_item % cutlass.Int32(mma_n) + qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) + qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) + if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): + qz_i = cutlass.select_( + qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) + ) + qz_k_global = rank + qz_i * cutlass.Int32(split) + if cutlass.const_expr(contig): + qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i + qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + qz_x_base = ( + qz_t_ld * cutlass.Int32(k_in if x_swiglu else k_in // 2) + + qz_k_global * cutlass.Int32(CTA_K // 2) + + qz_kb * cutlass.Int32(MMA_INST_K // 2) + ) + qz_wvecs = None + if cutlass.const_expr(din): + # The draft input's raw values (table row of the id / hidden row) and its norm weights. + qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) + qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) + qz_id = din_i.load(idx=qz_t_ld) + qz_hrow = qz_t_ld + if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): + qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) + qz_row = cutlass.select_( + din_hidden_half == cutlass.Int32(1), + qz_hrow, + cutlass.select_( + (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), + qz_id, + cutlass.Int32(0), + ), + ) + qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) + qz_vecs = [ + _din_load4( + din_table, + din_hidden, + din_hidden_half, + qz_src_w + cutlass.Int32(qz_v * 4), + ) + for qz_v in range(4) + ] + qz_wvecs = [ + din_w.load( + idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + else: + qz_vecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 + ) + for qz_v in range(4) + ] + if cutlass.const_expr(x_swiglu): + # The up half of the same columns, K values after the gate in the row. + qz_wvecs = [ + x_words.load( + idx=qz_x_base + cutlass.Int32(k_in // 2 + qz_v * 4), + vector_size=4, + alignment=16, + ) + for qz_v in range(4) + ] + qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) + if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): + for qz_g in cutlass.range_constexpr(len(qz_loaded)): + qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] + qz_vals = [] + qz_amax = cutlass.Float32(0.0) + if cutlass.const_expr(din): + qz_rs = din_rs.load( + idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) + ) + for qz_v in cutlass.range_constexpr(4): + for qz_e4 in cutlass.range_constexpr(4): + qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) + if cutlass.const_expr(x_swiglu): + qz_word = _swiglu_oai_bf16x2( + qz_word, cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + ) + qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) + qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) + if cutlass.const_expr(din): + # bf16(x * rs * w), the draft-input kernel's rounding. + qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) + qz_lo = _bf16_round( + qz_lo + * qz_rs + * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) + ) + qz_hi = _bf16_round( + qz_hi + * qz_rs + * (qz_wword & cutlass.Int32(-65536)).bitcast( + cutlass.Float32 + ) + ) + qz_vals.append(qz_lo) + qz_vals.append(qz_hi) + qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) + qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( + cutlass.Int32 + ) + qz_e = ( + (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) + ) + cutlass.select_( + (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + ) + qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) + qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( + cutlass.Float32 + ) + qz_scale = cutlass.select_( + qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) + ) + qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, + qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) + for qz_j in range(8)] # fmt: skip + if qz_item < cutlass.Int32(qz_items): + if qz_t < num_tokens: + qz_row_w = qz_i * cutlass.Int32( + x_tile_bytes // 4 + ) + qz_t * cutlass.Int32(CTA_K // 4) + qz_c0 = qz_kb * cutlass.Int32(2) + qz_b_words.store( + (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), + idx=qz_row_w + + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), + alignment=16, + ) + qz_b_words.store( + (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), + idx=qz_row_w + + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) + * cutlass.Int32(4), + alignment=16, + ) + smem_sfb.store( + cutlass.Uint8(qz_e), + idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) + + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) + + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) + + qz_kb, + ) + qz_loaded = [] + # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. + prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) + prims.barrier_cta_sync(2, thread_count=128) + if qz_tid == cutlass.Int32(0): + prims.mbarrier_arrive(act_full) + if cutlass.const_expr(st_async): + while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): + pass + else: + while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. + acc_all = [] + ld_col = 0 + for ld_w in (128, 64, 32, 16, 8): + if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): + piece = prims.tcgen05_ld( + "32x32b", + cutlass.inttoptr( + tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 + ), + num=ld_w, + ) + acc_all = acc_all + [piece[c] for c in range(ld_w)] + ld_col += ld_w + prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) + acc = [] + for t in cutlass.range_constexpr(mma_n): + v = cutlass.Float32(acc_all[t]) + for a in cutlass.range_constexpr(1, nacc): + v = v + cutlass.Float32(acc_all[a * mma_n + t]) + acc.append(v) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if tx == 128: + _stamp(timers, timed, bx, 4) + n = m_offset + row + if cutlass.const_expr(split == 1): + for t in cutlass.range_constexpr(mma_n): + if cutlass.Int32(t) < num_tokens: + y.store( + cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n + ) + if cutlass.const_expr(ar_push > 0): + prims.griddepcontrol(prims.GridDepAction.WAIT) + # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are + # issued before the first store so they overlap (one L2 round trip, not one per token). + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_v = cutlass.Float32(acc[t]) + ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) + ) + else: + owner = row // cutlass.Int32(rows_per_rank) + local = row % cutlass.Int32(rows_per_rank) + if cutlass.const_expr(rope): + # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. + pairs = rows_per_rank // 2 + owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) + local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( + row % cutlass.Int32(64) + ) % cutlass.Int32(pairs) + if owner != rank: + base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) + if cutlass.const_expr(st_async): + dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) + mb = _mapa_u32(reduce_ready.data_ptr(), owner) + words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] + for g in cutlass.range_constexpr(mma_n // 4): + _st_async_v4( + dst + cutlass.Int32(16 * g), + words[4 * g], + words[4 * g + 1], + words[4 * g + 2], + words[4 * g + 3], + mb, + ) + else: + for t in cutlass.range_constexpr(mma_n): + prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( + cutlass.Float32(acc[t]) + ) + prims.mbarrier_arrive( + prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER + ) + else: + if cutlass.const_expr(st_async): + while not _test_wait_cluster(reduce_ready.data_ptr(), 0): + pass + else: + while not prims.mbarrier_try_wait_parity( + reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER + ): + pass + if lane == 0: + if row // cutlass.Int32(rows_per_rank) == rank: + if (row % cutlass.Int32(rows_per_rank)) == 0: + _stamp(timers, timed, bx, 5) + totals = [] + for g in cutlass.range_constexpr(mma_n // 4): + sums = [cutlass.Float32(0.0)] * 4 + for q in cutlass.range_constexpr(split): + parts = _ld_shared_v4_f32( + mailbox.subview( + (cutlass.Int32(q * rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g) + ).data_ptr() + ) + sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), + parts[j]) for j in range(4)] # fmt: skip + for j in cutlass.range_constexpr(4): + totals.append(sums[j]) + if cutlass.const_expr(rope): + # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. + mailbox.store( + sums[j], + idx=(rank * cutlass.Int32(rows_per_rank) + local) + * cutlass.Int32(mb_stride) + + cutlass.Int32(4 * g + j), + ) + elif cutlass.Int32(4 * g + j) < num_tokens: + y.store( + cutlass.BFloat16(sums[j]), + idx=cutlass.Int32((4 * g + j) * n_out) + n, + ) + if cutlass.const_expr(ar_push > 0): + # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. + prims.griddepcontrol(prims.GridDepAction.WAIT) + ar_par = [ + _ld_volatile_s32( + ar_flags + + cutlass.Int64(4) + * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) + ) + & cutlass.Int32(1) + for t in range(mma_n) + ] + for t in cutlass.range_constexpr(mma_n): + ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) + if (lane % cutlass.Int32(2)) == cutlass.Int32(0): + if cutlass.Int32(t) < num_tokens: + ar_off = ( + cutlass.Int64(ar_par[t]) + * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) + + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) + * cutlass.Int64(n_out // 2) + + (cutlass.Int64(n) // cutlass.Int64(2)) + ) + _st_global_b32( + ar_ws + ar_off * cutlass.Int64(4), + _bf16x2_word(totals[t], ar_hi), + ) + prims.barrier_cta_sync(1, thread_count=128) + if cutlass.const_expr(rope): + # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over + # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to + # rope_q, K / V into the page of that position. + pairs = rows_per_rank // 2 + rp_e = tx - cutlass.Int32(128) + rp_p = rp_e % cutlass.Int32(pairs) + rp_t = rp_e // cutlass.Int32(pairs) + if rp_t < num_tokens: + rp_max_blocks = rope_dims.load(idx=0) + rp_tpb_log2 = rope_dims.load(idx=1) + rp_hq = rope_dims.load(idx=2) + rp_hkv = rope_dims.load(idx=3) + rp_T = rope_dims.load(idx=4) + rp_d = rank * cutlass.Int32(pairs) + rp_p + rp_a = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_b = _bf16_round( + mailbox.load( + idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) + * cutlass.Int32(mb_stride) + + rp_t + ) + ) + rp_req = rp_t // rp_T + rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) + rp_head = n_tile + if rp_head < rp_hq + rp_hkv: + rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) + rp_s = rope_cs.load( + idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) + + cutlass.Int32(1) + ) + rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) + rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) + rp_a = rp_ra + rp_b = rp_rb + rp_qs = rope_scale.load(idx=0) + rp_w = _e4m3x4( + rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) + ) + rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) + rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) + if rp_head < rp_hq: + rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d + rope_q.store(rp_lo, idx=rp_q) + rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) + else: + rp_kv = cutlass.select_( + rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) + ) + rp_kh = rp_head - rp_hq - rp_kv * rp_hkv + rp_page = rope_bt.load( + idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks + + (rp_tok >> rp_tpb_log2) + ) + # 64-bit byte offset: the pool view can exceed 2^31 bytes. + rp_dst = ( + (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) + << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) + ) + cutlass.Int64( + (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) + * cutlass.Int32(128) + + rp_d + ) + rp_addr = rope_pool.data_ptr(0).toint() + rp_dst + cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) + cutlass.inttoptr( + rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 + ).store(rp_hi) + if tx == 128: + _stamp(timers, timed, bx, 6) + if cutlass.const_expr(ar_push == 2): + # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. + prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) + if warp_id == 4: + prims.tcgen05_dealloc( + cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), + tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), + ) + + +def _bytes_map(t, cols, rows, box_rows, swizzle): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint8, + global_dims=[cols, rows], + global_strides=[cols // 16], + box_dims=[CTA_K, box_rows], + swizzle=swizzle, + ) + + +def _atom_map(t, atoms): + return cuda.create_tensor_map_tiled( + global_address=t.iterator.toint(), + dtype=cutlass.Uint16, + global_dims=[SF_ATOM_U16, atoms], + global_strides=[SF_ATOM_BYTES // 16], + box_dims=[SF_ATOM_U16, 1], + swizzle=cuda.TensorMapSwizzle.none, + ) + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[bool] = False, +) -> None: + # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. + assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( + "x_swiglu needs x_bf16 and excludes din" + ) + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, x_swiglu, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py new file mode 100644 index 000000000000..041dc12b1ffd --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py @@ -0,0 +1,521 @@ +# 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. +"""MXFP8 x MXFP8 decode GEMV (tcgen05, block-scaled) for already-quantized activations of at most 8 tokens. + +Compiled on the first call per (N, K, split, ring, early trigger, PDL, push, MMA width), which must happen outside +CUDA-graph capture, or ahead of any call by ``precompile``; once ``freeze`` closes compiles, a call whose kernel is not +compiled declines instead. +""" + +from __future__ import annotations + +import threading +from typing import Dict, Optional, Tuple + +import torch + +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +MAX_TOKENS = 48 +MMA_WIDTHS = ( + 8, + 16, + 32, + 48, +) # token columns of the compiled MMA; a call uses the smallest that holds M +_SMEM_LIMIT = 220 * 1024 + + +def _mma_n(m: int) -> int: + return next(w for w in MMA_WIDTHS if m <= w) + + +def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +_lock = threading.Lock() +_compiled: Dict[tuple, object] = {} + +# (N, K) -> (split, ring): GB300 exposed time behind a 4.25 us PDL predecessor at M 1 / 8, vs the fused low-M GEMM: +# QKV 2.78 / 3.16 vs 4.58 / 6.21 us, o_proj 2.23 / 2.41 vs 4.40 / 5.09, shared gate_up 2.43 / 2.74 vs 3.56 / 4.30, +# shared down 1.74 / 1.66 vs 2.92 / 3.38. A config is (split, ring[, contig, nacc, a_tmem, st_async]); the optional +# fields default to _CONTIG, _NACC, _A_TMEM, _ST_ASYNC. +_CONFIGS: Dict[Tuple[int, int], tuple] = { + (2560, 6144): (4, 12), + (6144, 2048): (2, 8), + # Shared-expert gate_up: split 2, ring 12, single accumulator, A operand from shared memory. + (1536, 6144): (2, 12, 1, 1, 0, 1), + # Shared-expert down: unsplit, ring 3, single accumulator, A operand from shared memory. + (6144, 768): (1, 3, 1, 1, 0, 1), + (2304, 6144): (4, 12), + (6144, 3072): (2, 12), + # Dense-layer gate_up: 6.4 / 6.3 us at M 1 / 8 vs FlashInfer 8.8. + (6144, 6144): (2, 12), +} +# Per-token-count configs (N, K, M): the grid that suits the concurrent routed-expert engine depends on how long +# that engine runs. At 2 tokens the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_BY_M: Dict[Tuple[int, int, int], tuple] = { + (1536, 6144, 2): (4, 12, 1, 1, 0, 1), +} +# Above 8 tokens (spec-decode verify rows): the shared-expert gate_up splits K over 4 CTAs. +_CONFIGS_WIDE: Dict[Tuple[int, int], tuple] = { + (1536, 6144): (4, 12, 1, 1, 0, 1), +} + +# Defaults of the optional config fields: contiguous k-tile ranges per split rank, one TMEM accumulator, the A +# operand staged in TMEM (when its columns fit), st.async split-K partial exchange. +_CONTIG = True +_NACC = 1 +_A_TMEM = True +_ST_ASYNC = True +_ensured = set() + + +_DUMMY: Dict[torch.device, torch.Tensor] = {} + + +def _dummy_timers(device) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = _DUMMY.get(device) + if t is None: + t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def _arg(t: torch.Tensor, align: int = 16): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: + cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and m > 8: + cfg = _CONFIGS_WIDE.get((n, k)) + cfg = cfg or _CONFIGS.get((n, k)) + # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. + if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: + return None + return cfg + + +def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + return True + tiles = k // 128 // split + mma_n = _mma_n(m) + # GB300 exposed time vs the FlashInfer prequantized GEMM (behind a PDL predecessor), at 24 / 48 tokens: QKV 3.76 / + # 4.63 vs 7.46 / 11.44 us, o_proj 2.74 / 3.33 vs 4.46 / 5.84, shared gate_up 4.57 / 5.71 vs 7.16 / 11.25. Behind + # only for K = 768 (unsplit shared down) at 17-32 tokens: 3.1 vs 2.9. + # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; + # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. + if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + + +def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: + """Whether the GEMV with its default config takes an [M, K] activation (bf16 or MXFP8) against an [N, K] weight.""" + cfg = get_config(n, k, m) + return cfg is not None and not _declines(n, k, m, x_bf16, cfg) + + +def _resolve( + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], + x_swiglu: bool = False, +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or get_config(n, k, m) + if cfg is None or _declines(n, k, m, x_bf16, cfg): + return None + split, ring = cfg[0], cfg[1] + contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else _NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC + if nacc != 1: + # More than one TMEM accumulator is broken on SM107 (783556 / 783557 / 783579: nondeterministic or garbage + # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = _mma_n(m) + while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: + ring -= 1 + if a_tmem: + from .m3_mx_gemv_kernel import tmem_cols + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into + # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 + if ar is not None and split <= 4 and n % 2 == 0: + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = get_env_enable_pdl() + # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) + return key, (ar_ws, ar_flags, ar_rank) + + +def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from . import m3_mx_gemv_kernel as kernel + + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with _lock: + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = cute.compile( + kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + x_swiglu, + ) + return fn + + +def _views( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight, + weight_scale, + y, + x_bf16: bool, + x_swiglu: bool = False, +) -> tuple: + # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. + k = x.shape[1] // 2 if x_swiglu else x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = _placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = _placeholder_words(x.device) + return ( + _arg(weight.view(torch.uint8)), + _arg(weight_scale.view(torch.uint8).reshape(-1)), + _arg(x_bytes), + _arg(xsf.view(torch.uint8).reshape(-1)), + _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), + _arg(x_words.reshape(-1)), + _arg(_dummy_timers(x.device)), + ) + + +def _run( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, + x_swiglu: bool = False, +) -> Optional[torch.Tensor]: + global last_pushed, last_uncompiled + last_pushed = False + last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + if x_swiglu: + k //= 2 + n = weight.shape[0] + if weight.shape[1] != k: + return None + resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu) + if resolved is None: + return None + key, runtime_ar = resolved + fn = _compiled.get(key) + if fn is None: + if _frozen[0]: + # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller + # reports the miss. + last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + extra = _extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = _compile(key, args, m, runtime_ar, stream, extra) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) + last_pushed = key[11] > 0 + return y + + +def precompile( + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, + x_swiglu: bool = False, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu) + if resolved is None: + return None + key, runtime_ar = resolved + if key in _compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros( + m, + 2 * k if x_swiglu else k, + dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, + device=dev, + ) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) + return key + + +def freeze(frozen: bool = True) -> None: + """Close (or reopen) compiles: while frozen, a call whose kernel is not compiled yet declines (returns None and + sets ``last_uncompiled``) instead of compiling. ``precompile`` and ``ensure_compiled`` still compile.""" + _frozen[0] = bool(frozen) + + +def frozen() -> bool: + return _frozen[0] + + +def is_compiled(key: tuple) -> bool: + return key in _compiled + + +def _extra_args(device) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = _placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) + r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) + return d + r + + +def m3_mx_gemv( + xq: torch.Tensor, + xsf: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``xq`` E4M3 [M <= 8, K] with ``xsf`` its 128x4-swizzled UE8M0 scales; ``weight`` E4M3 [N, K] with + ``weight_scale`` swizzled. Returns bf16 [M, N], or None when the shape is not served.""" + return _run(xq, xsf, weight, weight_scale, False, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_bf16( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """The same GEMV for a bf16 ``x`` [M <= 8, K] (contiguous), quantized to MXFP8 inside the kernel exactly as + torch.ops.trtllm.mxfp8_quantize(x, True) does.""" + if x.dtype != torch.bfloat16 or x.dim() != 2 or not x.is_contiguous() or x.data_ptr() % 16 != 0: + return None + return _run(x, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar) + + +def m3_mx_gemv_swiglu( + gate_up: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + trigger_early: bool = True, + cfg: Optional[tuple] = None, + use_pdl: Optional[bool] = None, + ar: Optional[tuple] = None, +) -> Optional[torch.Tensor]: + """``m3_mx_gemv_bf16(swiglu(gate_up), ...)`` in one launch: ``gate_up`` bf16 [M <= 8, 2K] (contiguous, gate then + up) is activated in the kernel's quantize prologue with the swigluoai arithmetic of the Triton silu_and_mul kernel + (alpha 1.702, limit 7, the (up + 1) offset; bit for bit), then quantized as the bf16 path does.""" + if ( + gate_up.dtype != torch.bfloat16 + or gate_up.dim() != 2 + or not gate_up.is_contiguous() + or gate_up.data_ptr() % 16 != 0 + or gate_up.shape[1] % 2 != 0 + ): + return None + return _run(gate_up, None, weight, weight_scale, True, trigger_early, cfg, use_pdl, ar, True) + + +_PLACEHOLDERS: Dict[tuple, torch.Tensor] = {} +# Whether the last call pushed its output into an all-reduce (see _run's ``ar``). +last_pushed = False +# The key of the last call declined because compiles are frozen and its kernel was not compiled (see ``freeze``). +last_uncompiled: Optional[tuple] = None +_frozen = [False] + + +def _placeholder_scales(k: int, device) -> torch.Tensor: + key = ("sf", k, device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def _placeholder_words(device) -> torch.Tensor: + key = ("w", device) + t = _PLACEHOLDERS.get(key) + if t is None: + t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in _ensured: + return + _ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in MMA_WIDTHS[1:]: + m3_mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + m3_mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/late_compile.py b/tensorrt_llm/_torch/late_compile.py new file mode 100644 index 000000000000..97eb4bf210ae --- /dev/null +++ b/tensorrt_llm/_torch/late_compile.py @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Late-compile detector: report kernel compiles that happen after the executor finished its warmup. + +A compile inside a served request costs that request its whole compile time, and in-memory compiles (CuTe DSL +handles, Triton specializations) leave no file in the JIT cache directories, so cache-file checks cannot see them. +After ``mark_server_ready()`` every ``cutlass.cute.compile`` call and every Triton JIT compile logs one warning line +``[late-compile] op= key=``. The executor arms it only for a model whose RuntimeHooks set +late_compile_detect. Kill switch: TRTLLM_LATE_COMPILE_DETECT=0. + +Not covered here: the TRTLLM-Gen FMHA runtime NVRTC compile (C++, an in-memory module cache warmed for generation +kernels only), whose warmup and persistence fix is tracked with that kernel library, and the fmha_sm100 package's +on-disk AOT cache, whose post-warmup "[aot_cache] Saved" lines are its own signal. +""" + +import os +from typing import Any + +from tensorrt_llm.logger import logger + +_ENABLED = os.environ.get("TRTLLM_LATE_COMPILE_DETECT", "1") != "0" +_ready = [False] + + +def note(op: str, key: Any) -> None: + """Log a compile of ``op`` for ``key`` when it happens after warmup.""" + if _ready[0]: + logger.warning(f"[late-compile] op={op} key={key}") + + +class _CuteCompile: + """Stands in for ``cutlass.cute.compile``: notes each call, then forwards it (``cute.compile[...]`` included).""" + + def __init__(self, inner: Any) -> None: + self._inner = inner + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + target = args[0] if args else None + note("cute.compile", getattr(target, "__qualname__", type(target).__name__)) + return self._inner(*args, **kwargs) + + def __getitem__(self, item: Any) -> "_CuteCompile": + return _CuteCompile(self._inner[item]) + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + +def _hook_cute() -> None: + try: + import cutlass.cute as cute + except ImportError: + return + if not isinstance(cute.compile, _CuteCompile): + cute.compile = _CuteCompile(cute.compile) + + +def _hook_triton() -> None: + try: + import triton + except ImportError: + return + runtime = getattr(getattr(triton, "knobs", None), "runtime", None) + if runtime is None or not hasattr(runtime, "jit_post_compile_hook"): + logger.info( + "[late-compile] Triton post-compile hook unavailable; Triton compiles are not reported" + ) + return + previous = runtime.jit_post_compile_hook + + def hook(*args: Any, **kwargs: Any) -> Any: + try: + fn = kwargs.get("fn") + name = getattr(getattr(fn, "fn", fn), "__qualname__", "?") + note("triton", f"{name} {kwargs.get('key', '')}") + except Exception: # noqa: BLE001 (the detector must never break a compile) + pass + return previous(*args, **kwargs) if previous is not None else None + + runtime.jit_post_compile_hook = hook + + +def mark_server_ready() -> None: + """Arm the detector once the executor's warmup is done (idempotent).""" + if not _ENABLED or _ready[0]: + return + _hook_cute() + _hook_triton() + _ready[0] = True + logger.info("[late-compile] detector armed after warmup") diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 8dc42e21e777..9ce5767c4b45 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -4642,9 +4642,16 @@ def instantiate_sampler( if engine.spec_config is not None and engine.spec_config.spec_dec_mode.has_spec_decoder( ): sampler = get_spec_decoder(sampler_args, engine.spec_config) - spec_step_fold = runtime_hooks_of(engine.model).spec_step_fold - if spec_step_fold is not None and isinstance(sampler, SpecSampler): - sampler.bind_spec_step_fold(spec_step_fold) + runtime_hooks = runtime_hooks_of(engine.model) + if isinstance(sampler, SpecSampler): + if runtime_hooks.spec_step_fold is not None: + sampler.bind_spec_step_fold(runtime_hooks.spec_step_fold) + sampler.d2h_side_stream_copies = runtime_hooks.sampler_d2h_side_stream + sampler.static_slot_ids = runtime_hooks.spec_small_h2d + if runtime_hooks.sampler_d2h_side_stream: + logger.info( + "[M3 host] spec sampler store copies on the D2H side stream" + ) return sampler if mm_encoder_only: diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 1f1893025ca2..43bffb485255 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -284,6 +284,8 @@ def __init__(self, config: CUDAGraphRunnerConfig): self._pre_replay_position_ids: Optional[torch.Tensor] = None self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} + # Graphs captured with a step-input pack's prologue, and that pack (step_inputs.capture_prologue). + self._step_input_prologue_keys: Dict[KeyType, Any] = {} self.graph_outputs: Dict[KeyType, Callable[[], Optional[torch.Tensor]]] = {} self.graph_metadata: Dict[KeyType, Dict[str, Any]] = {} @@ -911,6 +913,11 @@ def run_warmup(): output = None graph_pool = self.memory_pool or torch.cuda.graph_pool_handle() with nccl_window_graph_capture(graph, graph_pool): + # A pre-replay hook reads the static inputs before the replay, so their copies must land before it. + prologue_pack = (None if self._pre_replay_hooks else + step_inputs.capture_prologue()) + if prologue_pack is not None: + self._step_input_prologue_keys[key] = prologue_pack output = forward_fn(capture_inputs) # The addresses the graph baked in, for the strict buffer check. self._record_strict_buffer_snapshot(key) @@ -1039,8 +1046,14 @@ def replay( static_encoder_hidden_states[actual_num_encoder_tokens:].zero_() # Issue the step's packed input copies (no-op without an open pack) before - # anything below reads the static inputs on the device. - step_inputs.flush() + # anything below reads the static inputs on the device; a graph captured + # with the pack's prologue issues them itself at the start of the replay. + prologue_pack = self._step_input_prologue_keys.get(key) + step_inputs.flush(in_graph=prologue_pack is not None) + if prologue_pack is not None and not step_inputs.is_active( + prologue_pack): + # Replayed outside the pack's context: nothing filled its window. + prologue_pack.quiesce_graph_window() if self._pre_replay_hooks: hook_tensors = static_tensors if position_id_offsets is not None: @@ -1579,6 +1592,7 @@ def clear(self, *, release_nccl_window_owners: bool = True): self.graph_outputs.clear() self.graph_metadata.clear() self._pre_replay_position_ids = None + self._step_input_prologue_keys.clear() self.reset_padding_dummies() del self.memory_pool self.memory_pool = None diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index a9fc0625740e..eb180f9b11ee 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -2899,6 +2899,20 @@ def _get_remote_tail_decode_rank_tokens( ) return gathered[:, 0].tolist() + def _kv_lens_corrected_in_preprocess( + self, inputs: dict[str, Any], enable_spec_decode: bool + ) -> bool: + """Whether preprocess_inputs applies the overlap scheduler's KV-length offsets to kv_lens_cuda and then calls + on_update_kv_lens (the branch below).""" + attn_meta = inputs.get("attn_metadata") + return ( + enable_spec_decode + and not self._config.disable_overlap_scheduler + and attn_meta is not None + and attn_meta.kv_cache_manager is not None + and hasattr(attn_meta, "kv_lens_cuda") + ) + def preprocess_inputs( self, inputs: dict[str, Any], *, enable_spec_decode: bool, runtime_draft_len: int ) -> dict[str, Any]: @@ -2916,8 +2930,14 @@ def preprocess_inputs( if self.guided_decoder is not None: self.guided_decoder.token_event.record() return inputs - # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. - if attn_meta is not None: + # Invalidate per-forward-pass caches so they are recomputed (and captured) on every _forward_step. With + # RuntimeHooks.kv_lens_update_once, a step that corrects the KV lengths below (and then calls it again) calls + # it only there: the second call rewrites everything the first one writes, and nothing reads the metadata + # in between. + if attn_meta is not None and not ( + self.runtime_hooks.kv_lens_update_once + and self._kv_lens_corrected_in_preprocess(inputs, enable_spec_decode) + ): attn_meta.on_update_kv_lens() if enable_spec_decode and not self._config.disable_overlap_scheduler: @@ -2934,7 +2954,25 @@ def preprocess_inputs( num_ctx_tokens = inputs["attn_metadata"].num_ctx_tokens num_chunked_ctx_requests = inputs["attn_metadata"].num_chunked_ctx_requests previous_batch_tokens = inputs["input_ids"].shape[0] - num_ctx_tokens - if inputs["position_ids"].ndim == 3: # mrope: [3, 1, N] + # RuntimeHooks.spec_offsets_op: both corrections of a generation-row step in one launch. + fused_offsets = ( + self.runtime_hooks.spec_offsets_op is not None + and inputs["position_ids"].ndim != 3 + and hasattr(inputs["attn_metadata"], "kv_lens_cuda") + and not ( + num_ctx_requests >= num_chunked_ctx_requests + and num_chunked_ctx_requests > 0 + ) + and self.runtime_hooks.spec_offsets_op( + inputs["position_ids"][0, num_ctx_tokens:], + self._buffers.previous_pos_id_offsets_cuda[:previous_batch_tokens], + inputs["attn_metadata"].kv_lens_cuda[num_ctx_requests:num_seqs], + self._buffers.previous_kv_lens_offsets_cuda[:num_gen_requests], + ) + ) + if fused_offsets: + pass + elif inputs["position_ids"].ndim == 3: # mrope: [3, 1, N] inputs["position_ids"][:, :, num_ctx_tokens:] += ( self._buffers.previous_pos_id_offsets_cuda[:previous_batch_tokens] ) @@ -2944,7 +2982,9 @@ def preprocess_inputs( ) if hasattr(inputs["attn_metadata"], "kv_lens_cuda"): - if ( + if fused_offsets: + pass + elif ( num_ctx_requests >= num_chunked_ctx_requests and num_chunked_ctx_requests > 0 ): diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py index bbe234b984ff..ef68d85d0767 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/metadata.py @@ -143,6 +143,7 @@ def _arm_msa_policy(self, kv_cache_manager: KVCacheManager | KVCacheManagerV2 | kv_cache_manager.configure_decode_shortcuts( skip_unchanged_block_offsets=msa.skip_unchanged_block_offsets, reuse_pool_views=msa.reuse_pool_views, + block_ids_from_views=msa.block_ids_from_views, ) def setup_speculative( @@ -154,6 +155,13 @@ def setup_speculative( self._config, self.model.config, spec_resource_manager, self.mapping ) self.spec_metadata.skip_greedy_rng_upload = self.runtime_hooks.skip_greedy_rng_upload + self.spec_metadata.pack_rng_uploads = self.runtime_hooks.pack_spec_rng_uploads + if self.runtime_hooks.pack_spec_rng_uploads: + logger.info("[M3 host] sampled-step Philox uploads staged into the step-input pack") + if self.runtime_hooks.spec_small_h2d: + logger.info( + "[M3 host] spec step small uploads off the boundary (accepted counts packed, static slot ids)" + ) return self.spec_metadata @property diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index e9f9fb9f8be0..8df0fbab7b45 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -201,12 +201,41 @@ def __init__( self._config = config self.runtime_hooks = runtime_hooks_of(model) - self._h2d_cache = H2DSkipCache() if self.runtime_hooks.skip_unchanged_h2d else None + self._h2d_cache = ( + H2DSkipCache(changed_span=self.runtime_hooks.h2d_changed_span) + if self.runtime_hooks.skip_unchanged_h2d + else None + ) self._step_pack = ( - StepInputPack(torch.device("cuda", torch.cuda.current_device())) + StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=self.runtime_hooks.pack_parallel_scatter, + in_graph=self.runtime_hooks.pack_in_graph, + early_copy=self.runtime_hooks.pack_early_copy, + step_epoch=self.runtime_hooks.pack_step_epoch, + prologue_op=self.runtime_hooks.pack_prologue_op, + ) if self.runtime_hooks.pack_step_inputs else None ) + if self._step_pack is not None and self._step_pack.parallel_scatter: + logger.info("[M3 host] step-input pack: parallel scatter for non-overlapping copies") + if self._step_pack is not None and self._step_pack.in_graph: + logger.info("[M3 host] step-input pack issued by the decode graph's prologue") + if self._step_pack is not None and self._step_pack.early_copy: + logger.info( + "[M3 host] step-input pack: window uploaded ahead of the replay on a side copy stream" + ) + if self._step_pack is not None and self._step_pack.epoch is not None: + logger.info( + "[M3 host] step epoch: one device counter bumped by the decode graph's prologue" + ) + if self._step_pack is not None and self._step_pack.prologue_op is not None: + logger.info( + "[M3 host] step-input pack: the decode graph's prologue is one launch (scatter, header, epoch)" + ) + if self._h2d_cache is not None and self._h2d_cache.changed_span: + logger.info("[M3 host] changed metadata uploads stage only each row's changed span") self.forward_pass_callable = None # Optional in-graph sampling hook, registered by PyExecutor. Called at @@ -499,7 +528,11 @@ def _short_long_seq_len_threshold(self) -> int | None: Set by a sparse attention config that asks for the split, or by a model whose MSA policy stages identity selections for short sequences: a sequence within the - top-k blocks attends every valid block. + top-k blocks attends every valid block. The short family runs the identity selection, + so every token it attends must fit in the top-k blocks. The graph key estimates a + sequence's length from host token counts, which under the overlap scheduler miss up to + the in-flight step's accepted tokens; the identity threshold leaves one generation + step's tokens of headroom for that. """ sparse_config = self._sparse_attention_config if ( @@ -513,7 +546,8 @@ def _short_long_seq_len_threshold(self) -> int | None: "MsaRuntimePolicy.short_sequence_identity requires a MiniMax-M3 sparse " f"attention config, got {type(sparse_config).__name__}" ) - return sparse_config.sparse_topk_blocks * sparse_config.sparse_block_size + topk_span = sparse_config.sparse_topk_blocks * sparse_config.sparse_block_size + return topk_span - self._get_runtime_tokens_per_gen_step(self._config.max_draft_len) return None def _initialize_breakable_cuda_graph_runner(self) -> BreakableCUDAGraphRunner | None: diff --git a/tensorrt_llm/_torch/pyexecutor/h2d_cache.py b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py index 2a7a4d715447..bcb9eb8ba117 100644 --- a/tensorrt_llm/_torch/pyexecutor/h2d_cache.py +++ b/tensorrt_llm/_torch/pyexecutor/h2d_cache.py @@ -14,10 +14,11 @@ cache travels in the context rather than as an argument. Route a copy through it only for a destination that no other code writes (no device kernel, no other copy path, no other -engine): the shadow describes the device contents only under that invariant. Pooled metadata buffers satisfy it because -an engine that arms the skip gives its CUDA graph metadata a buffer pool of its own -(``RuntimeHooks.skips_unchanged_metadata_uploads``); eager metadata allocates its own buffers. A write through a -different view of the same buffer replaces the shadow, so overlapping views never skip on stale state. +engine, no graph replay): the shadow describes the device contents only under that invariant. CUDA graph metadata +buffers do not satisfy it: every graph key's metadata takes them from one pool, which hands the same block to each key +that asks for the same name, and a captured graph's device writes into a block run again at each of its replays. Those +destinations use ``copy_h2d_pool_shared`` / ``fill_tail_pool_shared``, which never skip. A write through a different +view of the same buffer replaces the shadow, so overlapping views never skip on stale state. """ from __future__ import annotations @@ -34,9 +35,19 @@ class H2DSkipCache: - """One engine's host shadows of its skip-routed destinations (see the module docstring).""" + """One engine's host shadows of its skip-routed destinations (see the module docstring). - def __init__(self) -> None: + ``changed_span``: a packed copy into a 2-D view with contiguous rows whose shadow covers the same rows (same + first element and row stride) uploads, per row, only the span from its first to its last changed column, plus any + columns the shadow does not cover. A page table that gained one page then stages a few bytes instead of every + padded row, so the pack still fits the decode graph's prologue window. + """ + + # A changed-span upload is used only while it moves at most this fraction of the view's bytes. + CHANGED_SPAN_MAX_FRACTION = 0.5 + + def __init__(self, changed_span: bool = False) -> None: + self.changed_span = changed_span # Device storage address -> (the storage, kept alive so the address is not reused; the view key; the host # values). Keyed by storage rather than by view base: tensors made under inference mode do not record their # base. @@ -61,17 +72,63 @@ def copy(self, dst: torch.Tensor, src: torch.Tensor, pack: bool) -> bool: entry = self._shadow.get(addr) if entry is not None and entry[1] == key and entry[2].matches(src, exact=True): return False + old = entry[2].host_values if entry is not None else None buffer = entry[2] if entry is not None and entry[1] == key else DeviceCopyBuffer(dst) buffer.invalidate() self._shadow.pop(addr, None) - if pack: - step_inputs.copy_(dst, src) - else: - dst.copy_(src, non_blocking=True) + if not ( + pack + and self.changed_span + and old is not None + and self._copy_changed_span(dst, src, (entry[1], old)) + ): + if pack: + step_inputs.copy_(dst, src) + else: + dst.copy_(src, non_blocking=True) buffer.record(src) self._shadow[addr] = (storage, key, buffer) return True + def _copy_changed_span(self, dst: torch.Tensor, src: torch.Tensor, shadow: tuple) -> bool: + """Stage only the changed span of each row of ``dst`` (see the class docstring); False if not applicable.""" + (offset, shape, stride, dtype), old = shadow + if ( + dst.dim() != 2 + or dst.stride(1) != 1 + or offset != dst.storage_offset() + or len(shape) != 2 + or shape[0] != dst.shape[0] + or stride != tuple(dst.stride()) + or dtype != dst.dtype + or old.dtype != src.dtype + or tuple(src.shape) != tuple(dst.shape) + or not step_inputs.is_open() + ): + return False + cols = dst.shape[1] + known = min(shape[1], cols) + changed = src[:, :known] != old[:, :known] + if cols > known: + # Columns the shadow does not cover hold unknown device values. + changed = torch.cat([changed, changed.new_ones(dst.shape[0], cols - known)], dim=1) + rows = changed.any(dim=1).nonzero().flatten().tolist() + spans = [] + for row in rows: + columns = changed[row].nonzero().flatten() + spans.append((row, int(columns[0]), int(columns[-1]) + 1)) + if sum(hi - lo for _, lo, hi in spans) > self.CHANGED_SPAN_MAX_FRACTION * dst.numel(): + return False + for row, lo, hi in spans: + step_inputs.copy_(dst[row, lo:hi], src[row, lo:hi]) + return True + + def forget(self, dst: torch.Tensor) -> None: + """Drop the shadow and tail record of ``dst``'s storage (it is about to be written outside the cache).""" + addr = dst.untyped_storage().data_ptr() + self._shadow.pop(addr, None) + self._tails.pop(addr, None) + def fill_tail(self, buf: torch.Tensor, start: int, value: int) -> None: storage = buf.untyped_storage() addr = storage.data_ptr() @@ -93,6 +150,10 @@ def fill_tail(self, buf: torch.Tensor, start: int, value: int) -> None: _active: ContextVar[Optional[H2DSkipCache]] = ContextVar("trtllm_h2d_skip_cache", default=None) +# Kill switch for A/B only: True routes the pool-shared metadata destinations (``copy_h2d_pool_shared``, +# ``fill_tail_pool_shared``) through the skip cache again, as before. False keeps them always written. +SKIP_POOL_SHARED_UPLOADS = False + @contextlib.contextmanager def active(cache: Optional[H2DSkipCache]) -> Iterator[None]: @@ -135,3 +196,40 @@ def fill_tail_if_changed(buf: torch.Tensor, start: int, value: int) -> None: buf[start:].fill_(value) return cache.fill_tail(buf, start, value) + + +def copy_h2d_pool_shared(dst: torch.Tensor, src: torch.Tensor, pack: bool = False) -> bool: + """``copy_h2d_if_changed`` for a destination in a CUDA graph metadata buffer pool: always issued. + + Every graph key's metadata takes its buffers from one pool, which hands the same block to each key that asks + for the same buffer name, and the captured graphs write some of those blocks on the device at every replay. The + host shadow therefore does not describe the device contents, so such a copy is never skipped. ``pack``: as for + ``copy_h2d_if_changed``, the copy may join the open ``step_inputs`` pack when a cache is active. + """ + cache = _active.get() + if cache is None: + dst.copy_(src, non_blocking=True) + return True + if SKIP_POOL_SHARED_UPLOADS: + return cache.copy(dst, src, pack) + cache.forget(dst) + if pack: + step_inputs.copy_(dst, src) + else: + dst.copy_(src, non_blocking=True) + return True + + +def fill_tail_pool_shared(buf: torch.Tensor, start: int, value: int) -> None: + """``buf[start:].fill_(value)`` for a buffer in a CUDA graph metadata buffer pool: always issued. + + See ``copy_h2d_pool_shared``: another graph key's replay can write ``buf`` past ``start``. + """ + cache = _active.get() + if cache is not None and SKIP_POOL_SHARED_UPLOADS: + cache.fill_tail(buf, start, value) + return + if cache is not None: + cache.forget(buf) + if start < buf.shape[0]: + buf[start:].fill_(value) diff --git a/tensorrt_llm/_torch/pyexecutor/profiling.py b/tensorrt_llm/_torch/pyexecutor/profiling.py index 13fc26c67ca2..b816276ddf5d 100644 --- a/tensorrt_llm/_torch/pyexecutor/profiling.py +++ b/tensorrt_llm/_torch/pyexecutor/profiling.py @@ -55,6 +55,8 @@ from tensorrt_llm.tools.layer_wise_benchmarks import get_calibrator from tensorrt_llm.tools.profiler.host_profile_tools.host_profiler import get_global_profiler +from .runtime_hooks import runtime_hooks_of + # Environment-variable names used by both the env-var-driven legacy # path and the runtime HTTP-endpoint path. They are duplicated as # module-level constants here so callers that only need the manager @@ -612,6 +614,18 @@ def profile_step(self): start_event_2 = None end_event_2 = torch.cuda.Event(enable_timing=True) prev_device_step_time = None + # With RuntimeHooks.iteration_timing_one_event, boundary_events[k % 3] marks the boundary before + # iteration k, recorded on timing_stream, and replaces the start/end ping-pong above. + one_timing_event = runtime_hooks_of( + getattr(executor.model_engine, "model", None) + ).iteration_timing_one_event + boundary_events = ( + [torch.cuda.Event(enable_timing=True) for _ in range(3)] if one_timing_event else None + ) + boundary_count = 0 + timing_stream = torch.cuda.Stream() if one_timing_event else None + if one_timing_event: + logger.info("[M3 host] iteration timing: one side-stream event per iteration boundary") # Cumulative fetch counter as of the previous iter log line; the # per-iteration ``num_fetched`` on the log line is the delta. prev_fetch_requests_cur_rank = 0 @@ -670,6 +684,7 @@ def profile_step_fn(): nonlocal it, enabled, start_time nonlocal start_event_1, end_event_1, start_event_2, end_event_2 nonlocal prev_device_step_time, prev_fetch_requests_cur_rank + nonlocal boundary_count nonlocal kv_util_str, kv_util_iter nonlocal torch_profiler, active_torch_trace_path nonlocal active_enable_torch_trace @@ -726,9 +741,21 @@ def profile_step_fn(): # ``self`` lets the /metrics serializer pick up the # values without going through the log line. should_capture_timing = executor.print_log or executor.enable_iter_perf_stats + if should_capture_timing and one_timing_event: + # Iteration k - 2 spans the two previous boundaries: the step the ping-pong reports here. + timing_stream.wait_stream(torch.cuda.current_stream()) + boundary_events[boundary_count % 3].record(timing_stream) + if boundary_count >= 2: + boundary_events[(boundary_count - 1) % 3].synchronize() + prev_device_step_time = boundary_events[(boundary_count - 2) % 3].elapsed_time( + boundary_events[(boundary_count - 1) % 3] + ) + boundary_count += 1 if should_capture_timing and start_time is not None: end_time = time.time() - if it % 2 == 0: + if one_timing_event: + pass + elif it % 2 == 0: end_event_1.record() if start_event_2 is not None: end_event_2.synchronize() @@ -858,7 +885,7 @@ def profile_step_fn(): calibrator.pre_step(it) start_time = time.time() - if should_capture_timing: + if should_capture_timing and not one_timing_event: if it % 2 == 0: if start_event_1 is None: start_event_1 = torch.cuda.Event(enable_timing=True) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ae1c06313552..eecbc30ec8bd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -121,7 +121,7 @@ from .resource_manager import (NoFreeSlotsError, ResourceManager, ResourceManagerType, is_context_only_kv_cache, request_context) -from .runtime_hooks import ExecutorEvents, runtime_hooks_of +from .runtime_hooks import ExecutorEvents, RuntimeHooks, runtime_hooks_of from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors) from .scheduler import (RequestScheduler, ScheduledRequests, @@ -424,6 +424,21 @@ class PendingEncoderStep: result: Optional[EncoderStepResult] = None +def _arm_runtime_opt_ins(hooks: RuntimeHooks) -> bool: + """Arm the process-wide features ``hooks`` opt into; return whether to arm + the late-compile detector after warmup. + + Both are off for a model that declares neither: the SM107 MXFP8 GEMM + tactic table (the GEMM runner does not know the model, so it is armed + here) and the late-compile detector. + """ + if hooks.mxfp8_pinned_tactics: + # Imported on use: only an opted-in model needs the table. + from tensorrt_llm._torch.custom_ops import mxfp8_cute_pinned_tactics + mxfp8_cute_pinned_tactics.arm() + return hooks.late_compile_detect + + def _request_flag(request, name: str) -> bool: """Read a request state predicate that is a nanobind read-only property on the C++ ``LlmRequest`` (``is_generation_only_request`` etc.) but a plain @@ -779,6 +794,20 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None: self._pending_num_new_active_requests = 0 self._pending_num_completed_requests = 0 self.stream_interval = self.llm_args.stream_interval + # Streaming: also emit on decoding iterations 2..N when the model's + # runtime hooks ask for it (RuntimeHooks.first_content_emit_iters; 1, + # the default, keeps the stock schedule). + model_hooks = runtime_hooks_of(model_engine.model) + self.first_content_emit_iters = max( + 1, model_hooks.first_content_emit_iters) + if self.stream_interval > 1 and self.first_content_emit_iters > 1: + logger.info( + f"first-content emission engaged: emit on decoding iters " + f"<= {self.first_content_emit_iters} (stream_interval " + f"{self.stream_interval}), rank {self.dist.rank}") + self._late_compile_detect = False + if model_hooks.mxfp8_pinned_tactics or model_hooks.late_compile_detect: + self._late_compile_detect = _arm_runtime_opt_ins(model_hooks) self.perf_manager = PerfMetricsManager( enabled=getattr(self.llm_args, 'return_perf_metrics', False)) self.attention_dp_enable_balance = ( @@ -1136,6 +1165,10 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None: self.encoder_launch_executor.submit( self._warmup_encoder_cuda_graphs_enc_dec).result() + if self._late_compile_detect: + from tensorrt_llm._torch.late_compile import mark_server_ready + mark_server_ready() + # Snapshot some cumulative KV cache counters so that stats reported to # users exclude blocks reused and missed during warmup dummy requests. if hasattr(self.kv_cache_manager, 'snapshot_warmup_baseline'): @@ -10352,6 +10385,28 @@ def _emit_first_token_responses(self, prev_scheduled_requests): self._enqueue_responses(new_responses) + def _should_emit_response(self, + request: LlmRequest, + emit_first_iter: bool = True, + first_gen_emit: bool = False) -> bool: + """Whether ``_handle_responses`` emits a response for ``request`` now. + + Emits on decoding iteration 1, on finish, every ``stream_interval`` + iterations, on the first generation-only content, and on iterations + ``2..first_content_emit_iters``. + """ + decoding_iter = request.py_decoding_iter + if request.is_finished: + return True + if decoding_iter == 1: + # The early-emit prototype issues the (non-terminal) iter-1 + # response from `_emit_first_token_responses`; suppress it here. + return emit_first_iter + if (self.first_content_emit_iters > 1 + and 1 < decoding_iter <= self.first_content_emit_iters): + return True + return first_gen_emit or decoding_iter % self.stream_interval == 0 + @nvtx_range("_handle_responses") def _handle_responses(self, emit_first_iter: bool = True): new_responses = [] @@ -10445,20 +10500,12 @@ def _handle_responses(self, emit_first_iter: bool = True): "route_capture", None) if route_capture is not None: route_capture.attach_routes(request) # R3: append routes - # A generation-only request skipped its iteration-1 response (the - # CTX returned the first token); without this its first streamed - # content waits for decoding iteration stream_interval. + # A generation-only request's first streamed content must not wait + # for stream_interval after the context-side first token. first_gen_emit = (request.is_generation_only_request and not request.py_gen_first_emitted) - should_emit = (request.py_decoding_iter == 1 or request.is_finished - or first_gen_emit or - request.py_decoding_iter % self.stream_interval == 0) - # The early-emit prototype issues the (non-terminal) iter-1 - # response from `_emit_first_token_responses`; suppress it here. - if (not emit_first_iter and request.py_decoding_iter == 1 - and not request.is_finished): - should_emit = False - if should_emit: + if self._should_emit_response(request, emit_first_iter, + first_gen_emit): if request.return_perf_metrics: # Response creation may finalize and copy scalar ctx GPU totals. self.perf_manager.compute_batch_gpu_times( diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py index 3409d72cb772..67276df9cc73 100644 --- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py +++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py @@ -15,7 +15,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Optional, Protocol +from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol if TYPE_CHECKING: import torch @@ -144,6 +144,9 @@ class MsaRuntimePolicy: skip_unchanged_block_offsets: bool = False # Build each KV pool view once per pool geometry instead of on every attention call. reuse_pool_views: bool = False + # Fill the per-request block-id table by slice-copying each request's page-index view instead of + # building a Python list of the whole padded row (max_seq_len / tokens_per_block entries) first. + block_ids_from_views: bool = False class DecodePrologue(Protocol): @@ -181,21 +184,78 @@ class RuntimeHooks: # Pack a decode step's small host-to-device input copies into one upload and # one scatter launch, issued just before the CUDA graph replay. pack_step_inputs: bool = False + # Issue the step-input pack's host copies with one parallel scatter (a program per copy) instead of the + # serial one-CTA scatter, when their destinations do not overlap (pack_step_inputs). + pack_parallel_scatter: bool = False + # Issue the step-input pack inside the replayed graph (a prologue copies a pinned window and scatters it), so no + # operation is left between the decode graphs (needs pack_parallel_scatter). + pack_in_graph: bool = False + # With pack_in_graph: a side copy stream uploads the pinned window into the prologue's own device buffer while the + # previous graph still runs, so the prologue only waits on that copy and scatters device to device (no host copy + # at the head of the replay). + pack_early_copy: bool = False + # With pack_in_graph: the pack keeps a step epoch, one int64 device counter the graph's prologue increments once per + # replay and never resets, so in-graph per-step flags can be tagged with it instead of cleared between steps. + pack_step_epoch: bool = False + # With pack_early_copy: the model's one-launch prologue, called inside the capture as + # ``op(window, epoch_or_None, max_desc)``. It writes the window's host copies, then its device copies, zeroes the + # window's header and increments the epoch, which the default prologue does with two scatter launches and a fill. + pack_prologue_op: Optional[Callable[[torch.Tensor, Optional[torch.Tensor], int], None]] = None decode_prologue: Optional[DecodePrologue] = None executor_events: Optional[ExecutorEvents] = None # Skip host-to-device metadata copies whose values the destination already holds. skip_unchanged_h2d: bool = False + # With skip_unchanged_h2d: a changed packed 2-D upload stages only each row's changed column span (a page table that + # gained a page re-sends a few bytes, not every padded row), so the pack keeps fitting the in-graph window. + h2d_changed_span: bool = False + # A speculative step under the overlap scheduler re-derives its KV-length metadata (on_update_kv_lens) once, + # after the forward's KV-length correction, instead of also before it: the first pass's device writes are all + # rewritten by the second. + kv_lens_update_once: bool = False + # The overlap scheduler's speculative step correction (position_ids += position offsets, kv_lens += KV-length + # offsets) as the model's one launch, called as ``op(positions, position_offsets, kv_lens, kv_lens_offsets)`` on the + # generation rows; it returns False to decline (the two in-tree adds run). + spec_offsets_op: Optional[ + Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], bool] + ] = None # Run the overlap executor loop with the execution stream current, so the # loop's own stream work and the forward are ordered on one stream. loop_on_execution_stream: bool = False # After autotuning, keep the fastest profiled tactic per shape on every # tensor-parallel rank, so the ranks capture the same kernels. merge_tuned_tactics_across_ranks: bool = False + # Streaming: also emit a response on decoding iterations 2..N, so the first + # counted text reaches the client one step after the first token instead of + # one stream interval later (a reasoning parser holds a leading reasoning + # tag back, so the iteration-1 chunk can carry no content). Emission only; + # tokens are unchanged. 1 keeps the stock schedule. + first_content_emit_iters: int = 1 + # On SM107, run FlashInfer's CuTe DSL MXFP8 GEMM above 32 rows at the pinned tuned tactics of + # custom_ops/mxfp8_cute_pinned_tactics (bitwise the default tactic's output, faster). The GEMM runner does not + # know the model, so the executor arms the table process-wide when the model opts in. + mxfp8_pinned_tactics: bool = False + # After warmup, log every CuTe DSL / Triton compile as a late compile (tensorrt_llm/_torch/late_compile.py). + late_compile_detect: bool = False # Skip the per-step Philox seed and offset uploads on all-greedy # one-model speculative batches (the argmax path never reads them). skip_greedy_rng_upload: bool = False + # Stage the one-model speculative sampler's per-step Philox seed and offset + # uploads into the step-input pack (pack_step_inputs), so they ride its one + # upload instead of each being a copy the next graph waits for. + pack_spec_rng_uploads: bool = False + # The one-model speculative step's small per-step uploads: the accepted-draft + # counts join the step-input pack, and the spec sampler keeps its slot ids in + # a device buffer it re-uploads only when they change. + spec_small_h2d: bool = False # Speculative step-input staging folded into the forward (see SpecStepFold). spec_step_fold: Optional[SpecStepFold] = None + # Copy the one-model speculative sampler's per-slot store to the host on the sampler's D2H side + # stream, so those copies overlap the next step's input staging instead of preceding it. + sampler_d2h_side_stream: bool = False + # Time the iteration log's device step with one timing event per iteration boundary, recorded on a side + # stream that waits for the current stream, instead of an end and a start event back to back on the + # current stream (the same instant, with no GPU work between them). The reported step times are unchanged. + iteration_timing_one_event: bool = False # MiniMax-M3 MSA metadata preparation. msa: MsaRuntimePolicy = MsaRuntimePolicy() diff --git a/tensorrt_llm/_torch/pyexecutor/step_inputs.py b/tensorrt_llm/_torch/pyexecutor/step_inputs.py index c538896c04da..448064f4968c 100644 --- a/tensorrt_llm/_torch/pyexecutor/step_inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -19,19 +19,132 @@ destination between the copy site and the flush: the pack moves the copy to the flush point. A copy that cannot be packed (non-contiguous, device-side dtype change) flushes the pack first and then runs as a plain copy, which keeps the original order. A pack holds every tensor it staged until it is issued, so a temporary source outlives its copy. + +With ``parallel_scatter`` (``RuntimeHooks.pack_parallel_scatter``) a pack of host copies only, whose destinations do +not overlap, is issued with one program per copy (``_parallel_scatter_kernel``) instead of the serial scatter: the same +bytes land in the same places, and with no overlap their order cannot matter. Any other pack takes the serial scatter. + +With ``in_graph`` (``RuntimeHooks.pack_in_graph``, on top of ``parallel_scatter``) the issue itself moves into the +replayed graph: a graph captured with ``capture_prologue`` starts by copying a fixed pinned window (``ghost``) to the +device staging buffer, records ``ghost_free``, and runs the parallel scatter with the copy counts read from the staged +header (one program per descriptor slot; the slots past the counts write nothing). ``flush(in_graph=True)`` before +such a replay then only fills the window on the host, after ``ghost_free`` shows the previous replay has copied it: no +operation is left between the graphs. A pack that does not fit the window or is not order-free is issued eagerly as +before, with a zero-count window so the graph's prologue writes nothing. + +With ``early_copy`` (``RuntimeHooks.pack_early_copy``, on top of ``in_graph``) the window reaches the device before the +replay: ``flush(in_graph=True)`` fills the pinned window and enqueues its upload on a side copy stream, into a device +buffer of the prologue's own (``gdev``), while the previous graph may still run. The prologue then waits for that upload +(``staged_ready``) and scatters device to device; it zeroes the window's header after the scatter and records +``prologue_done``, which the next upload waits for before it overwrites ``gdev``. The next fill of the pinned window +waits for ``ghost_free``, which the upload records once it has read the window. + +With ``step_epoch`` (``RuntimeHooks.pack_step_epoch``, on top of ``in_graph``) the pack keeps one int64 device counter, +``StepInputPack.epoch`` (``step_epoch()`` for the active pack), which the prologue's first scatter launch increments +once per replay (program 0, before any other kernel of the graph runs). It is never reset: it starts at 0, advances once +when the prologue kernel is compiled and then once per replay of any graph with the prologue, with or without a new pack +to scatter, and not on an eager forward. A kernel later in the graph reads it after its programmatic-dependent-launch +wait, so per-step device flags can be tagged with it instead of being cleared between steps. + +With ``prologue_op`` (``RuntimeHooks.pack_prologue_op``, on top of ``early_copy``) the prologue is the model's one +launch, ``prologue_op(gdev, epoch, MAX_DESC)``, in place of the two scatter launches and the header fill: it writes the +same bytes, zeroes the same header and increments the epoch the same once per replay (and once at warm-up). """ from __future__ import annotations import contextlib from contextvars import ContextVar -from typing import Iterator, Optional +from typing import Callable, Iterator, Optional import torch +import triton +import triton.language as tl from tensorrt_llm._utils import prefer_pinned +# One compiled variant for every descriptor offset: specializing ``first`` (== 1, divisible by 16) would compile a +# new kernel when a step's pack first holds a new number of copies, after the warm-up. +@triton.jit(do_not_specialize=["first"]) +def _parallel_scatter_kernel( + stage_u8, stage_i64, first, DEVICE_SRC: tl.constexpr, BLOCK: tl.constexpr +): + """Program i writes copy ``first + i`` ([dst, src, bytes] descriptors after the 2-word header). ``src`` is an offset + into the staged buffer for a host copy, a device address for a device copy (``DEVICE_SRC``).""" + d = 2 + 3 * (first + tl.program_id(0)) + dst = tl.load(stage_i64 + d).to(tl.pointer_type(tl.uint8)) + src_word = tl.load(stage_i64 + d + 1) + if DEVICE_SRC: + src = src_word.to(tl.pointer_type(tl.uint8)) + else: + src = stage_u8 + src_word + nbytes = tl.load(stage_i64 + d + 2) + for start in range(0, nbytes, BLOCK): + o = start + tl.arange(0, BLOCK) + tl.store(dst + o, tl.load(src + o, mask=o < nbytes), mask=o < nbytes) + + +@triton.jit +def _graph_scatter_kernel( + stage_u8, + stage_i64, + epoch, + DEVICE_SRC: tl.constexpr, + BLOCK: tl.constexpr, + BUMP_EPOCH: tl.constexpr, +): + """Descriptor slot i of the host copies (``DEVICE_SRC`` False) or of the device copies, with the counts read from + the staged header; a slot past its count writes nothing. ``BUMP_EPOCH``: program 0 also increments ``epoch[0]``.""" + i = tl.program_id(0) + if BUMP_EPOCH: + if i == 0: + tl.store(epoch, tl.load(epoch) + 1) + num_h2d = tl.load(stage_i64) + if DEVICE_SRC: + count = tl.load(stage_i64 + 1) + d = 2 + 3 * (num_h2d + i) + else: + count = num_h2d + d = 2 + 3 * i + if i < count: + dst = tl.load(stage_i64 + d).to(tl.pointer_type(tl.uint8)) + src_word = tl.load(stage_i64 + d + 1) + if DEVICE_SRC: + src = src_word.to(tl.pointer_type(tl.uint8)) + else: + src = stage_u8 + src_word + nbytes = tl.load(stage_i64 + d + 2) + for start in range(0, nbytes, BLOCK): + o = start + tl.arange(0, BLOCK) + tl.store(dst + o, tl.load(src + o, mask=o < nbytes), mask=o < nbytes) + + +def _disjoint(ranges: list[tuple[int, int]]) -> bool: + """Whether the [start, start + bytes) ranges are pairwise disjoint.""" + end = -1 + for start, nbytes in sorted(ranges): + if start < end: + return False + end = start + nbytes + return True + + +def _parallel_ok(h2d: list[tuple[int, int, int]], d2d: list[tuple[int, int, int]]) -> bool: + """Whether issuing the host copies in parallel and then the device copies in parallel writes what the serial + scatter (every copy in order) writes: the host copies' destinations are disjoint, and the device copies' + destinations are disjoint and none of them is a device copy's source (a host copy may still feed a device copy, + as the launches run in order).""" + if not _disjoint([(dst, n) for dst, _, n in h2d]): + return False + if not d2d: + return True + dsts = [(dst, n) for dst, _, n in d2d] + if not _disjoint(dsts): + return False + return all(_disjoint(dsts + [(src, n)]) for _, src, n in d2d) + + @torch.library.register_fake("trtllm::step_inputs_scatter") def _step_inputs_scatter_fake(stage: torch.Tensor) -> None: # Writes the staged rows into their destinations; returns nothing. @@ -48,8 +161,47 @@ class StepInputPack: HEADER_BYTES = 16 DESC_BYTES = 24 ALIGN = 16 - - def __init__(self, device: torch.device) -> None: + # The pinned window a graph's prologue copies (in_graph): the header, the descriptors and up to ~10 KB of data. + GRAPH_WINDOW = 1 << 14 + + def __init__( + self, + device: torch.device, + parallel_scatter: bool = False, + in_graph: bool = False, + early_copy: bool = False, + step_epoch: bool = False, + prologue_op: Optional[Callable[[torch.Tensor, Optional[torch.Tensor], int], None]] = None, + ) -> None: + self.parallel_scatter = parallel_scatter + self.in_graph = in_graph and parallel_scatter + self.early_copy = early_copy and self.in_graph + # The model's one-launch prologue (see the module docstring); None: the scatter launches and the header fill. + self.prologue_op = prologue_op if self.early_copy else None + # The step epoch (see the module docstring); None without step_epoch. + self.epoch: Optional[torch.Tensor] = ( + torch.zeros(1, dtype=torch.int64, device=device) + if step_epoch and self.in_graph + else None + ) + if self.early_copy: + # The prologue's own device window: an eager flush (self.dev) can never leave a header in it. + self.gdev = torch.zeros(self.GRAPH_WINDOW, dtype=torch.uint8, device=device) + self.copy_stream = torch.cuda.Stream(device=device) + # Recorded by the side upload once it landed (the prologue waits on it) ... + self.staged_ready = torch.cuda.Event(external=True) + # ... and by every replayed prologue after its scatter (the next upload waits on it). + self.prologue_done = torch.cuda.Event(external=True) + # A CUDA event is created on its first record, and a wait on one never created is skipped: a capture + # would then leave the prologue's wait out of the graph. Create both now. + self.staged_ready.record(self.copy_stream) + self.prologue_done.record(self.copy_stream) + if self.in_graph: + self.ghost = torch.zeros( + self.GRAPH_WINDOW, dtype=torch.uint8, pin_memory=prefer_pinned() + ) + # Recorded by every replayed prologue right after it copied the window (an event node). + self.ghost_free = torch.cuda.Event(external=True) self.host = [ torch.empty(self.CAPACITY, dtype=torch.uint8, pin_memory=prefer_pinned()) for _ in range(self.RING) @@ -64,6 +216,86 @@ def __init__(self, device: torch.device) -> None: # Every tensor a staged copy reads or writes, kept alive until the pack is issued. self.staged: list[torch.Tensor] = [] self.used = self.data_start + if parallel_scatter: + self._warm_up(device) + if self.in_graph: + self._warm_up_graph_kernel() + + def _warm_up(self, device: torch.device) -> None: + """Compile the parallel scatter now, not on a serving step's flush.""" + dst = torch.zeros(8, dtype=torch.int32, device=device) + self.begin() + self.stage(dst[:4], torch.arange(4, dtype=torch.int32)) + self.stage(dst[4:], dst[:4]) + self.flush() + torch.cuda.synchronize(device) + + def _warm_up_graph_kernel(self) -> None: + """Compile the prologue's scatter now (a zero-count header: it writes nothing).""" + self.dev[: self.HEADER_BYTES].zero_() + if self.prologue_op is not None: + # The window it is captured with; its header is zero until the first upload. + self.prologue_op(self.gdev, self.epoch, self.MAX_DESC) + else: + self._launch_graph_scatter(self.dev) + torch.cuda.synchronize(self.dev.device) + + def _launch_graph_scatter(self, window: torch.Tensor) -> None: + stage_i64 = window.view(torch.int64) + for device_src in (False, True): + bump = self.epoch is not None and not device_src + _graph_scatter_kernel[(self.MAX_DESC,)]( + window, + stage_i64, + self.epoch if bump else stage_i64, + DEVICE_SRC=device_src, + BLOCK=256, + BUMP_EPOCH=bump, + num_warps=1, + ) + + def capture_prologue(self) -> bool: + """Inside a graph capture: the prologue that issues the window ``flush(in_graph=True)`` fills (see the module + docstring). Returns whether the graph has it.""" + if not self.in_graph: + return False + if self.early_copy: + torch.cuda.current_stream().wait_event(self.staged_ready) + if self.prologue_op is not None: + self.prologue_op(self.gdev, self.epoch, self.MAX_DESC) + else: + self._launch_graph_scatter(self.gdev) + # Consumed: a replay with no new upload scatters nothing. + self.gdev[: self.HEADER_BYTES].zero_() + self.prologue_done.record() + return True + self.dev[: self.GRAPH_WINDOW].copy_(self.ghost, non_blocking=True) + self.ghost_free.record() + self._launch_graph_scatter(self.dev) + return True + + def _upload_window(self, nbytes: int) -> None: + """early_copy: upload the filled pinned window into gdev on the side stream, once the last replayed prologue + has consumed gdev; the window is fully written on the host before this is enqueued.""" + stream = self.copy_stream + stream.wait_event(self.prologue_done) + with torch.cuda.stream(stream): + self.gdev[:nbytes].copy_(self.ghost[:nbytes], non_blocking=True) + self.ghost_free.record(stream) + self.staged_ready.record(stream) + + def _ghost_writable(self) -> None: + # The last replayed prologue has copied the window (a no-op before the first replay). + self.ghost_free.synchronize() + + def quiesce_graph_window(self) -> None: + """Before replaying a graph with the prologue without a flush: a zero-count window, so it writes nothing.""" + if self.in_graph and self.ghost[: self.HEADER_BYTES].any(): + self._ghost_writable() + self.ghost[: self.HEADER_BYTES].zero_() + if self.early_copy: + # A window uploaded but never replayed would still sit in gdev: overwrite its header. + self._upload_window(self.HEADER_BYTES) def begin(self) -> None: if self.open: @@ -119,10 +351,20 @@ def stage(self, dst: torch.Tensor, src: torch.Tensor) -> bool: return True return False - def flush(self) -> None: + def flush(self, in_graph: bool = False) -> None: + """Issue the pack. ``in_graph``: the next replay is of a graph captured with ``capture_prologue``.""" if not self.open: + if in_graph: + self.quiesce_graph_window() return self.open = False + if in_graph and self.in_graph: + if self._fill_graph_window(): + if self.early_copy: + self._upload_window(self.used) + self._clear() + return + self.quiesce_graph_window() if not self.h2d and not self.d2d: return host = self.host[self.parity] @@ -132,7 +374,18 @@ def flush(self) -> None: host[: 8 * len(desc)].view(torch.int64).copy_(torch.tensor(desc, dtype=torch.int64)) # Descriptors and data in one upload; the gap between them is unused padding. self.dev[: self.used].copy_(host[: self.used], non_blocking=True) - torch.ops.trtllm.step_inputs_scatter(self.dev) + if self.parallel_scatter and _parallel_ok(self.h2d, self.d2d): + stage_i64 = self.dev.view(torch.int64) + if self.h2d: + _parallel_scatter_kernel[(len(self.h2d),)]( + self.dev, stage_i64, 0, DEVICE_SRC=False, BLOCK=256, num_warps=2 + ) + if self.d2d: + _parallel_scatter_kernel[(len(self.d2d),)]( + self.dev, stage_i64, len(self.h2d), DEVICE_SRC=True, BLOCK=256, num_warps=2 + ) + else: + torch.ops.trtllm.step_inputs_scatter(self.dev) event = self.events[self.parity] if event is None: event = self.events[self.parity] = torch.cuda.Event() @@ -141,6 +394,20 @@ def flush(self) -> None: # The scatter is queued on the current stream: the caching allocator keeps the staged tensors' memory for it. self._clear() + def _fill_graph_window(self) -> bool: + """Write the pack into the pinned window the replay's prologue issues; False when it does not fit or is not + order-free (the caller issues it eagerly).""" + if self.used > self.GRAPH_WINDOW or not _parallel_ok(self.h2d, self.d2d): + return False + desc = [len(self.h2d), len(self.d2d)] + [ + value for copy in self.h2d + self.d2d for value in copy + ] + self._ghost_writable() + host = self.host[self.parity] + self.ghost[: self.used].copy_(host[: self.used]) + self.ghost[: 8 * len(desc)].view(torch.int64).copy_(torch.tensor(desc, dtype=torch.int64)) + return True + def discard(self) -> None: """Close the pack without issuing it (its step failed).""" self.open = False @@ -167,6 +434,11 @@ def begin() -> None: pack.begin() +def is_active(pack: StepInputPack) -> bool: + """Whether ``pack`` is the calling context's pack.""" + return _active.get() is pack + + def is_open() -> bool: pack = _active.get() return pack is not None and pack.open @@ -193,8 +465,24 @@ def discard() -> None: pack.discard() -def flush() -> None: - """Issue the active open pack (one upload and one scatter launch on the current stream) and close it.""" +def flush(in_graph: bool = False) -> None: + """Issue the active open pack (one upload and one scatter launch on the current stream) and close it. + + ``in_graph``: the next replay is of a graph whose prologue issues the pack (``capture_prologue``).""" pack = _active.get() if pack is not None: - pack.flush() + pack.flush(in_graph=in_graph) + + +def step_epoch() -> Optional[torch.Tensor]: + """The active pack's step epoch (an int64 [1] device tensor), or None when there is none (see the module + docstring).""" + pack = _active.get() + return pack.epoch if pack is not None else None + + +def capture_prologue() -> Optional[StepInputPack]: + """Inside a graph capture: add the active pack's in-graph prologue; the pack whose window the graph issues, or + None when the graph has no prologue.""" + pack = _active.get() + return pack if pack is not None and pack.capture_prologue() else None diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 3ce6a6a38116..b090a2fd0201 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -32,6 +32,7 @@ from ..attention.backends.trtllm import (AttentionBackend, TrtllmAttention, TrtllmAttentionMetadata) from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE +from ..pyexecutor import step_inputs from ..pyexecutor.resource_manager import ResourceManagerType if TYPE_CHECKING: @@ -642,6 +643,9 @@ def is_ragged_verify(self) -> bool: # from the model's RuntimeHooks). The argmax path never reads them, and the # next sampled batch rewrites every row it reads. skip_greedy_rng_upload: bool = False + # Stage the per-step Philox seed/offset uploads into the engine's step-input + # pack (RuntimeHooks.pack_spec_rng_uploads) instead of copying each one. + pack_rng_uploads: bool = False # Advanced-sampling specialization (deploy-time; from DecodingBaseConfig.advanced_sampling_mode). advanced_sampling_mode: AdvancedSamplingMode = AdvancedSamplingMode.FULL # Whether the occurrence penalties are enabled (deploy-time; from @@ -955,10 +959,15 @@ def _populate_request_rng_state( device='cuda') def _upload(dst: torch.Tensor, values: list[int]) -> None: - dst[:len(values)].copy_(torch.tensor(values, - dtype=torch.int64, - pin_memory=prefer_pinned()), - non_blocking=True) + src = torch.tensor(values, + dtype=torch.int64, + pin_memory=prefer_pinned()) + if self.pack_rng_uploads: + # The forward is the first reader of these buffers, so the copy + # may move to the pack's flush (a plain copy when no pack is open). + step_inputs.copy_(dst[:len(values)], src) + else: + dst[:len(values)].copy_(src, non_blocking=True) # An all-greedy batch runs the argmax graph, which reads none of these # buffers. The offset windows above still advance, so skipping the diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 8e71b8fe089c..e60168a814b2 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -76,6 +76,11 @@ "Drop top_p_decay / top_p_min from the request, or disable " "speculative decoding." ) +UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG = ( + "logits_processor (a logits post-processor callback) is not supported with one-model " + "speculative decoding. Drop logits_processor from the request, or disable speculative " + "decoding." +) UNSUPPORTED_MIN_P_MSG = ( "min_p requires 'advanced_sampling_mode: full' in the speculative decoding " "config when using one-model speculative decoding. Restore that default, " @@ -116,6 +121,10 @@ def one_model_sampling_rejection_reason( return UNSUPPORTED_NO_REPEAT_NGRAM_MSG if sampling_params.embedding_bias is not None: return UNSUPPORTED_EMBEDDING_BIAS_MSG + # A per-request callback becomes py_logits_post_processors; the one-model path never + # calls it (a fused verify may not even materialize the full logits). + if sampling_params.logits_processor: + return UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG # The same predicate top_p_decay_active() delegates to, so "active" cannot # mean one thing here and another in the sampler. if SamplingParams.params_imply_top_p_decay_active(sampling_params.top_p_decay): @@ -171,6 +180,14 @@ class SpecSampler(Sampler[SampleStateSpec], AsyncWorkerMixin): # The model's speculative step-input staging, when it opted into one (bind_spec_step_fold). spec_step_fold: Optional[SpecStepFold] = None + # Copy the store on the D2H side stream also on the steps a step fold's forward stored the rows, when the fold + # awaits the copies through its store_copies_done (RuntimeHooks.sampler_d2h_side_stream). + d2h_side_stream_copies: bool = False + # Keep the stored rows' slot ids in a device buffer re-uploaded only when they change + # (RuntimeHooks.spec_small_h2d). + static_slot_ids: bool = False + _slot_ids_host: Optional[list[int]] = None + _slot_ids_cuda: Optional[torch.Tensor] = None SampleState = SampleStateSpec @@ -235,8 +252,8 @@ def _validate_unsupported_logits_processors(request: LlmRequest) -> None: """Reject logits-side sampling features the one-model path cannot apply. TorchSampler implements min_length, bad_words, no_repeat_ngram_size, - embedding_bias and top_p_decay by editing the target logits before - sampling. The one-model path has no logits-editing hook, so each is + embedding_bias, top_p_decay and logits post-processor callbacks by editing + the target logits before sampling. The one-model path has no logits-editing hook, so each is rejected here instead of being silently dropped. Each check gates on a non-neutral value rather than on presence, since a @@ -255,6 +272,11 @@ def _validate_unsupported_logits_processors(request: LlmRequest) -> None: raise ValueError(UNSUPPORTED_NO_REPEAT_NGRAM_MSG) if getattr(request, "py_embedding_bias", None) is not None: raise ValueError(UNSUPPORTED_EMBEDDING_BIAS_MSG) + # The one-model sampler never calls a request's logits post-processors, and a fused + # verify (e.g. the MiniMax-M3 shard verify) does not even gather the full logits, so + # the callback would be skipped silently. + if getattr(request, "py_logits_post_processors", None): + raise ValueError(UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG) # Reuse the handler's own predicate so "active" cannot drift between paths. if top_p_decay_active(_request_get_sampling_params(request)): raise ValueError(UNSUPPORTED_TOP_P_DECAY_MSG) @@ -331,6 +353,67 @@ def __init__( new_tokens_lens=int_tensor((seq_slots,)), ) + def _device_slot_ids(self, seq_slots: list[int]) -> torch.Tensor: + """The slot ids on the device, uploaded only when they differ from the last upload. + + The buffer is this sampler's own: only this method writes it, so the last uploaded + list describes its contents, and the stream orders each upload after the scatter that + read the previous ids. + """ + n = len(seq_slots) + if self._slot_ids_cuda is None or self._slot_ids_cuda.numel() < n: + self._slot_ids_cuda = torch.empty(max(n, 8), dtype=torch.long, device="cuda") + self._slot_ids_host = None + if self._slot_ids_host != list(seq_slots): + self._slot_ids_cuda[:n].copy_( + torch.tensor(seq_slots, dtype=torch.long, pin_memory=prefer_pinned()), + non_blocking=True, + ) + self._slot_ids_host = list(seq_slots) + return self._slot_ids_cuda[:n] + + def _copy_outputs_to_host( + self, + o_verify_lens: Optional[torch.Tensor], + forward_stored_rows: bool, + store_copies_done: Optional[torch.cuda.Event] = None, + ) -> tuple[SampleStateTensorsSpec, Optional[torch.cuda.Event]]: + """Host copies of the step's stored outputs, and the side-stream event to await them by (or None). + + The per-slot store is copied on the D2H side stream, so the copies overlap the next step's input staging + instead of preceding it on the main stream. That is sound while nothing rewrites the store before the next + ``sample_async`` waits on ``_store_copies_done``, which holds when this sampler stored the rows itself. + ``verify_lens`` is a forward output buffer the next replay rewrites, so it is always copied on the main + stream. + + When the forward stored the rows (a step fold), the next replay rewrites the store, so the copies stay on + the main stream unless the fold awaits them: with ``d2h_side_stream_copies`` and the fold's + ``store_copies_done``, that event is recorded on the side stream after the copies and the forward waits + on it before it writes the store. + """ + verify_lens = self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None + fold_awaits_copies = self.d2h_side_stream_copies and store_copies_done is not None + if self._async_worker_active() or (forward_stored_rows and not fold_awaits_copies): + host = SampleStateTensorsSpec( + new_tokens=self._copy_to_host(self.store.new_tokens), + new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), + next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), + verify_lens=verify_lens, + ) + return host, None + with self._make_side_stream_copier() as copier: + host = SampleStateTensorsSpec( + new_tokens=copier.stage_copy_to_host(self.store.new_tokens), + new_tokens_lens=copier.stage_copy_to_host(self.store.new_tokens_lens), + next_draft_tokens=copier.stage_copy_to_host(self.store.next_draft_tokens), + verify_lens=verify_lens, + ) + self._store_copies_done = copier.event + if store_copies_done is not None and copier.event is not None: + # The fold's forward waits on it before every store write, whoever stored this step's rows. + store_copies_done.record(self._d2h_side_stream) + return host, copier.event + def bind_spec_step_fold(self, spec_step_fold: SpecStepFold) -> None: """Let the model's forward store into this sampler's per-slot store (see SpecStepFold).""" self.spec_step_fold = spec_step_fold @@ -424,8 +507,11 @@ def _scatter_to_stores( ) -> None: """Move rows ``num_skip`` onward of the forward's outputs into the stores at ``slots``.""" num_sampling_requests = len(slots) - slots_device = torch.as_tensor(slots, dtype=torch.long) - slots_device = slots_device.to(device="cuda", non_blocking=True) + if self.static_slot_ids: + slots_device = self._device_slot_ids(slots) + else: + slots_device = torch.as_tensor(slots, dtype=torch.long) + slots_device = slots_device.to(device="cuda", non_blocking=True) o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests] o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests] @@ -559,9 +645,10 @@ def sample_async( torch.cuda.current_stream().wait_event(self._store_copies_done) slots = [r.py_seq_slot for r in sampling_requests] step_fold = self.spec_step_fold - scattered = step_fold is not None and step_fold.sampler_rows_stored( + forward_stored_rows = step_fold is not None and step_fold.sampler_rows_stored( outputs, num_skip, slots ) + scattered = forward_stored_rows # A model's step fold scatters the stores itself, in _scatter_to_stores. if not scattered and step_fold is None and self._store_scatter is not None: # A step whose batch is unchanged leaves the slot table as it is. @@ -586,33 +673,11 @@ def sample_async( next_draft_tokens=self.store.next_draft_tokens, ) - if self._async_worker_active(): - host_tensors = SampleStateTensorsSpec( - new_tokens=self._copy_to_host(self.store.new_tokens), - new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), - next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), - verify_lens=( - self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None - ), - ) - sampler_event = self._record_sampler_event() - else: - # The next step's input preparation waits on this stream and only - # needs the stores, so the host copies run on the side stream. - # update_requests syncs their event before reading them, and the - # next step's store update waits for it on this stream. - host_verify_lens = ( - self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None - ) - with self._make_side_stream_copier() as copier: - host_tensors = SampleStateTensorsSpec( - new_tokens=copier.stage_copy_to_host(self.store.new_tokens), - new_tokens_lens=copier.stage_copy_to_host(self.store.new_tokens_lens), - next_draft_tokens=copier.stage_copy_to_host(self.store.next_draft_tokens), - verify_lens=host_verify_lens, - ) - self._store_copies_done = copier.event - sampler_event = self._record_sampler_event(side_stream_event=copier.event) + store_copies_done = getattr(step_fold, "store_copies_done", None) if step_fold else None + host_tensors, side_stream_event = self._copy_outputs_to_host( + o_verify_lens, forward_stored_rows, store_copies_done + ) + sampler_event = self._record_sampler_event(side_stream_event=side_stream_event) # Add dummy draft tokens to context requests for KV cache preparation for request in finished_context_requests: diff --git a/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py index a5c777b5cf7c..7c3ff30b3c30 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py @@ -18,11 +18,11 @@ criteria would quietly fall back to the built-in implementation, pass, and report the built-in's numbers as the target's. -The deployment is the one the targets route (``models/minimax_m3/routing.py``): the NVFP4 checkpoint on sm_103, -tensor-parallel attention and experts on 4 or 8 ranks, the MiniMax-M3 MSA sparse attention with the fused five-way -QKV/index projection, an FP8 indexer and an NVFP4 KV cache. Each target runs GSM8K without speculation and with -one-model Eagle3, where a chat-GSM8K acceptance probe also guards the drafter (accuracy alone does not notice a -corrupted draft KV). +The deployment is the one the targets route (``models/minimax_m3/routing.py``): the NVFP4 checkpoint on sm_103 +(GB300, 4 or 8 ranks) or sm_107 (Rubin, 2, 4 or 8 ranks), tensor-parallel attention and experts, the MiniMax-M3 MSA +sparse attention with the fused five-way QKV/index projection, an FP8 indexer and an NVFP4 KV cache. Each target runs +GSM8K without speculation and with one-model Eagle3, where a chat-GSM8K acceptance probe also guards the drafter +(accuracy alone does not notice a corrupted draft KV). """ import json @@ -48,7 +48,11 @@ # The targets assert their own SM at construction: certification is per GPU architecture. skip_not_sm103 = pytest.mark.skipif( get_sm_version() != 103, - reason="the MiniMax-M3 modeling_v2 targets are certified on sm_103 only", + reason="these MiniMax-M3 modeling_v2 targets are certified on sm_103 only", +) +skip_not_sm107 = pytest.mark.skipif( + get_sm_version() != 107, + reason="these MiniMax-M3 modeling_v2 targets are certified on sm_107 only", ) # The chat-GSM8K acceptance floors of the in-tree test on an NVFP4 KV cache (test_llm_api_pytorch.TestMiniMaxM3). @@ -166,3 +170,51 @@ def test_gsm8k(self): @pytest.mark.skip_less_mpi_world_size(8) def test_eagle3_gsm8k_and_acceptance(self): self._gsm8k(eagle3=True) + + +class TestModelingV2MiniMaxM3Nvfp4Sm107Tp2(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_107 / tp2: two Rubin GPUs.""" + + TP = 2 + + @skip_not_sm107 + @pytest.mark.skip_less_device(2) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm107 + @pytest.mark.skip_less_device(2) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) + + +class TestModelingV2MiniMaxM3Nvfp4Sm107Tp4(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_107 / tp4: four Rubin GPUs.""" + + TP = 4 + + @skip_not_sm107 + @pytest.mark.skip_less_device(4) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm107 + @pytest.mark.skip_less_device(4) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) + + +class TestModelingV2MiniMaxM3Nvfp4Sm107Tp8(_MiniMaxM3Target): + """MiniMax-M3 NVFP4 / sm_107 / tp8: eight Rubin GPUs over two trays (an 8-rank MPI world across the two nodes).""" + + TP = 8 + + @skip_not_sm107 + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k(self): + self._gsm8k(eagle3=False) + + @skip_not_sm107 + @pytest.mark.skip_less_mpi_world_size(8) + def test_eagle3_gsm8k_and_acceptance(self): + self._gsm8k(eagle3=True) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index cd5a6fd51581..5d44c296060d 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7634,6 +7634,7 @@ def _run_nvfp4(self, torch_compile_config=TorchCompileConfig() if piecewise else None, max_seq_len=4096, + modeling_v2="off", trust_remote_code=True) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION if kv_dtype == "nvfp4": @@ -7757,6 +7758,7 @@ def _run_nvfp4_eagle3(self, # Keep the whole acceptance probe: the default buffer holds # only the latest 1000 iterations. max_stats_len=5000, + modeling_v2="off", trust_remote_code=True, **piecewise_kwargs) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 044a6b612c56..18d596032c05 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -142,7 +142,7 @@ l0_b300: # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU # stage instead), so 30 is 2.2x headroom. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm --ignore=unittest/_torch/modeling_v2/models/kimi_k3_vl --ignore=unittest/_torch/modeling_v2/models/minimax_m3 TIMEOUT (30) - - unittest/_torch/modeling_v2/models/minimax_m3 --import-mode=importlib --ignore-glob=unittest/_torch/modeling_v2/models/minimax_m3/**/comm/* TIMEOUT (30) + - unittest/_torch/modeling_v2/models/minimax_m3 --import-mode=importlib --ignore-glob=unittest/_torch/modeling_v2/models/minimax_m3/**/comm/* --ignore-glob=unittest/_torch/modeling_v2/models/minimax_m3/targets/sm107* TIMEOUT (30) # K3 single-GPU kernel and wrapper checks. - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 31ba50c6488d..1b89fddf0061 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -12,6 +12,8 @@ l0_cpu: backend: generic orchestrator: mpi tests: + # MiniMax-M3 Rubin target unit tests. + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py # MiniMax-M3 target unit tests. - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_103__tp8/test_modeling_v2_m3_tp4_drift.py - unittest/_torch/modules/test_linear_output_slice.py::test_fp8_weight_slice_cache @@ -94,6 +96,12 @@ l0_cpu: - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_kimi_k3_spec_worker_gate.py - unittest/_torch/modeling_v2/models/deepseek_v3/test_modeling_v2_deepseek_default_shards.py - unittest/_torch/modeling_v2/models/gpt_oss/test_modeling_v2_gpt_oss_default_shards.py + # MiniMax-M3 sm_107 tp2 target: its device-free (cpu_only) tests. + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py + - unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_dense_backend_library_order - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_multimodal_embeds.py - unittest/_torch/modeling_v2/models/qwen4_exp/test_modeling_v2_qwen4_exp_moe_scalars.py diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml index a9cbbc03f821..94bc49c05f3e 100644 --- a/tests/integration/test_lists/test-db/l0_gr100.yml +++ b/tests/integration/test_lists/test-db/l0_gr100.yml @@ -13,6 +13,40 @@ l0_gr100: stage: post_merge backend: pytorch tests: + # MiniMax-M3 Rubin target unit tests. + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py + - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True] - unittest/_torch/attention/test_flashinfer_trtllm_gen_sm107_p128.py - unittest/_torch/modeling/test_kimi_k3_mla.py::test_k3_mla_nvfp4_kv_with_fp8_kv_b_matches_huggingface diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py new file mode 100644 index 000000000000..fac388a85ec1 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 V2 cache manager: the SM107 per-token estimate adds what the generic sum leaves out.""" + +import types + +import pytest + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import cache_manager + +KVM = cache_manager.MiniMaxM3KVCacheManagerV2 + + +def _model_config(nvfp4=True, index_dtype="fp8"): + quant_mode = types.SimpleNamespace( + has_fp4_kv_cache=lambda: nvfp4, has_fp8_kv_cache=lambda: not nvfp4 + ) + return types.SimpleNamespace( + pretrained_config=types.SimpleNamespace( + num_hidden_layers=60, + num_key_value_heads=4, + head_dim=128, + hidden_size=6144, + num_attention_heads=64, + ), + sparse_attention_config=types.SimpleNamespace( + sparse_index_dim=128, indexer_kv_dtype=index_dtype + ), + quant_config=types.SimpleNamespace(quant_mode=quant_mode), + ) + + +def _mapping(tp_size, attention_dp): + return types.SimpleNamespace(tp_size=tp_size, pp_size=1, enable_attention_dp=attention_dp) + + +@pytest.fixture +def one_draft_layer(monkeypatch): + monkeypatch.setattr( + cache_manager, "shared_draft_layer_count", lambda spec_config, layer_mask: 1 + ) + + +@pytest.mark.parametrize( + "attention_dp,expected_extra", + [ + # TP2 attention, 2 KV heads per rank: index-K 57 x 128 + dense 3 x (512 - 288) + draft 512. + (False, 57 * 128 + 3 * (512 - 288) + 512), + # Attention DP, 4 KV heads per rank: index-K 57 x 128 + dense 3 x (1024 - 576) + draft 1024. + (True, 57 * 128 + 3 * (1024 - 576) + 1024), + ], +) +def test_extra_bytes_per_token_nvfp4(one_draft_layer, attention_dp, expected_extra): + extra = KVM._extra_bytes_per_token( + _model_config(), _mapping(2, attention_dp), None, {"spec_config": object()} + ) + assert extra == expected_extra + + +def test_separate_draft_cache_not_double_counted(one_draft_layer): + kwargs = {"spec_config": object(), "use_separate_draft_kv_cache": True} + extra = KVM._extra_bytes_per_token(_model_config(), _mapping(2, False), None, kwargs) + assert extra == 57 * 128 + 3 * (512 - 288) + + +def test_estimate_gate(monkeypatch): + calls = [] + + def base(cls, model_config, mapping, num_layers=None, **kwargs): + calls.append(kwargs) + return (17280, 0) + + monkeypatch.setattr( + cache_manager.KVCacheManagerV2, "get_cache_size_per_token", classmethod(base) + ) + monkeypatch.setattr(KVM, "_extra_bytes_per_token", classmethod(lambda cls, *a: 8480)) + monkeypatch.setattr(cache_manager, "get_sm_version", lambda: 107) + assert KVM.get_cache_size_per_token(None, None) == (25760, 0) + assert KVM.get_cache_size_per_token(None, None, is_draft=True) == (17280, 0) + monkeypatch.setattr(cache_manager, "get_sm_version", lambda: 103) + assert KVM.get_cache_size_per_token(None, None) == (17280, 0) + monkeypatch.setattr(cache_manager, "get_sm_version", lambda: 107) + monkeypatch.setattr(cache_manager, "_SM107_EXACT_KV_SIZE_ESTIMATE", False) + assert KVM.get_cache_size_per_token(None, None) == (17280, 0) diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py index 31912c481f51..8e23a16240a1 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py @@ -48,14 +48,18 @@ @pytest.mark.cpu_only +@pytest.mark.parametrize( + "tail_in_forward", [True, False], ids=["tail_in_forward", "tail_in_prepare"] +) @pytest.mark.parametrize("empty_state", ["empty_batch", "no_requests", "no_lengths"]) def test_msa_metadata_clears_padded_cache_slot_tail( - monkeypatch: pytest.MonkeyPatch, empty_state: str + monkeypatch: pytest.MonkeyPatch, empty_state: str, tail_in_forward: bool ) -> None: """A smaller replay must not reuse the previous step's live cache slots.""" from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend monkeypatch.setattr(msa_backend, "maybe_pin_memory", lambda tensor: tensor) + monkeypatch.setattr(msa_backend, "TAIL_FILL_IN_FORWARD", tail_in_forward) metadata_cls = MiniMaxM3MsaSparseAttention.Metadata metadata = metadata_cls.__new__(metadata_cls) metadata._msa_buffers_ready = True @@ -81,6 +85,8 @@ def test_msa_metadata_clears_padded_cache_slot_tail( # live count off it. metadata._num_tokens = count metadata._build_msa_fields() + # TAIL_FILL_IN_FORWARD: the forward (on_update_kv_lens) fills the tail before anything reads it. + metadata._fill_tail_in_forward() assert metadata.msa_out_cache_loc.tolist() == list(range(12, 12 + count)) + [-1] * ( 4 - count ) @@ -518,6 +524,45 @@ def test_msa_buffers_stage_local_cache_views(monkeypatch: pytest.MonkeyPatch, nv torch.testing.assert_close(index_cache, torch.full_like(index_cache, 3)) +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("kv_dtype", "expect_lengths"), + [ + (DataType.FP8, False), + (DataType.NVFP4, True), + ], +) +def test_msa_buffers_reserve_cumulative_lengths_for_csr_prefill( + monkeypatch: pytest.MonkeyPatch, + kv_dtype: DataType, + expect_lengths: bool, +) -> None: + """The NVFP4 CSR kernel reads graph-stable cumulative lengths, so only an + NVFP4 cache reserves them.""" + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend + + metadata = _buffer_metadata(dtype=kv_dtype, sparse_layer_ids=[], layer_offsets={}) + reserved = {} + + def get_empty(buffers, shape, **kwargs): + reserved[kwargs["cache_name"]] = tuple(shape) + return torch.empty(shape, dtype=kwargs["dtype"]) + + monkeypatch.setattr(metadata, "get_empty", get_empty) + monkeypatch.setattr(msa_backend, "uniform_subpages_per_slot", lambda manager: 0) + metadata._create_msa_buffers() + + names = ("msa_qo_lens_cuda", "msa_cu_q_lens", "msa_cu_kv_lens") + if expect_lengths: + assert reserved["msa_qo_lens_cuda"] == (MAX_NUM_SEQUENCES,) + assert reserved["msa_cu_q_lens"] == (MAX_NUM_SEQUENCES + 1,) + assert reserved["msa_cu_kv_lens"] == (MAX_NUM_SEQUENCES + 1,) + assert all(getattr(metadata, name) is not None for name in names) + else: + assert not set(names) & set(reserved) + assert all(getattr(metadata, name) is None for name in names) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize( ("factors", "expected"), diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py index 0e022aaa28f2..8ec99712dee6 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_runtime_policy.py @@ -9,16 +9,23 @@ from unittest.mock import create_autospec +import numpy as np import pytest import torch -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import MiniMaxM3KVCacheManagerV2 +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + msa_backend, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.common import ( MiniMaxM3SparseMetadataParams, ) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( MSA_REQUIRED_TOPK, ) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( + write_subpage_block_table, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( MiniMaxM3MsaSparseAttentionMetadata, _fill_unused_slots, @@ -30,6 +37,7 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.llmapi.llm_args import MiniMaxM3SparseAttentionConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -48,8 +56,8 @@ def fake_base_init(self, *args, **kwargs) -> None: self.is_disagg = False self.dtype = DataType.BF16 self.layer_offsets = {3: 0} - # _prepare_page_table_tensor's result for the two-pool host table the copy tests build. - self.num_attention_op_pools = 2 + # Set by the base's page-table setup, which this fake base skips. + self.num_pools = self.num_attention_op_pools = 1 def no_index_k_buffer(self, layer_idx, **kwargs): del self, layer_idx, kwargs @@ -58,10 +66,10 @@ def no_index_k_buffer(self, layer_idx, **kwargs): monkeypatch.setattr(KVCacheManagerV2, "__init__", fake_base_init) monkeypatch.setattr(KVCacheManagerV2, "get_index_k_buffer", no_index_k_buffer) monkeypatch.setattr(MiniMaxM3KVCacheManagerV2, "_compute_num_total_slots", lambda self: 0) - # The constructor resolves every layer's K/V view, which needs a real pool. Stub it there - # only: the pool-view tests below call get_buffers against their own geometry. - with monkeypatch.context() as construction: - construction.setattr(MiniMaxM3KVCacheManagerV2, "get_buffers", lambda self, *args: None) + # The constructor resolves each layer's K/V views; the bare base has no pools to resolve them + # from. Stub that only while constructing, so the pool-view tests below build real views. + with monkeypatch.context() as construct: + construct.setattr(MiniMaxM3KVCacheManagerV2, "get_buffers", lambda self, *args: None) return MiniMaxM3KVCacheManagerV2( num_layers=4, sparse_attention_config=MiniMaxM3SparseAttentionConfig() ) @@ -397,3 +405,206 @@ def test_unused_slots_are_filled_after_a_full_buffer_step(): buf[:3] = 0 _fill_unused_slots(buf, 3) assert buf.tolist() == [0, 0, 0, -1, -1, -1, -1, -1] + + +class _PageIndexView: + """A request's KV cache as get_block_ids_per_seq reads it: an int32 page-index view per pool.""" + + def __init__(self, view: np.ndarray): + self._view = view + + def get_base_page_indices(self, pool_id: int) -> np.ndarray: + assert pool_id == 0 + return self._view + + +def _page_index_views(seed: int): + rng = np.random.default_rng(seed) + batch = int(rng.integers(0, 5)) + same_width = rng.random() < 0.5 + width = int(rng.integers(1, 9000)) + views = {} + for rid in range(batch): + w = width if same_width else int(rng.integers(0, 9000)) + view = rng.integers(0, 2**31 - 1, w, dtype=np.int64).astype(np.int32) + view[int(rng.integers(0, w + 1)) :] = BAD_PAGE_INDEX + view[rng.random(w) < float(rng.choice([0.0, 0.01, 0.5]))] = BAD_PAGE_INDEX + view.setflags(write=False) + views[100 + rid] = _PageIndexView(view) + return views + + +def _bare_manager() -> MiniMaxM3KVCacheManagerV2: + """A manager with only the state get_block_ids_per_seq reads (no pools, no V2 base).""" + manager = object.__new__(MiniMaxM3KVCacheManagerV2) + manager.layer_specs = None + manager.kv_factor = 2 + return manager + + +def test_block_ids_from_views_is_off_by_default(): + assert not _bare_manager().block_ids_from_views + + +@pytest.mark.parametrize("seed", range(24)) +def test_block_ids_from_views_matches_the_list_build(seed): + manager = _bare_manager() + manager.kv_cache_map = _page_index_views(seed) + request_ids = list(manager.kv_cache_map) + + manager.configure_decode_shortcuts(skip_unchanged_block_offsets=False, reuse_pool_views=False) + built_from_lists = manager.get_block_ids_per_seq(request_ids) + manager.configure_decode_shortcuts( + skip_unchanged_block_offsets=False, reuse_pool_views=False, block_ids_from_views=True + ) + copied_from_views = manager.get_block_ids_per_seq(request_ids) + + assert copied_from_views.dtype == built_from_lists.dtype == torch.int32 + assert copied_from_views.is_pinned() == built_from_lists.is_pinned() + assert torch.equal(copied_from_views, built_from_lists) + + +# --------------------------------------------------------------------------- +# Sub-page block table expanded inside the forward +# --------------------------------------------------------------------------- + + +def _subpage_metadata( + rows: int, max_blocks: int, factor: int +) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + metadata._msa_subpage_rows = rows # as prepare() stages them + metadata.msa_block_table = torch.randint( + 0, 1 << 20, (rows + 2, max_blocks), dtype=torch.int32, device="cuda" + ) + metadata.msa_subpage_block_table = torch.full( + (rows + 2, 2, max_blocks), -7, dtype=torch.int32, device="cuda" + ) + metadata._msa_subpages_per_slot = factor + metadata._msa_subpage_expanded = False + return metadata + + +def _expected_subpages(metadata, rows: int) -> torch.Tensor: + out = torch.empty_like(metadata.msa_subpage_block_table[:rows]) + write_subpage_block_table(metadata.msa_block_table[:rows], metadata._msa_subpages_per_slot, out) + return out + + +@requires_cuda +def test_the_first_reader_of_a_forward_expands_the_subpage_table(monkeypatch): + monkeypatch.setattr(msa_backend, "SUBPAGE_IN_GRAPH", True) + metadata = _subpage_metadata(rows=3, max_blocks=40, factor=2) + calls = [] + real = msa_backend.write_subpage_block_table + monkeypatch.setattr( + msa_backend, "write_subpage_block_table", lambda *a: (calls.append(1), real(*a)) + ) + + table, factor = metadata.msa_subpage_rows(0, 3) + again, _ = metadata.msa_subpage_rows(1, 3) + + assert factor == 2 and len(calls) == 1 + assert torch.equal(table, _expected_subpages(metadata, 3)) + assert torch.equal(again, table[1:3]) + # Rows past the step's requests are left alone. + assert metadata.msa_subpage_block_table[3:].eq(-7).all() + + +@requires_cuda +def test_each_forward_expands_the_current_block_table(monkeypatch): + monkeypatch.setattr(msa_backend, "SUBPAGE_IN_GRAPH", True) + metadata = _subpage_metadata(rows=2, max_blocks=16, factor=4) + metadata.msa_subpage_rows(0, 2) + metadata.msa_block_table[:2, :5] += 11 + # The next forward (on_update_kv_lens opens it). + metadata._msa_subpage_expanded = False + + table, _ = metadata.msa_subpage_rows(0, 2) + + assert torch.equal(table, _expected_subpages(metadata, 2)) + + +@requires_cuda +def test_a_captured_forward_expands_the_table_at_every_replay(monkeypatch): + monkeypatch.setattr(msa_backend, "SUBPAGE_IN_GRAPH", True) + metadata = _subpage_metadata(rows=4, max_blocks=32, factor=2) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + metadata._msa_subpage_expanded = False + metadata.msa_subpage_rows(0, 4) + metadata.msa_subpage_rows(2, 4) + torch.cuda.current_stream().wait_stream(stream) + + for step in range(3): + # The pack lands this step's block table; then the replay. + metadata.msa_block_table[:4].copy_( + torch.randint(0, 1 << 20, (4, 32), dtype=torch.int32, device="cuda") + ) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(metadata.msa_subpage_block_table[:4], _expected_subpages(metadata, 4)), ( + step + ) + + +def test_nothing_is_expanded_without_a_prepared_step(): + metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + metadata.msa_subpage_block_table = torch.arange(12, dtype=torch.int32).reshape(2, 2, 3) + metadata._msa_subpages_per_slot = 2 + + table, factor = metadata.msa_subpage_rows(0, 2) + + assert factor == 2 and torch.equal(table, torch.arange(12, dtype=torch.int32).reshape(2, 2, 3)) + + +# --------------------------------------------------------------------------- +# msa_out_cache_loc tail filled inside the forward +# --------------------------------------------------------------------------- + + +def _tail_metadata(cap: int, start) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = object.__new__(MiniMaxM3MsaSparseAttentionMetadata) + metadata.msa_out_cache_loc = torch.arange(100, 100 + cap, dtype=torch.int32, device="cuda") + metadata._msa_tail_start = start + return metadata + + +@requires_cuda +def test_the_forward_fills_only_past_the_live_slots(): + metadata = _tail_metadata(10, 4) + metadata._fill_tail_in_forward() + assert metadata.msa_out_cache_loc.tolist() == [100, 101, 102, 103] + [-1] * 6 + + untouched = _tail_metadata(6, None) + untouched._fill_tail_in_forward() + assert untouched.msa_out_cache_loc.tolist() == list(range(100, 106)) + + full = _tail_metadata(6, 6) + full._fill_tail_in_forward() + assert full.msa_out_cache_loc.tolist() == list(range(100, 106)) + + +@requires_cuda +def test_a_captured_forward_refills_its_tail_at_every_replay(): + """Another graph key sharing the pool block writes live slots past this key's tail between replays; each replay of + this key's forward must invalidate them again before its own readers run.""" + metadata = _tail_metadata(12, 5) + buf = metadata.msa_out_cache_loc + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + metadata._fill_tail_in_forward() + torch.cuda.current_stream().wait_stream(stream) + + for step in range(3): + buf[:5] = ( + torch.arange(5, dtype=torch.int32, device="cuda") + 10 * step + ) # this step's live slots + buf[5:9] = 777 # another key's in-graph writes into the shared block + graph.replay() + torch.cuda.synchronize() + assert buf.tolist() == [10 * step + i for i in range(5)] + [-1] * 7, step diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_skip_cache_routing.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_skip_cache_routing.py new file mode 100644 index 000000000000..9e50093d41ff --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_skip_cache_routing.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""No buffer a captured graph writes on the device is routed through the H2D skip cache (K3 audit #150, check #6). + +The skip cache (``pyexecutor/h2d_cache.py``) skips an upload when its host shadow says the destination already holds +the values. That is sound only while nothing else writes the destination: a graph replay that rewrites it on the device +(``on_update_kv_lens`` / ``msa.kv_update`` for the MSA lengths, slots and valid-block counts, ``_preprocess_inputs`` for +``kv_lens_cuda``) leaves the shadow describing values the device no longer holds. This test scans +``tensorrt_llm/_torch`` and fails if any of those buffers is the destination of a skipping helper; they have to go +through ``step_inputs.copy_`` or the never-skipping ``copy_h2d_pool_shared`` / ``fill_tail_pool_shared``. + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +# Buffers a captured forward writes on the device. +_DEVICE_WRITTEN = { + "msa_seq_lens_cuda", + "msa_cu_kv_lens", + "msa_out_cache_loc", + "msa_n_valid_blocks", + "_msa_prefill_n_valid_blocks", + "kv_lens_cuda", +} +# Helpers that may skip the write; their first argument is the destination. +_SKIPPING = {"copy_h2d_if_changed", "fill_tail_if_changed"} + + +def _receiver_names(node) -> set: + """Attribute / variable names along the receiver chain of an expression (``self.a[:n].view(...)`` -> {"a"}).""" + names = set() + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute): + names.add(node.attr) + node = node.func if isinstance(node, ast.Call) else node.value + if isinstance(node, ast.Name): + names.add(node.id) + return names + + +def _skipped_device_buffers(tree: ast.AST): + hits = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not node.args: + continue + f = node.func + name = f.id if isinstance(f, ast.Name) else (f.attr if isinstance(f, ast.Attribute) else "") + if name in _SKIPPING: + bad = _receiver_names(node.args[0]) & _DEVICE_WRITTEN + if bad: + hits.append((node.lineno, name, sorted(bad))) + return hits + + +def test_no_device_written_buffer_goes_through_the_skip_cache() -> None: + found = [] + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if not any(helper in src for helper in _SKIPPING): + continue + for line, helper, bad in _skipped_device_buffers(ast.parse(src)): + found.append(f"{path.relative_to(_ROOT).as_posix()}:{line} {helper}({', '.join(bad)})") + assert not found, "device-written buffers routed through the H2D skip cache:\n" + "\n".join( + found + ) + + +def test_the_scan_sees_a_skipped_device_buffer() -> None: + """The scan itself: a skipping helper writing a device-written buffer is reported, a never-skipping one is not.""" + bad = ast.parse( + "copy_h2d_if_changed(self.msa_seq_lens_cuda[:b], host)\nfill_tail_if_changed(m.msa_out_cache_loc, n, -1)" + ) + ok = ast.parse( + "copy_h2d_pool_shared(self.msa_seq_lens_cuda[:b], host)\ncopy_h2d_if_changed(self.msa_block_table[:b], h)" + ) + assert [h[1] for h in _skipped_device_buffers(bad)] == [ + "copy_h2d_if_changed", + "fill_tail_if_changed", + ] + assert _skipped_device_buffers(ok) == [] diff --git a/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py b/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py index a9e1c62b1e0e..83ab912fdad6 100644 --- a/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py +++ b/tests/unittest/_torch/attention/test_context_fmha_kernel_presence.py @@ -35,6 +35,10 @@ # A head width the SM100-family kernel table does carry for 16-bit paged-context # attention. _SUPPORTED_HEAD_SIZE = 64 +# The page size and head width of MiniMax-M3's dense and draft attention. The +# kernel table carries no 128-token-page kernels; SM107 generates them at runtime. +_NVRTC_TOKENS_PER_BLOCK = 128 +_NVRTC_HEAD_SIZE = 128 _NO_GPU = not torch.cuda.is_available() _NO_GPU_REASON = "needs a GPU; missing bindings on a GPU are a test failure" @@ -188,3 +192,30 @@ def test_lookup_is_page_size_sensitive(): ) assert present assert not odd_page + + +@pytest.mark.parametrize("kv_cache_dtype", [DataType.BF16, DataType.FP8], ids=["bf16", "fp8"]) +def test_lookup_reports_present_for_sm107_128_token_pages(kv_cache_dtype): + """SM107 serves dense paged context over 128-token pages with BF16 output. + + No precompiled kernel exists for that page size, so matched BF16/E4M3 Q + and KV with BF16 output must be generated at runtime rather than reported + absent, which would refuse paged-context attention for such models. The + runtime path is limited to that output precision: FP8 output at the same + page size has no kernel either way and must still report absent. + """ + if get_sm_version() != 107: + pytest.skip("runtime-generated 128-token-page context kernels are specific to SM107") + assert thop.fused_context_fmha_kernel_exists( + head_size=_NVRTC_HEAD_SIZE, + kv_cache_dtype=kv_cache_dtype, + tokens_per_block=_NVRTC_TOKENS_PER_BLOCK, + output_dtype=DataType.BF16, + ) + if kv_cache_dtype == DataType.FP8: + assert not thop.fused_context_fmha_kernel_exists( + head_size=_NVRTC_HEAD_SIZE, + kv_cache_dtype=kv_cache_dtype, + tokens_per_block=_NVRTC_TOKENS_PER_BLOCK, + output_dtype=DataType.FP8, + ) diff --git a/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py new file mode 100644 index 000000000000..479d5f279cfa --- /dev/null +++ b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""FallbackFmha's refusal of paged prefill at a KV manager's TRTLLM-Gen extra page sizes: kept on every SM but SM107, +where no TRTLLM-Gen FMHA exists and the thop paged-context kernel serves those pages (its SM107 P128 oracle is +test_modeling_v2_m3_draft_attention_paged_context.py). CPU only: the check runs before any kernel.""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.attention.backends.fmha import fallback +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase + +pytestmark = pytest.mark.cpu_only + + +class _PastRefusal(Exception): + """Raised by the stubbed paged-context validation that follows the refusal check.""" + + +def _fmha(): + fmha = object.__new__(fallback.FallbackFmha) + + def _validate(metadata): + raise _PastRefusal() + + fmha._validate_paged_context_fmha = _validate + return fmha + + +def _context_batch(tokens_per_block: int = 128, paged: bool = True): + return SimpleNamespace( + num_contexts=1, + kv_cache_manager=SimpleNamespace(trtllm_gen_extra_tokens_per_block=frozenset({128})), + use_paged_context_fmha=paged, + tokens_per_block=tokens_per_block, + ) + + +@pytest.mark.parametrize("sm", [90, 100, 103, 120]) +def test_extra_page_paged_prefill_is_refused_off_sm107(monkeypatch, sm) -> None: + monkeypatch.setattr(fallback, "get_sm_version", lambda: sm) + with pytest.raises(RuntimeError, match="cannot serve paged prefill with 128-token KV pages"): + _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT) + + +def test_extra_page_paged_prefill_is_served_on_sm107(monkeypatch) -> None: + monkeypatch.setattr(fallback, "get_sm_version", lambda: 107) + with pytest.raises(_PastRefusal): + _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT) + + +@pytest.mark.parametrize("sm", [100, 107]) +def test_other_page_sizes_and_unpaged_prefill_pass_the_check(monkeypatch, sm) -> None: + monkeypatch.setattr(fallback, "get_sm_version", lambda: sm) + for batch in (_context_batch(tokens_per_block=32), _context_batch(paged=False)): + with pytest.raises(_PastRefusal): + _fmha()._is_supported(None, None, None, batch, None, phase=FmhaPhase.CONTEXT) diff --git a/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py b/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py new file mode 100644 index 000000000000..319bb4abd8b6 --- /dev/null +++ b/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TRTLLM_AR_TUNE_NCCL_SYMMETRIC: the AllReduce tuner's candidate list with and without NCCL_SYMMETRIC.""" + +import pytest +import torch + +from tensorrt_llm._torch.autotuner import OptimizationProfile +from tensorrt_llm._torch.custom_ops import torch_custom_ops +from tensorrt_llm.functional import AllReduceStrategy + + +def _runner(): + runner = object.__new__(torch_custom_ops.AllReduceRunner) + runner.tp_size = 2 + return runner + + +@pytest.mark.parametrize("rows", [8, 4096, 16384]) +@pytest.mark.parametrize("tune_symmetric", [True, False]) +def test_nccl_symmetric_candidate_follows_knob( + monkeypatch, rows: int, tune_symmetric: bool +) -> None: + monkeypatch.setattr(torch_custom_ops, "_AR_TUNE_NCCL_SYMMETRIC", tune_symmetric) + inputs = [torch.empty(rows, 6144, dtype=torch.bfloat16)] + tactics = _runner().get_valid_tactics(inputs, OptimizationProfile()) + assert AllReduceStrategy.NCCL.value in tactics + assert (AllReduceStrategy.NCCL_SYMMETRIC.value in tactics) is tune_symmetric + # Off removes only NCCL_SYMMETRIC: the other candidates are unchanged. + monkeypatch.setattr(torch_custom_ops, "_AR_TUNE_NCCL_SYMMETRIC", True) + stock = _runner().get_valid_tactics(inputs, OptimizationProfile()) + assert [t for t in stock if t != AllReduceStrategy.NCCL_SYMMETRIC.value] == [ + t for t in tactics if t != AllReduceStrategy.NCCL_SYMMETRIC.value + ] diff --git a/tests/unittest/_torch/kernels/parallel/test_mxfp8_cute_pinned_tactics.py b/tests/unittest/_torch/kernels/parallel/test_mxfp8_cute_pinned_tactics.py new file mode 100644 index 000000000000..f9c38516d2fb --- /dev/null +++ b/tests/unittest/_torch/kernels/parallel/test_mxfp8_cute_pinned_tactics.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The pinned FlashInfer CuTe DSL MXFP8 GEMM tactics (custom_ops/mxfp8_cute_pinned_tactics): every entry is +non-split-K, each +pinned call equals FlashInfer's default (heuristic) CuTe DSL path bit for bit, the runner takes the pinned path above 32 +rows only, and the kill switch restores the default path.""" + +import hashlib + +import pytest +import torch + +from tensorrt_llm._torch.custom_ops import mxfp8_cute_pinned_tactics as pin +from tensorrt_llm._utils import get_sm_version + +sm107 = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 107, reason="SM107 (Rubin) only" +) + + +@pytest.fixture(autouse=True) +def _armed(monkeypatch): + """The table as an opted-in model's executor arms it (RuntimeHooks.mxfp8_pinned_tactics).""" + monkeypatch.setattr(pin, "_ARMED", [True]) + + +def _table_digest() -> str: + return hashlib.md5(repr(sorted(pin.TACTICS.items())).encode()).hexdigest()[:8] + + +def test_table_is_non_split_k_and_pinned(): + assert pin.TACTICS, "empty table" + for (n, k, m), (tile, cluster, swap, prefetch, split_k) in pin.TACTICS.items(): + assert split_k == 1, ((n, k, m), "a split-K tactic changes the K order (not bitwise)") + assert m >= 64 and m & (m - 1) == 0 and m <= pin.MAX_BUCKET, (n, k, m) + assert ( + len(tile) == 2 + and len(cluster) == 2 + and isinstance(swap, bool) + and isinstance(prefetch, bool) + ) + # Drift guard: a re-tuned table is a new measured artifact (bitwise + timing), not an edit in place. + assert _table_digest() == "d43d4a0d", _table_digest() + + +def test_bucket_and_row_gate(): + assert ( + pin._bucket(33) == 64 + and pin._bucket(64) == 64 + and pin._bucket(65) == 128 + and pin._bucket(184) == 256 + ) + assert pin.tactic_for(32, 4992, 6144) is None # 1..32 rows keep FlashInfer's own low-M path + + +def _operands(m, n, k, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, device="cuda", generator=g) * 0.5).to(torch.bfloat16) + w = (torch.randn(n, k, device="cuda", generator=g) * 0.05).to(torch.bfloat16) + q, s = torch.ops.trtllm.mxfp8_quantize(x, True) + wq, ws = torch.ops.trtllm.mxfp8_quantize(w, True) + return q, s, wq, ws + + +def _default(q, s, wq, ws): + from flashinfer import autotune, mm_mxfp8 + + with autotune(tune_mode=False, skip_ops="mxfp8_gemm"): + return mm_mxfp8( + q, wq.t(), s, ws, out_dtype=torch.bfloat16, use_8x4_sf_layout=False, backend="cute-dsl" + ) + + +@sm107 +@pytest.mark.parametrize("key", sorted(pin.TACTICS), ids=lambda k: f"N{k[0]}_K{k[1]}_M{k[2]}") +def test_pinned_equals_default_bitwise(key): + n, k, bucket = key + for m in sorted({bucket // 2 + 1 if bucket > 64 else 33, bucket}): # both ends of the bucket + q, s, wq, ws = _operands(m, n, k, seed=hash(key) % 1000 + m) + assert pin.tactic_for(m, n, k) == pin.TACTICS[key] + out = pin.run(q, s, wq, ws, torch.bfloat16) + assert out is not None, "the pinned path declined a table entry" + assert torch.equal(out, _default(q, s, wq, ws)), (key, m) + + +@sm107 +def test_runner_takes_pinned_above_32_rows_and_kill(monkeypatch): + from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + calls = [] + real = pin.run + monkeypatch.setattr(pin, "run", lambda *a, **kw: calls.append(a[0].shape[0]) or real(*a, **kw)) + for m in (16, 80): + q, s, wq, ws = _operands(m, 4992, 6144, seed=m) + out = runner.forward([q, s, wq, ws], tactic=FlashInferMXFP8GemmRunner.CUTE_DSL) + assert torch.equal(out, _default(q, s, wq, ws)), m + assert calls == [16, 80] + monkeypatch.setattr(pin, "ENABLED", False) + assert pin.tactic_for(80, 4992, 6144) is None diff --git a/tests/unittest/_torch/kernels/parallel/test_mxfp8_mxfp8_gemm.py b/tests/unittest/_torch/kernels/parallel/test_mxfp8_mxfp8_gemm.py index d6b489dad7b1..13b2e935d4c0 100644 --- a/tests/unittest/_torch/kernels/parallel/test_mxfp8_mxfp8_gemm.py +++ b/tests/unittest/_torch/kernels/parallel/test_mxfp8_mxfp8_gemm.py @@ -23,8 +23,9 @@ @pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="MXFP8 GEMM is supported on SM100 and SM103 only. Current SM is %d." % getSMVersion(), + getSMVersion() not in (100, 103, 107), + reason="MXFP8 GEMM is supported on SM100, SM103 and SM107 only. Current SM is %d." + % getSMVersion(), ) @pytest.mark.parametrize( "m,n,k", @@ -62,8 +63,8 @@ def test_mxfp8_mxfp8_gemm_large_m(m: int, n: int, k: int): @pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="MXFP8 tactic runner requires SM100 or SM103. Current SM is %d." % getSMVersion(), + getSMVersion() not in (100, 103, 107), + reason="MXFP8 tactic runner requires SM100, SM103 or SM107. Current SM is %d." % getSMVersion(), ) def test_mxfp8_mxfp8_runner_tactics(): """Every exposed tactic and the generic fallback produce aligned output.""" @@ -94,8 +95,8 @@ def test_mxfp8_mxfp8_runner_tactics(): @pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="MXFP8 tactic cache requires SM100 or SM103. Current SM is %d." % getSMVersion(), + getSMVersion() not in (100, 103, 107), + reason="MXFP8 tactic cache requires SM100, SM103 or SM107. Current SM is %d." % getSMVersion(), ) def test_mxfp8_mxfp8_native_tactic_cache(): """The direct op uses cached tactics and preserves the generic fallback.""" @@ -156,8 +157,8 @@ def test_mxfp8_mxfp8_native_tactic_cache(): @pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="MXFP8 tactic cache requires SM100 or SM103. Current SM is %d." % getSMVersion(), + getSMVersion() not in (100, 103, 107), + reason="MXFP8 tactic cache requires SM100, SM103 or SM107. Current SM is %d." % getSMVersion(), ) @pytest.mark.parametrize( "lower_bound,upper_bound,bucket", 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__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py index d3ac1acda903..980cd15a09ef 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__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py @@ -7,6 +7,7 @@ need the 4-node launch its docstring describes. """ +__extra_import_path__ = ["."] import _glm_rank_job import pytest import torch 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__tp16/spec/test_modeling_v2_draft_ehproj_table.py index b536ec7d560e..af00750b1940 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__tp16/spec/test_modeling_v2_draft_ehproj_table.py @@ -7,6 +7,7 @@ 16-rank run over four nodes. """ +__extra_import_path__ = ["."] import _draft_step_rank_job import pytest 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__tp16/spec/test_modeling_v2_draft_sample_glue.py index faf00183c753..d784f9480daf 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__tp16/spec/test_modeling_v2_draft_sample_glue.py @@ -7,6 +7,7 @@ 16-rank run over four nodes. """ +__extra_import_path__ = ["."] import _draft_step_rank_job import pytest diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..4052191ded74 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_silu_and_mul entry of the MiniMax-M3 tp2 target: SwiGLU-OAI and plain SwiGLU.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.activation.m3_silu_and_mul import ( # noqa: E501 + m3_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + # Gate and up spread past the clamp (|x| up to ~12) so both clamps engage. + x = (torch.randn(spec["m"], 2 * spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + return (x, *OAI) if spec["oai"] else (x, None, None, None) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few + fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); + measured on GB300: worst element 0.5 unit, relative RMS 0.01.""" + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(*args) + assert out.shape == (cell.spec["m"], cell.spec["inter"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=1.0, rms_ulp=0.1) + + +def test_the_gate_sees_a_missing_clamp() -> None: + """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" + spec = op.CELLS[0].spec + x, *oai = _build(dict(spec, m=64), 3) + out = op(x, *oai) + with pytest.raises(AssertionError): + assert_within_ulp(out, op.reference(x, None, oai[1], oai[2]), element_ulp=1.0, rms_ulp=0.1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_attn_runtime_tp2.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_attn_runtime_tp2.py new file mode 100644 index 000000000000..137230cd0345 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_attn_runtime_tp2.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real MiniMax-M3 attention runtime state at the TP2 per-rank geometry, for the backend-step entries' GPU tests. + +One hybrid NVFP4 KV cache manager (``MiniMaxM3KVCacheManagerV2``) laid out as the target's: a dense layer in the FP8 +half, a sparse MSA layer in the NVFP4 half with its FP8 index-K cache, and (``draft=True``) the one-model Eagle3 +drafter's shared layer, read through the manager's FP8 draft view as the runtime's draft swap reads it. Metadata is +the runtime's own type, filled by its own ``prepare()``; the backends are built the way the target builds them. + +The geometry is one rank's: 32 query heads over two KV heads and two index heads of dim 128 (one index-K head). +The mapping is a single +rank holding exactly those heads, so nothing is sharded again. + +History rows are written by the in-tree scatters (``fused_write_layer_caches`` for the FP8 layers, +``fused_write_layer_caches_nvfp4`` for the NVFP4 layer and its index-K), at the slots the manager assigned, which are +not in request order. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + MiniMaxM3MsaSparseAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, + fused_write_layer_caches_nvfp4, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HEADS = 32 +KV_HEADS = 2 +INDEX_HEADS = 2 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +DENSE, SPARSE, DRAFT = 0, 1, 2 +MAX_BATCH = 8 +MAX_SEQ_LEN = 8192 +MAX_NUM_TOKENS = 8192 +#: The Eagle3 drafter checkpoint's RoPE: NeoX over the whole head, theta 5e6. +DRAFT_ROPE_THETA = 5.0e6 + +#: E4M3-exact planted index levels, 1/8 apart in relative terms (0.25 .. 15): 48 distinct block scores. +LEVELS = tuple(2.0**e * (1 + m / 8) for e in range(-2, 4) for m in range(8)) +#: The index value of a block's non-peak rows; its score (8) is below every level's (>= 16). +LOW = 0.125 + + +def mapping() -> Mapping: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + +def sparse_attention_config() -> MiniMaxM3SparseAttentionConfig: + """The target's runtime: MSA, the FP8 indexer, the fused five-way projection (M3's checkpoint selection rule: + no forced initial blocks, the query's own block forced).""" + return MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_num_index_heads=INDEX_HEADS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + ) + + +def make_manager(*, draft: bool) -> MiniMaxM3KVCacheManagerV2: + torch.cuda.init() + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_tokens=4 * MAX_SEQ_LEN, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + event_buffer_max_size=0, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=[KV_HEADS] * (3 if draft else 2), + head_dim=HEAD_DIM, + tokens_per_block=PAGE, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + mapping=mapping(), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") if draft else None + ), + sparse_layer_ids=[SPARSE], + sparse_attention_config=sparse_attention_config(), + ) + + +def make_msa_backend(layer: int) -> MiniMaxM3MsaSparseAttention: + """One layer's MSA backend as ``MiniMaxM3Tp2ModelingV2`` builds it (NVFP4 KV quant config on every layer).""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + backend = create_attention( + "TRTLLM", + layer, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=None, + quant_config=quant, + q_scaling=1.0, + sparse_params=sparse_attention_config().to_sparse_params(layer_idx=layer), + ) + backend.update_quant_config(quant) + assert isinstance(backend, MiniMaxM3MsaSparseAttention), type(backend).__name__ + return backend + + +def make_draft_backend() -> TrtllmAttention: + """The drafter layer's attention as ``M3Eagle3Draft`` builds it: fused NeoX RoPE, an FP8 KV cache (the target + overrides the drafter's KV dtype to FP8), no other quantization.""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + backend = create_attention( + "TRTLLM", + DRAFT, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=HEAD_DIM, theta=DRAFT_ROPE_THETA, max_positions=MAX_SEQ_LEN), + ), + quant_config=quant, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + backend.update_quant_config(quant) + assert type(backend) is TrtllmAttention, type(backend).__name__ + return backend + + +def make_msa_metadata(manager) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = MiniMaxM3MsaSparseAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + max_num_sequences=MAX_BATCH, + kv_cache_manager=manager, + mapping=mapping(), + sparse_metadata_params=sparse_attention_config().to_sparse_metadata_params(), + ) + # The target's MSA preparation choice that changes what prepare() stages (the identity tables of short steps). + metadata.msa_policy = MsaRuntimePolicy(short_sequence_identity=True) + return metadata + + +def make_draft_metadata(manager) -> TrtllmAttentionMetadata: + """Drafter metadata over the shared manager's FP8 draft view, the manager the runtime's draft swap installs.""" + view = manager.get_draft_kv_cache_view() + assert view is not None and view.dtype == DataType.FP8 + return TrtllmAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=mapping(), + ) + + +def allocate(manager, request_ids: Sequence[int], qlens: Sequence[int], cached: Sequence[int]): + manager.add_dummy_requests(list(request_ids), token_nums=[c + n for c, n in zip(cached, qlens)]) + + +def prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + c + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +def flat_slots(manager, layer: int, request_id: int, positions: torch.Tensor) -> torch.Tensor: + """The cache slots (page slot id * PAGE + offset) of a request's positions in ``layer``.""" + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0], + dtype=torch.int64, + device="cuda", + ) + positions = positions.to(device="cuda", dtype=torch.int64) + return (pages[positions // PAGE] * PAGE + positions % PAGE).to(torch.int32) + + +def write_fp8_rows(manager, layer, request_id, positions, k, v) -> None: + """K / V rows [n, 128] (E4M3) into an FP8 layer's cache at ``positions``.""" + view = manager.get_buffers(layer, "HND") + assert view.dtype == torch.float8_e4m3fn + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches(view[:, 0], view[:, 1], None, slots, k, v, None) + + +def write_nvfp4_rows(manager, metadata, layer, request_id, positions, k, v, idx_k, quant) -> None: + """BF16 K / V rows quantized into the NVFP4 layer's cache, and index-K rows into its FP8 index cache.""" + data = manager.get_buffers(layer, "HND") + scales = manager.get_block_scale_buffers(layer, "HND") + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches_nvfp4( + data[:, 0], + data[:, 1], + scales[:, 0], + scales[:, 1], + metadata.msa_idx_k_cache(layer), + slots, + k, + v, + idx_k, + quant, + ) + + +def index_directions(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """Two orthogonal ±1 index directions (the halves of the 128 dims), E4M3-exact.""" + signs = torch.randint(0, 2, (HEAD_DIM,), generator=gen, device="cuda").float() * 2 - 1 + u0 = torch.zeros(HEAD_DIM, device="cuda") + u1 = torch.zeros(HEAD_DIM, device="cuda") + u0[: HEAD_DIM // 2] = signs[: HEAD_DIM // 2] + u1[HEAD_DIM // 2 :] = signs[HEAD_DIM // 2 :] + return u0, u1 + + +def planted_index_keys(length: int, u0, u1, gen: torch.Generator) -> torch.Tensor: + """Index-K rows [length, 128] (bf16, E4M3-exact) whose per-block maximum score is unambiguous. + + Every block has one peak row at a random position, ``a_b u0 + c_b u1`` with ``a`` and ``c`` two independent + permutations of ``LEVELS``; its other rows are ``LOW (u0 + u1)``. A query along ``u0`` then ranks the blocks by + ``a``, one along ``u1`` by ``c``, every score 1/8 or more apart from the next.""" + blocks = (length + PAGE - 1) // PAGE + assert blocks <= len(LEVELS) + levels = torch.tensor(LEVELS, device="cuda") + a = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + c = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + rows = (LOW * (u0 + u1)).expand(length, HEAD_DIM).clone() + for b in range(blocks): + span = min(PAGE, length - b * PAGE) + peak = b * PAGE + int(torch.randint(0, span, (1,), generator=gen, device="cuda")) + rows[peak] = a[b] * u0 + c[b] * u1 + return rows.to(torch.bfloat16) + + +def nvfp4_quant(*values: torch.Tensor) -> torch.Tensor: + """The (Q, K, V) quantization scales ``fused_write_layer_caches_nvfp4`` takes: each role's amax at the top of the + E2M1 x E4M3 range.""" + quant = torch.ones(3, dtype=torch.float32, device="cuda") + for role, x in enumerate(values, start=1): + quant[role] = (448.0 * 6.0) / x.abs().max().float() + return quant diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_nvfp4_cache.py new file mode 100644 index 000000000000..61886497be86 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/_m3_nvfp4_cache.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A MiniMax-M3 TP2 sparse-layer decode step over a packed NVFP4 cache, for the attention entries' tests. + +The cache is written by the in-tree scatter (``fused_write_layer_caches_nvfp4``, +which quantizes with ``torch.ops.trtllm.fp4_quantize``), into one +``[pages, 2 (K, V), 2 (KV heads), 128, .]`` pool per role exactly as ``msa_paged_kv`` and +``get_block_scale_buffers`` hand the K and V views over. Pages are shuffled so +a kernel that ignored the block table would read the wrong ones. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +NUM_HEADS = 32 +NUM_KV_HEADS = 2 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 + + +def token_kv_len(seq_len: int, decode_query_len: int, offset: int) -> int: + """The KV extent query token ``offset`` of a request attends, as every decode kernel derives it.""" + return max(seq_len - decode_query_len + offset + 1, 0) + + +def build_step( + seq_lens: Sequence[int], decode_query_len: int, *, identity: bool, seed: int +) -> SimpleNamespace: + """One decode step: E4M3 q [T, 32, 128], the packed NVFP4 K/V views and their scales, a [2, T, 16] + head-major top-k table (the identity table, or a random ascending selection with a -1 tail per KV head), the + block table and per-request lengths.""" + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = max(1, max((s + PAGE - 1) // PAGE for s in seq_lens)) + num_pages = batch * max_blocks + num_slots = num_pages * PAGE + block_table = ( + torch.randperm(num_pages, device="cuda", generator=g) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + data_pool = torch.zeros( + num_pages, 2, NUM_KV_HEADS, PAGE, HEAD_DIM // 2, dtype=torch.uint8, device="cuda" + ) + scale_pool = torch.zeros( + num_pages, 2, NUM_KV_HEADS, PAGE, HEAD_DIM // 16, dtype=torch.uint8, device="cuda" + ) + quant = torch.ones(3, dtype=torch.float32, device="cuda") + values = [] + for role in (0, 1): + x = torch.randn( + num_slots, NUM_KV_HEADS * HEAD_DIM, device="cuda", generator=g, dtype=torch.float32 + ).to(torch.bfloat16) + quant[role + 1] = (448.0 * 6.0) / x.abs().max().float().item() + values.append(x) + assert fused_write_layer_caches_nvfp4( + data_pool[:, 0], + data_pool[:, 1], + scale_pool[:, 0], + scale_pool[:, 1], + None, + torch.arange(num_slots, device="cuda", dtype=torch.int32), + values[0], + values[1], + None, + quant, + ) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) + topk = torch.full((NUM_KV_HEADS, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + nblocks = (token_kv_len(seq_lens[req], decode_query_len, offset) + PAGE - 1) // PAGE + for head in range(NUM_KV_HEADS): + if identity: + assert nblocks <= TOPK, "an identity step has every request within 16 blocks" + chosen = torch.arange(nblocks, device="cuda") + else: + chosen = torch.randperm(nblocks, device="cuda", generator=g)[:TOPK].sort().values + topk[head, tok, : chosen.numel()] = chosen.to(torch.int32) + return SimpleNamespace( + q=q, + k_paged=data_pool[:, 0], + v_paged=data_pool[:, 1], + k_block_scale=scale_pool[:, 0], + v_block_scale=scale_pool[:, 1], + k_global_scale=(1.0 / quant[1:2]).contiguous(), + v_global_scale=(1.0 / quant[2:3]).contiguous(), + topk=topk, + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device="cuda"), + decode_query_len=decode_query_len, + sm_scale=HEAD_DIM**-0.5, + source=values, + quant=quant, + num_slots=num_slots, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py new file mode 100644 index 000000000000..2dfd7aa71ebb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_block_table_contract.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The MSA page table is written on the host side only, before the graph replay. + +The index scorer's early issue (``ops/m3_index_score.py``, ``early_issue``) reads ``block_table[b, split]`` BEFORE +its PDL grid-dependency wait and loads that page's K tile right after it. That is exact only while no kernel of the +step writes ``MsaMetadata.msa_block_table``: today its one writer is the host-to-device copy in +``_build_msa_fields`` (called from ``prepare()``, before the replay), besides the allocation. This test scans +``tensorrt_llm/_torch`` and fails on any other write site, so a new in-graph writer has to revisit the early issue. + +(``msa_seq_lens_cuda`` IS patched in-graph by ``on_update_kv_lens``; the scorer reads it before the wait only as a hint +and bounds every load and store by the value it reads after the wait.) + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +_NAME = "msa_block_table" +# In-place tensor methods that write their receiver. +_MUTATORS = { + "copy_", + "fill_", + "zero_", + "index_put_", + "index_copy_", + "scatter_", + "masked_fill_", + "add_", + "sub_", +} +# Helpers whose first argument is the destination they write. +_COPY_HELPERS = {"copy_h2d_if_changed", "copy_h2d_pool_shared"} +# (file relative to tensorrt_llm/_torch, enclosing function) -> what it may do +_ALLOWED = { + ( + "attention/backends/sparse/minimax_m3/msa_backend.py", + "_create_msa_buffers", + ), # the allocation + ( + "attention/backends/sparse/minimax_m3/msa_backend.py", + "_build_msa_fields", + ), # the H2D copy in prepare() +} + + +def _mentions(node) -> bool: + """Whether the expression's receiver chain (attributes / subscripts / slices) is ``*.msa_block_table``.""" + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute) and node.attr == _NAME: + return True + node = node.func if isinstance(node, ast.Call) else node.value + return isinstance(node, ast.Name) and node.id == _NAME + + +def _write_sites(tree: ast.AST): + sites = [] + + def visit(node, func): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + func = node.name + if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + if any(_mentions(t) for t in targets) and not isinstance(node, ast.AnnAssign): + sites.append((func, node.lineno)) + if isinstance(node, ast.Call): + f = node.func + if isinstance(f, ast.Attribute) and f.attr in _MUTATORS and _mentions(f.value): + sites.append((func, node.lineno)) + name = ( + f.id + if isinstance(f, ast.Name) + else (f.attr if isinstance(f, ast.Attribute) else "") + ) + if name in _COPY_HELPERS and node.args and _mentions(node.args[0]): + sites.append((func, node.lineno)) + for child in ast.iter_child_nodes(node): + visit(child, func) + + visit(tree, None) + return sites + + +def test_only_prepare_writes_the_msa_page_table() -> None: + found = set() + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if _NAME not in src: + continue + rel = path.relative_to(_ROOT).as_posix() + for func, line in _write_sites(ast.parse(src)): + assert (rel, func) in _ALLOWED, ( + f"{rel}:{line} ({func}) writes msa_block_table; the index scorer's early issue reads it before its " + "PDL grid wait (ops/m3_index_score.py), so an in-graph writer must be ordered before it or the " + "early page read moved after the wait" + ) + found.add((rel, func)) + assert found == _ALLOWED, f"expected write sites {_ALLOWED}, found {found}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..0ebd8285aa9e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7): the TP2 dense layers' FP8 q/k/v (minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8) carries the bits +of the stock path, the BF16 fused_qk_norm_rope output cast to E4M3, for every norm / RoPE variant: Q and K are rounded +to BF16 before E4M3 (K is stored in the FP8 cache), V is cast.""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers the trtllm ops) + +assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) + +HQ, HK, HV, D, ROT = 32, 2, 2, 128, 64 + + +@pytest.mark.parametrize("tokens", [1, 4, 8, 4096]) +@pytest.mark.parametrize("use_gemma", [False, True], ids=["rmsnorm", "gemma"]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_dense_fp8_qkv_is_the_bf16_kernel_cast(tokens, use_gemma, is_neox) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8") + g = torch.Generator(device="cuda").manual_seed(tokens * 4 + int(use_gemma) * 2 + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + ours = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8( + qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, use_gemma + ) + stock = qkv.clone() + torch.ops.trtllm.fused_qk_norm_rope( + stock, HQ, HK, HV, D, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, + 1.0, 0.0, 0.0, 1.0, True, use_gemma, False, 0, 0, + ) # fmt: skip + assert torch.equal(ours.view(torch.uint8), stock.to(torch.float8_e4m3fn).view(torch.uint8)) + + +@pytest.mark.parametrize("tokens,live", [(1, 1), (4, 4), (8, 5), (64, 64)]) +@pytest.mark.parametrize("tokens_per_block", [32, 128]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_fused_kv_insert_stores_the_bytes_of_write_kv_slots_and_the_compact_q( + tokens, live, tokens_per_block, is_neox +) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + g = torch.Generator(device="cuda").manual_seed(tokens * 1000 + tokens_per_block + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + pages, layers, layer = 16, 3, 1 + # The layer's HND view of a multi-layer pool, strided as get_buffers returns it. + raw = torch.randint( + 0, + 256, + (pages, layers, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.randperm(pages * tokens_per_block, device="cuda", generator=g)[:tokens].to( + torch.int32 + ) + slots[live:] = -1 # padded rows own no slot + + def view(t): + return t.view(torch.float8_e4m3fn)[:, layer].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, False) + q_compact = torch.empty(tokens, HQ, D, device="cuda", dtype=torch.float8_e4m3fn) + ours = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, live, q_compact + ) + plain = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + assert torch.equal(ours.view(torch.uint8), plain.view(torch.uint8)) + assert torch.equal( + q_compact.view(tokens, HQ * D).view(torch.uint8), plain[:, : HQ * D].view(torch.uint8) + ) + + ref_raw = raw.clone() + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D) + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D) + ref = view(ref_raw) + write_kv_slots(ref[:, 0], slots, k, live, layout="HND") + write_kv_slots(ref[:, 1], slots, v, live, layout="HND") + assert torch.equal(ours_raw, ref_raw) + + +def test_the_fused_kv_insert_skips_a_negative_slot() -> None: + """A live row whose slot is the -1 sentinel stores nothing (no write before the pool row).""" + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + tokens, tokens_per_block, pages = 4, 32, 8 + g = torch.Generator(device="cuda").manual_seed(7) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + raw = torch.randint( + 0, + 256, + (pages, 3, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.tensor([5, -1, 40, 77], device="cuda", dtype=torch.int32) + + def view(t): + return t.view(torch.float8_e4m3fn)[:, 1].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, True, pos, False) + torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, tokens + ) + plain = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + keep = torch.tensor([0, 2, 3], device="cuda") + ref_raw = raw.clone() + ref = view(ref_raw) + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D)[keep] + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D)[keep] + write_kv_slots(ref[:, 0], slots[keep], k, 3, layout="HND") + write_kv_slots(ref[:, 1], slots[keep], v, 3, layout="HND") + assert torch.equal(ours_raw, ref_raw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..e3de00587e67 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_attention entry: the Eagle3 drafter layer's attention against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp2``): the target's hybrid NVFP4 +``MiniMaxM3KVCacheManagerV2`` with its one-model Eagle3 draft layer, the requests allocated in it, drafter +``TrtllmAttentionMetadata`` over the manager's FP8 draft view (the manager the runtime's draft swap installs) filled +by its ``prepare()``, and the drafter's ``TrtllmAttention`` built as ``M3Eagle3Draft`` builds it (fused NeoX RoPE at +the drafter checkpoint's theta, FP8 KV). Earlier rows go into the draft layer's cache through the in-tree scatter at +the slots the manager assigned; the call rotates and appends the step's own. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. +# The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: +# it takes the rotated Q in E4M3 (rounding the reference's Q to E4M3 takes the RMS gap from ~9 to ~7 ulps) and rounds +# the softmax probabilities before the PV product. Measured on GB300 over these cells (max element / RMS): 13.6 / 8.6 +# (4 x 4 rows), 9.9 / 7.2 (8 x 1), 20.4 / 7.7 (512 context rows). The band leaves about 1.5x; a reference shifted +# by one row is far outside it. +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 + +# Cached tokens per request of each cell; the new rows per request come from the cell. +_CACHED = { + (2, 4): (200, 1000), + (2, 1): (127, 2047), + (1, None): (0,), +} +#: Rows past the live ones in the fused QKV, as a padded step hands them over. +_PAD = 3 + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + +# Non-paged context FMHA on SM107 does not apply RoPE (this fixture's path; the served drafter runs the paged context +# FMHA, which rotates): the context cell is a strict expected failure until the runtime rotates there too. +_UNROTATED_CONTEXT = pytest.mark.xfail( + strict=True, + reason="non-paged context FMHA on SM107 does not apply RoPE (shared runtime)", +) +_CELLS = [ + pytest.param(c, marks=_UNROTATED_CONTEXT) if c.spec.get("context") else c for c in op.CELLS +] + + +@pytest.mark.parametrize("cell", _CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + context = spec.get("context") + cached = _CACHED[spec["reqs"], spec.get("qlen")] + qlens = (context,) if context else (spec["qlen"],) * spec["reqs"] + assert len(cached) == spec["reqs"] + contexts = 1 if context else 0 + rids = list(range(spec["reqs"])) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = rt.make_draft_metadata(manager) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + assert int(metadata.num_tokens) == total + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + certifying(op) + out = op(backend, qkv, metadata) + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{cell.why}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + group = rt.HEADS // rt.KV_HEADS + for h in range(rt.KV_HEADS): + cols = slice(h * group * rt.HEAD_DIM, (h + 1) * group * rt.HEAD_DIM) + element, rms = _ulp_stats(out[:, cols], ref[:, cols]) + print(f"[{cell.why}] KV head {h}: max element {element:.2f} ulp, rms {rms:.3f} ulp") + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + # The call's append: each live row's rotated K and its V at the row's position, E4M3 (unit KV scale), to + # within one E4M3 step of the float64 rotation (the kernel rotates in FP32). + pool = manager.get_buffers(rt.DRAFT, "HND") + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + for h in range(rt.KV_HEADS): + k_cols = q_size + h * rt.HEAD_DIM + v_cols = q_size + (rt.KV_HEADS + h) * rt.HEAD_DIM + want_k = _rope(rows[:, k_cols : k_cols + rt.HEAD_DIM].float(), positions) + want_v = rows[:, v_cols : v_cols + rt.HEAD_DIM].to(torch.float8_e4m3fn) + got_k = pool[slots // rt.PAGE, 0, h, slots % rt.PAGE].to(torch.float64) + got_v = pool[slots // rt.PAGE, 1, h, slots % rt.PAGE] + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), ( + f"K head {h} is appended rotated, in E4M3" + ) + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)), f"V head {h}" + first += n + + # The band discriminates: a reference whose rows are shifted by one token is outside it. + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py new file mode 100644 index 000000000000..d401a9d532e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's context attention on SM107 as the server runs it: paged context FMHA (chunked prefill +and speculative draft tokens set ``use_paged_context_fmha``) over the shared manager's FP8 draft view at its +128-token pages, which SM107 serves through the thop paged-context kernel (no TRTLLM-Gen FMHA there). + +Each case is checked against the exact float64 reference of ``m3_draft_attention`` (``test_modeling_v2_m3_draft_ +attention.py``'s band), with fresh contexts, chunked appends over a cached prefix (page-aligned or not), two contexts +in one step and a mixed step with generation rows beside the context. The appended K / V rows and a shifted-reference +discriminator are checked as there. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._utils import get_sm_version + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp # noqa: E402 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" +pytestmark = pytest.mark.skipif( + get_sm_version() != 107, reason="the SM107 thop paged-context path (no TRTLLM-Gen FMHA)" +) + +# test_modeling_v2_m3_draft_attention.py's band (BF16 ulps of each row's largest magnitude / of the relative RMS). +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 +_PAD = 3 + +# (case id, cached tokens per request, new rows per request, context requests first). +_CASES = [ + ("fresh_ctx_512", (0,), (512,), 1), + ("append_1000_unaligned", (1000,), (700,), 1), + ("append_1024_aligned", (1024,), (300,), 1), + ("two_ctx_prefix_384", (0, 384), (257, 129), 2), + ("prefix_ctx_beside_verify", (500, 200, 1000), (200, 4, 4), 1), + ("fresh_ctx_beside_decode", (0, 200, 1000), (200, 1, 1), 1), + ("fresh_ctx_beside_verify", (0, 200, 1000), (200, 4, 4), 1), +] + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _paged_draft_metadata(manager) -> TrtllmAttentionMetadata: + """``rt.make_draft_metadata`` with the server's runtime features (chunked prefill, Eagle3 draft tokens).""" + view = manager.get_draft_kv_cache_view() + metadata = TrtllmAttentionMetadata( + max_num_requests=rt.MAX_BATCH, + max_num_tokens=rt.MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=rt.mapping(), + runtime_features=AttentionRuntimeFeatures( + chunked_prefill=True, has_speculative_draft_tokens=True + ), + ) + assert metadata.use_paged_context_fmha + return metadata + + +def _prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + """``rt.prepare`` as the serving runner prepares a step: a context request's prompt length is its CURRENT CHUNK's + token count (``prompt_lengths.append(len(prompt_tokens))``) with the prefix in ``num_cached_tokens_per_seq``, so + the chunk is appended at the cached length; generation rows as ``rt.prepare``.""" + from tensorrt_llm._torch.metadata import KVCacheParams + + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +@pytest.mark.parametrize("case", _CASES, ids=[c[0] for c in _CASES]) +def test_paged_context_at_128_token_pages(case) -> None: + name, cached, qlens, contexts = case + rids = list(range(len(qlens))) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(name.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = _paged_draft_metadata(manager) + assert metadata.tokens_per_block == rt.PAGE == 128 + assert 128 in manager.get_draft_kv_cache_view().trtllm_gen_extra_tokens_per_block + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + pool = manager.get_buffers(rt.DRAFT, "HND") + prefix_slots = { + rid: rt.flat_slots(manager, rt.DRAFT, rid, torch.arange(c, device="cuda")).long() + for rid, c in zip(rids, cached) + if c + } + prefix_before = { + rid: pool[sl // rt.PAGE, :, :, sl % rt.PAGE].clone() for rid, sl in prefix_slots.items() + } + _prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + out = op(backend, qkv, metadata) + # The cached prefix's K / V rows are left as they were (the call appends only the step's rows). + torch.cuda.synchronize() + changed = { + rid: int( + ( + pool[sl // rt.PAGE, :, :, sl % rt.PAGE].view(torch.uint8) + != prefix_before[rid].view(torch.uint8) + ) + .any(dim=-1) + .sum() + ) + for rid, sl in prefix_slots.items() + } + print(f"\n[{name}] cached prefix rows x heads x K/V changed by the call: {changed}") + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{name}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + # Every live row's K / V landed in its page and slot (the append of a chunk past a cached prefix). + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + for h in range(rt.KV_HEADS): + v_cols = q_size + (rt.KV_HEADS + h) * rt.HEAD_DIM + want_v = rows[:, v_cols : v_cols + rt.HEAD_DIM].to(torch.float8_e4m3fn) + got_v = pool[slots // rt.PAGE, 1, h, slots % rt.PAGE] + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)), f"V head {h}" + first += n + + assert not any(changed.values()), f"the call overwrote cached prefix K / V rows: {changed}" + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py new file mode 100644 index 000000000000..6a42e13000de --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for ops/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V +append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row +first draft pass and the 1-row later passes of 1-2 requests. + +Lengths sweep the KV tile edges each query row in turn ends on (the rows of a Q CTA with no key in its trailing +splits), with mixed per-request lengths; every call is held to the fp32 reference over the cache after the append +(E4M3 query, unit KV scales, no RoPE) within the FP8-KV band, finite, for both split-KV reducers. +""" + +import math +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) + +assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +FP8_TOL = 6e-2 +# Persistent like the FMHA manager's workspace and the spec-decoding metadata of a real step: the in-tree preprocess +# kernels must not see these move between calls. +_PERSISTENT = {} + + +def _persistent(name: str, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _edge_lens(plan, qlen: int): + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +def _library(qlen: int, nreq: int, own_reducer: bool): + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = own_reducer + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = torch.cuda.get_device_properties(0).multi_processor_count + lib._on_taken = lambda rows: None + # The in-tree preprocess pass (its RoPE-free cache append); the one-launch preprocess has its own bitwise test. + lib._fused_rope_mode = None + lib._pdl = False + lib._on_rope_taken = None + return lib + + +def _attn(qlen: int): + return SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=0, theta=10000.0, scale_type=0, scale=1.0, max_positions=1024 + ), + ) + + +def _run(lib, qlen: int, kv_lens, seed: int, spec: bool = True): + """One drafter generation pass over a fresh FP8 cache; returns (output, fp32 reference). ``spec`` passes the + spec-decoding lengths/positions (a multi-token generation step); without them a multi-row pass is the drafter's + plain generation step of input_seq_length rows, as it runs in-model.""" + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g).clamp_(-4, 4) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + for r, n in enumerate( + kv_lens + ): # shift the tail third of V, so a missed trailing split fails the band + start = 2 * n // 3 + for pg in range(start // PAGE, math.ceil(n / PAGE)): + pool[int(perm[r, pg]), 1, :, (start % PAGE if pg == start // PAGE else 0) :] += 1.0 + pool = pool.to(torch.float8_e4m3fn) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + pointers = torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64) + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=dev, generator=g) + seq = torch.tensor(kv_lens, dtype=torch.int32, device=dev) + one = torch.ones(1, dtype=torch.float32, device=dev) + multi = qlen > 1 and spec + out = torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=_attn(qlen), + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=pointers, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=seq, + spec_decoding_generation_lengths=( + _persistent( + f"gen_lengths{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=( + _persistent( + f"position_offsets{qlen}x{nreq}", + lambda: ( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous() + ), + ) + if multi + else None + ), + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=out, + ) + lib.run_generation(params) + torch.cuda.synchronize() + # fp32 reference over the cache after the append: E4M3 query, unit scales, causal over the step's rows. + q = qkv[:, : H * D].float().view(nreq, qlen, H, D).to(torch.float8_e4m3fn).float() + flat = pool.float() + ref = torch.empty(nreq, qlen, H, D, device=dev) + for r, n in enumerate(kv_lens): + blocks = perm[r, : math.ceil(n / PAGE)] + k = flat[blocks, 0].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + v = flat[blocks, 1].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + k = k.repeat_interleave(H // KVH, dim=0) + v = v.repeat_interleave(H // KVH, dim=0) + s = torch.einsum("qhd,hkd->hqk", q[r], k) * D**-0.5 + visible = ( + torch.arange(n, device=dev)[None, :] + < (n - qlen + 1 + torch.arange(qlen, device=dev))[:, None] + ) + s = s.masked_fill(~visible[None], float("-inf")) + ref[r] = torch.einsum("hqk,hkd->qhd", torch.softmax(s, dim=-1), v) + return out.float().view(nreq, qlen, H, D), ref + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_the_drafter_pass_at_tile_edges(qlen, nreq) -> None: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if fullctx.has_own_reducer() else (False,) + for own in reducers: + fullctx.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_LEN, own, get_env_enable_pdl() + ) + lens = _edge_lens(fullctx.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + lens = [n for n in lens if n > qlen] + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for i, kv_lens in enumerate(batches): + # The first draft pass runs in-model as a plain multi-row generation step; every other batch also takes the + # spec-decoding form. + spec = qlen > 1 and i % 2 == 1 + for own in reducers: + got, ref = _run(_library(qlen, nreq, own), qlen, kv_lens, sum(kv_lens), spec=spec) + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +def _drafter_model(scale: float): + """A drafter whose one attention module holds a TrtllmAttention backend as ``attn`` (the backend is no module).""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + backend = TrtllmAttention.__new__(TrtllmAttention) + backend.kv_cache_scaling_factor = torch.full((1,), scale, device="cuda") + backend.is_mla_enable, backend.kv_lora_rank, backend.v_head_dim, backend.head_dim = ( + False, + None, + None, + D, + ) + backend._fmha_manager = SimpleNamespace(fmha_libs=[]) + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.zeros(1, device="cuda")) + model.self_attn = torch.nn.Module() + model.self_attn.attn = backend + return model, backend + + +def test_install_declines_a_scaled_kv_cache() -> None: + """At unit KV scales the arithmetic is the stock decode's; a scaled cache keeps the stock path.""" + reasons = [] + model, layer = _drafter_model(0.5) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 0 + assert layer._fmha_manager.fmha_libs == [] and reasons == ["a scaled FP8 KV cache"] + model, layer = _drafter_model(1.0) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 1 + assert isinstance(layer._fmha_manager.fmha_libs[0], drafter.M3DrafterFullCtxFmha) + + +def test_a_cache_bound_past_the_plan_capacity_keeps_the_stock_path() -> None: + """The plans are compiled for ``max_kv_len`` of K/V: a step whose cache length bound is longer falls through.""" + + from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase + + class _Backend: + num_heads, num_kv_heads, head_dim, quant_mode = H, KVH, D, FP8_KV_QUANT_MODE + kv_scale_orig_quant = kv_scale_quant_orig = None + + backend = _Backend() + backend.kv_scale_orig_quant, backend.kv_scale_quant_orig = object(), object() + reasons = [] + lib = _library(1, 1, False) + lib._on_uncovered = reasons.append + lib._attn_ref = weakref.ref(backend) + fwd = SimpleNamespace( + attention_input_type=None, + kv_scale_orig_quant=backend.kv_scale_orig_quant, + kv_scale_quant_orig=backend.kv_scale_quant_orig, + ) + meta = SimpleNamespace(num_contexts=0, num_generations=1, max_seq_len=MAX_LEN + 8) + q = torch.empty(1, H * D, device="cuda") + assert not lib._is_supported(q, None, None, meta, fwd, phase=FmhaPhase.GENERATION) + assert reasons == ["the cache length bound is beyond the plans' K/V capacity"] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py new file mode 100644 index 000000000000..3d77b932b3e7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -0,0 +1,406 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for minimax_m3_sm107tp2_drafter_rope_fp8_kv, the TP2 drafter layer's one-launch generation +preprocess: bitwise against the in-tree ``thop.trtllm_gen_generation_preprocess`` followed by the torch E4M3 cast +of its BF16 query, on the query bytes and on every byte of the paged cache, with GPT-NeoX RoPE from a cos/sin table +(full and partial rotary width), multi-token (spec-decoding positions) and plain passes, one and two requests, and +query values past the E4M3 range. The library's fused pass is also held bitwise to its preprocess pass end to end. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp2_drafter_rope_fp8_kv"): + pytest.skip("minimax_m3_sm107tp2_drafter_rope_fp8_kv is not built", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +ROPE_GPT_NEOX = 2 +_PERSISTENT = {} + + +def _persistent(name, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _case(qlen, kv_lens, rotary_dim, multi, seed): + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = (torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g) * 2).to( + torch.float8_e4m3fn + ) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, device=dev, generator=g) * 4 + # A few query / key values near and past the E4M3 range (both casts saturate at 448). + hot = torch.rand(qkv.shape, device=dev, generator=g) < 0.01 + qkv = torch.where( + hot, qkv.sign() * (430 + 80 * torch.rand(qkv.shape, device=dev, generator=g)), qkv + ) + qkv = qkv.to(torch.bfloat16).contiguous() + max_pos = max(kv_lens) + 64 + cos_sin = torch.randn(max_pos, rotary_dim // 2, 2, device=dev, generator=g).float().contiguous() + attn = SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=ROPE_GPT_NEOX, + rotary_inv_freq=None, + rotary_cos_sin=cos_sin, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=rotary_dim, theta=10000.0, scale_type=0, scale=1.0, max_positions=max_pos + ), + ) + one = torch.ones(1, dtype=torch.float32, device=dev) + spec_offsets = None + if multi: + spec_offsets = _persistent( + f"pos{qlen}x{nreq}", + lambda: torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous(), + ) + # Spec positions are not the plain row index: shift them per request. + spec_offsets.copy_( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1) + + torch.arange(nreq, dtype=torch.int32, device=dev)[:, None] + ) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=attn, + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=torch.tensor(kv_lens, dtype=torch.int32, device=dev), + spec_decoding_generation_lengths=( + _persistent( + f"len{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=spec_offsets, + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev), + ) + return params, pool + + +def _preprocess(params, mpc): + """run_generation's in-tree preprocess call and the E4M3 cast of its query.""" + attn, meta, fwd, rope = params.attn, params.meta, params.fwd, params.attn.rope_params + q_processed = thop.trtllm_gen_generation_preprocess( + params.qkv_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope.dim, + rope.theta, + int(rope.scale_type), + rope.scale, + rope.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + False, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + mpc, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + )[0] + rows = params.num_tokens * attn.num_heads * attn.head_dim + return ( + q_processed.flatten()[:rows] + .to(torch.float8_e4m3fn) + .view(params.num_tokens, attn.num_heads, attn.head_dim) + ) + + +def _op(params, pool, mode): + attn = params.attn + nreq, qlen = params.batch_size, params.input_seq_length + kv_pages = pool.view(-1, KVH, PAGE, D) + tables = params.meta.kv_cache_block_offsets[0, :nreq] + q = torch.empty( + params.num_tokens, + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=pool.device, + ) + multi = params.spec_decoding_generation_lengths is not None and qlen > 1 + torch.ops.trtllm.minimax_m3_sm107tp2_drafter_rope_fp8_kv( + params.qkv_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + attn.rope_params.dim, + PAGE, + mode, + True, + ) + return q + + +# One-request multi-token passes keep the in-tree preprocess (it reads an unbuilt token map), so the spec cases hold +# two requests. +CASES = [ + (qlen, kv_lens, rotary, multi) + for qlen, multi in ((1, False), (4, True), (4, False)) + for kv_lens in (((4, 130) if multi else (4,)) if qlen > 1 else (130,), (8191, 131073)) + for rotary in (128, 64) +] + + +def _bytes(t): + return t.contiguous().view(torch.uint8) + + +def _matches(case, seed): + """Per rope mode: (query bytes equal, cache bytes equal) against the in-tree preprocess + cast.""" + qlen, kv_lens, rotary, multi = case + kv_lens = [max(n, qlen) for n in kv_lens] + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(qlen, kv_lens, rotary, multi, seed) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + pool_ref = pool.clone() + out = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + out[mode] = ( + bool(torch.equal(_bytes(q), _bytes(q_ref))), + bool(torch.equal(_bytes(pool), _bytes(pool_ref))), + ) + return out + + +@pytest.mark.parametrize( + "case", + CASES, + ids=lambda c: f"q{c[0]}-{'x'.join(map(str, c[1]))}-r{c[2]}-{'spec' if c[3] else 'plain'}", +) +def test_the_pinned_rope_mode_is_bitwise(case) -> None: + result = _matches(case, seed=7) + print(f"[rope modes] {case}: {result}", flush=True) + assert result[modeling.DRAFTER_ROPE_MODE] == (True, True), result + + +def test_the_fused_pass_matches_the_preprocess_pass() -> None: + mpc = torch.cuda.get_device_properties(0).multi_processor_count + for qlen, nreq in sorted(modeling.FULLCTX_DECODE_RECIPES): + for multi in (False, True) if qlen > 1 else (False,): + kv_lens = [70001 + 517 * r for r in range(nreq)] + outs = [] + for fused in (False, True): + params, pool = _case(qlen, kv_lens, 128, multi, seed=11) + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = True + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = mpc + lib._fused_rope_mode = modeling.DRAFTER_ROPE_MODE if fused else None + lib._pdl = True + lib._one_request_multi = drafter.ROPE_ONE_REQUEST_MULTI + taken = [] + lib._on_rope_taken = lambda rows, what: taken.append(what) + lib._on_taken = lambda rows: None + fullctx.prebuild(torch.device("cuda"), lib._recipes, MAX_LEN, True, True) + lib.run_generation(params) + torch.cuda.synchronize() + assert bool(taken) == ( + fused and (lib._one_request_multi or not (multi and nreq == 1)) + ) + outs.append((params.output.clone(), pool.clone())) + assert torch.equal(_bytes(outs[0][1]), _bytes(outs[1][1])), (qlen, nreq, multi, "cache") + assert torch.equal(_bytes(outs[0][0]), _bytes(outs[1][0])), ( + qlen, + nreq, + multi, + "output", + ) + + +def _round_bf16(x): + u = x.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000 + return u.to(torch.int64).to(torch.int32).view(torch.float32) + + +def _discriminating(n, seed): + """(cos, sin, x, pair) quadruples whose rotated first-half output's E4M3 byte differs between the rope modes + (fp32 CPU emulation; the GPU comparison is the verdict).""" + g = torch.Generator().manual_seed(seed) + keep = [] + while sum(k.shape[0] for k in keep) < n: + m = 1 << 20 + c, s = torch.randn(m, generator=g), torch.randn(m, generator=g) + x = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + xp = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + sn = -s + r0 = (c * x) + (sn * xp) + r1 = ((c.double() * x.double()) + (sn * xp).double()).float() + r2 = ((sn.double() * xp.double()) + (c * x).double()).float() + b = [ + _round_bf16(r.contiguous()).to(torch.float8_e4m3fn).view(torch.uint8) + for r in (r0, r1, r2) + ] + diff = (b[0] != b[1]) | (b[0] != b[2]) | (b[1] != b[2]) + diff &= (r0.abs() < 400) & (r1.abs() < 400) + keep.append(torch.stack([c, s, x, xp], 1)[diff]) + return torch.cat(keep)[:n] + + +def test_crafted_inputs_pin_the_rope_mode() -> None: + """Rotations whose E4M3 byte depends on how the two products are summed: the in-tree preprocess and exactly the + pinned mode agree on every one.""" + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(1, [130], 128, False, seed=3) + half = 64 + quads = _discriminating(H * half, seed=5).cuda() + pos = 129 # the one row's rope position: cache length 130, one query row + table = params.attn.rotary_cos_sin.view(-1, half, 2) + # One (cos, sin) per dim: every head shares the position's table row, so the heads take the same dims' quads with + # their own (x, pair). + table[pos, :, 0] = quads[:half, 0] + table[pos, :, 1] = quads[:half, 1] + qkv = params.qkv_input.view(-1, H + 2 * KVH, D) + for h in range(H): + qkv[0, h, :half] = quads[:half, 2].to(torch.bfloat16) + qkv[0, h, half:] = quads[:half, 3].to(torch.bfloat16) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + agree = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + agree[mode] = int((_bytes(q)[0, :, :half] == _bytes(q_ref)[0, :, :half]).sum()) + print( + f"[rope modes crafted] agreeing first-half query bytes of {H * half}: {agree}", flush=True + ) + assert agree[modeling.DRAFTER_ROPE_MODE] == H * half, agree + + +def test_one_request_multi_token_matches_the_guarded_preprocess() -> None: + """The drafter's first pass at one request (several tokens): the in-tree preprocess places those rows by a + workspace token map it does not build for one request, which the TP2 token-map guard (ops/m3_tokmap_guard) + writes right before the call. After POISONING the map with a two-request call, the fused op must be bitwise the + guarded preprocess + cast; this is what drafter.ROPE_ONE_REQUEST_MULTI rests on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_tokmap_guard as guard, + ) + + guard.install(torch.device("cuda"), 64, lambda rows: None) + mpc = torch.cuda.get_device_properties(0).multi_processor_count + results = {} + for kv in (130, 8191, 131073, 165521): + poison, _ = _case(4, [131072, 4000], 128, True, seed=11) + _preprocess(poison, mpc) + params, pool = _case(4, [kv], 128, True, seed=5) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + pool_ref = pool.clone() + pool.copy_(start) + q = _op(params, pool, modeling.DRAFTER_ROPE_MODE) + torch.cuda.synchronize() + results[kv] = ( + bool(torch.equal(_bytes(q), _bytes(q_ref))), + bool(torch.equal(_bytes(pool), _bytes(pool_ref))), + ) + print(f"[one-request multi-token vs guarded preprocess] {results}", flush=True) + if drafter.ROPE_ONE_REQUEST_MULTI: + assert all(v == (True, True) for v in results.values()), results diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..918174c1295f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp8_indexer_qk_norm_rope entry of the MiniMax-M3 tp2 target: index-Q out and index-K in the +cache against the in-place fused norm + RoPE followed by the E4M3 cast, byte for byte up to the 0.1% budget the +in-tree test of the kernel allows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 + m3_fp8_indexer_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" + +DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 + + +def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: + assert a.shape == b.shape and a.dtype == b.dtype == torch.float8_e4m3fn + differ = int((a.view(torch.uint8) != b.view(torch.uint8)).sum()) + assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + m, heads = cell.spec["m"], cell.spec["heads"] + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + qk = torch.randn(m, (heads + 1) * DIM, generator=g, device="cuda").to(torch.bfloat16) + qw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (m,), generator=g, device="cuda", dtype=torch.int32) + pages = (m + PAGE - 1) // PAGE + 1 + cache = torch.zeros(pages, 1, PAGE, DIM, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(pages * PAGE, generator=g, device="cuda")[:m].to(torch.int32) + before = qk.clone() + args = (qk, cache, slots, heads, DIM, ROTARY, EPS, qw, kw, THETA, pos) + with validating(op): + q = op(*args) + assert torch.equal(qk, before), "the BF16 input is left untouched" + q_ref, k_ref = op.reference(*args) + _bytes_close(q.reshape(m, heads, DIM), q_ref) + flat = cache.view(-1, DIM) + _bytes_close(flat[slots.long()], k_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py new file mode 100644 index 000000000000..adfd40f1d941 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_saturate.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""paged_cache.fp8_e4m3_saturate (FP8_SATURATE): the E4M3 cache writes and the transient FP8 query casts saturate to ++-448 like the CUDA writers (__NV_SATFINITE); every finite in-range value keeps the plain cast's bits, values already +in E4M3 pass through, and the kill switch is exactly the plain cast. Where torch's cast already saturates (this image) +the conversion is the cast alone, one kernel; otherwise it clamps first, the same bytes.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import paged_cache + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fp8_e4m3_saturate_matches_satfinite(dtype, monkeypatch) -> None: + g = torch.Generator(device="cuda").manual_seed(0) + inside = (torch.randn(4096, generator=g, device="cuda") * 100).clamp(-448, 448).to(dtype) + plain = inside.to(torch.float8_e4m3fn) + assert torch.equal( + paged_cache.fp8_e4m3_saturate(inside).view(torch.uint8), plain.view(torch.uint8) + ) + edge = torch.tensor( + [449.0, 463.0, 464.0, 1e4, float("inf"), -500.0, -float("inf")], dtype=dtype, device="cuda" + ) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0], device="cuda") + assert torch.equal(paged_cache.fp8_e4m3_saturate(edge).float(), want) + nan = paged_cache.fp8_e4m3_saturate(torch.tensor([float("nan")], dtype=dtype, device="cuda")) + assert torch.isnan(nan.float()).all() + already = edge.clamp(-448, 448).to(torch.float8_e4m3fn) + assert paged_cache.fp8_e4m3_saturate(already) is already + monkeypatch.setattr(paged_cache, "FP8_SATURATE", False) + killed = paged_cache.fp8_e4m3_saturate(edge) + assert torch.equal(killed.view(torch.uint8), edge.to(torch.float8_e4m3fn).view(torch.uint8)) + + +def test_cache_write_saturates() -> None: + """The cache-dtype conversion every write_kv_slots branch applies (compared as bytes: torch has no float8 + index_copy_ kernel), and an E4M3 source passing through untouched.""" + values = torch.tensor([[1.0, 500.0, -1e5, float("inf")]], dtype=torch.bfloat16, device="cuda") + out = paged_cache._to_cache_dtype(values, torch.float8_e4m3fn) + assert out.dtype == torch.float8_e4m3fn + want = torch.tensor([[1.0, 448.0, -448.0, 448.0]], device="cuda").to(torch.float8_e4m3fn) + assert torch.equal(out.view(torch.uint8), want.view(torch.uint8)) + assert paged_cache._to_cache_dtype(out, torch.float8_e4m3fn) is out + bf = paged_cache._to_cache_dtype(values, torch.bfloat16) + assert bf.dtype == torch.bfloat16 and torch.equal(bf, values) + + +_EDGES = [ + 448.0, + 449.0, + 463.0, + 464.0, + 1e4, + float("inf"), + -448.0, + -449.0, + -464.0, + -1e4, + -float("inf"), + float("nan"), + 0.0, +] + + +def _cuda_kernels(fn) -> int: + from torch.profiler import ProfilerActivity, profile + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + return sum(1 for e in prof.events() if e.device_type.name == "CUDA") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("cast_saturates", [True, False], ids=["cast_only", "clamp_fallback"]) +def test_fp8_e4m3_saturate_bytes_equal_clamp_then_cast(dtype, cast_saturates, monkeypatch) -> None: + """Both paths give the bytes of clamp(+-448) followed by the cast, edge values included.""" + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", cast_saturates) + x = torch.tensor(_EDGES, dtype=dtype, device="cuda") + want = x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + assert torch.equal(paged_cache.fp8_e4m3_saturate(x).view(torch.uint8), want.view(torch.uint8)) + + +def test_fp8_e4m3_saturate_takes_the_cast_only_path_here(monkeypatch) -> None: + """This image's torch cast saturates, on the CPU (the import-time probe) and on CUDA, so the conversion is one + kernel; with the probe failing it falls back to clamp + cast, two kernels.""" + assert paged_cache._CAST_SATURATES + x = torch.tensor(_EDGES, dtype=torch.bfloat16, device="cuda") + assert torch.equal( + x.to(torch.float8_e4m3fn).view(torch.uint8), + x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).view(torch.uint8), + ) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 1 + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", False) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 2 + + +@pytest.mark.parametrize("device", ["meta", "cuda"]) +def test_cast_probe_ignores_the_default_device(device) -> None: + """The import-time probe runs on the CPU even under a default-device context, as when the module is first + imported while a model is built on the meta device.""" + with torch.device(device): + assert paged_cache._cast_saturates() == paged_cache._CAST_SATURATES diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py new file mode 100644 index 000000000000..6a21693bcf9a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for the attention/m3_fullctx_decode entry: the dense layers' full-context FP8 decode on PrimTS +split-KV decode at the recipe the target pins per (query rows per request, requests). + +Every cell runs against the fp32 reference (q rounded to E4M3 first, as the kernel runs it), eager and replayed from a +CUDA graph (the same bits), and twice (the same bits). The K/V pages are read through the strided views of a flat +sub-page pool, as the KV cache manager hands them out. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_fullctx_decode import ( # noqa: E501 + HEAD_DIM, + HEADS, + KV_HEADS, + MAX_KV_LEN, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_fullctx_decode import ( # noqa: E501 + m3_fullctx_decode as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) + +PAGE = 128 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the entry asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor, subpages: int): + self.pool = pool + self.subpages = subpages + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, self.subpages + + +def _inputs(kv_lens, qlen: int, seed: int, subpages: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * subpages, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1, so a split that missed the tail fails the band. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + for page in range(start // PAGE, (n + PAGE - 1) // PAGE): + v = int(block_table[b, page]) * subpages + 1 + pool[v, :, (start % PAGE if page == start // PAGE else 0) :] += 1.0 + # q as the layer hands it over on decode steps: the FP8 query columns of the fused norm/RoPE kernel's + # [q | k | v] output (row stride wider than q), so no cast makes it compact. + qkv = torch.randn( + (batch * qlen, (HEADS + 2 * KV_HEADS) * HEAD_DIM), device="cuda", generator=g + ).to(torch.float8_e4m3fn) + q = qkv[:, : HEADS * HEAD_DIM].view(batch * qlen, HEADS, HEAD_DIM) + seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32) + return q, _Pool(pool.to(torch.float8_e4m3fn), subpages), block_table, seq_lens + + +@pytest.mark.parametrize("subpages", [2, 3], ids=["kv", "kv+index"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, subpages, certifying) -> None: # noqa: F811 + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], qlen, zlib.crc32(cell.why.encode()) % 2**31, subpages + ) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + certifying(op) + ref = op.reference(*args) + eager = op(*args).clone() + torch.testing.assert_close(eager.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + again = op(*args).clone() + assert torch.equal(again, eager), "not deterministic" + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args) + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager), "the CUDA-graph replay differs from the eager call" + + +def test_the_cells_cover_the_target_rows() -> None: + certified = {(c.spec["qlen"], len(c.spec["kv_lens"])) for c in op.CELLS} + assert set(modeling.FULLCTX_DECODE.keys) <= certified + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: + """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once + in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + + if not kernel.has_own_reducer(): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_fullctx_reduce") + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs(spec["kv_lens"], qlen, 11, 2) + outs = [] + for own in (True, False): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + op(q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe, own_reducer=own) + outs.append(out.float()) + ours, primts = outs + _, exp = torch.frexp(primts.abs()) + spacing = torch.where( + primts == 0, + torch.full_like(primts, 2.0**-133), + torch.ldexp(torch.ones_like(primts), exp - 8), + ) + assert bool(((ours - primts).abs() <= spacing).all()) + + +def _tile_edge_kv_lens(plan, qlen: int): + """KV lengths around every kind of boundary the split ranges have, so each query row of the Q CTA in turn ends + its keys on a KV tile edge (k * tile + d, d in -1 .. qlen - 1): the first tiles, the reducer's 8-group stride, the + lengths where a split's range grows (splits * insts tiles) and a split edge past it.""" + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: + """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys + (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode + step first hit (NaN rows from exp2(-inf - -inf)).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, False, get_env_enable_pdl() + ) + lens = _tile_edge_kv_lens(kernel.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, sum(kv_lens), 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: + """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just + past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + lengths = [(1 << 20) + 1, (1 << 20) + 8] + ([1_500_001] if nreq == 1 else []) + batches = [[n] * nreq if nreq == 1 else [n, (1 << 20) - 3] for n in lengths] + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, 13, 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, own, get_env_enable_pdl() + ) + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_identity_selection.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_identity_selection.py new file mode 100644 index 000000000000..0c2af37bf5ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_identity_selection.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the short-sequence identity selection (runtime policy ``short_sequence_identity``) is the stock table. + +On a step whose every query token has at most 16 valid blocks, this target skips the index scorer and the select and +stages the table (i < n_valid ? i : -1). Stock runs the in-tree CuTe DSL scorer and ``select_blocks_from_maxscore``, +which returns ascending ids and drops a valid block only when its max score is -inf. The scorer's scores are FP8 +products summed in FP32 (|score| <= 448^2 * 128), so -inf comes only from the causal mask, which n_valid excludes. +This test drives the real scorer with extreme FP8 operands at the TP2 geometry and checks both claims. +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) +from tensorrt_llm._torch.custom_ops import ( + cute_dsl_custom_ops, # noqa: F401 — registers the scorer op +) + +PAGE, DIM, TOPK, HEADS, QLEN = 128, 128, 16, 2, 4 +E4M3_MAX = 448.0 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + + +def _fill(kind, shape, gen): + if kind == "random": + x = torch.randn(shape, generator=gen, device="cuda") * 64 + elif kind == "max": + x = torch.full(shape, E4M3_MAX, device="cuda") + elif kind == "negmax": + x = torch.full(shape, -E4M3_MAX, device="cuda") + elif kind == "alternating": + x = torch.full(shape, E4M3_MAX, device="cuda") + x[..., 1::2] = -E4M3_MAX + elif kind == "zero": + x = torch.zeros(shape, device="cuda") + else: # the smallest E4M3 subnormal + x = torch.full(shape, 2.0**-9, device="cuda") + return x.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + + +_KINDS = ("random", "max", "negmax", "alternating", "zero", "tiny") + + +@pytest.mark.parametrize( + "seq_lens", [(4,), (129,), (2048,), (700, 2048), (5, 1500, 2047, 300), (1, 128, 256, 1024)] +) +@pytest.mark.parametrize("q_kind", _KINDS) +@pytest.mark.parametrize("k_kind", _KINDS) +def test_stock_selection_of_a_short_step_is_the_identity_table(seq_lens, q_kind, k_kind) -> None: + gen = torch.Generator(device="cuda").manual_seed(0) + reqs, total_q = len(seq_lens), len(seq_lens) * QLEN + nblk = [(s + PAGE - 1) // PAGE for s in seq_lens] + pages = sum(nblk) + 4 + perm = torch.randperm(pages, generator=gen, device="cuda").int() + block_table = torch.zeros((reqs, max(nblk)), dtype=torch.int32, device="cuda") + start = 0 + for r, n in enumerate(nblk): + block_table[r, :n] = perm[start : start + n] + start += n + n_valid = torch.tensor( + [(max(s - QLEN + o + 1, 0) + PAGE - 1) // PAGE for s in seq_lens for o in range(QLEN)], + dtype=torch.int32, + device="cuda", + ) + idx_q = _fill(q_kind, (total_q, HEADS, DIM), gen) + idx_k = _fill(k_kind, (pages, PAGE, DIM), gen) + max_score = torch.full((HEADS, max(nblk), total_q), float("-inf"), device="cuda") + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, + idx_k, + block_table, + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + max_score.transpose(1, 2), + QLEN, + ) + for tok, nv in enumerate(n_valid.tolist()): + assert bool(torch.isfinite(max_score[:, :nv, tok]).all()), ( + f"token {tok}: a causally valid block scored non-finite" + ) + table = select_blocks_from_maxscore( + max_score, + topk=TOPK, + n_valid_blocks=n_valid, + init_blocks=1, + local_blocks=4, + head_major_output=True, + ).to(torch.int32) + # The selector returns [tokens, kv_heads, topk]; the decode reads it as [kv_heads, tokens, topk] (layers.py). + table = table.permute(1, 0, 2) + i = torch.arange(TOPK, dtype=torch.int32, device="cuda") + identity = torch.where( + i[None, :] < n_valid[:, None], i[None, :], torch.full_like(i[None, :], -1) + ) + assert torch.equal(table, identity[None].expand(HEADS, total_q, TOPK)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..85c5e6b746c1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score entry: the MSA index scorer with its pre-wait schedule. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score`` at its own 256-CTA split) bit for bit, which the +pre-wait prefetch, the early tile issue and the cell's split must not change. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_index_score import ( # noqa: E501 + PAGE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 2 # index heads per rank at TP2 +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + assert spec["heads"] == HEADS + args = (idx_q, index_k, block_table, seq, score, qlen, spec["split_k"], spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, 1536, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4 + ) + assert torch.equal(score, stock) + + +@pytest.mark.parametrize("qlen", [1, 2, 3, 4]) +def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> None: + """Lengths around the edge between pre-staged (all-old) and post-wait first blocks, 1 and 2 requests.""" + for k in (1, 2, 1294): + for length in (k * PAGE - qlen - 1, k * PAGE - qlen, k * PAGE - qlen + 1, k * PAGE + 1): + for seq_lens, split in (((length,), 1536), ((length, max(qlen, length // 3)), 768)): + idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) + score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) + op(idx_q, index_k, block_table, seq, score, 4, split, True) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + assert torch.equal(score, stock), (qlen, seq_lens) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, 1536, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..4f7e66dd8910 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_msa_attention entry: one layer's MSA backend step against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp2``): a hybrid NVFP4 ``MiniMaxM3KVCacheManagerV2`` +with the requests allocated in it, ``MiniMaxM3MsaSparseAttentionMetadata`` filled by its ``prepare()``, and the layer's +``MiniMaxM3MsaSparseAttention`` built as the target builds it. Earlier rows (and, on a sparse layer, the step's own +rows, which the target's producer writes before the call) go into the cache through the in-tree scatters, at the +slots the manager assigned. The sparse layers' index-K rows are planted so every block's score is 1/8 or more apart +from every other's: the top-k selection the reference recomputes is then unambiguous. + +The cells are the target's calls (``_attention_core`` in modeling.py): a dense layer with live bf16 K / V in a +prefill chunk and live FP8 K / V in a decode step (the dense layers sit in the hybrid cache's FP8 half), a sparse +layer whose producer already wrote the caches (``k = v = None``) in a prefill chunk beyond 16 blocks and in a mixed +step, both with the NVFP4 scales the target passes. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_msa_attention import ( # noqa: E501 + m3_msa_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" + +# Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, +# per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 +# FMHA does), so what is left is each FMHA's own arithmetic, measured on GB300 over these cells (max element / RMS): +# * fmha_sm100 over the FP8 pool (dense context rows): 10.1 / 6.1. It runs E4M3 x E4M3 and its softmax +# probabilities are rounded before the PV product. +# * trtllm-gen over the FP8 pool (dense generation rows): 7.0 / 6.1. It rounds the probabilities to E4M3 (scaled by +# 448 against the running maximum): emulating that brings a single-tile request to 0.6 RMS. +# * the NVFP4 CSR prefill (sparse context rows): 34.4 / 11.8 at 4096 rows, 26.8 / 9.1 at 512. It stages the +# reconstructed K / V in E4M3 (the in-tree A/B test, test_minimax_m3_sparse_attn_decode, notes it and holds it +# to 3e-2). +# * the Triton NVFP4 sparse decode (sparse generation rows): 1.0 / 0.25 -- the band of its own entry +# (m3_sparse_decode): exact operands, FP16 probabilities, BF16 rounding. +# Each band leaves about 1.5x the measured maximum; a shifted or wrongly selected reference is far outside them. +_FP8_CONTEXT = dict(element_ulp=16.0, rms_ulp=9.0) +_FP8_DECODE = dict(element_ulp=12.0, rms_ulp=9.0) +_NVFP4_CONTEXT = dict(element_ulp=48.0, rms_ulp=16.0) +_NVFP4_DECODE = dict(element_ulp=2.0, rms_ulp=0.5) +_BANDS = { + ("dense", "context"): (_FP8_CONTEXT, None), + ("dense", "generation"): (None, _FP8_DECODE), + ("sparse", "context"): (_NVFP4_CONTEXT, None), + ("sparse", "mixed"): (_NVFP4_CONTEXT, _NVFP4_DECODE), +} + + +def _assert_in_bands(output, ref, bands, context_rows: int) -> None: + """Gate the context rows and the generation rows each on their own kernel's band.""" + context_band, decode_band = bands + if context_rows: + assert_within_ulp(output[:context_rows], ref[:context_rows], **context_band) + if context_rows < output.shape[0]: + assert_within_ulp(output[context_rows:], ref[context_rows:], **decode_band) + + +# (request rows, cached tokens per request, new tokens per request, context rows) per cell. +_SHAPES = { + ("dense", "context"): ((384,), (1024,), 1), + ("dense", "generation"): ((128, 2500), (1, 1), 0), + ("sparse", "context"): ((0,), (4096,), 1), + ("sparse", "mixed"): ((0, 2999, 4500), (512, 1, 1), 1), +} + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """The two quantities ``assert_within_ulp`` gates, in BF16 ulps.""" + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _shift(x: torch.Tensor) -> torch.Tensor: + """``x``'s rows moved down by one (byte-wise, so E4M3 rows move too).""" + return x.view(torch.uint8).roll(1, 0).view(x.dtype) if x.element_size() == 1 else x.roll(1, 0) + + +def _shape(cell): + spec = cell.spec + cached, qlens, contexts = _SHAPES[spec["layer"], spec["phase"]] + if spec["phase"] == "context": + assert sum(qlens) == spec["tokens"] + elif spec["phase"] == "generation": + assert len(qlens) == spec["reqs"] + else: + assert sum(qlens[:contexts]) == spec["tokens"] and len(qlens) - contexts == spec["reqs"] + return cached, qlens, contexts + + +def _run_dense(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + rt.write_fp8_rows(manager, rt.DENSE, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + # A prefill chunk arrives in BF16 (the hybrid cache's dense layers take FP8 Q / K / V only on decode steps). + dtype = torch.bfloat16 if contexts else torch.float8_e4m3fn + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + k = torch.randn(total, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + v = torch.randn(total, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + # The NVFP4 target passes its layer scales on every layer; an FP8 layer's write and FMHA take none. + inv = torch.tensor([1.0, 0.5, 2.0], device="cuda") + fwd = 1.0 / inv + return (q, k, v, None, inv, fwd), dict(rids=rids, cached=cached, qlens=qlens, contexts=contexts) + + +def _run_sparse(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + u0, u1 = rt.index_directions(gen) + lengths = [c + n for c, n in zip(cached, qlens)] + keys = [ + torch.randn(n, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths + ] + values = [ + torch.randn(n, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths + ] + quant = rt.nvfp4_quant(torch.cat(keys), torch.cat(values)) + for rid, n, kk, vv in zip(rids, lengths, keys, values): + # Every row, the step's own included: the producer writes a sparse layer's caches before the call. + idx_k = rt.planted_index_keys(n, u0, u1, gen) + rt.write_nvfp4_rows( + manager, + metadata, + rt.SPARSE, + rid, + torch.arange(n), + kk.to(torch.bfloat16), + vv.to(torch.bfloat16), + idx_k, + quant, + ) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + # Query tokens alternate between the two planted directions, so neighbouring tokens select different blocks; the + # rank's index heads take alternating directions, so with more than one KV head they select different blocks too. + turn = torch.arange(total, device="cuda") % 2 + directions = torch.stack([u0, u1]) + idx_q = torch.cat([directions[(turn + h) % 2] for h in range(rt.INDEX_HEADS)], dim=1) + idx_q = idx_q.to(torch.float8_e4m3fn) + return (q, None, None, idx_q, quant, 1.0 / quant), dict( + rids=rids, cached=cached, qlens=qlens, contexts=contexts + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + sparse = spec["layer"] == "sparse" + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=False) + try: + backend = rt.make_msa_backend(rt.SPARSE if sparse else rt.DENSE) + metadata = rt.make_msa_metadata(manager) + run = _run_sparse if sparse else _run_dense + (q, k, v, idx_q, inv, fwd), step = run(cell, gen, manager, backend, metadata) + output = torch.full( + (q.shape[0], rt.HEADS * rt.HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + certifying(op) + result = op(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert result is output + ref = op.reference(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert torch.isfinite(output.float()).all(), "every live row is written" + bands = _BANDS[spec["layer"], spec["phase"]] + context_rows = sum(step["qlens"][: step["contexts"]]) + for name, rows in ( + ("context", slice(0, context_rows)), + ("generation", slice(context_rows, None)), + ): + if output[rows].numel(): + element, rms = _ulp_stats(output[rows], ref[rows]) + print( + f"\n[{cell.why}] {name} rows: max element {element:.2f} ulp, rms {rms:.3f} ulp" + ) + _assert_in_bands(output, ref, bands, context_rows) + + if k is not None: + # The call's own cache write: the live K / V rows land, as a plain E4M3 cast, at their positions. + pool = manager.get_buffers(rt.DENSE, "HND") + first = 0 + for rid, c, n in zip(step["rids"], step["cached"], step["qlens"]): + slots = rt.flat_slots(manager, rt.DENSE, rid, torch.arange(c, c + n)).long() + for role, src in ((0, k), (1, v)): + for h in range(rt.KV_HEADS): + got = pool[slots // rt.PAGE, role, h, slots % rt.PAGE] + want = src[first : first + n, h * rt.HEAD_DIM : (h + 1) * rt.HEAD_DIM] + want = want.to(torch.float8_e4m3fn) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose query rows are shifted by one token is outside it. + shifted = op.reference(backend, _shift(q), k, v, idx_q, metadata, output, sparse, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, shifted, bands, context_rows) + if sparse: + # And the selection is load-bearing: attending every causal block instead is outside it too. + dense = op.reference(backend, q, k, v, idx_q, metadata, output, False, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, dense, bands, context_rows) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py new file mode 100644 index 000000000000..b883a15acd38 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_n_valid_contract.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-token valid-block count is written only before the step's first layer. + +The long-row select (``ops/m3_select_long_dsl.py``, ``SELECT_N_VALID_PREWAIT``) reads its row's entry of +``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the +buffer happens before the step's first layer: + +- the allocation (``_alloc_msa_proxy_scratch``); +- the host-to-device copy in ``prepare()`` (``_build_step_plans``); +- ``on_update_kv_lens``, which the runner calls before the forward. It writes through the plain + ``<<<>>>`` launch of ``msa_kv_update_kernel`` or through torch copies, so any PDL grid of the step starts only + after it completes. + +This test scans ``tensorrt_llm/_torch`` and fails on any other write site (including writes through a local alias of +the buffer), so a new writer has to revisit the pre-wait read. A buffer handed to an opaque kernel launch is not +tracked; such a launch inside a layer must be reviewed by hand. + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +_NAME = "msa_n_valid_blocks" +# In-place tensor methods that write their receiver. +_MUTATORS = { + "copy_", + "fill_", + "zero_", + "index_put_", + "index_copy_", + "scatter_", + "masked_fill_", + "add_", + "sub_", +} +# Calls whose first argument is the destination they write: ``step_inputs.copy_(dst, src)`` and the H2D helpers. +_COPY_HELPERS = {"copy_", "copy_h2d_if_changed", "copy_h2d_pool_shared"} +_BACKEND = "attention/backends/sparse/minimax_m3/msa_backend.py" +# (file relative to tensorrt_llm/_torch, enclosing function) -> writes allowed there +_ALLOWED = { + (_BACKEND, "_alloc_msa_proxy_scratch"), # the allocation + (_BACKEND, "_build_step_plans"), # the H2D copy in prepare() + (_BACKEND, "on_update_kv_lens"), # the device-side correction before the forward +} + + +def _base(node): + """The innermost receiver of an attribute / subscript / call chain.""" + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute) and node.attr == _NAME: + return node + node = node.func if isinstance(node, ast.Call) else node.value + return node + + +def _mentions(node, aliases) -> bool: + """Whether the expression's receiver chain is ``*.msa_n_valid_blocks`` or a local alias of it.""" + b = _base(node) + if isinstance(b, ast.Attribute): + return b.attr == _NAME + return isinstance(b, ast.Name) and (b.id == _NAME or b.id in aliases) + + +def _refers(node, aliases) -> bool: + """Whether any sub-expression (e.g. a branch of a conditional) names the buffer or an alias.""" + return any( + (isinstance(n, ast.Attribute) and n.attr == _NAME) + or (isinstance(n, ast.Name) and (n.id == _NAME or n.id in aliases)) + for n in ast.walk(node) + ) + + +def _write_sites(tree: ast.AST): + sites = [] + + def visit(node, func, aliases): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + func, aliases = node.name, set() + if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + value = node.value + for t in targets: + if isinstance(t, ast.Name) and value is not None and _refers(value, aliases): + aliases.add(t.id) # a local view of the buffer + elif ( + not isinstance(t, ast.Name) + and _mentions(t, aliases) + and not isinstance(node, ast.AnnAssign) + ): + sites.append((func, node.lineno)) + elif ( + isinstance(t, ast.Attribute) + and t.attr == _NAME + and isinstance(node, ast.AnnAssign) + ): + pass # a dataclass field declaration + if isinstance(node, ast.Call): + f = node.func + if isinstance(f, ast.Attribute) and f.attr in _MUTATORS and _mentions(f.value, aliases): + sites.append((func, node.lineno)) + name = ( + f.id + if isinstance(f, ast.Name) + else (f.attr if isinstance(f, ast.Attribute) else "") + ) + if name in _COPY_HELPERS and node.args and _mentions(node.args[0], aliases): + sites.append((func, node.lineno)) + for child in ast.iter_child_nodes(node): + visit(child, func, aliases) + + visit(tree, None, set()) + return sites + + +def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: + found = set() + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if _NAME not in src: + continue + rel = path.relative_to(_ROOT).as_posix() + for func, line in _write_sites(ast.parse(src)): + assert (rel, func) in _ALLOWED, ( + f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " + "wait (ops/m3_select_long_dsl.py, SELECT_N_VALID_PREWAIT), so a writer inside the " + "step must be ordered before the select's grid or the read moved after the wait" + ) + found.add((rel, func)) + assert found == _ALLOWED, f"expected write sites {_ALLOWED}, found {found}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..dcdfc973c8ba --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py @@ -0,0 +1,218 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_producer entry: the pre-wait-staged NVFP4 QKV + indexer producer at the TP2 heads. + +Every cell is held to the native producer (``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert``) bit for bit on q, +index-q and every K / V / scale / index-K cache byte, and to the fp32 torch reference: the FP8 outputs within one E4M3 +step, the decoded NVFP4 K / V within one E2M1 step of their block scale. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_producer import ( # noqa: E501 + FP4_BLOCK, + HEADS_INDEX, + HEADS_KV, + HEADS_PACKED, + HEADS_Q, + PAGE, + decode_nvfp4_token, + scale_index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_producer import ( # noqa: E501 + m3_producer as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_producer requires a CUDA device" + +PAGES = 1024 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _caches(): + return ( + torch.zeros(PAGES, 2, HEADS_KV, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, HEADS_KV, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _inputs(tokens: int, scale: float, bad_slots: bool, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + weights = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + packed = ( + torch.randn(tokens, HEADS_PACKED * 128, device="cuda", generator=g) * scale + ).bfloat16() + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return packed, slots, quant, weights, rotary, positions + + +def _args(inputs, caches, early: bool = True): + packed, slots, quant, weights, rotary, positions = inputs + return (packed, *caches, slots, quant, EPS, *weights, rotary, positions, HEADS_Q, True, early) + + +def _native(inputs, caches): + packed, slots, quant, weights, rotary, positions = inputs + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, + *caches, + slots, + quant, + HEADS_Q, + HEADS_KV, + HEADS_INDEX, + 128, + 64, + EPS, + *weights, + rotary, + positions, + ) + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + early = spec.get("early", True) + inputs = _inputs( + spec["tokens"], spec["scale"], spec.get("bad_slots", False), zlib.crc32(cell.why.encode()) + ) + stock_caches, caches = _caches(), _caches() + stock = _native(inputs, stock_caches) + args = _args(inputs, caches, early) + certifying(op) + ref = op.reference(*args) + out = op(*args) + torch.cuda.synchronize() + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, s, o in zip(names, (*stock, *stock_caches), (*out, *caches)): + assert torch.equal(s.view(torch.uint8), o.view(torch.uint8)), ( + f"{name} differs from the native producer" + ) + + _within_fp8(out[0], ref["q"], "q") + _within_fp8(out[1], ref["index_q"], "index_q") + kv_data, kv_scale, index_k = caches + slots = inputs[1] + checked = 0 + for t in range(spec["tokens"]): + slot = int(slots[t]) + if slot < 0 or slot >= PAGES * PAGE: + continue + page, within = divmod(slot, PAGE) + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") + for (plane, name), head in ((p, h) for p in ((0, "k"), (1, "v")) for h in range(HEADS_KV)): + scales = kv_scale[page, plane, head].reshape(-1)[scale_index(within, plane).cuda()] + dq = decode_nvfp4_token(kv_data[page, plane, head, within], scales, KV_QUANT[1 + plane]) + r = ref[name][t, head].bfloat16().float() + # Round to nearest E2M1 under the block's E4M3-rounded scale s (in value units s / quant_scale): at most + # half the widest E2M1 step (one unit), plus the clipping at 6 units when s rounded down; a scale that + # underflows E4M3 zeroes the block. A slack of one E4M3 step covers the kernel's approximate reciprocals. + qs = KV_QUANT[1 + plane] + amax = r.view(-1, FP4_BLOCK).abs().amax(dim=1) + unit = (qs * amax / 6.0).to(torch.float8_e4m3fn).float() / qs + bound = torch.where( + unit > 0, 1.15 * unit + (amax - 6.0 * unit / 1.15).clamp_min(0.0), amax + ) + bound = bound.repeat_interleave(FP4_BLOCK) + 1e-12 + assert bool(((dq - r).abs() <= bound).all()), (name, t, head) + checked += 1 + assert checked >= spec["tokens"] - 2 + + +def test_early_trigger_does_not_change_the_bits() -> None: + inputs = _inputs(48, 1.0, False, 7) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=True)) + out_b = op(*_args(inputs, b, early=False)) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +@pytest.mark.parametrize("tokens", [1, 4, 8, 48]) +def test_prewait_step_inputs_do_not_change_the_bits(tokens: int) -> None: + """Reading the position ids, slots and RoPE rows before the grid wait leaves every output and cache byte as is, + padded rows (slot -1) and a slot past the pool included.""" + inputs = _inputs(tokens, 1.0, tokens >= 3, 11) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a), prewait_step_inputs=True) + out_b = op(*_args(inputs, b), prewait_step_inputs=False) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +def test_guard_refuses_a_repeated_slot() -> None: + inputs = list(_inputs(4, 1.0, False, 3)) + inputs[1] = inputs[1].clone() + inputs[1][1] = inputs[1][0] + with pytest.raises(AssertionError, match="must not repeat a slot"): + with validating(op): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_operand_off_the_device() -> None: + inputs = list(_inputs(4, 1.0, False, 9)) + inputs[5] = inputs[5].cpu() # position ids on the host + with pytest.raises(RuntimeError, match="packed's CUDA device"): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_unaligned_index_k_page_stride() -> None: + """The index-K rows are written as 32-bit words: a page stride that is not a multiple of 4 is refused.""" + kv_data, kv_scale, _ = _caches() + backing = torch.zeros(PAGES * (128 * 128 + 1), dtype=torch.float8_e4m3fn, device="cuda") + index_k = backing.as_strided((PAGES, 1, PAGE, 128), (128 * 128 + 1, 128 * 128, 128, 1)) + with pytest.raises(RuntimeError, match="index_k"): + op(*_args(_inputs(4, 1.0, False, 13), (kv_data, kv_scale, index_k))) + + +def test_op_refuses_unaligned_weights_or_rope_table() -> None: + """The norm weights are read as 8-byte words and the RoPE rows as float4: offset views are refused.""" + packed, slots, quant, weights, rotary, positions = _inputs(4, 1.0, False, 17) + offset_weight = torch.zeros(129, dtype=torch.bfloat16, device="cuda")[1:] + with pytest.raises(RuntimeError, match="norm weights"): + op( + *_args( + (packed, slots, quant, [offset_weight, *weights[1:]], rotary, positions), _caches() + ) + ) + flat = torch.zeros(rotary.numel() + 1, device="cuda") + offset_rope = flat[1:].view(rotary.shape) + with pytest.raises(RuntimeError, match="rotary_cos_sin"): + op(*_args((packed, slots, quant, weights, offset_rope, positions), _caches())) + + +def test_each_kv_head_lands_in_its_own_block() -> None: + """K / V of KV head 1 go to head block 1 of the HND pools and differ from head 0's (distinct packed rows).""" + inputs = _inputs(4, 1.0, False, 21) + caches = _caches() + op(*_args(inputs, caches)) + torch.cuda.synchronize() + kv_data = caches[0] + page, within = divmod(int(inputs[1][0]), PAGE) + for plane in (0, 1): + assert not torch.equal(kv_data[page, plane, 0, within], kv_data[page, plane, 1, within]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..bcc0ad67c843 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_qk_norm_rope_fp8 entry of the MiniMax-M3 tp2 target: the out-of-place FP8 form of the fused +per-head Gemma RMSNorm + partial NeoX RoPE against the in-place kernel followed by the E4M3 cast. V is a cast, so it +matches bit for bit; Q and K round once from fp32 where the composition rounds through bf16 first, so each element is +within one E4M3 step of it.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 + m3_qk_norm_rope_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" + +HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 + + +def _args(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q, kv = spec["q"], spec["kv"] + qkv = torch.randn(spec["m"], (q + 2 * kv) * HEAD_DIM, generator=g, device="cuda").to( + torch.bfloat16 + ) + qw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (spec["m"],), generator=g, device="cuda", dtype=torch.int32) + return ( + qkv, + q, + kv, + kv, + HEAD_DIM, + ROTARY_DIM, + EPS, + qw, + kw, + THETA, + True, + pos, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = args[0].clone() + with validating(op): + out = op(*args) + assert torch.equal(args[0], before), "the FP8 form must leave qkv untouched" + ref = op.reference(*args) + assert out.dtype == torch.float8_e4m3fn and out.shape == ref.shape + v_start = (cell.spec["q"] + cell.spec["kv"]) * HEAD_DIM + assert torch.equal(out[:, v_start:].view(torch.uint8), ref[:, v_start:].view(torch.uint8)) + assert _e4m3_steps(out[:, :v_start], ref[:, :v_start]) <= 1 + + +def _e4m3_steps(a, b) -> int: + """The largest distance between two E4M3 tensors in representable values (E4M3 codes order by magnitude).""" + + def signed(t): + c = t.view(torch.uint8).int() + # Sign-magnitude to a line: -0 and +0 both map to 0. + return torch.where(c >= 128, 128 - c, c) + + return int((signed(a) - signed(b)).abs().max()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..4be27dc1a2ff --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_qkv_insert entry: the generic horizontal producer on an FP8 main cache. + +Q, index-Q and every written K / V / index-K cache row are held within one E4M3 step of the fp32 reference. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_qkv_insert import ( # noqa: E501 + m3_qkv_insert as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" + +PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 32, 2, 2, 128, 64 +MAX_POSITIONS = 262144 +EPS = 1e-6 + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +def _args(tokens: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + tokens, (HEADS + 2 * KV + INDEX + 1) * D, device="cuda", generator=g + ).bfloat16() + kv = torch.zeros(PAGES, 2, KV, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + index_k = torch.zeros(PAGES, 1, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + weights = [(torch.randn(D, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + rotary = torch.randn(MAX_POSITIONS, 2, ROTARY // 2, device="cuda", generator=g) + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + return ( + packed, + kv, + None, + index_k, + slots, + None, + HEADS, + KV, + INDEX, + D, + ROTARY, + EPS, + *weights, + rotary, + positions, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _args(cell.spec["m"], zlib.crc32(cell.why.encode())) + certifying(op) + ref = op.reference(*args) + q, index_q = op(*args) + torch.cuda.synchronize() + _within_fp8(q.view(ref["q"].shape), ref["q"], "q") + _within_fp8(index_q.view(ref["index_q"].shape), ref["index_q"], "index_q") + kv, index_k, slots = args[1], args[3], args[4] + for t in range(cell.spec["m"]): + page, within = divmod(int(slots[t]), PAGE) + _within_fp8(kv[page, 0, :, within], ref["k"][t], f"k token {t}") + _within_fp8(kv[page, 1, :, within], ref["v"][t], f"v token {t}") + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..19123d839dea --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_blocks entry of the MiniMax-M3 tp2 target: MSA top-16 block selection with the forced +init / local blocks and per-row valid-block masks.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_blocks import ( # noqa: E501 + m3_select_blocks as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" + +TOPK, INIT, LOCAL = 16, 0, 1 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + n_blocks, total_q = spec["n_blocks"], spec["total_q"] + # Distinct scores (a permutation) so the selection has no ties to break. + scores = ( + torch.randperm(n_blocks * total_q, generator=g, device="cuda") + .float() + .view(1, n_blocks, total_q) + ) + if spec["mixed"]: + n_valid = torch.randint( + 17, n_blocks + 1, (total_q,), generator=g, device="cuda", dtype=torch.int32 + ) + else: + n_valid = torch.full((total_q,), n_blocks, dtype=torch.int32, device="cuda") + with validating(op): + out = op(scores, TOPK, n_valid, INIT, LOCAL, True) + ref = op.reference(scores, TOPK, n_valid, INIT, LOCAL, True) + assert out.shape == ref.shape + assert torch.equal(out, ref) + assert out.permute(1, 0, 2).is_contiguous(), "head-major backing" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..9fc6e361a03e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp2 (Rubin) m3_select_long entry (top-16 MSA block selection for long decode rows, 2 KV heads).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long import ( # noqa: E501 + m3_select_long as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key +# widths, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 2048, 2049, 4096, 4097, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # re-read path). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + # Every row: 10 winners above, ~1500 blocks tied at the 16th score, the rest below; the in-tree op stages up to + # 2048 boundary candidates, so it must break this tie to the smaller block like this op. + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + # Odd tokens: almost everything -inf, so picks run past the finite scores. + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +def _mass_tie_rows(scores, nvb, init_blocks, local_blocks): + """[total_q, num_kv_heads] mask of rows where more than 2048 valid blocks share the 16th-largest score.""" + num_kv_heads, num_blocks, total_q = scores.shape + blocks = torch.arange(num_blocks, device=scores.device) + raw = nvb[:total_q].long() + valid = raw.clamp(0, num_blocks) + forced = scores.permute(2, 0, 1).clone() # [total_q, heads, blocks] + forced = torch.where(blocks < init_blocks, 1.0e30, forced) + forced = torch.where(blocks >= (raw - local_blocks).clamp_min(0)[:, None, None], 1.0e29, forced) + in_row = (blocks[None, :] < valid[:, None])[:, None, :] + ranked = torch.where(in_row, forced, float("-inf")).sort(dim=-1, descending=True).values + kth = ranked[..., min(TOPK, num_blocks) - 1 : min(TOPK, num_blocks)] + return ((forced == kth) & in_row).sum(dim=-1) > 2048 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + + # Selection is an exact claim: the same blocks or a different attention. + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # And it is the in-tree selector's answer, bit for bit (the op the target falls back to) -- except on rows + # where more than 2048 blocks tie at the 16th-largest score: there the in-tree op does not break the tie to + # the smallest blocks, while this op and the reference do. + scores, nvb, _, init_blocks, local_blocks, _ = args + in_tree = select_blocks_from_maxscore( + scores, + topk=TOPK, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + mass_tie = _mass_tie_rows(scores, nvb, init_blocks, local_blocks) + assert torch.equal(out[~mass_tie], in_tree[~mass_tie]), ( + "selected blocks differ from the in-tree minimax_m3_select_blocks" + ) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + """Releasing the dependent grid right after the select's own wait changes only when the next kernel may launch.""" + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + """Releasing the dependent grid once the row's scores are loaded changes only when the next kernel may launch.""" + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) + + +def test_guard_refuses_out_aliasing_scores() -> None: + scores, nvb, _, init_blocks, local_blocks, _ = _case(op.CELLS[0].spec, seed=3) + aliased = ( + scores.view(torch.int32) + .reshape(-1)[: scores.shape[2] * TOPK] + .view(scores.shape[2], 1, TOPK) + ) + with pytest.raises(AssertionError, match="alias"): + with validating(op): + op(scores, nvb, aliased, init_blocks, local_blocks, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py new file mode 100644 index 000000000000..60b04b9dde33 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp2 (Rubin) m3_select_long_dsl entry: the long-row top-16 select from one 384-thread CTA per row, +against the reference and bit for bit against the 1024-thread m3_select_long on every cell.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( + m3_select_long_dsl as select_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, + m3_select_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_long_dsl import ( # noqa: E501 + m3_select_long_dsl as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key +# widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 1536, 1537, 4224, 4225, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # whole-row re-rank). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # Every row, mass ties included, is the 1024-thread kernel's answer. + scores, nvb, _, init_blocks, local_blocks, _ = args + long_out = torch.empty_like(out) + m3_select_long(scores, nvb, long_out, init_blocks, local_blocks, False) + assert torch.equal(out, long_out), "selected blocks differ from m3_select_long" + + +@pytest.mark.parametrize("layout", ["kv_blocks_q", "kv_q_blocks"]) +@pytest.mark.parametrize("valid", [17, 1294, 2524, 3907, 9216]) +def test_layouts_and_widths_match_the_1024_thread_select(layout, valid) -> None: + """Both score layouts the scorer can hand over, at a valid count in each key width.""" + g = torch.Generator(device="cuda").manual_seed(valid) + nb = max(valid, 129) + s = torch.randn(1, nb, 8, generator=g, device="cuda") + if layout == "kv_q_blocks": + s = s.permute(0, 2, 1).contiguous().permute(0, 2, 1) + nvb = torch.tensor([valid - (t % 2) for t in range(8)], dtype=torch.int32, device="cuda") + out = torch.empty(1, 8, TOPK, dtype=torch.int32, device="cuda").permute(1, 0, 2) + ref = torch.empty_like(out) + op(s, nvb, out, 0, 1, False) + m3_select_long(s, nvb, ref, 0, 1, False) + assert torch.equal(out, ref) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +@pytest.mark.parametrize("early_trigger", [False, True]) +@pytest.mark.parametrize("cell", range(len(op.CELLS))) +def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: + """ops/m3_select_long_dsl.SELECT_N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the + grid wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and + tied rows).""" + phase = 0 if early_trigger else 1 + spec = op.CELLS[cell].spec + scores, nvb, before, init_blocks, local_blocks, _ = _case(spec, seed=11 + cell) + after = torch.empty_like(before) + select_kernels._compile((True, early_trigger, phase, True))( + scores, nvb, before, init_blocks, local_blocks + ) + select_kernels._compile((True, early_trigger, phase, False))( + scores, nvb, after, init_blocks, local_blocks + ) + assert torch.equal(before, after) + + +def test_graph_replay_selects_the_same() -> None: + spec = op.CELLS[1].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=9) + ref = op.reference(scores, nvb, out, init_blocks, local_blocks) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + for _ in range(3): + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..feb450a5d350 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,451 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode entry (the target's Triton block-sparse decode and its routes) at the TP2 heads. + +Besides the band of every route around a float64 reference, the Triton merge's MXFP8 o_proj hand-off (published to +``m3_prequant.SLOT``) is held to ``torch.ops.trtllm.mxfp8_quantize(output, True)`` bit for bit: on identity steps and on +selected steps with the long-context range killed (the routes the o_proj consumes when LC2 does not run). +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, # noqa: E501 +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_prequant, + triton_sparse_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + triton_sparse_decode as tsd, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode as stock_decode, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import ( # noqa: E402 — needs the path declared above + HEAD_DIM, + NUM_HEADS, + build_step, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + +# Band, in BF16 ulps (2^-8) of each [32, 128] row's largest magnitude, and of the relative RMS. Every route +# computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference +# is the BF16 rounding on both sides (one ulp of the row at most), the FP16 probabilities before PV, the split-K +# partials kept in FP32 and merged by their log2-sum-exp, and approximate exp2/log2: 2 ulps per element, the two +# roundings (each at most half an ulp) dominating the RMS, so half an ulp. (GB300 TP4: at most 1.80 ulps per element +# and 0.20 ulps RMS over the cells.) +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, output, identity): + return op( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + identity, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=spec["identity"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + certifying(op) + _call(case, output, spec["identity"]) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + spec["identity"], + ) + assert torch.isfinite(output.float()).all(), "padding rows must be zeros, not NaN" + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_flag_does_not_change_the_answer() -> None: + """On an identity table the identity route (the per-row producer) and the selected-table route (the + long-context kernel) compute the same attention.""" + case = build_step((2048, 130), 4, identity=True, seed=17) + outs = [] + for identity in (True, False): + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, identity) + outs.append(output) + assert_within_ulp(outs[0], outs[1], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_guard_refuses_a_padded_selection_slot() -> None: + """The Triton producers index the block table with every one of a token's first min(16, blocks) ids, + -1 included.""" + case = build_step((1500,), 1, identity=True, seed=23) + case.topk[0, 0, 3] = -1 + output = torch.empty(1, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="blocks of the request"): + with validating(op): + _call(case, output, True) + + +def test_guard_refuses_a_false_identity_promise() -> None: + """identity_blocks lets the multi-query producer read one table for the whole request.""" + case = build_step((1500, 1600, 1700, 1800, 1500, 1600, 1700, 1800), 4, identity=True, seed=29) + case.topk[0, 5, :2] = case.topk[0, 5, :2].flip(0) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="identity_blocks promises"): + with validating(op): + _call(case, output, True) + + +def _swizzled_scale_offsets(rows: int, width: int) -> torch.Tensor: + """Offsets of rows [0, rows)'s UE8M0 scales in the 128x4-swizzled layout of mxfp8_quantize(x, True).""" + k_sf = width // 32 + k_tiles = (k_sf + 3) // 4 + r = torch.arange(rows, device="cuda").view(-1, 1) + k = torch.arange(k_sf, device="cuda").view(1, -1) + offsets = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + return offsets.reshape(-1) + + +def _mxfp8_handoff_matches_stock(output) -> None: + """The real rows' E4M3 bytes and UE8M0 scales equal mxfp8_quantize's. The merge writes only the step's rows into + its persistent 128-row scale tile; the padding rows' scales are never read by the o_proj and are not compared.""" + rows = output.shape[0] + flat = output.view(rows, -1) + taken = m3_prequant.SLOT.take(flat) + assert taken is not None, "the route published no MXFP8 hand-off for its output" + q8, sf = taken + stock_q, stock_sf = torch.ops.trtllm.mxfp8_quantize(flat, True) + assert torch.equal( + q8.view(torch.uint8).reshape(rows, -1), stock_q.view(torch.uint8).reshape(rows, -1) + ) + at = _swizzled_scale_offsets(rows, flat.shape[1]) + assert torch.equal( + sf.view(torch.uint8).reshape(-1)[at], stock_sf.view(torch.uint8).reshape(-1)[at] + ) + + +@pytest.mark.parametrize("seq_lens", [(1900,), (2048, 0), (700, 1500)]) +def test_identity_merge_mxfp8_handoff_is_stock_mxfp8_quantize(seq_lens) -> None: + """The Triton merge writes the o_proj hand-off (E4M3 rows, 128x4-swizzled UE8M0 scales) as mxfp8_quantize.""" + case = build_step(seq_lens, 4, identity=True, seed=31 + len(seq_lens)) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(case, output, True) + _mxfp8_handoff_matches_stock(output) + + +@pytest.mark.parametrize("seq_lens", [(40000,), (9000, 40000)]) +def test_killed_long_context_route(seq_lens) -> None: + """With SPARSE_DECODE_LC killed a selected table runs the Triton producer + merge: in the band of the reference, + and its MXFP8 hand-off is mxfp8_quantize of the rows it wrote.""" + case = build_step(seq_lens, 4, identity=False, seed=41 + len(seq_lens)) + landed = triton_sparse_decode.LC_KEYS + triton_sparse_decode.LC_KEYS = frozenset() + try: + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, False) + _mxfp8_handoff_matches_stock(output) + finally: + triton_sparse_decode.LC_KEYS = landed + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_route_is_the_certified_producer_range() -> None: + """Every identity step at the TP2 heads takes the per-row producer, the key modeling.SPARSE_DECODE_PRODUCER + covers (the multi-query producer needs one KV head).""" + assert modeling.SPARSE_DECODE_PRODUCER.covers("per_row") + assert not modeling.SPARSE_DECODE_PRODUCER.covers("multi_query") + for cell in op.CELLS: + keys = tuple(modeling._sparse_decode_producer(cell.spec)) + assert keys == (("per_row",) if cell.spec["identity"] else ()), (cell.why, keys) + + +_IDENTITY_LENS = (1900, 2048, 700, 1500, 1200, 2000, 300, 1024) +_SELECTED_LENS = (9000, 40000, 20000, 5000, 7000, 12000, 30000, 3000) + + +def _stock(case, output): + stock_decode( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + sm_scale=case.sm_scale, + output=output, + decode_query_len=case.decode_query_len, + k_block_scale=case.k_block_scale, + v_block_scale=case.v_block_scale, + k_global_scale=case.k_global_scale, + v_global_scale=case.v_global_scale, + ) + + +@pytest.mark.parametrize("identity", [True, False], ids=["identity", "selected_lc_killed"]) +@pytest.mark.parametrize("requests", [1, 2, 4, 8]) +def test_triton_route_is_bitwise_the_stock_decode(requests, identity) -> None: + """The target's Triton producer + split-K merge writes the stock in-tree MSA decode's rows bit for bit: on SM107 it + takes stock's split-K factor at every row count (4 / 8 / 16 / 32 verify rows here).""" + lens = (_IDENTITY_LENS if identity else _SELECTED_LENS)[:requests] + case = build_step(lens, 4, identity=identity, seed=50 + requests + 100 * identity) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC_KEYS + triton_sparse_decode.LC_KEYS = frozenset() + try: + _call(case, target, identity) + finally: + triton_sparse_decode.LC_KEYS = landed + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)) + + +@pytest.mark.parametrize("requests", [33, 46, 48], ids=["132_rows", "184_rows", "192_rows"]) +def test_plain_route_is_bitwise_the_stock_decode(requests) -> None: + """Above MergeMxfp8.MAX_ROWS (128 rows) a selected table is outside SPARSE_DECODE_LC ("plain") and runs the + Triton producer + merge: the stock decode's rows bit for bit, and no MXFP8 hand-off is published (the o_proj + quantizes its BF16 input itself). 33..48 requests x 4 verify rows: the wide target's steps past one scale tile, + up to WIDE_MAX_TOKENS.""" + case = build_step((3000,) * requests, 4, identity=False, seed=77) + rows = case.q.shape[0] + assert rows > tsd.MergeMxfp8.MAX_ROWS + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _call(case, target, False) + assert m3_prequant.SLOT.take(target.view(rows, -1)) is None + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)) + + +@pytest.mark.parametrize("identity", [True, False], ids=["triton_identity", "lc2_selected"]) +def test_handoff_kill_publishes_nothing_and_keeps_the_rows(identity) -> None: + """MERGE_MXFP8_HANDOFF False: no route publishes the MXFP8 hand-off (the o_proj then quantizes its BF16 input), and + the attention rows are the same bits as with the hand-off.""" + case = build_step((1900,) if identity else (40000,), 4, identity=identity, seed=91 + identity) + on = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + off = torch.empty_like(on) + _call(case, on, identity) + assert m3_prequant.SLOT.take(on.view(on.shape[0], -1)) is not None + triton_sparse_decode.MERGE_MXFP8_HANDOFF = False + try: + _call(case, off, identity) + assert m3_prequant.SLOT.take(off.view(off.shape[0], -1)) is None + finally: + triton_sparse_decode.MERGE_MXFP8_HANDOFF = True + assert torch.equal(on.view(torch.int16), off.view(torch.int16)) + + +def _short_lens(requests: int, g: torch.Generator) -> tuple: + """Short contexts of 1..16 blocks of 128 tokens (the GSM / identity-route regime), a few exactly on a block edge.""" + lens = torch.randint(4, 16 * 128 + 1, (requests,), generator=g).tolist() + if requests > 1: + lens[0] = 128 * int(torch.randint(1, 17, (1,), generator=g)) + return tuple(lens) + + +@pytest.mark.parametrize("seed", range(12)) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_short_table_stress_is_bitwise_the_stock_decode(requests, seed) -> None: + """Short contexts (1..16 blocks) at 1..4 requests x 4 verify rows, the identity route and a random selected route + with -1 tails past each token's blocks: the target's Triton route writes the stock decode's rows bit for bit.""" + g = torch.Generator().manual_seed(1000 * requests + seed) + for identity in (True, False): + case = build_step(_short_lens(requests, g), 4, identity=identity, seed=7 * seed + requests) + target = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC_KEYS + triton_sparse_decode.LC_KEYS = frozenset() + try: + _call(case, target, identity) + finally: + triton_sparse_decode.LC_KEYS = landed + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + identity, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("fault", ["neg_in_window", "id_past_row", "page_out_of_pool"]) +def test_invalid_slot_fails_closed(fault) -> None: + """A corrupt decode input (a -1 inside a token's block window, an id past the block-table row, or a block-table + page outside the pool) is read as an absent block instead of a wild KV read: the rows equal the reference over the + remaining blocks.""" + case = build_step((1500, 2048, 700), 4, identity=False, seed=303) + tok, slot = 5, 2 + if fault == "neg_in_window": + case.topk[1, tok, slot] = -1 + elif fault == "id_past_row": + case.topk[1, tok, slot] = case.block_table.shape[1] + 7 + else: + req = tok // 4 + blk = int(case.topk[1, tok, slot]) + bad_page = case.k_paged.shape[0] + 1000 + case.block_table = case.block_table.clone() + case.block_table[req, blk] = bad_page + # Every token of the request that selected that block loses it. + hits = (case.topk[:, req * 4 : req * 4 + 4] == blk).nonzero().tolist() + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + landed = triton_sparse_decode.LC_KEYS + triton_sparse_decode.LC_KEYS = frozenset() + try: + _call(case, output, False) + finally: + triton_sparse_decode.LC_KEYS = landed + torch.cuda.synchronize() + # The reference walks each token's first min(16, blocks) entries and drops -1: the corrupt slots as -1. + ref_topk = case.topk.clone() + if fault == "page_out_of_pool": + for h, t, k in hits: + ref_topk[h, req * 4 + t, k] = -1 + else: + ref_topk[1, tok, slot] = -1 + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + ref_topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert torch.isfinite(output.float()).all() + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def _identity_table( + seq_lens: torch.Tensor, decode_query_len: int, heads: int, lag: int = 0 +) -> torch.Tensor: + """The identity selection of every verify token, staged from lengths `lag` tokens behind seq_lens.""" + seq = (seq_lens.cpu().to(torch.int64) - lag).clamp(min=0).repeat_interleave(decode_query_len) + intra = torch.arange(seq.shape[0]) % decode_query_len + n_valid = ((seq - decode_query_len + intra + 1).clamp(min=0) + 127) // 128 + ar = torch.arange(16) + table = torch.where(ar[None, :] < n_valid[:, None], ar[None, :], -1).to(torch.int32) + return table.unsqueeze(0).expand(heads, -1, -1).contiguous().cuda() + + +# Each crosses a 128-token edge within its four verify tokens or between the stale and the true lengths. +_EDGE_LENS = (1284, 1287, 1281, 129, 132, 2048, 1028, 5, 643, 1920, 256, 1) + + +@pytest.mark.parametrize("fault", ["stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_identity_route_reads_its_own_lengths_not_the_table(requests, fault) -> None: + """The identity route derives block i for slot i from the kv lengths the kernel reads: a staged table that lags + the lengths by a verify step (a 128-token edge crossed), or is garbage, gives the stock decode's rows over the + true identity selection bit for bit.""" + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=400 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + assert torch.equal(case.topk, true_table), ( + "build_step's identity table is the identity selection" + ) + if fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _stock(case, stock) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + _call(case_staged, target, True) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py new file mode 100644 index 000000000000..477b7ef94ebf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode_lc.py @@ -0,0 +1,490 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU tests for trtllm::minimax_m3_sm107_tp2_sparse_decode_lc (one-launch long-context sparse decode, TP2). + +Two contracts: +- the BF16 rows sit in the same band around a float64 reference as the stock Triton sparse decode, and every variant + (linear softmax, pre-wait block-table staging, MXFP8 hand-off on or off) writes the same rows as the stock + log-domain recurrence where the arithmetic is shared; +- the MXFP8 hand-off (q8, sf) is bitwise torch.ops.trtllm.mxfp8_quantize(out, True) of the rows the kernel wrote: the + E4M3 bytes row-major and the UE8M0 scales in its 128x4-swizzled layout. The o_proj consumes this as its + prequantized activation, so this test is the interface contract. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + m3_sparse_decode_lc as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + mxfp8_handoff_reference, +) + +__extra_import_path__ = ["../../../.."] +from _validating import ( # noqa: E402, F401 — needs the path declared above + assert_within_ulp, + certifying, +) + +PAGE, DIM, TOPK, GROUP, KV_HEADS = 128, 128, 16, 16, 2 +WIDTH = KV_HEADS * GROUP * DIM +SM_SCALE = DIM**-0.5 +# Band, in BF16 ulps (2^-8) of each output row's largest magnitude, and of the relative RMS. Same as the sm_103 TP4 +# kernel; on SM107 the stock Triton decode and this kernel both measure at most 1.03 ulps and 0.26 RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + +_E2M1 = [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6] + + +def _op(): + return torch.ops.trtllm.minimax_m3_sm107_tp2_sparse_decode_lc + + +def _e4m3_bytes(shape, gen, lo=0x28, hi=0x48): + mag = torch.randint(lo, hi, shape, dtype=torch.int32, device="cuda", generator=gen) + sign = torch.randint(0, 2, shape, dtype=torch.int32, device="cuda", generator=gen) << 7 + return (mag | sign).to(torch.uint8) + + +def _case(seq_lens, seed, dql=4, overlap=14, q_hi=0x48, tie=False): + """A packed-NVFP4 paged cache, ascending top-16 tables with a -1 tail, and FP8 queries for TP2 geometry.""" + gen = torch.Generator(device="cuda").manual_seed(seed) + reqs, total_q = len(seq_lens), len(seq_lens) * dql + nblk = [(s + PAGE - 1) // PAGE for s in seq_lens] + pages = sum(nblk) + 8 + k = torch.randint( + 0, 256, (pages, KV_HEADS, PAGE, DIM // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + v = torch.randint( + 0, 256, (pages, KV_HEADS, PAGE, DIM // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + if tie: + k[:] = k[:1, :, :1, :] + ksc = _e4m3_bytes((pages, KV_HEADS, PAGE, DIM // 16), gen) + vsc = _e4m3_bytes((pages, KV_HEADS, PAGE, DIM // 16), gen) + perm = torch.randperm(pages, device="cuda", generator=gen).int() + bt = torch.full((reqs, max(nblk)), -1, dtype=torch.int32, device="cuda") + start = 0 + for r, n in enumerate(nblk): + bt[r, :n] = perm[start : start + n] + start += n + topk = torch.full((KV_HEADS, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for r, s in enumerate(seq_lens): + for h in range(KV_HEADS): + last = (s + PAGE - 1) // PAGE + base = torch.randperm(max(last - 1, 1), device="cuda", generator=gen)[ + : TOPK - 1 + ].tolist() + for qo in range(dql): + nb = (s - dql + qo + 1 + PAGE - 1) // PAGE + sel = [b for b in base[:overlap] if b < nb - 1] + for b in torch.randperm(max(nb - 1, 1), device="cuda", generator=gen).tolist(): + if len(sel) >= min(TOPK, nb) - 1: + break + if b not in sel and b < nb - 1: + sel.append(b) + sel = sorted(set(sel + [nb - 1]))[: min(TOPK, nb)] + topk[h, r * dql + qo, : len(sel)] = torch.tensor( + sel, dtype=torch.int32, device="cuda" + ) + q = _e4m3_bytes((total_q, KV_HEADS * GROUP, DIM), gen, lo=0x20, hi=q_hi).view( + torch.float8_e4m3fn + ) + return dict( + q=q, + k=k, + v=v, + ksc=ksc, + vsc=vsc, + kgs=torch.tensor([0.03], device="cuda"), + vgs=torch.tensor([0.05], device="cuda"), + topk=topk, + bt=bt, + seq=torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + dql=dql, + ) + + +def _dequant(packed, sc, quad): + lut = torch.tensor(_E2M1, dtype=torch.float64, device="cuda") + vals = torch.stack([lut[(packed & 0xF).long()], lut[(packed >> 4).long()]], dim=-1).reshape( + *packed.shape[:3], DIM + ) + if ( + quad + ): # V scales: token t, column c at byte (t // 4) * 32 + t % 4 + 4 * c within the page/head + t = torch.arange(PAGE, device="cuda")[:, None] + idx = (t // 4) * 32 + t % 4 + 4 * torch.arange(8, device="cuda")[None, :] + sc = sc.reshape(*sc.shape[:2], PAGE * 8)[:, :, idx] + return vals * sc.contiguous().view(torch.float8_e4m3fn).to(torch.float64).repeat_interleave( + 16, dim=-1 + ) + + +def _reference(c): + q = c["q"].to(torch.float64) + out = torch.zeros(q.shape, dtype=torch.float64, device="cuda") + dql = c["dql"] + for tok in range(q.shape[0]): + r, qo = divmod(tok, dql) + kv_len = max(int(c["seq"][r]) - dql + qo + 1, 0) + for h in range(KV_HEADS): + blks = [b for b in c["topk"][h, tok].tolist() if b >= 0][ + : min(TOPK, (kv_len + PAGE - 1) // PAGE) + ] + if not blks: + continue + pg = c["bt"][r, blks].long() + kk = _dequant(c["k"][pg, h : h + 1], c["ksc"][pg, h : h + 1], False)[:, 0].reshape( + -1, DIM + ) + vv = _dequant(c["v"][pg, h : h + 1], c["vsc"][pg, h : h + 1], True)[:, 0].reshape( + -1, DIM + ) + pos = ( + torch.tensor(blks, device="cuda")[:, None] * PAGE + + torch.arange(PAGE, device="cuda") + ).reshape(-1) + s = (q[tok, h * GROUP : (h + 1) * GROUP] @ kk.T) * (SM_SCALE * c["kgs"].item()) + s[:, pos >= kv_len] = float("-inf") + out[tok, h * GROUP : (h + 1) * GROUP] = (torch.softmax(s, dim=-1) @ vv) * c[ + "vgs" + ].item() + return out.reshape(q.shape[0], WIDTH) + + +def _band(x, ref): + x = x.to(torch.float64).reshape(-1, DIM) + ref = ref.reshape(-1, DIM) + rowmax = ref.abs().amax(dim=-1, keepdim=True) + live = rowmax[:, 0] > 0 + unit = (rowmax * 2.0**-8).clamp_min(1e-300) + err = ((x - ref).abs() / unit)[live] + rms = (((x - ref) ** 2).mean(dim=-1).sqrt() / unit[:, 0])[live] + return (float(err.max()) if err.numel() else 0.0), (float(rms.max()) if rms.numel() else 0.0) + + +def _run(c, linear=False, prewait=False, mxfp8=False, n1=None, use_pdl=False, out=None): + total_q = c["q"].shape[0] + if out is None: + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + q8 = torch.zeros(total_q, WIDTH, dtype=torch.uint8, device="cuda") if mxfp8 else None + sf_bytes = 128 * (WIDTH // 32) * ((total_q + 127) // 128) + sf = torch.zeros(sf_bytes, dtype=torch.uint8, device="cuda") if mxfp8 else None + _op()( + c["q"].view(torch.uint8), + c["k"], + c["v"], + c["ksc"], + c["vsc"], + c["kgs"], + c["vgs"], + c["topk"], + c["bt"], + c["seq"], + out, + q8, + sf, + WIDTH // 128, + c["dql"], + SM_SCALE, + linear, + use_pdl, + prewait, + n1, + ) + return out, q8, sf + + +_CASES = { + "b1-165k": dict(seq_lens=[165521]), + "b2-mixed": dict(seq_lens=[165521, 98765]), + "b2-short": dict(seq_lens=[130, 5000]), + "b1-one-token-tail": dict(seq_lens=[128 * 40 + 1]), + "b1-padded-topk": dict(seq_lens=[700]), + "b1-tiny": dict(seq_lens=[4]), + "b2-large-logits": dict(seq_lens=[165521, 64000], q_hi=0x70), + "b1-ties": dict(seq_lens=[20000], tie=True), + "b2-disjoint": dict(seq_lens=[165521, 165521], overlap=0), + "b4": dict(seq_lens=[165521, 120000, 80000, 30000]), +} + + +@pytest.mark.parametrize("name", list(_CASES)) +@pytest.mark.parametrize("seed", [0, 1]) +def test_band_vs_float64(name, seed) -> None: + c = _case(seed=seed, **_CASES[name]) + out, _, _ = _run(c) + elem, rms = _band(out, _reference(c)) + assert torch.isfinite(out.float()).all() + assert elem <= _ELEMENT_ULP and rms <= _RMS_ULP, (elem, rms) + + +@pytest.mark.parametrize("name", ["b1-165k", "b2-mixed", "b1-padded-topk", "b4"]) +def test_variants_write_the_same_rows(name) -> None: + c = _case(seed=7, **_CASES[name]) + base, _, _ = _run(c) + for prewait in (False, True): + for mxfp8 in (False, True): + out, _, _ = _run(c, prewait=prewait, mxfp8=mxfp8) + assert torch.equal(out.view(torch.int16), base.view(torch.int16)), (prewait, mxfp8) + lin, _, _ = _run(c, linear=True) + elem, rms = _band(lin, _reference(c)) + assert elem <= _ELEMENT_ULP and rms <= _RMS_ULP, (elem, rms) + + +@pytest.mark.parametrize("name", list(_CASES)) +@pytest.mark.parametrize("linear", [False, True]) +def test_mxfp8_handoff_is_stock_mxfp8_quantize(name, linear) -> None: + """The o_proj interface: (q8, sf) == mxfp8_quantize(out, is_sf_swizzled_layout=True), bitwise.""" + c = _case(seed=3, **_CASES[name]) + out, q8, sf = _run(c, linear=linear, mxfp8=True) + ref_q, ref_sf = torch.ops.trtllm.mxfp8_quantize(out, True) + assert torch.equal(q8, ref_q.view(torch.uint8).reshape(q8.shape)) + ref_sf = ref_sf.view(torch.uint8).reshape(-1) + assert torch.equal(sf[: ref_sf.numel()], ref_sf) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The catalog entry's cells through its own __call__ (guarded by is_valid) and its own reference.""" + spec = cell.spec + c = _case( + list(spec["seq_lens"]), + seed=zlib.crc32(cell.why.encode()) % 2**31, + dql=spec["decode_query_len"], + ) + total_q = c["q"].shape[0] + # the hand-off buffers hold whole 128-row tiles (MergeMxfp8.MAX_ROWS) + alloc = (total_q + 127) // 128 * 128 + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + q8 = torch.zeros(alloc, WIDTH, dtype=torch.uint8, device="cuda") if spec["mxfp8"] else None + sf = ( + torch.zeros(alloc * (WIDTH // 32), dtype=torch.uint8, device="cuda") + if spec["mxfp8"] + else None + ) + args = ( + c["q"].view(torch.uint8), + c["k"], + c["v"], + c["ksc"], + c["vsc"], + c["kgs"], + c["vgs"], + c["topk"], + c["bt"], + c["seq"], + out, + q8, + sf, + WIDTH // 128, + c["dql"], + SM_SCALE, + spec["linear_softmax"], + False, + spec["prewait"], + ) + certifying(op) + op(*args) + ref = op.reference(*args) + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + assert torch.isfinite(out.float()).all() + if spec["mxfp8"]: + ref_q8, ref_sf = mxfp8_handoff_reference(out, alloc) + assert torch.equal(q8[:total_q], ref_q8) + assert torch.equal(sf, ref_sf) + stock_q, stock_sf = torch.ops.trtllm.mxfp8_quantize(out, True) + assert torch.equal(q8[:total_q], stock_q.view(torch.uint8).reshape(total_q, WIDTH)) + assert torch.equal(sf[: stock_sf.numel()], stock_sf.view(torch.uint8).reshape(-1)) + + +# LC on vs the killed range (SPARSE_DECODE_LC = frozenset(): the stock Triton producer + merge) through the target's +# minimax_m3_sparse_attn_decode on the same inputs. Not bitwise: the split-K merge sums the 16 partials in Triton's +# reduction order there and in chunk order here (the per-block math is the same), so the rows differ by summation +# order only. Band, in BF16 ulps of each row's largest magnitude: measured on SM107 at most 1.92 per element and +# 0.18 RMS (R4/R13 of the MSA lane), so 2 / 0.25. +_LC_VS_KILLED_ELEMENT_ULP = 2.0 +_LC_VS_KILLED_RMS_ULP = 0.25 + + +def _target_decode(c, killed): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_prequant, + triton_sparse_decode, + ) + + landed = triton_sparse_decode.LC_KEYS + if killed: + triton_sparse_decode.LC_KEYS = frozenset() + try: + total_q = c["q"].shape[0] + out = torch.empty(total_q, NUM_HEADS_TOTAL, DIM, dtype=torch.bfloat16, device="cuda") + triton_sparse_decode.minimax_m3_sparse_attn_decode( + c["q"], + c["k"], + c["v"], + c["topk"], + c["bt"], + c["seq"], + sm_scale=SM_SCALE, + output=out, + decode_query_len=c["dql"], + k_block_scale=c["ksc"], + v_block_scale=c["vsc"], + k_global_scale=c["kgs"], + v_global_scale=c["vgs"], + ) + torch.cuda.synchronize() + finally: + triton_sparse_decode.LC_KEYS = landed + m3_prequant.SLOT.clear() + return out.view(total_q, WIDTH) + + +NUM_HEADS_TOTAL = KV_HEADS * GROUP + + +@pytest.mark.parametrize("name", list(_CASES)) +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_lc_vs_killed_stock_path(name, seed) -> None: + """The landed cell and its kill switch on the same step: rows within summation-order distance of each other.""" + c = _case(seed=seed, **_CASES[name]) + lc = _target_decode(c, killed=False) + stock = _target_decode(c, killed=True) + assert torch.isfinite(lc.float()).all() and torch.isfinite(stock.float()).all() + elem, rms = _band(lc, stock.to(torch.float64)) + assert elem <= _LC_VS_KILLED_ELEMENT_ULP and rms <= _LC_VS_KILLED_RMS_ULP, (elem, rms) + ref = _reference(c) + for out in (lc, stock): + e, r = _band(out, ref) + assert e <= _ELEMENT_ULP and r <= _RMS_ULP, (e, r) + + +# N1 (n1_ids): the layer's previous-step block ids, prefetched into L2 before the grid wait and overwritten with this +# step's. A hint: the rows must be bitwise the rows without it, whatever the buffer holds. + + +def _n1_expected(c, rows, before=None): + """The buffer after a launch: row ``tok`` < T holds topk[h, tok, r] for r below the token's block count, else -1; + rows at or past T keep what they held (``before``, default -1).""" + total_q, dql = c["q"].shape[0], c["dql"] + exp = ( + before.clone() + if before is not None + else torch.full((KV_HEADS, rows, TOPK), -1, dtype=torch.int32, device="cuda") + ) + for tok in range(min(total_q, rows)): + r, qo = divmod(tok, dql) + kv_len = max(int(c["seq"][r]) - dql + qo + 1, 0) + real = min(TOPK, (kv_len + PAGE - 1) // PAGE) + exp[:, tok] = -1 + exp[:, tok, :real] = c["topk"][:, tok, :real] + return exp + + +def _n1_garbage(rows, seed): + gen = torch.Generator(device="cuda").manual_seed(seed) + ids = torch.randint( + -5, 1 << 20, (KV_HEADS, rows, TOPK), dtype=torch.int32, device="cuda", generator=gen + ) + ids[0, :, 0] = 2**31 - 1 + ids[1, :, 0] = -(2**31) + return ids + + +# In-model request counts (1, 2, 4, 5, 8 Eagle3 verify requests of 4 rows), a full 128-row step, a plain decode. +_N1_CASES = { + "m1": dict(seq_lens=[165521]), + "m2": dict(seq_lens=[165521, 98765]), + "m4": dict(seq_lens=[165521, 120000, 700, 30000]), + "m5": dict(seq_lens=[165521, 4, 130, 98765, 5000]), + "m8": dict(seq_lens=[165521, 98765, 64000, 30000, 20000, 5000, 700, 130]), + "m32": dict(seq_lens=[4096 * (i + 1) + 77 * i for i in range(32)]), + "m2-dql1": dict(seq_lens=[165521, 5000], dql=1), +} + + +@pytest.mark.parametrize("name", list(_N1_CASES)) +def test_n1_rows_bitwise_and_ids_recorded(name) -> None: + c = _case(seed=11, **_N1_CASES[name]) + total_q = c["q"].shape[0] + for prewait in (False, True): + for mxfp8 in (False, True): + base, bq8, bsf = _run(c, prewait=prewait, mxfp8=mxfp8) + ids = torch.full((KV_HEADS, 128, TOPK), -1, dtype=torch.int32, device="cuda") + # First step (all -1: no prefetch), second (last step's ids: prefetch), then stale garbage ids. + for step, init in enumerate((None, None, _n1_garbage(128, 3))): + if init is not None: + ids.copy_(init) + exp = _n1_expected(c, 128, ids) + out, q8, sf = _run(c, prewait=prewait, mxfp8=mxfp8, n1=ids) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), base.view(torch.int16)), ( + prewait, + mxfp8, + step, + ) + if mxfp8: + assert torch.equal(q8, bq8) and torch.equal(sf, bsf), (prewait, mxfp8, step) + assert torch.equal(ids, exp), (prewait, mxfp8, step) + if total_q > 1: + # A buffer with fewer rows than the step: those rows record, the rest skip both prefetch and record. + rows = total_q // 2 + ids = torch.full((KV_HEADS, rows, TOPK), -1, dtype=torch.int32, device="cuda") + out, _, _ = _run(c, n1=ids) + assert torch.equal(out.view(torch.int16), _run(c)[0].view(torch.int16)) + assert torch.equal(ids, _n1_expected(c, rows)) + + +def test_n1_pdl_stress_and_graph_replay() -> None: + """Back-to-back PDL launches of two layers (each with its own id buffer, alternating steps), eager and replayed + from a CUDA graph: every launch writes the rows of the launch without N1, and the buffers end at the last step's + ids.""" + cases = [_case(seed=21, **_N1_CASES["m2"]), _case(seed=22, **_N1_CASES["m5"])] + bases = [_run(c, mxfp8=True)[0] for c in cases] + ids = [torch.full((KV_HEADS, 128, TOPK), -1, dtype=torch.int32, device="cuda") for _ in cases] + iters = 200 + outs = [[torch.empty_like(b) for _ in range(iters)] for b in bases] + for i in range(iters): + for j, c in enumerate(cases): + _run(c, mxfp8=True, n1=ids[j], use_pdl=True, out=outs[j][i]) + torch.cuda.synchronize() + for j in range(len(cases)): + assert all(torch.equal(o.view(torch.int16), bases[j].view(torch.int16)) for o in outs[j]), j + assert torch.equal(ids[j], _n1_expected(cases[j], 128)), j + for t in ids: + t.fill_(-1) + gouts = [torch.empty_like(b) for b in bases] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream): + for j, c in enumerate(cases): # warm the launch configuration outside capture + _run(c, mxfp8=True, n1=ids[j], use_pdl=True, out=gouts[j]) + torch.cuda.synchronize() + with torch.cuda.graph(graph, stream=stream): + for _ in range(4): + for j, c in enumerate(cases): + _run(c, mxfp8=True, n1=ids[j], use_pdl=True, out=gouts[j]) + garbage = _n1_garbage(128, 5) + for t in ids: + t.copy_(garbage) + for _ in range(25): + for o in gouts: + o.zero_() + graph.replay() + torch.cuda.synchronize() + for j in range(len(cases)): + assert torch.equal(gouts[j].view(torch.int16), bases[j].view(torch.int16)), j + assert torch.equal(ids[j], _n1_expected(cases[j], 128, garbage)), j diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py new file mode 100644 index 000000000000..3a19b3bf5a7d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for ops/m3_tokmap_guard: a one-request multi-token call of +thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a +stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the +appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for +every length. Without the guard the stale map moves rows (the negative control), and a several-request call passes +through with its own map untouched. + +The fresh-map reference forces the op's own build (a scaled RoPE type with no inv-freq cache, at rotary dim 0, so no +row is rotated either way). +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_tokmap_guard as guard, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = 32, 2, 128, 128 # the TP2 drafter shard (64 / 2 heads, 4 / 2 KV heads) +BUCKET = 4 +_PEER_EXTRA = 777 # the second request's extra length behind the first +FP8_KV_QUANT_MODE = 128 +_DEV = torch.device("cuda") # unindexed on purpose: install() and the lookup normalize it +_TAKEN = [] + + +@pytest.fixture(scope="module", autouse=True) +def _installed(): + guard.install(_DEV, 48, _TAKEN.append) + yield + + +class _Step: + """One request (or ``nreq``) of BUCKET padded rows over a fresh FP8 paged cache, persistent across calls like a + real step's buffers.""" + + def __init__(self, kv_len: int, nreq: int = 1, seed: int = 0): + g = torch.Generator(device=_DEV).manual_seed(seed + kv_len) + self.nreq = nreq + # Per-request inputs (page tables, generation lengths, position offsets) for at least two requests, the page + # tables at the peer's longer length (_seq_with_peer): the stale map sends rows to request 1, which must have + # real entries behind it, also on the unguarded calls. + tables = max(nreq, 2) + pages = math.ceil((kv_len + _PEER_EXTRA + BUCKET) / PAGE) + self.nblocks = tables * pages + 2 + self.pool0 = ( + torch.randn(self.nblocks, 2, KVH, PAGE, D, device=_DEV, generator=g) + .clamp_(-4, 4) + .to(torch.float8_e4m3fn) + ) + self.pool = self.pool0.clone() + perm = torch.randperm(self.nblocks, device=_DEV, generator=g)[: tables * pages].view( + tables, pages + ) + self.offsets = torch.zeros(1, tables, 2, pages, dtype=torch.int32, device=_DEV) + self.offsets[0, :, 0] = 2 * perm + self.offsets[0, :, 1] = 2 * perm + 1 + self.pointers = torch.tensor([[self.pool.data_ptr(), 0]], dtype=torch.int64) + self.mapping = torch.zeros(1, 2, dtype=torch.int32) + self.qkv = torch.randn( + nreq * BUCKET, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=_DEV, generator=g + ) + # Request lengths after this step's append; generation lengths (accepted rows) set per case. + self.seq = torch.full((nreq,), kv_len, dtype=torch.int32, device=_DEV) + self.gen = torch.full((tables,), BUCKET, dtype=torch.int32, device=_DEV) + self.pos = ( + torch.arange(BUCKET, dtype=torch.int32, device=_DEV).repeat(tables, 1).contiguous() + ) + self.one = torch.ones(1, dtype=torch.float32, device=_DEV) + self.ws = torch.empty(64 << 20, dtype=torch.uint8, device=_DEV) + self.kv_len = kv_len + + def call(self, op, fresh_map: bool = False): + """One generation preprocess; returns a copy of the processed query.""" + out = op( + self.qkv, + self.ws, + self.seq, + self.gen, + self.pos, + self.offsets, + self.pointers, + self.mapping, + self.one, + self.one, + None, + None, + None, + None, + 0, + 0, + H, + KVH, + D, + PAGE, + FP8_KV_QUANT_MODE, + 1 << 20, + 1 << 20, + self.nreq * BUCKET, + self.nreq, + BUCKET, + self.kv_len, + 0, + 10000.0, + 1 if fresh_map else 0, + 1.0, + 1 << 20, + 0, + 1.0, + 1.0, + False, + BUCKET, + 0, + torch.cuda.get_device_properties(0).multi_processor_count, + self.nblocks, + 2, + True, + False, + ) + return out[0].clone() + + def poison_map(self) -> None: + """A stale multi-request map in the op's one-request map slice: rows 1.. on request 1.""" + layout = thop.get_trtllm_gen_generation_workspace_layout( + self.qkv.dtype, 1, self.nreq * BUCKET, H, D, 0, KVH, 0, False, False + ) + off = int(layout["tokens_info_offset"]) + stale = torch.tensor( + [[0, 0]] + [[1, t] for t in range(BUCKET - 1)], dtype=torch.int32, device=_DEV + ) + self.ws[off : off + 8 * BUCKET].view(torch.int32).view(BUCKET, 2).copy_(stale) + + def reset(self) -> None: + self.pool.copy_(self.pool0) + + +def _seq_with_peer(step: _Step) -> None: + # A second request's length behind the first, so a row the stale map moves to request 1 reads a real length. + step.seq = torch.tensor( + [step.kv_len, step.kv_len + _PEER_EXTRA], dtype=torch.int32, device=_DEV + ) + + +@pytest.mark.parametrize("kv_len", [130, 8191, 34117]) +@pytest.mark.parametrize("n", [1, 2, 3, 4]) +def test_guard_equals_a_fresh_map(kv_len, n) -> None: + step = _Step(kv_len) + _seq_with_peer(step) + step.gen.fill_(n) + step.poison_map() + step.reset() + q_ref = step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + + step.poison_map() + step.reset() + q = step.call(thop.trtllm_gen_generation_preprocess) + # The processed query's n live rows (rows past the request's length are padding the op leaves unwritten). + assert torch.equal(q[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)) + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + assert BUCKET in _TAKEN + + +def test_stale_map_moves_rows_without_the_guard() -> None: + step = _Step(34117) + _seq_with_peer(step) + step.gen.fill_(BUCKET) + step.poison_map() + step.reset() + step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + step.poison_map() + step.reset() + step.call(guard._original) + assert not torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + + +def test_one_captured_graph_serves_every_length() -> None: + step = _Step(34000) + _seq_with_peer(step) + refs = {} + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + q = step.call(guard._original, fresh_map=True) + refs[n] = (q, step.pool.clone()) + step.gen.fill_(BUCKET) + step.call(thop.trtllm_gen_generation_preprocess) # eager first, as the warmup runs it + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = thop.trtllm_gen_generation_preprocess + q_static = step.call(out) + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + g.replay() + torch.cuda.synchronize() + q_ref, pool_ref = refs[n] + assert torch.equal(q_static[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)), n + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)), n + + +def test_several_requests_pass_through() -> None: + step = _Step(8191, nreq=2) + layout = thop.get_trtllm_gen_generation_workspace_layout( + step.qkv.dtype, 2, 2 * BUCKET, H, D, 0, KVH, 0, False, False + ) + step.reset() + step.call(thop.trtllm_gen_generation_preprocess) + pool_guarded = step.pool.clone() + step.reset() + step.call(guard._original) + assert torch.equal(step.pool.view(torch.uint8), pool_guarded.view(torch.uint8)) + assert int(layout["tokens_info_size"]) >= 8 * 2 * BUCKET diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..d35556356264 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the msa_kv_update entry (the one-launch MSA KV-length correction).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.msa_kv_update import ( # noqa: E501 + TOKENS_PER_BLOCK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" + +#: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. +_SPARE = 3 + + +def _case(spec, seed): + """The metadata's buffers for one step, with stale values in every output so a skipped write shows.""" + g = torch.Generator(device="cpu").manual_seed(seed) + batch, qlen = spec["batch"], spec["query_len"] + total_q = batch * qlen + staged = torch.randint(qlen, 5000, (batch,), generator=g, dtype=torch.int32) + accept = spec["accept"] + if accept == "full": + kv = staged.clone() + elif accept == "page_edge": + staged = torch.tensor([130, 257, 2 * TOKENS_PER_BLOCK + qlen][:batch], dtype=torch.int32) + kv = staged - torch.tensor([0, 1, qlen - 1][:batch], dtype=torch.int32) + else: + kv = staged - torch.randint(0, qlen, (batch,), generator=g, dtype=torch.int32) + if accept == "padded": + kv[batch // 2 :] = 0 + staged[batch // 2 :] = 0 + ncols = (5000 + qlen) // TOKENS_PER_BLOCK + 2 + table = ( + torch.randperm(batch * ncols + 7, generator=g)[: batch * ncols] + .to(torch.int32) + .reshape(batch, ncols) + ) + dev = dict(device="cuda", dtype=torch.int32) + stale = -12345 + return ( + torch.cat([kv, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.cat([staged, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.full((batch + _SPARE,), qlen, **dev), + torch.arange(batch, **dev).repeat_interleave(qlen), + torch.arange(qlen, **dev).repeat(batch), + table.cuda(), + TOKENS_PER_BLOCK, + batch, + total_q, + torch.full((batch + _SPARE,), stale, **dev), + torch.full((batch + 1 + _SPARE,), stale, **dev) if spec["cu_kv"] else None, + torch.full((total_q + _SPARE,), stale, **dev), + torch.full((total_q + _SPARE,), stale, **dev) if spec["n_valid"] else None, + ) + + +def _check(args, ref): + seq, cu_kv, loc, n_valid = args[9], args[10], args[11], args[12] + batch, total_q = args[7], args[8] + ref_seq, ref_cu, ref_loc, ref_nv = ref + # Integer bookkeeping: exact, and nothing past the live rows is touched. + assert torch.equal(seq[:batch], ref_seq) + assert torch.equal(loc[:total_q], ref_loc) + assert bool((seq[batch:] == -12345).all()) and bool((loc[total_q:] == -12345).all()) + assert (cu_kv is None) == (ref_cu is None) and (n_valid is None) == (ref_nv is None) + if cu_kv is not None: + assert torch.equal(cu_kv[: batch + 1], ref_cu) + assert bool((cu_kv[batch + 1 :] == -12345).all()) + if n_valid is not None: + assert torch.equal(n_valid[:total_q], ref_nv) + assert bool((n_valid[total_q:] == -12345).all()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + op(*args) + _check(args, op.reference(*args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_hook_takes_every_cell(cell) -> None: + """The target's hook (what the metadata calls) takes each certified step and writes what the op does.""" + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + 1) + assert hook(*args) is True + _check(args, op.reference(*args)) + + +def test_the_hook_declines_what_the_op_would_misread() -> None: + """A strided operand is read wrongly by the op, so the hook declines it (the metadata then runs its own + torch update) and writes nothing.""" + args = list(_case(op.CELLS[2].spec, seed=5)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + before = args[11].clone() + assert hook(*args) is False + assert torch.equal(args[11], before) + + +def test_guard_refuses_a_strided_operand() -> None: + args = list(_case(op.CELLS[2].spec, seed=6)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_index_past_the_batch() -> None: + args = list(_case(op.CELLS[2].spec, seed=7)) + args[3] = args[3].clone() + args[3][-1] = args[7] + with pytest.raises(AssertionError, match="q_batch_row"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py new file mode 100644 index 000000000000..822a3b2a0ffc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the ``m3_allreduce`` entry, 2 ranks (SM107): its cells against the entry's reference, +and the fused draft loop's layer boundaries built on it (``M3Eagle3Draft.ar_norm``: the plain all-reduce, then +FlashInfer's fused add + RMSNorm) against stock Eagle3's, bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy the single-node drafter can run (AUTO, +ONESHOT, TWOSHOT, NCCL; MNNVL is multi-node only), at tp2 shapes (H = 6144, eps = 1e-6, 1 to 32 rows). The inputs are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own 2-rank launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_allreduce.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py +""" + +import math +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 4, 8, 32) +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +RANK = 0 +WORLD = 1 + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on each rank. + parts = [x / 2, x / 2] + return ( + [p.bfloat16() for p in parts[:WORLD]], + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A + B, -A - B + d) onto a unit residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a + b, (-a - b + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return [p.bfloat16() for p in parts[:WORLD]], residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +def _ulps_at(out: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the magnitude ``scale`` (the operands'), not at the result's.""" + return ((out.double() - ref.double()).abs() / _ulp_at(scale.abs())).max().item() + + +def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary (``M3Eagle3Draft.ar_norm``): ``m3_allreduce`` plain, then the catalog's + ``flashinfer_fused_add_rmsnorm``.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + + return flashinfer_fused_add_rmsnorm(m3_allreduce(x, ar, None), residual.clone(), w, EPS) + + +def _check_cells(mapping: Any) -> None: + """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum + to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude + u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final + roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). + The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm.functional import AllReduceStrategy + + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating + + ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + for i, cell in enumerate(m3_allreduce.CELLS): + spec = cell.spec + xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) + params = ( + None + if spec.get("plain") + else AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual.clone(), + norm_weight=w, + eps=EPS, + trigger_completion_at_end=False, + ) + ) + # The reference first: the fused path may write the updated residual into ``params.residual``. + ref = m3_allreduce.reference(xs, ar, params) + with validating(m3_allreduce): + out = m3_allreduce(xs[RANK], ar, params) + torch.cuda.synchronize() + total = sum(x.double() for x in xs) + if params is None: + err = _ulps_at(out, ref, total) + assert err <= 1.0, (cell.why, err) + else: + err = _ulps_at(out[1], ref[1], torch.maximum(total.abs(), residual.double().abs())) + assert err <= 2.5, (cell.why, "residual", err) + h = out[1].double() + own = h * torch.rsqrt(h.square().mean(-1, keepdim=True) + EPS) * w.double() + assert _ulps(out[0], own) <= 2.0, (cell.why, "normed", _ulps(out[0], own)) + dist.barrier() + if RANK == 0: + print(f"{RESULT_TAG} cell {cell.why} passed", flush=True) + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + _check_cells(mapping) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + ulp, res_ulp = 0.0, 0.0 + for m in TOKENS: + xs, residual, w = build(m) + target_out, target_res = _target(ar, xs[RANK], residual, w) + stock_out, stock_res = _stock(ar, xs[RANK], residual, w) + torch.cuda.synchronize() + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {name} m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {name:34s} | target = stock vs fp64: normed {ulp:.3f}" + f" | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 2, f"the drafter runs at tp2, got {WORLD} ranks" + try: + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + ) + 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", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..8e0e595bdc1c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,755 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``, 2 ranks (SM107). + +Its own launcher, as the catalog's collective matrices are: run plainly it spawns +one rank process per device in CUDA_VISIBLE_DEVICES and kills them all on the +first failure or on the deadline. A rank whose peers never push polls forever, +so the deadline is what keeps a broken exchange from wedging the calling run. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_mc_allreduce_norm_op_matrix.py + +Each rank builds the three multicast workspaces (8 / 16 / 48 tokens) the target +builds, over the torch.distributed group (the non-MPI path of +``_make_mnnvl_mcast_buffer``). Every rank can regenerate every rank's inputs from +the shared seed convention, which is what keeps each reference arithmetic rather +than a second collective. Producer pushes (``prepushed`` 2 / 3) are emulated by +storing the canonicalized rows through the multicast mapping at the call parity +the flags select, which is exactly what the producing kernels store. + +Not a pytest module, despite the ``check_*`` bodies: they are one fixed sequence +inside a single 2-rank job whose workspace state (call parities) carries from +one to the next. The collected entry point is +``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +finalize_entry: Any = None +ar_module: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +# ---------------------------------------------------------------- inputs + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + """One rank's [m, H] BF16 rows; ``rank`` -1 for tensors every rank holds alike.""" + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return _rows(seed + 7, -1, m, 4.0) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _finalize_inputs(seed: int, rank: int, m: int, topk: int, dropped: int): + """One rank's unfinalized MoE outputs: fc2 [m * topk, H], perm [m, topk], fp32 scales, shared [m, H].""" + g = _gen(seed + 23, rank) + rows = m * topk + fc2 = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + perm = torch.randperm(rows, generator=g, device="cuda").to(torch.int32) + if dropped: + drop = torch.randperm(rows, generator=g, device="cuda")[:dropped] + perm[drop] = -1 + perm = perm.view(m, topk).contiguous() + scales = torch.rand(m, topk, generator=g, device="cuda", dtype=torch.float32) + shared = _rows(seed + 29, rank, m, 0.5) + return fc2, perm, scales, shared + + +# ---------------------------------------------------------------- workspace / producer emulation + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor, exchange_set: int) -> None: + """Store this rank's rows into every rank's exchange, as a producer kernel does before a prepushed call. + + Slot of (token t, cluster CTA c): set * 2 * half + parity(t, c) * half + (t * world + rank) * H / 2 + c * H / 16 + words, parity read from the CTA's flag; -0.0 halves are stored as +0.0 (0x80000000 is the empty word). + """ + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = ( + exchange_set * 2 * half + + parity.view(m, CLUSTER, 1) * half + + (t * WORLD + RANK) * (HIDDEN // 2) + + c * chunk + + w + ) + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +# ---------------------------------------------------------------- gates + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +#: Share of normed elements one BF16 step off the float64 reference: measured at most 1.1e-4 per call over +#: every call of this matrix on GB300 (see ``_assert_normed``). The gate sits ~18x above that and far below what a +#: wrong statistic produces (most elements move). +NORMED_STEP_FRACTION = 0.002 +NORMED_STEP_SEEN: List[float] = [] + + +def _assert_normed(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """At most one BF16 rounding step from the float64 reference, and rarely that. + + The kernel's RMS statistic is a shuffle / warp / cluster tree of FP32 partial sums and its ``rsqrtf`` is + approximate (~2^-22 relative): the scaled value moves by ~1e-6 relative, which changes the BF16 rounding only + for values within that distance of a rounding midpoint -- one step, never two. The fraction bound catches a + statistic that is systematically off (a missing rank, a wrong H), which moves most elements. + """ + assert out.dtype == ref.dtype == torch.bfloat16 and out.shape == ref.shape, where + o, r = out.float(), ref.float() + _, exponent = torch.frexp(torch.maximum(o.abs(), r.abs())) + step = torch.ldexp(torch.ones_like(o), exponent - 8) + diff = (o - r).abs() + assert bool((diff <= step).all()), f"{where}: max diff {diff.max().item()} beyond one BF16 step" + frac = (diff != 0).float().mean().item() + NORMED_STEP_SEEN.append(frac) + assert frac <= NORMED_STEP_FRACTION, f"{where}: {frac:.4f} of normed elements off by one step" + + +def _assert_quant(outs, ref, m: int, where: str) -> None: + """The MXFP8 copy: bit-exact against the recipe applied to the op's own normed rows (a pure function of + them), and equal to the reference's wherever the reference's normed 32-column block is the op's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + mxfp8_rows, + sf_written, + ) + + q, sf = outs[2], outs[3] + own_q, own_sf = mxfp8_rows(outs[0]) + written = sf_written(m, HIDDEN // 32, sf.device) + _assert_bitwise(q, own_q, f"{where} e4m3 rows") + _assert_bitwise(sf[written], own_sf[written], f"{where} UE8M0 scales") + same_block = (outs[0].view(m, -1, 32) == ref[0].view(m, -1, 32)).all(-1) + qb = q.view(torch.uint8).view(m, -1, 32) + rqb = ref[2].view(torch.uint8).view(m, -1, 32) + assert bool((qb == rqb).all(-1)[same_block].all()), ( + f"{where}: e4m3 differs on a block whose normed rows agree" + ) + + +def _check_outputs(outs, ref, m: int, quant: bool, where: str) -> None: + assert len(outs) == (4 if quant else 2), (where, len(outs)) + _assert_bitwise(outs[1], ref[1], f"{where} updated residual") + _assert_normed(outs[0], ref[0], f"{where} normed") + if quant: + _assert_quant(outs, ref, m, where) + + +# ---------------------------------------------------------------- calls + + +def _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, opts): + return allreduce_entry.m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + ws.max_tokens, + quant, + prepushed, + opts, + ) + + +def _plain_case(spec: Dict[str, Any], seed: int, where: str) -> None: + m, maxt, quant, prepushed = spec["m"], spec["maxt"], spec["quant"], spec["prepushed"] + ws = WORKSPACES[maxt] + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] if spec["x2"] else None + residual, norm_w = _residual(seed, m), _norm_w(seed) + x, x2 = xs[RANK], (None if x2s is None else x2s[RANK]) + if prepushed >= 2: + _push(ws, x, 0) + if prepushed == 3: + _push(ws, x2, 1) + with validating(allreduce_entry.m3_mc_allreduce_norm): + outs = _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, spec["opts"]) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, residual, norm_w, EPS, quant, spec.get("drafter", False) + ) + _check_outputs(outs, ref, m, quant, where) + dist.barrier() + + +# ---------------------------------------------------------------- checks + + +def check_plain_cells() -> None: + """Every cell of ``m3_mc_allreduce_norm.CELLS`` through ``__call__`` under ``validating``, gated on its + ``reference``.""" + for i, cell in enumerate(allreduce_entry.m3_mc_allreduce_norm.CELLS): + _plain_case(cell.spec, 1000 + 17 * i, cell.why) + + +def check_finalize_cells() -> None: + """Every cell of ``m3_mc_finalize_allreduce_norm.CELLS``, gated on its ``reference``.""" + op = finalize_entry.m3_mc_finalize_allreduce_norm + for i, cell in enumerate(op.CELLS): + spec = cell.spec + m, maxt, quant, topk = spec["m"], spec["maxt"], spec["quant"], spec["topk"] + seed = 2000 + 17 * i + ws = WORKSPACES[maxt] + per_rank = [_finalize_inputs(seed, r, m, topk, spec["dropped"]) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + with validating(op): + outs = op( + fc2, + perm, + scales, + shared, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + maxt, + quant, + spec["opts"], + ) + torch.cuda.synchronize() + ref = op.reference( + [p[0] for p in per_rank], + [p[1] for p in per_rank], + [p[2] for p in per_rank], + [p[3] for p in per_rank], + residual, + norm_w, + EPS, + quant, + ) + _check_outputs(outs, ref, m, quant, cell.why) + dist.barrier() + + +def _capture_reference( + updated: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor +) -> None: + """The capture the boundary replaces: m3_spec_tail.capture_norm of the updated residual (no second addend).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_spec_tail, + ) + + m3_spec_tail.capture_norm(updated, None, weight, eps, out) + + +def check_capture_norm_in_boundary() -> None: + """CAPTURE_NORM_IN_BOUNDARY: with cap_out the boundary also writes the capture, BITWISE equal to capture_norm of + its own updated residual, into a strided [M, H] block of a [rows, 3H] capture buffer (other bytes untouched); its + normed / updated / MXFP8 outputs are bitwise the ones of the same call without cap_out. The row kernel (the plain + entry at 2 ranks: x, x + x2, prepushed) and the cluster kernel (MoE finalize), 1 / 4 / 8 rows, residual scales + from 1e-3 to 1e3 (the statistic's range), the fc_norm eps 1e-6 and 1e-5.""" + cap_eps_values = (1e-6, 1e-5) + for case, (m, scale, kind) in enumerate( + ( + (1, 1.0, "plain"), + (4, 1e3, "plain_x2"), + (8, 1e-3, "prepushed"), + (4, 1.0, "prepushed"), + (1, 4.0, "finalize"), + (4, 1e3, "finalize"), + (8, 1e-3, "finalize"), + ) + ): + seed = 5000 + 31 * case + cap_eps = cap_eps_values[case % 2] + ws = WORKSPACES[8] + residual = (_residual(seed, m).float() * scale).to(torch.bfloat16) + norm_w = _norm_w(seed) + cap_w = _norm_w(seed + 101) + quant = case % 3 == 1 + outs_by = [] + for with_cap in (False, True): + buf = torch.full((m + 3, 3 * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + cap = buf[:m, HIDDEN : 2 * HIDDEN] if with_cap else None + if kind == "finalize": + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, m, 4, 1) + fc2 = (fc2.float() * scale).to(torch.bfloat16) + outs = finalize_entry.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), + ws.flags_for(CLUSTER), RANK, WORLD, CLUSTER, False, 8, quant, 7, + cap, cap_w if with_cap else None, cap_eps, + ) # fmt: skip + else: + x = (_rows(seed, RANK, m).float() * scale).to(torch.bfloat16) + x2 = ( + (_rows(seed + 3, RANK, m).float() * scale).to(torch.bfloat16) + if kind == "plain_x2" + else None + ) + prepushed = 2 if kind == "prepushed" else 0 + if prepushed: + _push(ws, x, 0) + outs = allreduce_entry.m3_mc_allreduce_norm( + x, x2, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), ws.flags_for(CLUSTER), RANK, WORLD, + CLUSTER, False, 8, quant, prepushed, (7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + cap, cap_w if with_cap else None, cap_eps, + ) # fmt: skip + torch.cuda.synchronize() + outs_by.append((outs, buf)) + dist.barrier() + (plain, _), (capped, buf) = outs_by + where = f"capture in boundary, {kind}, {m} rows, scale {scale}" + assert len(plain) == len(capped), where + for i, (a, b) in enumerate(zip(plain, capped)): + if i == 3: + # The scale buffer is torch.empty: compare the bytes the kernel writes (the valid rows'). + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written, + ) + + written = sf_written(m, HIDDEN // 32, a.device) + a, b = a[written], b[written] + _assert_bitwise(b, a, f"{where}: output {i} with vs without cap_out") + ref = torch.empty((m, HIDDEN), dtype=torch.bfloat16, device="cuda") + _capture_reference(capped[1], cap_w, cap_eps, ref) + torch.cuda.synchronize() + _assert_bitwise(buf[:m, HIDDEN : 2 * HIDDEN], ref, f"{where}: capture vs capture_norm") + untouched = buf.clone() + untouched[:m, HIDDEN : 2 * HIDDEN] = 7.0 + assert bool((untouched == 7.0).all()), f"{where}: bytes outside the capture block changed" + dist.barrier() + + +def check_the_target_wrappers_select_bucket_and_options() -> None: + """The target's own wrappers (bucket by token count, option table) on top of the entries, every token count of + the 8-token bucket plus one of each larger bucket, and the finalize wrapper; the MXFP8 publication is the + prequant slot's business and is left off.""" + buckets = ar_module.McBuckets([WORKSPACES[b] for b in BUCKETS]) + for m in (*range(1, 9), 11, 40): + seed = 3000 + m + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + normed, updated = ar_module.m3_mc_allreduce_norm( + xs[RANK], residual, norm_w, EPS, buckets, x2=x2s[RANK] + ) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, False) + _check_outputs([normed, updated], ref, m, False, f"wrapper m={m}") + dist.barrier() + for m in (9, 33): + seed = 3100 + m + per_rank = [_finalize_inputs(seed, r, m, 4, 2) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + normed, updated = ar_module.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, buckets + ) + torch.cuda.synchronize() + ref = finalize_entry.m3_mc_finalize_allreduce_norm.reference( + *[[p[j] for p in per_rank] for j in range(4)], residual, norm_w, EPS, False + ) + _check_outputs([normed, updated], ref, m, False, f"finalize wrapper m={m}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """Two boundaries captured in one graph (8- and 16-token buckets), replayed with fresh inputs. + + The decode graph holds every boundary call; each replay rotates the same per-CTA parities the eager calls do, + so a replay that reused a stale half or skipped a re-arm reads the previous replay's rows and fails the + bitwise residual gate (every replay gets new inputs). + """ + shapes = ((3, 8, 15), (12, 16, 15)) + norm_w = _norm_w(4000) + bufs = [] + for m, maxt, opts in shapes: + bufs.append( + dict( + m=m, + maxt=maxt, + opts=opts, + x=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + x2=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + residual=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + ) + + def body(): + return [ + _call_plain( + WORKSPACES[b["maxt"]], b["x"], b["x2"], b["residual"], norm_w, False, 0, b["opts"] + ) + for b in bufs + ] + + def fill(seed: int): + refs = [] + for j, b in enumerate(bufs): + s = seed + 31 * j + xs = [_rows(s, r, b["m"]) for r in range(WORLD)] + x2s = [_rows(s + 3, r, b["m"]) for r in range(WORLD)] + b["x"].copy_(xs[RANK]) + b["x2"].copy_(x2s[RANK]) + b["residual"].copy_(_residual(s, b["m"])) + refs.append( + allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, b["residual"].clone(), norm_w, EPS, False + ) + ) + torch.cuda.synchronize() + dist.barrier() + return refs + + refs = fill(4100) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm = body() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(warm[j], refs[j], b["m"], False, f"graph warm-up m={b['m']}") + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = body() + for rep in range(4): + refs = fill(4200 + 100 * rep) + graph.replay() + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(captured[j], refs[j], b["m"], False, f"replay {rep} m={b['m']}") + dist.barrier() + del graph + dist.barrier() + + +def check_the_gate_sees_a_missing_rank() -> None: + """The residual gate is not blind to the reduction: dropping any one rank's partial from the reference + moves most elements of the updated residual (each rank's rows are distinct by construction).""" + m, seed = 4, 5000 + xs = [_rows(seed, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + full = allreduce_entry.m3_mc_allreduce_norm.reference(xs, None, residual, norm_w, EPS, False) + for r in range(WORLD): + partial = [x if i != r else torch.zeros_like(x) for i, x in enumerate(xs)] + wrong = allreduce_entry.m3_mc_allreduce_norm.reference( + partial, None, residual, norm_w, EPS, False + ) + frac = (wrong[1] != full[1]).float().mean().item() + assert frac > 0.9, f"dropping rank {r} changes only {frac:.3f} of the residual" + dist.barrier() + + +def _expect_rejected(fn, match: str, where: str) -> None: + try: + fn() + except AssertionError as exc: + assert match in str(exc), f"{where}: rejected for the wrong reason: {exc}" + else: + raise AssertionError(f"{where}: the guard admitted it") + + +def check_guards() -> None: + """``is_valid`` refuses what the op accepts and answers wrongly. Nothing here reaches the kernel: the guard + raises first, on every rank alike, so no rank is left polling.""" + m, seed = 4, 6000 + ws8, ws16 = WORKSPACES[8], WORKSPACES[16] + x, residual, norm_w = _rows(seed, RANK, m), _residual(seed, m), _norm_w(seed) + op = allreduce_entry.m3_mc_allreduce_norm + + def call(ws_uc, flags, norm): + with validating(op): + op( + x, + None, + residual, + norm, + EPS, + ws_uc, + ws8.mc.data_ptr(), + flags, + RANK, + WORLD, + CLUSTER, + False, + 8, + False, + 0, + 15, + ) + + strided = torch.empty(2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[::2] + strided.copy_(norm_w) + _expect_rejected( + lambda: call(ws8.uc, ws8.flags_for(CLUSTER), strided), "contiguous", "strided norm_w" + ) + _expect_rejected( + lambda: call(ws8.uc[: ws8.uc.numel() // 2], ws8.flags_for(CLUSTER), norm_w), + "words", + "short workspace", + ) + _expect_rejected( + lambda: call(ws8.uc, ws16.flags_for(CLUSTER), norm_w), "counters", "another bucket's flags" + ) + + fin = finalize_entry.m3_mc_finalize_allreduce_norm + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, 9, 4, 0) + bad = perm.clone() + bad[0, 0] = fc2.shape[0] + + def call_fin(): + with validating(fin): + fin( + fc2, + bad, + scales, + shared, + _residual(seed, 9), + norm_w, + EPS, + ws16.uc, + ws16.mc.data_ptr(), + ws16.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + 16, + False, + 15, + ) + + _expect_rejected(call_fin, "perm spans", "perm past the FC2 rows") + dist.barrier() + + +CHECKS = ( + check_plain_cells, + check_finalize_cells, + check_capture_norm_in_boundary, + check_the_target_wrappers_select_bucket_and_options, + check_cuda_graph_replays, + check_the_gate_sees_a_missing_rank, + check_guards, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, allreduce_entry, finalize_entry, ar_module, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 2, f"the entries are certified at 2 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as _finalize, + ) + + allreduce_entry, finalize_entry, ar_module = _plain, _finalize, _ar + # Spawned by file path, not collected by pytest: reach the sibling helper the plain way (confined to the + # rank processes). /comm -> tests/unittest/_torch/modeling_v2. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peers' polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print( + f"[rank {RANK}] {len(CHECKS)} checks passed; normed one-step fraction max " + f"{max(NORMED_STEP_SEEN):.5f} over {len(NORMED_STEP_SEEN)} calls", + flush=True, + ) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + ) + 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", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + 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), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py new file mode 100644 index 000000000000..c6baaa215341 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py @@ -0,0 +1,411 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification of the MoE engine's all-reduce push (``m3_moe_engine`` with ``arp``), 2 ranks (SM107). + +The decode forward's MoE boundary (MOE_PUSH rows): the engine pushes its bf16 sum (routed + shared) into set A of +the boundary's multicast exchange while it writes it, and the boundary all-reduce runs with ``prepushed`` 2, reading +the pushed words instead of loading and pushing the rows itself. Certified against the unpushed boundary the target +runs otherwise: the same engine call without ``arp``, then ``m3_mc_allreduce_norm`` with ``prepushed`` 0. Every +output must be bit for bit the same: the engine's own output tensor, the updated residual, the normed rows and, with +``quant``, the MXFP8 hand-off it publishes. Both read the same bf16 words, so the reduction and the norm see equal +inputs; any difference is a push that stored the wrong word, slot or parity. + +Each rank holds its own expert shard (operand seed per rank) and the same hidden rows, residual and norm weight, as +the TP2 decode step does. Eager calls alternate pushed and unpushed in both orders so both call parities of the +exchange are exercised; a CUDA graph of the pushed engine + boundary is replayed on fresh rows against eager +unpushed references. The engine's early PDL release (EARLY_AR, the target's default) is checked against the release +after the push: scheduling only, the same bits. + +Its own launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is: run plainly it spawns one rank process per device +in CUDA_VISIBLE_DEVICES and kills them all on the first failure or on the deadline (a rank whose peer never pushes +polls forever). The collected entry point is ``test_modeling_v2_m3_moe_push_op_matrix.py``. + + CUDA_VISIBLE_DEVICES=0,1 python _m3_moe_push_op_matrix.py +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_moe_engine requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1200 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +MAX_TOKENS = 8 +EMPTY_WORD = -2147483648 +RESERVED_SMS = 24 +OPERAND_SEED = 20261013 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +ar = None +sf_written: Any = None +engine_op: Any = None +ar_push_desc: Any = None +moe_state: Any = None +operands: Any = None +prequant: Any = None +RANK = 0 +WORLD = 1 +WS: Any = None + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return (torch.randn(m, HIDDEN, generator=_gen(seed + 7), device="cuda") * 4.0).to( + torch.bfloat16 + ) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _make_workspace(): + """The 8-token bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed + group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * MAX_TOKENS * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=MAX_TOKENS, + ) + ws.flags_for(CLUSTER) + return ws + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +_STATES = {} +_LAYERS = {} + + +def _layer(m: int, case: str, early_release: bool = True): + """This rank's engine layer for ``m`` tokens: its own expert shard; ``concentrated`` routes every token to the + same four experts. ``early_release``: the state's EARLY_AR build (the target's default) or the release after the + push.""" + if (m, early_release) not in _STATES: + sms = torch.cuda.get_device_properties(0).multi_processor_count + _STATES[m, early_release] = moe_state( + torch.device("cuda"), + num_tokens=m, + num_ctas=sms - RESERVED_SMS, + shared=True, + early_release=early_release, + ) + key = (m, case, early_release) + if key not in _LAYERS: + ops = operands.experts(OPERAND_SEED + RANK) + bias = ops.bias.clone() + if case == "concentrated": + bias[list(operands.CONCENTRATED)] += 1.0 + _LAYERS[key] = _STATES[m, early_release].layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=ops.shared, + ) + return _LAYERS[key] + + +def _desc(): + flags = WS.flags_for(CLUSTER) + return ar_push_desc( + WS.mc.data_ptr(), flags.data_ptr(), RANK, WORLD, WS.max_tokens, flags.device + ) + + +def _boundary(x, residual, norm_w, quant, prepushed): + prequant.SLOT.clear() + outs = ar.m3_mc_allreduce_norm( + x, residual, norm_w, EPS, WS, use_pdl=True, quant=quant, prepushed=prepushed + ) + assert outs is not None, "the boundary all-reduce declined its operands" + pub = prequant.SLOT.pending + prequant.SLOT.clear() + return outs, (None if pub is None else (pub[3].clone(), pub[4].clone())) + + +def _unpushed(hidden, layer, residual, norm_w, quant): + out = engine_op(hidden, layer) + (normed, res), pub = _boundary(out, residual.clone(), norm_w, quant, 0) + return out.clone(), normed.clone(), res.clone(), pub + + +def _pushed(hidden, layer, residual, norm_w, quant): + out = engine_op(hidden, layer, arp=_desc()) + (normed, res), pub = _boundary(out, residual.clone(), norm_w, quant, 2) + return out.clone(), normed.clone(), res.clone(), pub + + +def _compare(got, ref, quant: bool, where: str) -> None: + _assert_bitwise(got[0], ref[0], f"{where} engine output") + _assert_bitwise(got[2], ref[2], f"{where} updated residual") + _assert_bitwise(got[1], ref[1], f"{where} normed") + if quant: + assert got[3] is not None and ref[3] is not None, f"{where}: no MXFP8 hand-off published" + _assert_bitwise(got[3][0], ref[3][0], f"{where} e4m3 rows") + # Only the scale atoms of the m rows are written; the swizzled buffer's padding is not. + written = sf_written(got[0].shape[0], HIDDEN // 32, got[3][1].device) + _assert_bitwise(got[3][1][written], ref[3][1][written], f"{where} UE8M0 scales") + + +# ---------------------------------------------------------------- checks + + +def check_pushed_boundary_is_the_unpushed_one() -> None: + """Every token count, both routing cases, with and without the MXFP8 hand-off, both call orders.""" + for m in range(1, MAX_TOKENS + 1): + for case in ("random", "concentrated"): + layer = _layer(m, case) + for quant in (False, True): + for seed in (1, 2): + s = ( + 1000 * m + + 10 * seed + + (5 if quant else 0) + + (500 if case == "concentrated" else 0) + ) + hidden = operands.hidden_states(m, s) + residual, norm_w = _residual(s, m), _norm_w(s) + where = f"m={m} {case} quant={quant} seed={seed}" + if seed == 1: + ref = _unpushed(hidden, layer, residual, norm_w, quant) + got = _pushed(hidden, layer, residual, norm_w, quant) + else: + got = _pushed(hidden, layer, residual, norm_w, quant) + ref = _unpushed(hidden, layer, residual, norm_w, quant) + torch.cuda.synchronize() + _compare(got, ref, quant, where) + + +def check_graph_replays() -> None: + """A captured pushed engine + prepushed boundary, replayed on fresh rows, equals the eager unpushed path.""" + for m in (1, 4, 8): + layer = _layer(m, "random") + hidden = operands.hidden_states(m, 90000 + m).clone() + residual = _residual(90000 + m, m).clone() + norm_w = _norm_w(90000 + m) + desc = _desc() + for _ in range( + 2 + ): # warm both parities eagerly; every push is consumed by its prepushed boundary + _pushed(hidden, layer, residual, norm_w, False) + _unpushed(hidden, layer, residual, norm_w, False) + torch.cuda.synchronize() + dist.barrier() + res_in = residual.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + g_out = engine_op(hidden, layer, arp=desc) + g_normed, g_res = ar.m3_mc_allreduce_norm( + g_out, res_in, norm_w, EPS, WS, use_pdl=True, prepushed=2 + ) + torch.cuda.synchronize() + dist.barrier() + for rep in range(6): + s = 91000 + 100 * m + rep + fresh_h, fresh_r = operands.hidden_states(m, s), _residual(s, m) + hidden.copy_(fresh_h) + res_in.copy_(fresh_r) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + got = (g_out.clone(), g_normed.clone(), g_res.clone(), None) + ref = _unpushed(fresh_h, layer, fresh_r, norm_w, False) + torch.cuda.synchronize() + _compare(got, ref, False, f"graph m={m} replay {rep}") + dist.barrier() + + +def check_early_release_is_scheduling_only() -> None: + """The pushed engine + boundary with the early PDL release and with the release after the push: the same bits.""" + for m in range(1, MAX_TOKENS + 1): + for quant in (False, True): + s = 50000 + 10 * m + (5 if quant else 0) + hidden = operands.hidden_states(m, s) + residual, norm_w = _residual(s, m), _norm_w(s) + early = _pushed(hidden, _layer(m, "random", True), residual, norm_w, quant) + late = _pushed(hidden, _layer(m, "random", False), residual, norm_w, quant) + torch.cuda.synchronize() + _compare(early, late, quant, f"m={m} quant={quant} early vs late release") + + +CHECKS = [ + check_pushed_boundary_is_the_unpushed_one, + check_early_release_is_scheduling_only, + check_graph_replays, +] + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, ar, sf_written, engine_op, ar_push_desc, moe_state, operands, prequant, WS + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 2, f"the push is certified at 2 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_prequant as _pq # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written as _sf_written, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + M3MoeState, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + ar_push_desc as _desc_fn, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as _op, + ) + + ar, prequant, moe_state, ar_push_desc, engine_op = _ar, _pq, M3MoeState, _desc_fn, _op + sf_written = _sf_written + # Spawned by file path, not collected by pytest: the sibling moe/ operands the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "moe")) + import _m3_moe_operands as _operands + + operands = _operands + WS = _make_workspace() + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peer's polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + ) + 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", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + 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), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_rank_job.py new file mode 100644 index 000000000000..2d024dc2d3e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_rank_job.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run this target's collective rank job from pytest, and report what it did. + +The same contract as ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``: the +launcher beside this file owns the per-device rank processes and the deadline +that kills a wedged collective; this only selects the devices, starts it in a +fresh interpreter (by file path: this tree is not a package) and turns its exit +code into an assertion. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +WORLD_SIZE = 2 +"""tp2's world size -- the topology these entries are certified for.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the launcher's own deadline, so its message wins the race.""" + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + if len(devices) < WORLD_SIZE: + pytest.skip( + f"this entry is certified at world size {WORLD_SIZE}; only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def _load_launcher_constants(launcher: Path): + """Read the launcher's module-level budget without importing it by name (it imports only torch at + module scope).""" + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str) -> str: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices; returns its stdout.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + + # The ranks import tensorrt_llm absolutely. From a source checkout that is not installed, put the repo root + # (tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root) on their path; from a tests + # tree run against an installed package, they import the installed one. + repo_root = Path(__file__).resolve().parents[8] + if (repo_root / "tensorrt_llm").is_dir(): + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p + ) + else: + assert importlib.util.find_spec("tensorrt_llm") is not None, ( + f"no tensorrt_llm/ at the repo root {repo_root} and no installed tensorrt_llm" + ) + + constants = _load_launcher_constants(launcher) + timeout = constants.DEADLINE_S + _LAUNCHER_GRACE_S + + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) + return completed.stdout diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_allreduce.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_allreduce.py new file mode 100644 index 000000000000..2e3d69a51316 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_allreduce.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the ``m3_allreduce`` certification matrix. + +The matrix is ``_m3_allreduce_op_matrix.py`` beside this file, its own 2-rank launcher; see ``_rank_job``. It drives +every cell of the entry against its reference, then requires the fused draft loop's boundaries built on it (plain +all-reduce, then the fused add + RMSNorm) to equal stock Eagle3's bit for bit under every strategy, and reports both +against an fp64 reference as a warning, one line per strategy and case. +""" + +import warnings + +import _rank_job +import pytest +import torch + +pytestmark = [ + pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _rank_job.WORLD_SIZE, + reason=f"certified at world size {_rank_job.WORLD_SIZE}", + ), +] + +_RESULT_TAG = "[ulp]" + + +# Spawns 2 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_allreduce_op_matrix() -> None: + out = _rank_job.run("m3_allreduce") + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + "draft boundary (= stock Eagle3) vs fp64, max BF16 ulps\n" + "\n".join(lines), stacklevel=1 + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..75f31cdb994b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the multicast all-reduce + RMSNorm entries' certification matrix. + +Covers both ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm`` (one +kernel, one exchange protocol; the finalize entry's reference has its own test in +``test_modeling_v2_m3_mc_finalize_allreduce_norm.py``). The matrix is ``_m3_mc_allreduce_norm_op_matrix.py`` +beside this file, its own 2-rank launcher; see ``_rank_job`` for why. +""" + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + + +# Spawns 2 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_mc_allreduce_norm_op_matrix() -> None: + _rank_job.run("m3_mc_allreduce_norm") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..984ed13ced0e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``m3_mc_finalize_allreduce_norm``: its reference's rounding. The entry's GPU cells (9..48 rows, every one gated on +that reference) run in the multicast all-reduce matrix, ``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``, which +launches both entries of the one kernel at 2 ranks.""" + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + finalize_partial, + m3_mc_finalize_allreduce_norm, +) + + +def test_finalize_partial_rounds_the_routed_sum_before_the_shared_expert() -> None: + """As stock: the finalized routed output is BF16, and the shared expert is added to it (shared.add_(routed)). + Four routed rows of 1.0 at FP32 routing weight 0.2509765625 sum to 1.00390625, which rounds to BF16 1.0; plus the + shared -1 that is 0, where adding before the rounding would leave 0.00390625.""" + fc2 = torch.ones((4, 8), dtype=torch.bfloat16) + perm = torch.arange(4, dtype=torch.int32).view(1, 4) + scales = torch.full((1, 4), 0.2509765625, dtype=torch.float32) + shared = torch.full((1, 8), -1.0, dtype=torch.bfloat16) + out = finalize_partial(fc2, perm, scales, shared) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(out, shared.clone().add_(routed)) + assert torch.equal(out, torch.zeros(1, 8, dtype=torch.bfloat16)) + + +def test_the_cells_cover_every_row_the_target_sends_here() -> None: + """The decode layer's MC finalize (MC_FINALIZE_MIN_TOKENS 9 .. the b32 plan's 32 rows) and the generic layer's + (past 32, up to GENERIC_MC_FINALIZE_MAX_TOKENS 48): every one of those row counts is a certified cell.""" + rows = {cell.spec["m"] for cell in m3_mc_finalize_allreduce_norm.CELLS} + assert set(range(9, 49)) <= rows diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_moe_push_op_matrix.py new file mode 100644 index 000000000000..c548a51e3060 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_moe_push_op_matrix.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the MoE engine's all-reduce push certification (``_m3_moe_push_op_matrix.py`` beside +this file, its own 2-rank launcher; see ``_rank_job``).""" + +import _rank_job +import pytest +import torch + +pytestmark = [ + pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _rank_job.WORLD_SIZE, + reason=f"certified at world size {_rank_job.WORLD_SIZE}", + ), +] + + +# Spawns 2 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_moe_push_op_matrix() -> None: + _rank_job.run("m3_moe_push") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..b77aabe64e20 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_embedding entry of the MiniMax-M3 tp2 target: an exact gather with zero rows outside the +table.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.embedding.m3_embedding import ( # noqa: E501 + m3_embedding as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" + +VOCAB, HIDDEN = 4096, 6144 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + ids = torch.randint( + 0, VOCAB, (spec["rows"],), generator=g, device="cuda", dtype=spec["id_dtype"] + ) + if spec["pad"]: + ids[1::3] = -1 + with validating(op): + out = op(ids, table) + assert torch.equal(out, op.reference(ids, table)) + if spec["pad"]: + assert not out[1::3].any(), "padded ids must embed to zero rows" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py new file mode 100644 index 000000000000..e44ad56d743d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" + +import torch + + +def swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): + """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" + v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) + return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..b455f6f2a049 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_bf16_linear entry of the MiniMax-M3 tp2 target (the drafter's projections).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_bf16_linear import ( # noqa: E501 + m3_bf16_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the + row, RMS well under one; measured on GB300: worst 0.5 unit.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w = (torch.randn(spec["n"], spec["k"], generator=g, device="cuda") * 0.02).to(torch.bfloat16) + with validating(op): + out = op(x, w) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(x, w), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py new file mode 100644 index 000000000000..f1d036a31909 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the dense MLP's activation seam (``modeling._dense_swiglu``, the ``m3_dense_swiglu`` C++ op): between +the dense MLP's two decode GEMVs, the seam's down-GEMV output is bit for bit the Triton swigluoai's +(``m3_silu_and_mul``) at every decode row count, eager and in a CUDA graph, with PDL; and the seam declines (to the +Triton kernel) what it does not cover.""" + +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.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + DENSE_GATE_UP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.mx_gemv import ( # noqa: E501 + op as gemv, +) + +assert torch.cuda.is_available(), "the dense seam test requires a CUDA device" + +ALPHA, LIMIT, BETA = 1.702, 7.0, 1.0 +ROWS = list(range(1, modeling.TINY_MAX_TOKENS + 1)) +CORE = SimpleNamespace(swiglu_alpha=ALPHA, swiglu_limit=LIMIT, swiglu_beta=BETA) + + +def _weights(n, k, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return wq.view(n, k), wsf + + +GATE_UP = None +DOWN = None + + +@pytest.fixture(scope="module", autouse=True) +def _mlp_weights(): + global GATE_UP, DOWN + GATE_UP = _weights(*DENSE_GATE_UP, 11) + DOWN = _weights(*DENSE_DOWN, 13) + + +def _gemv(x, wts): + out = gemv.m3_mx_gemv_bf16(x.contiguous(), *wts, use_pdl=True) + assert out is not None, ( + "the target's bf16 GEMV serves every decode row of the dense projections" + ) + return out + + +def _mlp(x, seam: bool): + """gate_up GEMV, the activation (the seam, or the Triton kernel), down GEMV.""" + gu = _gemv(x, GATE_UP) + if seam: + act = modeling._dense_swiglu(CORE, gu) + else: + act = modeling.m3_silu_and_mul(gu, LIMIT, ALPHA, BETA) + return _gemv(act, DOWN) + + +def _x(rows, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + return (torch.randn(rows, DENSE_GATE_UP[1], generator=g, device="cuda") * 2).to(torch.bfloat16) + + +@pytest.fixture +def op_calls(monkeypatch): + """Records each call that reaches the C++ op through the seam.""" + calls = [] + real = modeling.m3_dense_swiglu + + class _Recording: + CELLS = real.CELLS + + def __call__(self, gu, *args): + calls.append(gu.shape[0]) + return real(gu, *args) + + monkeypatch.setattr(modeling, "m3_dense_swiglu", _Recording()) + return calls + + +@pytest.mark.parametrize("rows", ROWS) +def test_seam_on_is_bitwise_seam_off(op_calls, rows) -> None: + for trial in range(3): + x = _x(rows, 100 * rows + trial) + on = _mlp(x, True) + off = _mlp(x, False) + torch.cuda.synchronize() + assert torch.equal(on.view(torch.int16), off.view(torch.int16)), f"trial {trial}" + assert op_calls == [rows] * 3, "the seam runs the C++ op at every decode row count" + + +@pytest.mark.parametrize("rows", [1, 4, modeling.TINY_MAX_TOKENS]) +def test_seam_in_a_graph_is_bitwise_seam_off(rows) -> None: + """Both paths captured; replays on new inputs (PDL edges inside the graph).""" + x = _x(rows, 7) + _mlp(x, True) # compile outside the capture + graph_on = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph_on): + on = _mlp(x, True) + _mlp(x, False) + graph_off = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph_off): + off = _mlp(x, False) + for step in range(3): + x.copy_(_x(rows, 31 + step)) + graph_on.replay() + graph_off.replay() + torch.cuda.synchronize() + assert torch.equal(on.view(torch.int16), off.view(torch.int16)), f"replay {step}" + + +def test_seam_declines_what_it_does_not_cover(op_calls, monkeypatch) -> None: + """Rows past the op's certified rows and other activation parameters take the Triton kernel (the same bits).""" + max_rows = max(cell.spec["rows"] for cell in modeling.m3_dense_swiglu.CELLS) + # Past the decode rows the GEMV declines too; the seam only reads the gate_up rows, so any BF16 rows will do. + g = torch.Generator(device="cuda").manual_seed(5) + gu = torch.randn(max_rows + 1, DENSE_GATE_UP[0], generator=g, device="cuda").to(torch.bfloat16) + got = modeling._dense_swiglu(CORE, gu) + assert op_calls == [], "rows past the op's certified rows" + assert torch.equal(got, modeling.m3_silu_and_mul(gu, LIMIT, ALPHA, BETA)) + other = SimpleNamespace(swiglu_alpha=1.0, swiglu_limit=LIMIT, swiglu_beta=BETA) + gu = _gemv(_x(4, 5), GATE_UP) + got = modeling._dense_swiglu(other, gu) + assert op_calls == [], "other activation parameters" + assert torch.equal(got, modeling.m3_silu_and_mul(gu, LIMIT, 1.0, BETA)) + + +def test_seam_needs_no_capture_scope(op_calls, monkeypatch) -> None: + """The op replaces only the activation kernel: the small-batch decode graphs are captured outside the FlashInfer + MXFP8 capture scope (r4.10.1 never engaged at 1-8 rows).""" + monkeypatch.setattr(modeling, "flashinfer_mxfp8_decode_graph_capture_active", lambda: False) + modeling._dense_swiglu(CORE, _gemv(_x(4, 5), GATE_UP)) + assert op_calls == [4], "outside the FlashInfer capture scope" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py new file mode 100644 index 000000000000..2b8b3a4f2e0c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise oracle of the dense MLP's swigluoai C++ op (ops/gemm/m3_dense_swiglu, +trtllm::minimax_m3_sm107_tp2_swiglu_oai) against the Triton silu_and_mul kernel the dense MLP runs today +(modules/swiglu.py at limit 7, alpha 1.702, beta 1): every bf16 gate bit pattern against special and random up values, +every up bit pattern against special and random gates, random pairs, and the dense MLP's shapes at every decode row +count, eager and replayed from a CUDA graph, with and without the programmatic launch.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.silu_and_mul +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_dense_swiglu import ( # noqa: E501 + ALPHA, + BETA, + LIMIT, + m3_dense_swiglu, +) + +assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_swiglu_oai"): + pytest.skip( + "needs the TP2 op library with minimax_m3_sm107_tp2_swiglu_oai", allow_module_level=True + ) + +# The TP2 dense MLP's activation width (gate_up 3072 -> d 1536). +DENSE_D = modeling.DENSE_DOWN[1] + + +def _triton(x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=LIMIT, swiglu_alpha=ALPHA, swiglu_beta=BETA + ) + + +def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" + return out + + +def _row(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + return torch.cat([gate.view(1, -1), up.view(1, -1)], dim=1).contiguous() + + +def _all_bf16() -> torch.Tensor: + return ( + torch.arange(1 << 16, dtype=torch.int32, device="cuda").to(torch.int16).view(torch.bfloat16) + ) + + +def _specials(n: int, seed: int) -> list: + g = torch.Generator(device="cuda").manual_seed(seed) + vals = [0.0, -0.0, 1.0, -1.0, 6.99, 7.0, 7.01, -7.0, -7.01, 1e-38, -1e-38, 1e-40, 88.0, -88.0, 3e38, -3e38, + float("inf"), float("-inf"), float("nan")] # fmt: skip + rows = [torch.full((n,), v, dtype=torch.bfloat16, device="cuda") for v in vals] + for scale in (0.05, 1.0, 4.0, 30.0): + rows.append((scale * torch.randn(n, generator=g, device="cuda")).to(torch.bfloat16)) + rows.append( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + return rows + + +def _assert_bitwise(got: torch.Tensor, ref: torch.Tensor, what: str) -> None: + bad = (got.view(torch.int16) != ref.view(torch.int16)).nonzero().flatten() + assert bad.numel() == 0, f"{what}: {bad.numel()} mismatches, first at {bad[:4].tolist()}" + + +def test_every_gate_pattern_is_bitwise_the_triton_kernel() -> None: + gate = _all_bf16() + for j, up in enumerate(_specials(gate.numel(), 1)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"up set {j}") + + +def test_every_up_pattern_is_bitwise_the_triton_kernel() -> None: + up = _all_bf16() + for j, gate in enumerate(_specials(up.numel(), 2)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"gate set {j}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: + g = torch.Generator(device="cuda").manual_seed(100 + seed) + n = 1 << 22 + + def pattern() -> torch.Tensor: + bits = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + return bits.to(torch.int16).view(torch.bfloat16) + + x = _row(pattern(), pattern()) + _assert_bitwise(_op(x), _triton(x), f"seed {seed}") + + +@pytest.mark.parametrize("pdl", [True, False], ids=["pdl", "plain"]) +@pytest.mark.parametrize("rows", range(1, modeling.TINY_MAX_TOKENS + 1)) +def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pdl: bool) -> None: + """The TP2 dense MLP's gate_up rows (a GEMV output's range plus values past the clamps), eager and replayed from a + CUDA graph after fresh rows are written in place, against the Triton kernel on the same rows.""" + g = torch.Generator(device="cuda").manual_seed(2000 + rows) + x = (3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, DENSE_D : DENSE_D + 8] = x[:, :8] + _assert_bitwise(_op(x, pdl), _triton(x), f"eager rows {rows}") + out = torch.empty(rows, DENSE_D, dtype=torch.bfloat16, device="cuda") + _op(x, pdl) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out.copy_(_op(x, pdl)) + x.copy_((3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(out, _triton(x), f"graph rows {rows}") + + +def test_the_op_declines_other_parameters_and_layouts() -> None: + x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + + +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: + """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton + kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" + rows, d = cell.spec["rows"], cell.spec["d"] + g = torch.Generator(device="cuda").manual_seed(3000 + rows) + x = (3.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, d : d + 8] = x[:, :8] + _assert_bitwise(_op(x), _triton(x), cell.why) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py new file mode 100644 index 000000000000..b44554a8d717 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_draft_lowm entry: the Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M +kernels (direct SIMT and split-K tensor-core), at the TP2 drafter shards, at the rows its table certifies. + +Every cell is held to the fp64 reference within the derived band. The drafter's projection call +(``M3Eagle3Draft.linear``) at a covered row count is the entry's call bit for bit, at an uncovered one the in-tree +GEMM's. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + TACTICS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 + m3_draft_lowm as op, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" + +# Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). +# +# Kernel and reference multiply the same BF16 operands (every product exact in fp32); the kernel sums in fp32 (the +# direct kernel one dot product in K order, the split-K kernel tensor-core blocks and then its split partials), the +# reference in fp64, and both round once to BF16. With w ~ N(0, 1 / K) and x ~ N(0, 1) the outputs are O(1) and the +# fp32 order difference is ~sqrt(K) * 2^-24 of the summed magnitudes, a few 1e-6 at K = 18432: far below half a BF16 +# ulp, so the two roundings disagree by at most one ulp of the element -- up to two units of the row scale for an +# element in its row maximum's binade -- and only where the sum sits that close to a rounding boundary, a few in ten +# thousand elements, hence 0.1 ulp RMS. The same derivation as the decode GEMV's band (test_modeling_v2_m3_mx_gemv), +# with no slack under it: a second ulp would mean a wrong sum. +ELEMENT_ULP = 2.0 +RMS_ULP = 0.1 + +#: The drafter's projections (N, K) and their weight keys on ``M3Eagle3Draft.w``. +DRAFTER_LINEARS = { + (4608, 12288): "qkv", + (6144, 4096): "o", + (18432, 6144): "gu", + (6144, 9216): "dn", + (6144, 18432): "fc", +} + + +@pytest.fixture(autouse=True) +def _policy(): + modeling.set_require(False) + modeling.reset_logged() + yield + modeling.set_require(False) + modeling.reset_logged() + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + return x, w + + +def test_the_table_is_the_tp2_drafter() -> None: + assert set(TACTICS) == set(DRAFTER_LINEARS) + assert {c.spec["nk"] for c in op.CELLS} == set(TACTICS) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + kind, tactic = tactic_for(n, k, m) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + args = ( + x, + w, + out, + kind, + tactic, + True, + ) # the target launches with PDL when the environment enables it + certifying(op) + assert op(*args) is out + ref = op.reference(*args) + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=RMS_ULP) + # PDL only changes the launch attribute. + plain = torch.empty_like(out) + op(x, w, plain, kind, tactic, False) + assert torch.equal(plain.view(torch.int16), out.view(torch.int16)) + + +def _draft(linears=DRAFTER_LINEARS): + """A drafter holding only its projection weights and the low-M bookkeeping (the module is not built).""" + draft = object.__new__(modeling.M3Eagle3Draft) + draft.w = {key: build(*nk, 1, seed=nk[0] + nk[1])[1] for nk, key in linears.items()} + draft._lowm_built, draft._lowm_engaged = set(), set() + draft._push_arm, draft._pushed = None, 0 + return draft + + +@pytest.fixture(scope="module") +def prebuilt(): + draft = _draft() + assert draft.prebuild_linears() == len(op.CELLS) + assert draft.prebuild_linears() == 0, "a second prebuild finds nothing left to compile" + return draft + + +@pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) +def test_the_projection_is_the_entry_and_falls_back_past_its_rows(prebuilt, nk) -> None: + n, k = nk + key = DRAFTER_LINEARS[nk] + w = prebuilt.w[key] + use_pdl = get_env_enable_pdl() + modeling.set_require(True) + # The certified rows the target enables (DRAFT_LOWM_MAX_ROWS caps the table: the batch-4 / batch-8 cells are off + # by default since they read flat in-model). + for m in sorted(modeling.DRAFT_LOWM_ROWS[nk]): + assert modeling.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt.linear(x, key) + kind, tactic = tactic_for(n, k, m) + want = op(x, w, torch.empty(m, n, dtype=x.dtype, device="cuda"), kind, tactic, use_pdl) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)), m + # Rows the table does not certify (the b2 rows, and past it) or that DRAFT_LOWM_MAX_ROWS leaves off: the in-tree + # GEMM, silently, even under require. + for m in sorted( + (set(range(1, modeling.TINY_MAX_TOKENS + 2)) | set(ROWS[nk])) + - set(modeling.DRAFT_LOWM_ROWS[nk]) + ): + assert not modeling.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt.linear(x, key) + assert torch.equal(got.view(torch.int16), F.linear(x, w).view(torch.int16)), m + + +def test_a_projection_outside_the_table_is_the_in_tree_gemm() -> None: + draft = _draft({(256, 512): "o"}) + x, _ = build(256, 512, 1, seed=3) + assert draft._lowm_linear(draft.w["o"], x) is None + got = draft.linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_a_covered_call_not_prebuilt_raises_under_require() -> None: + draft = _draft({(6144, 4096): "o"}) + x, _ = build(6144, 4096, 1, seed=5) + modeling.set_require(True) + with pytest.raises(modeling.FastPathUnavailable, match="not compiled before warmup"): + draft.linear(x, "o") + # Outside require it falls back to the in-tree GEMM. + modeling.set_require(False) + got = draft.linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_guard_refuses_a_strided_activation() -> None: + x, w = build(6144, 4096, 2, seed=6) + wide = torch.zeros(2, 8192, dtype=torch.bfloat16, device="cuda") + wide[:, :4096] = x + kind, tactic = tactic_for(6144, 4096, 2) + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op( + wide[:, :4096], + w, + torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda"), + kind, + tactic, + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..ebfee3b50ee0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp32_linear entry of the MiniMax-M3 tp2 target (the router beyond the decode routers' rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """torch's default fp32 matmul precision on this hardware may be TF32 (~2^-11 relative per product): the band is + stated in fp32 ulps of each row's largest logit and covers it (K = 6144 TF32 products: ~2^13 ulps worst element); + measured on GB300 with the engine's defaults: worst element ~1700 ulps, RMS ~300.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(spec["e"], spec["k"], generator=g, device="cuda") * 0.02 + with validating(op): + out = op(x, w) + assert out.dtype == torch.float32 and out.shape == (spec["m"], spec["e"]) + assert_within_ulp(out, op.reference(x, w), element_ulp=8192.0, rms_ulp=2048.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..9057207107af --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv entry: the decode MXFP8 GEMV on a producer-quantized activation. + +Operands as the target holds them: the weight through mxfp8_quantize with the swizzled scale layout (the MXFP8 linear +method's storage), the activation through mxfp8_quantize with the swizzled layout (what the boundary all-reduce and +the attention merge publish). +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" + + +# Element band, in ULP units of the row scale: see compare. +ELEMENT_ULP = 1.5 + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Element ELEMENT_ULP of the row scale, relative RMS 0.1 ulp. + + Kernel and reference multiply bit-identical dequantized operands; the kernel sums in fp32 (tensor-core blocks, + then the split-K partials), the reference in fp64. Over K <= 6144 that order difference is ~1e-6 relative, so the + two bf16 roundings disagree only where the sum sits within that of a rounding boundary: one bf16 ulp of that + element, which for an element in its row maximum's binade is up to two of ``assert_within_ulp``'s units of the row + scale, and a few in ten thousand elements in aggregate. Any fp32 sum shows it: torch's IEEE fp32 matmul of the + same dequantized operands measured at most 1.22 units against the fp64 reference over every row count 1..48 on + the o_proj and the QKV (6 seeds each, GB300), the kernel at most 1.25, both at most 0.009 ulp RMS. The band + leaves the margin between that and a full extra ulp. The kernel is the same on SM107 (the tp2 table only picks + split and ring), and its TP2 projections' K is 1536..6144. + """ + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=0.1) + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + return x, xq.view(m, k), xsf, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + _, xq, xsf, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(xq, xsf, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + compare(out, op.reference(xq, xsf, w, wsf)) + + +def test_an_unserved_shape_returns_none() -> None: + """Shapes outside the tuned table fall back at the call site; the op says so rather than computing.""" + _, xq, xsf, w, wsf = build(256, 512, 1, seed=1) + assert op(xq, xsf, w, wsf) is None + + +def test_guard_refuses_linear_activation_scales() -> None: + """A linear [M, K / 32] scale buffer is read as the swizzled one.""" + x, _, _, w, wsf = build(6144, 4096, 4, seed=2) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, False) + with pytest.raises(AssertionError, match="swizzled"): + with validating(op): + op(xq.view(4, 4096), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + _, xq, xsf, w, wsf = build(6144, 4096, 1, seed=3) + wide = torch.zeros(6144, 8192, dtype=w.dtype, device="cuda") + wide[:, :4096] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :4096], wsf) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..4055f2d8db4e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv_bf16 entry: the decode MXFP8 GEMV quantizing a bf16 activation in its prologue. + +The weight as the MXFP8 linear method stores it (mxfp8_quantize, swizzled scales). Besides the reference band, the +prologue's quantization is held to mxfp8_quantize itself: the GEMV on a bf16 x must equal, bit for bit, the +prequantized GEMV on mxfp8_quantize(x) -- the claim the target relies on when it mixes the two paths. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv_bf16 import ( # noqa: E501 + m3_mx_gemv_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" + +# The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. +ELEMENT_ULP = 1.5 + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return x, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(x, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + # The band of test_modeling_v2_m3_mx_gemv.compare: the quantized operands are the reference's exactly (checked + # below), the sums differ only in order. + assert_within_ulp(out, op.reference(x, w, wsf), element_ulp=ELEMENT_ULP, rms_ulp=0.1) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + pre = m3_mx_gemv(xq.view(m, k), xsf, w, wsf, use_pdl=True) + assert pre is not None and torch.equal(out.view(torch.int16), pre.view(torch.int16)), ( + "the prologue's MXFP8 quantization differs from mxfp8_quantize" + ) + + +def test_a_strided_activation_is_declined() -> None: + """The prologue reads x densely; the op declines a strided x (None) rather than computing on it.""" + x, w, wsf = build(6144, 4096, 2, seed=1) + wide = torch.zeros(2, 8192, dtype=torch.bfloat16, device="cuda") + wide[:, :4096] = x + assert op(wide[:, :4096], w, wsf) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..f59ca4d385b4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_gemm entry of the MiniMax-M3 tp2 target (the decode-graph GEMM on a prequantized +activation).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most + 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, False), + element_ulp=2.0, + rms_ulp=0.25, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..5d272004698e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_linear entry of the MiniMax-M3 tp2 target: MXFP8 quantize + GEMM on each backend the +engine selects (native CUTLASS, FlashInfer, the per-bucket tuned quantize + GEMM).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" + +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and + one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" + spec = cell.spec + if ( + spec["path"] == "flashinfer" + and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None + ): + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + args = (x, w, w_sf, alpha, torch.bfloat16, *_FLAGS[spec["path"]]) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py new file mode 100644 index 000000000000..e685bc5df277 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for MX-TAC: the pinned decode-graph MXFP8 decision (CuTe DSL quantize + CuTe DSL GEMM) gives the same bits +as every other candidate stock's AutoTuner can pick (TRT-LLM quantize, CUTLASS GEMM), the same on every call.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.modeling import ( # noqa: E501 + _pinned_mxfp8_gemm, + _pinned_mxfp8_quantize, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + +assert torch.cuda.is_available(), "MX-TAC requires a CUDA device" + +# The TP2 projections at decode rows the GEMVs do not serve (12, 16) and a few beyond. +CASES = [ + ((4992, 6144), 16), + ((6144, 4096), 12), + ((3072, 6144), 24), + ((12288, 6144), 32), + ((6144, 1536), 48), +] + + +@pytest.mark.parametrize("nk,m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_pinned_decision_matches_every_stock_candidate(nk, m, scale) -> None: + n, k = nk + g = torch.Generator(device="cuda").manual_seed(n + k + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = (torch.randn(m, k, generator=g, device="cuda") * scale).to(torch.bfloat16) + x[0, ::32] *= 1e3 # one outlier per 32-element block on the first row + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + wq = wq.view(n, k) + xq, xsf = _pinned_mxfp8_quantize(x) + tq, tsf = torch.ops.trtllm.mxfp8_quantize(x, True) + assert torch.equal(xq.view(torch.uint8).view(m, k), tq.view(torch.uint8).view(m, k)) + assert torch.equal(xsf.view(torch.uint8).flatten(), tsf.view(torch.uint8).flatten()) + xq = xq.view(m, k) + out = _pinned_mxfp8_gemm(xq, xsf, wq, wsf, torch.bfloat16) + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + for tactic in (FlashInferMXFP8GemmRunner.CUTE_DSL, FlashInferMXFP8GemmRunner.CUTLASS): + ref = runner.forward([xq, xsf, wq, wsf], tactic=tactic) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), tactic + again = _pinned_mxfp8_gemm(xq, xsf, wq, wsf, torch.bfloat16) + assert torch.equal(out.view(torch.int16), again.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py new file mode 100644 index 000000000000..46548f91c1d7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the uGPU target head (ops/gemm/m3_ugpu_head): the die-placed copy of the head shard is the weight bit +for bit with its halves in the two locality domains; the module's own weight (what every other reader of the head +uses) is never placed; the per-die split GEMM gives the in-tree head's logits bit for bit at the certified row counts, +eager and in a CUDA graph; an in-place change of the weight makes the copy stale.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm import ( # noqa: E501 + m3_ugpu_head as op, +) +from tensorrt_llm._torch.modules.linear import Linear + +assert torch.cuda.is_available(), "the uGPU head requires a CUDA device" + + +def _locality(): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_locality import ( # noqa: E501 + LocalityMap, + ) + + try: + return LocalityMap.get(torch.cuda.current_device()) + except RuntimeError as err: # one locality domain (not a two-die GPU) + pytest.skip(f"no two locality domains: {err}") + + +@pytest.fixture(scope="module") +def placed(): + lm = _locality() + head = Linear(op.HIDDEN, op.TARGET_SHARD, bias=False, dtype=torch.bfloat16) + head.to("cuda") + g = torch.Generator(device="cuda").manual_seed(20261015) + with torch.no_grad(): + head.weight.copy_( + (0.02 * torch.randn(op.TARGET_SHARD, op.HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ) + ) + before = head.weight.detach().clone() + x = torch.randn(16, op.HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + stock = {m: head.apply_linear(x[:m], None) for m in (4, 8, 12, 16)} + ugpu = op.prepare_head(head) + assert ugpu is not None and ugpu.source is head.weight and op.current(ugpu, head.weight) + return lm, head, ugpu, before, x, stock + + +def test_no_reader_of_the_module_sees_placed_memory(placed) -> None: + """A kernel reading die-placed memory from all SMs runs 1.5-1.7x slower: every tensor of the head module (what + the in-tree GEMM, prefill logits, warmup and fallbacks read) must stay in ordinary memory; only the op's copy is + placed.""" + lm, head, ugpu, before, _, _ = placed + for name, t in list(head.named_parameters()) + list(head.named_buffers()): + end = t.data_ptr() + t.numel() * t.element_size() - 1 + assert lm.domain_of(t.data_ptr()) == -1 and lm.domain_of(end) == -1, name + assert t.data_ptr() != ugpu.weight.data_ptr(), name + assert torch.equal(head.weight, before) + + +def test_placement_is_exact_and_per_die(placed) -> None: + lm, head, ugpu, before, _, _ = placed + w = ugpu.weight + assert torch.equal(w, before) + row = w.shape[1] * w.element_size() + assert lm.domain_of(w.data_ptr()) == 0 + assert lm.domain_of(w.data_ptr() + ugpu.split * row - 1) == 0 + assert lm.domain_of(w.data_ptr() + ugpu.split * row) == 1 + assert lm.domain_of(w.data_ptr() + w.numel() * w.element_size() - 1) == 1 + assert (ugpu.split * row) % lm.granularity == 0 and abs(2 * ugpu.split - w.shape[0]) < 2 * 512 + + +@pytest.mark.parametrize( + "cell", op.m3_ugpu_head.CELLS, ids=[c.why[:44] for c in op.m3_ugpu_head.CELLS] +) +def test_split_head_is_the_in_tree_head_bit_for_bit(placed, cell) -> None: + m = cell.spec["m"] + _, head, ugpu, _, x, stock = placed + got = op.m3_ugpu_head(x[:m], ugpu) + torch.cuda.synchronize() + assert torch.equal(got, stock[m]) + assert torch.equal(got, op.m3_ugpu_head.reference(x[:m], ugpu)) + assert torch.equal( + head.apply_linear(x[:m], None), stock[m] + ) # stock readers of the placed weight + + +@pytest.mark.parametrize( + "cell", op.m3_ugpu_head.CELLS, ids=[c.why[:44] for c in op.m3_ugpu_head.CELLS] +) +def test_split_head_in_a_cuda_graph(placed, cell) -> None: + m = cell.spec["m"] + _, _, ugpu, _, x, stock = placed + xs = x[:m].clone() + box = {} + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + box["y"] = op.m3_ugpu_head(xs, ugpu) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + box["y"] = op.m3_ugpu_head(xs, ugpu) + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + assert torch.equal(box["y"], stock[m]) + + +def test_an_in_place_weight_change_makes_the_copy_stale(placed) -> None: + _, head, ugpu, _, _, _ = placed + assert op.current(ugpu, head.weight) + with torch.no_grad(): + head.weight.add_(0) # an in-place write bumps the version counter + assert not op.current(ugpu, head.weight) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py new file mode 100644 index 000000000000..a3509b9ab5a9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_r1b1_gemv_order_gate_a.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rule 1b gate a for the TP2 decode GEMV order override (ops/mx_gemv/op.STRIDED_K_TILES: strided k-tiles): +for every overridden (projection, rows) the output is the same on 3 eager calls and on a CUDA graph replay +(deterministic), and its max abs / max rel error vs the fp64 product of the dequantized operands stays within the +shipped order's band (the shared table's config on the same operands, over several seeds).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + mxfp8_dequant, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.mx_gemv import ( # noqa: E501 + op, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + +CASES = [(nk, m) for nk in (QKV_SPARSE, QKV, O_PROJ, DENSE_GATE_UP) for m in (1, 2, 4, 8)] + + +def _err(y, ref): + d = (y.double() - ref).abs() + return float(d.max()), float((d / ref.abs().amax(dim=1, keepdim=True).clamp_min(1e-30)).max()) + + +@pytest.mark.parametrize("nk, m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +def test_gate_a(nk, m) -> None: + n, k = nk + cfg = op.get_config(n, k, m) + assert cfg is not None and len(cfg) > 2 and cfg[2] == 0, ( + "the override is off or does not cover this call" + ) + shipped = cfg[:2] # the contiguous order + worst_new, worst_ship = [0.0, 0.0], [0.0, 0.0] + for seed in range(4): + g = torch.Generator(device="cuda").manual_seed(1000 * seed + m) + w = (torch.randn(n, k, device="cuda", generator=g) * 2).to(torch.float8_e4m3fn) + ws = torch.randint(118, 124, (n * k // 32,), dtype=torch.uint8, device="cuda", generator=g) + x = (torch.randn(m, k, device="cuda", generator=g) * (0.5 + seed)).to(torch.bfloat16) + q, s = torch.ops.trtllm.mxfp8_quantize(x, True) + ref = ( + mxfp8_dequant(q.view(torch.uint8), s).double() + @ mxfp8_dequant(w.view(torch.uint8), ws).double().t() + ) + ys = [m3_mx_gemv(q, s, w, ws, True, cfg, True, None) for _ in range(3)] + assert all(y is not None for y in ys), "the GEMV declined the override" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + out = [] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out.append(m3_mx_gemv(q, s, w, ws, True, cfg, True, None)) + graph.replay() + torch.cuda.synchronize() + for y in ys[1:] + [out[0]]: + assert torch.equal(ys[0].view(torch.int16), y.view(torch.int16)), ( + "not deterministic (eager / graph)" + ) + ship = m3_mx_gemv(q, s, w, ws, True, shipped, True, None) + torch.cuda.synchronize() + for acc, y in ((worst_new, ys[0]), (worst_ship, ship)): + e = _err(y, ref) + acc[0], acc[1] = max(acc[0], e[0]), max(acc[1], e[1]) + print( + f"\nR1B1 tp2 {n}x{k} m={m} cfg {cfg}: max abs / rel vs fp64 {worst_new[0]:.3e} / {worst_new[1]:.3e} | shipped " + f"order {worst_ship[0]:.3e} / {worst_ship[1]:.3e}", + flush=True, + ) + # The band: the shipped order's own error (bf16 output rounding dominates), with 25 % headroom. + assert worst_new[0] <= 1.25 * worst_ship[0] and worst_new[1] <= 1.25 * worst_ship[1] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..b5cb1d2e9109 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_moe_operands.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One tp2 rank's MiniMax-M3 MoE layer operands for the decode-engine entries' tests. + +The routed experts are random checkpoint-format NVFP4 tensors ([w3 (up); w1 (gate)] rows, e4m3 block scales) put +through TRT-LLM's own TRTLLM-Gen NVFP4 loader shuffle (``NVFP4TRTLLMGenFusedMoEMethod._shuffle_all_experts``), which +writes the buffers the engines read in place. The per-expert scalars follow the loader's conventions (act_beta and +act_clamp divided by fc31_alpha, fc31_scale_c = fc2_input_scale x fc31_alpha) and are sized so the FC1 sums land +around the SwiGLU-OAI clamp and the outputs near unit scale. The shared expert's MXFP8 weights come from +mxfp8_quantize with the swizzled scale layout, as the MXFP8 linear method stores them. +""" + +import functools +from types import SimpleNamespace + +import torch + +# Hidden size, the routed experts' intermediate per rank at TP2 (3072 / 2: every expert on every rank), experts, +# top-k, and the shared expert's intermediate per rank. +H, I_TP, NUM_EXPERTS, TOP_K = 6144, 1536, 128, 4 +SHARED_I = 1536 +SWIGLU_ALPHA, SWIGLU_LIMIT = 1.702, 7.0 +ROUTED_SCALING = 2.0 +MINIMAX2_ROUTING = 5 # RoutingMethodType.MiniMax2 +# NVFP4 global scales: 448 x 6 over the amax each side is sized for. +FC31_INPUT_SCALE = 448.0 * 6.0 / 8.0 +FC2_INPUT_SCALE = 448.0 * 6.0 / 24.0 + + +def _rand_e4m3_scales(shape, gen) -> torch.Tensor: + """E4M3 block scales in [0.5, 2): biased exponents 6 and 7, any mantissa.""" + return torch.randint(0x30, 0x40, shape, dtype=torch.uint8, device="cuda", generator=gen).view( + torch.float8_e4m3fn + ) + + +@functools.lru_cache(maxsize=None) +def experts(seed: int = 20261013): + """The routed experts and the router of one MoE layer, in the layout and conventions of the TRTLLM-Gen loader. + + Returns a SimpleNamespace: ``tensors`` (the ten expert operands in ``M3MoeState.layer``'s order), ``bias``, + ``router_weight``, ``fc31_input_scale`` and ``shared`` (gate_up, its scales, down, its scales). + """ + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + gen = torch.Generator(device="cuda").manual_seed(seed) + e = NUM_EXPERTS + w31 = torch.randint( + 0, 256, (e, 2 * I_TP, H // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + w2 = torch.randint(0, 256, (e, H, I_TP // 2), dtype=torch.uint8, device="cuda", generator=gen) + w31s = _rand_e4m3_scales((e, 2 * I_TP, H // 16), gen) + w2s = _rand_e4m3_scales((e, H, I_TP // 16), gen) + module = SimpleNamespace( + w3_w1_weight=w31, + w2_weight=w2, + w3_w1_weight_scale=w31s, + w2_weight_scale=w2s, + is_gated_activation=True, + bias=False, + ) + NVFP4TRTLLMGenFusedMoEMethod()._shuffle_all_experts(module) + torch.cuda.synchronize() + + # Raw FC1 sums: rms(x) x FC31_INPUT_SCALE x rms(E2M1 value) x rms(block scale) x sqrt(H) ~ 1.0e5 at rms(x) = 1; + # fc31_alpha brings them to ~1.5 in the dequantized domain, the FC2 alpha the outputs to ~1. + base = 1.5 / (FC31_INPUT_SCALE * 2.93 * 1.3 * H**0.5) + fc31_alpha = base * (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) + fc2_alpha = (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) / ( + FC2_INPUT_SCALE * 2.93 * 1.3 * I_TP**0.5 + ) + tensors = ( + w31, + w31s, + w2, + w2s, + fc31_alpha, + FC2_INPUT_SCALE * fc31_alpha, # fc31_scale_c + torch.full((e,), SWIGLU_ALPHA, device="cuda"), # act_alpha + 1.0 / fc31_alpha, # act_beta: the (up + 1) offset over fc31_alpha + SWIGLU_LIMIT / fc31_alpha, # act_clamp over fc31_alpha + fc2_alpha, + ) + bias = (0.02 * torch.randn(e, device="cuda", generator=gen)).float() + router_weight = (torch.randn(e, H, device="cuda", generator=gen) / H**0.5).float().contiguous() + fc31_input_scale = torch.tensor([FC31_INPUT_SCALE], dtype=torch.float32, device="cuda") + + gu = (0.02 * torch.randn(2 * SHARED_I, H, device="cuda", generator=gen)).to(torch.bfloat16) + dn = (0.05 * torch.randn(H, SHARED_I, device="cuda", generator=gen)).to(torch.bfloat16) + gu_q, gu_sf = torch.ops.trtllm.mxfp8_quantize(gu, True) + dn_q, dn_sf = torch.ops.trtllm.mxfp8_quantize(dn, True) + shared = (gu_q.view(2 * SHARED_I, H), gu_sf, dn_q.view(H, SHARED_I), dn_sf) + return SimpleNamespace( + tensors=tensors, + bias=bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared=shared, + ) + + +# Routing cases: random (up to 4 M distinct experts), and every token on the same four experts (the bias outweighs +# any sigmoid difference), which leaves later slot groups empty. +CONCENTRATED = (5, 40, 77, 126) + + +def layer_for(state, case: str, shared: bool = True): + """One layer on ``state`` over :func:`experts` for a routing case. ``shared=False``: the shared expert's weights + are zero, so the engine returns the routed part alone (bf16(0 + bf16(routed))).""" + ops = experts() + bias = ops.bias.clone() + if case == "concentrated": + bias[list(CONCENTRATED)] += 1.0 + shared_weights = ops.shared + if not shared: + gu, gu_sf, dn, dn_sf = ops.shared + shared_weights = (torch.zeros_like(gu), gu_sf, torch.zeros_like(dn), dn_sf) + return state.layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=shared_weights, + ) + + +def stock_routed( + hidden: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """The stock TRTLLM-Gen NVFP4 MoE (fp4_quantize, then fp4_block_scale_moe_runner with SwiGLU-OAI) on the given + routing, over :func:`experts`: what the built-in MiniMaxM3MoE computes for the routed part.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + + ops = experts() + w31, w31s, w2, w2s, fc31_alpha, scale_c, act_alpha, act_beta, act_clamp, fc2_alpha = ops.tensors + xq, xsf = torch.ops.trtllm.fp4_quantize(hidden, ops.fc31_input_scale, 16, False, False) + out = torch.ops.trtllm.fp4_block_scale_moe_runner( + None, None, xq, xsf.view(torch.float8_e4m3fn).flatten(), + w31, w31s, None, act_alpha, act_beta, act_clamp, w2, w2s, None, + scale_c, fc31_alpha, fc2_alpha, + NUM_EXPERTS, TOP_K, None, None, I_TP, 0, NUM_EXPERTS, ROUTED_SCALING, MINIMAX2_ROUTING, True, 0, + topk_weights=topk_weights.to(torch.bfloat16), topk_ids=topk_ids.to(torch.int32), + ) # fmt: skip + return out[0] if isinstance(out, (list, tuple)) else out + + +def hidden_states(m: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(m, H, device="cuda", generator=gen).to(torch.bfloat16) + + +def check_routed_part(op, state, m: int, seed: int) -> None: + """Exact, independent of the reference: the engine's routing (its route_dbg record: ids, then the bf16 weights' + fp32 bits) is the reference routing bit for bit, and its routed part is the stock TRTLLM-Gen MoE's output on that + routing bit for bit -- the claim the engines are built on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=False) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + stock = stock_routed(hidden, ids, weights) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)), ( + "routed part differs from TRTLLM-Gen's" + ) + + +def stock_shared(hidden: torch.Tensor) -> torch.Tensor: + """The stock MXFP8 shared expert as the built-in GatedMLP runs it (MXFP8LinearMethod's native path): mxfp8_quantize + + mxfp8_mxfp8_gemm, the fused swiglu (limit 7, alpha 1.702, beta 1), mxfp8_quantize + mxfp8_mxfp8_gemm.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + from tensorrt_llm._torch.modules.swiglu import swiglu + + gu, gu_sf, dn, dn_sf = experts().shared + one = torch.ones(1, dtype=torch.float32, device=hidden.device) + a, s = torch.ops.trtllm.mxfp8_quantize(hidden.contiguous(), True) + y = torch.ops.trtllm.mxfp8_mxfp8_gemm(a, s, gu, gu_sf, one, torch.bfloat16) + h = swiglu(y, swiglu_limit=SWIGLU_LIMIT, swiglu_alpha=SWIGLU_ALPHA, swiglu_beta=1.0) + a2, s2 = torch.ops.trtllm.mxfp8_quantize(h.contiguous(), True) + return torch.ops.trtllm.mxfp8_mxfp8_gemm(a2, s2, dn, dn_sf, one, torch.bfloat16) + + +def check_full_output(op, state, m: int, seed: int) -> None: + """Exact: the engine's bf16(shared + routed) is the stock MoE block's shared_output.add_(routed_output) bit for bit + (stock TRTLLM-Gen routed part on the engine's routing, stock MXFP8 shared expert).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=True) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + ref = stock_shared(hidden).add_(stock_routed(hidden, ids, weights)) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), ( + "bf16(shared + routed) differs from stock" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py new file mode 100644 index 000000000000..6d8fe90ad9d7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared operands and gates of the router entries' tests (m3_router_gemv, m3_router_gemv_fp4, m3_router_tc: one +kernel source).""" + +import zlib + +import torch + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv import ( # noqa: E501 + m3_router_gemv as gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv import ( # noqa: E501 + near_tie, + nvfp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 + m3_router_gemv_fp4 as gemv_fp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__all__ = [ + "SF_SCALE", + "_build", + "_compare_fp4", + "_compare_logits", + "_seed", + "gemv", + "gemv_fp4", + "tc", +] + +#: The experts' NVFP4 input scale the tc cells use: 448 * 6 / 8, block scales well inside E4M3. +SF_SCALE = 336.0 + + +def _build(m: int, e: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(e, k, generator=g, device="cuda") * 0.02 + return x, w + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +def _compare_logits(out: torch.Tensor, ref: torch.Tensor) -> None: + """fp32 logits that differ from a true fp32 product only in summation order. + + Over K = 6144 products a reordered fp32 sum moves by a few ulps of the sum of magnitudes, which near + cancellation is ~sqrt(K) times the row's largest logit's ulp at most. Gated in fp32 ulps of each row's largest + logit: 64 element-wise, 8 in relative RMS. Measured on GB300: GEMV 1.6 / 1.4 ulp, tensor-core router 3.8 / 2.0 + ulp at 2..48 tokens; the tensor-core logits against a true fp32 product (the TF32 rounding the band must see) + sit at 2200+ / 1700 ulp. + """ + assert out.dtype == ref.dtype == torch.float32 and out.shape == ref.shape + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=8.0) + + +def _codes(packed: torch.Tensor) -> torch.Tensor: + """[M, K / 2] packed e2m1 bytes -> [M, K] 4-bit codes, low nibble first.""" + lo, hi = packed & 0xF, packed >> 4 + return torch.stack((lo, hi), dim=-1).reshape(packed.shape[0], -1) + + +def _compare_fp4( + q: torch.Tensor, sf: torch.Tensor, x: torch.Tensor, sf_scale: torch.Tensor +) -> None: + """Scales bit-exact; codes bit-exact except within ``TIE_WINDOW`` of an e2m1 midpoint, where the kernel's + approximate reciprocals choose the side (and then only by one code).""" + ref_q, ref_sf, scaled = nvfp4_quantize(x, sf_scale) + assert q.dtype == torch.uint8 and q.shape == ref_q.shape + assert torch.equal(sf, ref_sf), f"{(sf != ref_sf).sum().item()} block scales differ" + got, want = _codes(q), _codes(ref_q) + differ = got != want + assert not bool((differ & ~near_tie(scaled)).any()), ( + f"{(differ & ~near_tie(scaled)).sum().item()} e2m1 codes differ away from a rounding tie" + ) + magnitude = ((got & 7).int() - (want & 7).int()).abs() + assert bool((magnitude[differ] <= 1).all()), "a tie resolved more than one code away" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py new file mode 100644 index 000000000000..f4116a9c71b2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The router of an unfused MoE past the tensor-core router's 48 rows: the fp32 router (``m3_fp32_linear``) is bit for +bit the built-in gate's generation-step arithmetic there, ``F.linear(h.to(float32), w)`` (the gate's TF32 fixed tile +is off in generation steps, its tensor-core router covers 9..16 rows, and below 2048 rows its fp32 cast is the plain +``.to(torch.float32)``), so a 49..192-row decode step routes as the agg canon does.""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear, +) +from tensorrt_llm._utils import get_sm_version + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 107, reason="SM107 (Rubin) only" +) + +HIDDEN, EXPERTS = 6144, 128 + + +@pytest.mark.parametrize("rows", [49, 64, 128, 192]) +def test_the_fp32_router_is_the_gates_generation_arithmetic(rows: int) -> None: + assert not modeling.ROUTER_TC_GEN.covers(rows), "the tensor-core router stops at 48 rows" + g = torch.Generator(device="cuda").manual_seed(rows) + h = torch.randn(rows, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + w = torch.randn(EXPERTS, HIDDEN, generator=g, device="cuda") / HIDDEN**0.5 + got = m3_fp32_linear(h, w) + want = F.linear(h.to(torch.float32), w) + assert got.dtype == torch.float32 + assert torch.equal(got.view(torch.int32), want.view(torch.int32)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..1d6de982a67a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe_engine entry: the 1..8-token MoE engine on its caller-owned state. + +One Rubin (SM107) rank of the tp2 deployment: 128 experts (all local, routed and shared intermediate 1536 per rank), +top 4, hidden 6144, the shared expert folded in, ``num_ctas`` = SMs - 24 as the target builds it (212 - 24 = 188). +Operands: ``_m3_moe_operands``. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + HIDDEN, + INTER, + SHARED_INTER, + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import ( # noqa: E402 + I_TP, + SHARED_I, + H, + check_full_output, + check_routed_part, + hidden_states, + layer_for, +) +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +_STATES = {} + + +def test_the_operands_are_the_tp2_engine_geometry() -> None: + assert (H, I_TP, SHARED_I) == (HIDDEN, INTER, SHARED_INTER) == (6144, 1536, 1536) + + +def test_the_engine_refuses_the_tp4_width() -> None: + with pytest.raises(ValueError, match="1536"): + M3MoeState(torch.device("cuda"), num_tokens=1, inter=768) + + +def test_the_engine_refuses_more_routed_ctas_than_fc2_blocks() -> None: + """A routed CTA owning none of the 192 FC2 output blocks would deadlock the hand-off (routed-only, 212 SMs).""" + with pytest.raises(ValueError, match="192"): + M3MoeState(torch.device("cuda"), num_tokens=4, num_ctas=193, shared=False) + + +def _state(m: int) -> M3MoeState: + if m not in _STATES: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + _STATES[m] = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + return _STATES[m] + + +def test_a_module_built_with_other_trace_config_compiles_its_own_kernel(monkeypatch) -> None: + """The compile cache keys on the config the kernel module was built with: a state whose module carries an extra + trace-time option (``timing``, which only adds per-CTA stamps) compiles its own kernel, never the plain build's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_moe import ( # noqa: E501 + op as kop, + ) + + hidden = hidden_states(1, seed=5) + op(hidden, layer_for(_state(1), "random")) # the plain build is compiled + n0 = len(kop._compiled) + build = kop._kernel_module + monkeypatch.setattr( + kop, "_kernel_module", lambda config, path: build({**config, "timing": 1}, path) + ) + sms = torch.cuda.get_device_properties(0).multi_processor_count + state = M3MoeState(torch.device("cuda"), num_tokens=1, num_ctas=sms - 24, shared=True) + op(hidden, layer_for(state, "random")) + assert state.mod.TIMING and not _state(1).mod.TIMING + assert len(kop._compiled) == n0 + 1, "the timing build reused another build's compiled kernel" + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """The engine against the plain-torch MoE: BAND_ELEMENT bf16 ulp of the row's largest magnitude per element, + BAND_RMS ulp relative RMS. + + Both sides quantize the same operands with the same recipes (the activation's and the intermediate's NVFP4 codes + and block scales, the shared expert's MXFP8); they differ in summation order (tensor-core fp32 against fp64) and in + the PTX approximations (rcp / ex2 / tanh). Such a difference changes the output only where it moves a value across + a rounding boundary: one bf16 ulp of a slot output or of the sum, or, rarely, one intermediate NVFP4 code, which + reaches the output as one weight column's share. Measured over every cell, both routing cases and 3 seeds each + on the tp4 engine (GB300, intermediate 768): at most 1.48 ulp element, 0.18 ulp RMS; the tp2 engine runs the + same recipes over twice the intermediate. The band leaves ~3x on both; a misread recipe (the activation's + 1 / (sf x (1 / sc)) spelling of the reciprocal on ~0.06% of its codes) measured 12 ulp element, 7 ulp RMS. + """ + assert_within_ulp(out, ref, element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + # The workspace is tagged per call and never reset: a repeated call is the same call. + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + compare(out, op.reference(hidden, layer)) + + +@pytest.mark.parametrize("m", (1, 2, 4, 8)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +@pytest.mark.parametrize("seed", (1, 2, 3)) +@pytest.mark.parametrize("m", (4, 8)) +def test_full_output_is_the_stock_block_bit_for_bit(m, seed) -> None: + """The TP2 operating points (one / two requests verifying 3 drafts): the whole layer output, shared expert + included, against the stock block (TRTLLM-Gen routed + MXFP8 GatedMLP shared, shared.add_(routed)).""" + check_full_output(op, _state(m), m, seed=200 + 10 * m + seed) + + +def test_layers_on_one_state_do_not_leak_into_each_other() -> None: + """Two layers share one state's workspace, in stream order: A, B, A gives A's output twice, bit for bit, and the + caller's ``out`` receives the same rows as a fresh tensor.""" + state = _state(4) + a, b = layer_for(state, "random"), layer_for(state, "concentrated") + hidden = hidden_states(4, seed=7) + first = op(hidden, a) + op(hidden, b) + out = torch.empty(4, H, dtype=torch.bfloat16, device="cuda") + second = op(hidden, a, out=out) + assert second.data_ptr() == out.data_ptr() + assert torch.equal(first.view(torch.int16), out.view(torch.int16)) + + +@pytest.mark.parametrize( + "bad", + ("misaligned hidden", "fp32 out", "short out", "long out", "int32 arp"), +) +def test_malformed_buffers_are_rejected_by_the_call(bad) -> None: + """Why the entry has no is_valid: the compiled call checks each buffer against its signature and raises.""" + state = _state(2) + layer = layer_for(state, "random") + hidden, kw = hidden_states(2, seed=3), {} + if bad == "misaligned hidden": + hidden = torch.randn(2 * H + 8, device="cuda").to(torch.bfloat16)[2 : 2 + 2 * H].view(2, H) + elif bad == "fp32 out": + kw["out"] = torch.empty(2, H, dtype=torch.float32, device="cuda") + elif bad == "short out": + kw["out"] = torch.empty(1, H, dtype=torch.bfloat16, device="cuda") + elif bad == "long out": + kw["out"] = torch.empty(3, H, dtype=torch.bfloat16, device="cuda") + else: + kw["arp"] = torch.zeros(8, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + op(hidden, layer, **kw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py new file mode 100644 index 000000000000..59559dc01d89 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the m3_moe engine's per-call epoch wraps past 2^31 without changing a bit. + +The engine tags its hand-off records with the call epoch (masked low bits plus a marker bit) and counts the B2IMG group +counters up by CNT_PER_CALL per call, waiting for ``(epoch + 1) * CNT_PER_CALL`` in wrapping int32 arithmetic. A state +seeded at epoch 0x7FFFFFF0 (counters seeded consistently, as after that many calls) runs 32 calls across the int32 +wrap; every output equals a fresh state's, bit for bit. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import hidden_states, layer_for # noqa: E402 + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +EPOCH0 = 0x7FFFFFF0 +CALLS = 32 + + +def _int32(v: int) -> int: + v &= 0xFFFFFFFF + return v - (1 << 32) if v >= 1 << 31 else v + + +def _seed(state: M3MoeState, epoch: int) -> None: + """The state as after ``epoch`` calls: every CTA's epoch, and each B2IMG group counter at epoch x CNT_PER_CALL.""" + state.epochs.fill_(_int32(epoch)) + mod = state.mod + if mod.B2IMG: + value = _int32(epoch * mod.CNT_PER_CALL) + for g in range(mod.NG_CAP): + state.cnt[8 * g] = value + torch.cuda.synchronize() + + +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_the_epoch_wraps_bit_for_bit(m) -> None: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + fresh = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + wrapped = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + _seed(wrapped, EPOCH0) + lf, lw = layer_for(fresh, "random"), layer_for(wrapped, "random") + for i in range(CALLS): + x = hidden_states(m, 4000 + i) + a, b = op(x, lf), op(x, lw) + torch.cuda.synchronize() + assert torch.equal(a, b), f"call {i} (epoch {EPOCH0 + i:#x})" + assert int(wrapped.epochs[0]) == _int32(EPOCH0 + CALLS) < 0 # crossed 2^31 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py new file mode 100644 index 000000000000..700ce9af2be9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_finalize_rows.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Who finalizes the wide target's MoE at each generic row count: past the decode rows, up to +GENERIC_MC_FINALIZE_MAX_TOKENS, the finalize (+ the shared expert) folds into the boundary's multicast all-reduce +(M3DecoderLayer._mc_finalize_boundary); from MOE_TRITON_FINALIZE_MIN_TOKENS, ops/m3_moe_finalize finalizes it. +The two ranges meet with no gap and no overlap, and the fold declines from shapes alone past its cap. + +CPU only: nothing here launches a kernel. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_mc_allreduce, +) + +H = 6144 + + +def _fold_cap() -> int: + return min(dispatch.GENERIC_MC_FINALIZE_MAX_TOKENS, m3_mc_allreduce.FINALIZE_MAX_TOKENS) + + +def test_the_fold_stays_inside_the_finalize_kernel() -> None: + assert dispatch.GENERIC_MC_FINALIZE_MAX_TOKENS <= m3_mc_allreduce.FINALIZE_MAX_TOKENS + + +def test_the_fold_and_the_triton_finalize_meet() -> None: + """With the fold on (MC_AR_TP2), the Triton finalize starts on the first row past it; with it off, nothing folds.""" + if not dispatch.MC_AR_TP2: + assert dispatch.GENERIC_MC_FINALIZE_MAX_TOKENS == 0 + return + assert dispatch.DECODE_MAX_TOKENS < _fold_cap() + if dispatch.MOE_TRITON_FINALIZE: + assert dispatch.MOE_TRITON_FINALIZE_MIN_TOKENS == _fold_cap() + 1 + + +def _target(): + """A DecodeTarget whose bucket lookup and MC finalize must not be reached (the row gate declines first).""" + + def unreachable(*args, **kwargs): + raise AssertionError("the generic MC finalize ran past its row gate") + + return SimpleNamespace(_bucket_flags=unreachable, _mc_finalize=unreachable) + + +@pytest.mark.parametrize("past", [1, 2, 16]) +def test_the_fold_declines_past_its_cap(past: int) -> None: + hidden = torch.empty((_fold_cap() + past, H), dtype=torch.bfloat16) + assert dispatch.DecodeTarget._generic_mc_finalize(_target(), 3, hidden, hidden, None) is None + + +@pytest.mark.parametrize("rows", [1, 8, 32]) +def test_the_fold_leaves_the_decode_layers_rows(rows: int) -> None: + rows = min(rows, dispatch.DECODE_MAX_TOKENS) + hidden = torch.empty((rows, H), dtype=torch.bfloat16) + assert dispatch.DecodeTarget._generic_mc_finalize(_target(), 3, hidden, hidden, None) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py new file mode 100644 index 000000000000..bc95a2586d74 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_flat_gate.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The TP2 MoE engine's FLAT stack is applied only up to FLAT_MAX_TOKENS (4) tokens per call. + +At 8 tokens FLAT was +2.96 us/call slower than the per-m tables, so above the cap the engine keeps the per-m +configuration (the engine with the cap at 0, bitwise). CPU only: M3MoeState is constructed with its device queries +stubbed and stopped at the kernel build, so the test reads the engine config and the engagement lines without a GPU. +""" + +import types + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_moe import ( # noqa: E501 + op, +) + + +class _Built(Exception): + """Raised by the stubbed kernel build: the config is final by then.""" + + +def _engine(monkeypatch, num_tokens): + """(config, info lines) of M3MoeState(num_tokens) up to its kernel build.""" + lines = [] + seen = {} + + def build(config, path): + seen["config"] = dict(config) + raise _Built() + + monkeypatch.setattr(op, "is_supported", lambda device=None: (True, "")) + monkeypatch.setattr(op, "_kernel_module", build) + monkeypatch.setattr(op.torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr( + op.torch.cuda, + "get_device_properties", + lambda device: types.SimpleNamespace(multi_processor_count=op.HIDDEN // 32), + ) + from tensorrt_llm.logger import logger + + monkeypatch.setattr(logger, "info", lambda msg, *a, **k: lines.append(str(msg))) + with pytest.raises(_Built): + op.M3MoeState("cuda", num_tokens, num_ctas=op.HIDDEN // 32) + return seen["config"], [line for line in lines if "MiniMax-M3 MoE engine FLAT" in line] + + +def _has_flat(config): + return all(config.get(k) == v for k, v in op.FLAT_STACK.items()) + + +@pytest.mark.parametrize("num_tokens", range(1, 9)) +def test_flat_applies_only_up_to_the_cap(monkeypatch, num_tokens): + config, lines = _engine(monkeypatch, num_tokens) + assert op.FLAT_MAX_TOKENS == 4 + if num_tokens <= 4: + assert _has_flat(config), config + assert lines == [ + "MiniMax-M3 MoE engine FLAT: b2img=0 lg_repl=16 ar_pre=0 late_trigger=0 " + f"tokens={num_tokens}" + ] + else: + # above the cap: the per-m tables, i.e. exactly the cap-0 engine + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 0) + off_config, _ = _engine(monkeypatch, num_tokens) + assert config == off_config + assert lines == [ + f"MiniMax-M3 MoE engine FLAT off above 4 tokens (per-m tables): tokens={num_tokens}" + ] + + +@pytest.mark.parametrize("num_tokens", range(1, 9)) +def test_cap_zero_runs_the_per_m_tables(monkeypatch, num_tokens): + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 0) + config, lines = _engine(monkeypatch, num_tokens) + assert lines == [ + f"MiniMax-M3 MoE engine FLAT off above 0 tokens (per-m tables): tokens={num_tokens}" + ] + assert config.get("b2img") == op.B2IMG_BY_M[num_tokens - 1] + assert config.get("ar_pre") == op.AR_PRE_BY_M[num_tokens - 1] + + +@pytest.mark.parametrize("num_tokens", range(1, 9)) +def test_cap_of_eight_runs_flat_at_every_token_count(monkeypatch, num_tokens): + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 8) + config, lines = _engine(monkeypatch, num_tokens) + assert _has_flat(config), config + assert len(lines) == 1 and lines[0].endswith(f"tokens={num_tokens}") + + +def test_flat_applies_helper(monkeypatch): + """op.flat_applies reads FLAT_MAX_TOKENS per call.""" + assert [op.flat_applies(m) for m in range(1, 9)] == [True] * 4 + [False] * 4 + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 8) + assert all(op.flat_applies(m) for m in range(1, 9)) + monkeypatch.setattr(op, "FLAT_MAX_TOKENS", 0) + assert not any(op.flat_applies(m) for m in range(1, 9)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py new file mode 100644 index 000000000000..0530ad9cc24e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: SH_TAKE (the shared-expert CTAs take the extra FC2 output blocks) changes no output bit. + +With 188 routed CTAs, 192 FC2 output blocks leave four routed CTAs with two blocks. Under SH_TAKE the last four +shared-expert CTAs run, after their shared work, one of those blocks each as FC2-only CTAs, with the same tiles, +scales and finalize. ``SH_TAKE_BY_M`` turns it on per token count (shipped all off). For every token count the +engine's output under the gated table (on at 6 / 7 / 8 tokens) equals its output under the shipped table, bit for bit, +eager and replayed from a CUDA graph, for spread and concentrated routing. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_moe import ( # noqa: E501 + op as engine_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as run, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import hidden_states, layer_for # noqa: E402 + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +SHIPPED = engine_op.SH_TAKE_BY_M +GATED = (0, 0, 0, 0, 0, 1, 1, 1) + + +def _state(monkeypatch, m, on): + monkeypatch.setattr(engine_op, "SH_TAKE_BY_M", GATED if on else SHIPPED) + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + return M3MoeState(dev, num_tokens=m, num_ctas=sms - engine_op.SHARED_CTAS, shared=True) + + +@pytest.mark.parametrize("m", range(1, 9)) +def test_tables(monkeypatch, m) -> None: + assert not any(SHIPPED) + on = _state(monkeypatch, m, True) + off = _state(monkeypatch, m, False) + nheavy = (engine_op.HIDDEN // 32) % on.num_ctas + expect = bool(GATED[m - 1]) and 0 < nheavy <= engine_op.SHARED_CTAS + assert bool(on.config.get("sh_take")) == expect + assert "sh_take" not in off.config + + +@pytest.mark.parametrize("case", ["random", "concentrated"]) +@pytest.mark.parametrize("m", range(1, 9)) +def test_bitwise_against_the_shipped_table(monkeypatch, m, case) -> None: + states = {on: _state(monkeypatch, m, on) for on in (True, False)} + layers = {on: layer_for(st, case) for on, st in states.items()} + for i in range(8): + x = hidden_states(m, 6000 + i) + a, b = run(x, layers[True]), run(x, layers[False]) + torch.cuda.synchronize() + assert torch.equal(a, b), f"eager call {i}" + x = hidden_states(m, 6100) + outs = {} + for on in (True, False): + run(x, layers[on]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + y = run(x, layers[on]) + graph.replay() + torch.cuda.synchronize() + outs[on] = y.clone() + assert torch.equal(outs[True], outs[False]), "graph replay" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py new file mode 100644 index 000000000000..5523a391661e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_triton_finalize.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ops/m3_moe_finalize: the routed experts' finalize (+ the shared expert) from the TRTLLM-Gen runner's unfinalized +parts, bitwise to the runner's finalizeKernelVecLoad order (acc = 0; acc += w_k * fc2[idx_k] for k in order, skipping +-1; one BF16 rounding) and to ``shared.add_(routed)``, at the rows the wide target routes there (49..16K), with dropped +(-1) slots, padded FC2 rows, eager and CUDA-graph replay. (The runner itself was compared at 512 / 2K / 4K / 16K rows on +SM107, maximum difference 0.)""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_moe_finalize import ( # noqa: E501 + m3_moe_finalize, +) + +H, TOP_K = 6144, 4 + + +def _parts(rows: int, seed: int, pad: int = 128): + g = torch.Generator(device="cuda").manual_seed(seed) + n = rows * TOP_K + fc2 = torch.randn((n + pad, H), device="cuda", generator=g).to(torch.bfloat16) + idx = torch.randperm(n, device="cuda", generator=g).to(torch.int32).reshape(rows, TOP_K) + drop = torch.rand((rows, TOP_K), device="cuda", generator=g) < 0.02 + idx = torch.where(drop, torch.full_like(idx, -1), idx) + w = torch.rand((rows, TOP_K), device="cuda", generator=g).to(torch.bfloat16) + shared = torch.randn((rows, H), device="cuda", generator=g).to(torch.bfloat16) + return fc2, w, idx.contiguous(), shared + + +def _reference(fc2, w, idx): + acc = torch.zeros((w.shape[0], H), dtype=torch.float32, device=fc2.device) + for k in range(TOP_K): + valid = (idx[:, k] >= 0).unsqueeze(1) + x = fc2[idx[:, k].clamp(min=0).long()].float() + acc = torch.where(valid, acc + w[:, k : k + 1].float() * x, acc) + return acc.to(torch.bfloat16) + + +@pytest.mark.parametrize("rows", [49, 64, 65, 1000, 4096, 16384]) +def test_bitwise_to_the_runner_order(rows: int) -> None: + fc2, w, idx, shared = _parts(rows, rows) + ref = _reference(fc2, w, idx) + out = m3_moe_finalize(fc2, w, idx) + assert out is not None and torch.equal(out.view(torch.int16), ref.view(torch.int16)) + fused = m3_moe_finalize(fc2, w, idx, shared) + assert torch.equal(fused.view(torch.int16), shared.clone().add_(ref).view(torch.int16)) + + +def test_graph_replay_matches_eager() -> None: + fc2, w, idx, shared = _parts(512, 7) + eager = m3_moe_finalize(fc2, w, idx, shared) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = m3_moe_finalize(fc2, w, idx, shared) + for seed in (8, 9): + f2, w2, i2, s2 = _parts(512, seed) + fc2.copy_(f2) + w.copy_(w2) + idx.copy_(i2) + shared.copy_(s2) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, m3_moe_finalize(fc2, w, idx, shared)) + assert eager.shape == out.shape + + +def test_declines_other_layouts() -> None: + fc2, w, idx, shared = _parts(64, 3) + assert m3_moe_finalize(fc2.float(), w, idx) is None + assert m3_moe_finalize(fc2, w, idx.long()) is None + assert m3_moe_finalize(fc2, w, idx, shared[:, :128].contiguous()) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..61bef8aacee1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_tc entry of the MiniMax-M3 tp2 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP2 target runs on SM107 only", +) + + +@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +def test_tc_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) + sf_scale = torch.tensor([SF_SCALE], dtype=torch.float32, device="cuda") if spec["fp4"] else None + with validating(tc): + outs = tc(x, w, sf_scale, False) + ref = tc.reference(x, w, sf_scale, False) + assert len(outs) == len(ref) == (3 if spec["fp4"] else 1) + _compare_logits(outs[0], ref[0]) + if spec["fp4"]: + _compare_fp4(outs[1], outs[2], x, sf_scale) + + +def test_tc_reference_is_tf32_and_the_gate_sees_it() -> None: + """The tensor-core router rounds the weight to TF32: against a true fp32 product its logits sit far outside + the band, so the TF32 rounding in ``reference`` is load-bearing rather than decoration.""" + x, w = _build(16, 128, 6144, 5) + out = tc(x, w, None, False)[0] + with pytest.raises(AssertionError): + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +def test_guard_refuses_a_misaligned_activation() -> None: + """A row view starting 8 bytes into the buffer keeps a legal row stride but not a legal bulk-copy source.""" + x, w = _build(2, 128, 6144, 3) + flat = torch.empty(2 * 6144 + 4, dtype=torch.bfloat16, device="cuda") + shifted = flat[4:].view(2, 6144) + shifted.copy_(x) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(tc): + tc(shifted, w, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..59201301c623 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/m3_embedding_firstnorm entry: the decode embedding gather + the first layer's RMSNorm. + +The table is the model's full bf16 [200064, 6144] embedding shape (random values at the checkpoint's scale), eps the +checkpoint's 1e-6. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + m3_embedding_firstnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" + +H = 6144 +EPS = 1e-6 + + +@functools.lru_cache(maxsize=None) +def _operands(): + g = torch.Generator(device="cuda").manual_seed(20261013) + table = (0.02 * torch.randn(VOCAB, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return table, weight + + +def _buffers(rows): + return (torch.empty(rows, H, dtype=torch.bfloat16, device="cuda"), + torch.empty(rows, H, dtype=torch.bfloat16, device="cuda")) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + rows, id_dtype = cell.spec["rows"], cell.spec["id_dtype"] + table, weight = _operands() + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device="cuda").to(id_dtype) + raw, out = _buffers(rows) + with validating(op): + normed, residual = op(ids, table, weight, EPS, raw, out) + assert normed.data_ptr() == out.data_ptr() and residual.data_ptr() == raw.data_ptr() + ref_normed, ref_raw = op.reference(ids, table, weight, EPS, raw, out) + # The gather is a copy: exact. + assert torch.equal(raw.view(torch.int16), ref_raw.view(torch.int16)) + # The norm: the kernel's fp32 sum of squares (warp butterflies, then across 4 warps) and fast rsqrt against an + # fp64 reference, then x * rstd * w rounded once to bf16. The fp32 side is ~1e-6 relative off, so the two roundings + # disagree only next to a bf16 rounding boundary: one ulp at most per element, a small fraction in aggregate. + # Measured over these cells (3 seeds each, GB300): at most 0.51 ulp element, 0.026 ulp RMS. + assert_within_ulp(out, ref_normed, element_ulp=1.0, rms_ulp=0.1) + + +def test_ids_outside_the_table_give_zero_rows() -> None: + """Documented, not refused: an id outside [0, V) gathers a zero row whose norm is zero.""" + table, weight = _operands() + ids = torch.tensor([-1, 5, VOCAB], dtype=torch.int64, device="cuda") + raw, out = _buffers(3) + op(ids, table, weight, EPS, raw, out) + assert not raw[0].any() and not raw[2].any() and not out[0].any() and not out[2].any() + assert torch.equal(raw[1], table[5]) + + +@pytest.mark.parametrize("rows", range(1, 17)) +def test_against_the_stock_embedding_and_rmsnorm(rows) -> None: + """What this target certifies on SM107 against the built-in path it replaces (the model's Embedding gather, then + the first layer's RMSNorm module), including a constant row, large and tiny rows and both vocab ends: the gather + bit for bit, the norm within one bf16 ulp (the sum-of-squares tree differs from flashinfer's: measured ~5e-6 of + the elements one ulp off over 12000 random rows).""" + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + table, weight = _operands() + table = table.clone() + g = torch.Generator(device="cuda").manual_seed(rows) + table[1] = 3.0 + table[2] = (300.0 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + table[3] = (1e-4 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + norm = RMSNorm(hidden_size=H, eps=EPS, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + norm.weight.copy_(weight) + for ids in ( + torch.randint(0, VOCAB, (rows,), generator=g, device="cuda"), + # The special rows (constant, large, tiny, both vocab ends), repeated to the row count past 8. + torch.tensor(([1, 2, 3, 0, VOCAB - 1, 1, 2, 3] * 4)[:rows], device="cuda"), + ): + raw, out = _buffers(rows) + normed, residual = op(ids, table, weight, EPS, raw, out) + x = torch.nn.functional.embedding(ids, table) + assert torch.equal(residual.view(torch.int16), x.view(torch.int16)) + stock = norm(x) + off = normed != stock + assert int(off.sum()) <= 4, f"{int(off.sum())} normed elements differ from stock" + assert_within_ulp(normed, stock, element_ulp=1.0, rms_ulp=0.05) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..158aa26fc6e4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs shared by the tp4 target's spec entry tests: TP4 dimensions, single-device stand-ins for the cross-rank +exchanges, and the comparison bands.""" + +from __future__ import annotations + +import zlib + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) + +#: MiniMax-M3 at TP4: hidden size, vocabulary, the target head's shard and RMSNorm epsilon. +HIDDEN = 6144 +VOCAB = 200064 +WORLD = 2 +TARGET_SHARD = VOCAB // WORLD +EPS = 1e-6 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): + """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" + uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") + keys = torch.zeros(slot_rows, dtype=torch.int64, device="cuda") + counters = torch.tensor([0, call], dtype=torch.int32, device="cuda") + return uc, keys, counters + + +def plant_peers( + uc: torch.Tensor, + counters: torch.Tensor, + own_best: torch.Tensor, + rank: int, + world: int, + shard: int, + g: torch.Generator, + slot_rows: int = SHARD_ROWS, + ties: bool = False, +) -> None: + """Write the other ranks' keys for each row into this rank's slots, as their multicast pushes would. + + Each peer's maximum is the own row maximum shifted by +-0.25 (well clear of the accumulation-order band), so + some rows are won by a peer and some by this rank; with ``ties`` a peer on an even row reports exactly the own + maximum, and the lower global index must win. + """ + if world == 1: + return + m = own_best.shape[0] + half = int(counters[1].item()) & 1 + for r in range(world): + if r == rank: + continue + sign = torch.where(torch.rand(m, generator=g, device="cuda") < 0.5, -1.0, 1.0) + value = own_best.float() + 0.25 * sign + if ties: + value = torch.where(torch.arange(m, device="cuda") % 2 == 0, own_best.float(), value) + index = r * shard + torch.randint(0, shard, (m,), generator=g, device="cuda") + keys = make_keys(value, index) + for t in range(m): + uc[(half * slot_rows + t) * world + r] = keys[t] + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..15758e2645e2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_capture_norm entry (Eagle3 capture with fc_norm folded in).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_capture_norm import ( # noqa: E501 + m3_capture_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" + +LAYERS = 3 + + +def _case(spec, seed): + g = gen(seed) + m, slot = spec["m"], spec["slot"] + h = randn((m, HIDDEN), g, 2.0) + r = randn((m, HIDDEN), g, 2.0) if spec["residual"] else None + w = (1.0 + randn((HIDDEN,), g, 0.1).float()).to(torch.bfloat16) + buf = torch.full((m, LAYERS * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + dst = buf[:, slot * HIDDEN : (slot + 1) * HIDDEN] + return (h, r, w, EPS, dst), buf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer + untouched.""" + args, buf = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + slot = cell.spec["slot"] + others = torch.cat([buf[:, : slot * HIDDEN], buf[:, (slot + 1) * HIDDEN :]], dim=1) + assert bool((others == 7.0).all()) + + +def test_guard_refuses_a_strided_weight() -> None: + (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + w2 = torch.stack([w, w], dim=1)[:, 0] + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(h, r, w2, eps, dst) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py new file mode 100644 index 000000000000..6b2c334a7615 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_workspaces.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The lazily created exchange workspaces refuse to be created inside a CUDA-graph capture (K3 audit #150, check #10). + +The draft-step and sampled-accept workspaces are created collectively (both ranks, a multicast buffer) on the first +eager call. Creating one inside a capture would allocate and synchronize mid-capture, so a capture that finds it missing +reports a miss (``ops/m3_report.capture_missed``) and gets None (the caller then takes its stock path) instead. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( + m3_draft_step, + m3_report, + m3_sampled_accept, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU" +) + + +def _refuse(*args, **kwargs): + raise AssertionError("a workspace must not be created inside a CUDA-graph capture") + + +@pytest.mark.parametrize( + "module, creator", + [ + (m3_draft_step, "create_workspace"), + (m3_sampled_accept, "SampledAcceptWorkspace.create"), + ], + ids=["draft_step", "sampled_accept"], +) +def test_a_capture_without_the_workspace_misses_instead_of_creating_it( + monkeypatch, module, creator +): + monkeypatch.setattr(module, "_WS", None) + owner, _, attr = creator.rpartition(".") + monkeypatch.setattr(getattr(module, owner) if owner else module, attr, _refuse) + misses = [] + monkeypatch.setattr(m3_report, "capture_missed", lambda *args: misses.append(args)) + + device = torch.device("cuda", torch.cuda.current_device()) + found = [] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + found.append(module.workspace(None, device)) + torch.cuda.current_stream().wait_stream(stream) + + assert found == [None] + assert module._WS is None + assert len(misses) == 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..b255f786fc5e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_chunk_norm entry (the Eagle3 fc_norm over the captured layers).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_chunk_norm import ( # noqa: E501 + m3_chunk_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, c, h = spec["m"], spec["c"], spec["h"] + buf = randn((m, c * h + spec["extra_cols"]), g, 2.0) + src = buf[:, : c * h + spec["extra_cols"]] + w = (1.0 + randn((c, h), g, 0.1).float()).to(torch.bfloat16) + dst = torch.full((m, c * h), 7.0, dtype=torch.bfloat16, device="cuda") + return (src, w, EPS, dst, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + + +def test_guard_refuses_a_strided_weight() -> None: + src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() + wt.copy_(w) + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(src, wt, eps, dst, pdl) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py new file mode 100644 index 000000000000..deeb0b10c5d3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the drafter boundary push (modeling.DRAFT_AR_PUSH / DRAFT_AR_PUSH_DIRECT) on one rank: the push variants +of the split-K copy (ops/m3_lowm_splitk) and the direct copy (ops/m3_lowm_direct) write FlashInfer's output +bit for bit, and store exactly the boundary all-reduce's canonical exchange words of that output: word (t, column pair) += bf16 lo | bf16 hi << 16 with -0.0 halves as +0.0, at parity * maxt * world * H / 2 + (t * world + rank) * H / 2 + +column / 2, the parity being the (t, column / (H / 8)) flag's low bit; every other word stays untouched. Eager and in a +CUDA graph, both parities, at each certified push cell of the TP2 table (the 1-row direct and 4-row split-K o_proj / +down).""" + +import pytest +import torch +from flashinfer.gemm.kernels.dense_bf16_gemm_direct import DirectTactic as FiDirect +from flashinfer.gemm.kernels.dense_bf16_gemm_direct import run_direct_dense as fi_direct +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiSplitK +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_splitk + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_lowm_direct, + m3_lowm_splitk, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + tactic_for, +) + +assert torch.cuda.is_available(), "the push variants require a CUDA device" + +H, MAXT, EMPTY = 6144, 8, -2147483648 +# The boundary linears (output width H): o_proj and down, each certified row count on split-K or an even direct width, +# up to the 8-token bucket the push arms on (ops/m3_draft_lowm.arm_push); the batch-4/8 rows past it run unpushed. +CELLS = [ + (nk[1], m, *tactic_for(nk[0], nk[1], m)) + for nk, rows in sorted(ROWS.items()) + if nk[0] == H and nk[1] != 18432 # o_proj / down, not the fc + for m in rows + if m <= MAXT + and (tactic_for(nk[0], nk[1], m)[0] == "s" or tactic_for(nk[0], nk[1], m)[1][1] % 2 == 0) +] + + +def _expected_words(y: torch.Tensor, flags: torch.Tensor, world: int, rank: int) -> torch.Tensor: + m = y.shape[0] + half = MAXT * world * (H // 2) + ws = torch.full((2 * half,), EMPTY, dtype=torch.int32, device="cuda") + h16 = y.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, H // 2) + for t in range(m): + for c in range(8): + par = int(flags[t * 8 + c].item()) & 1 + lo = par * half + (t * world + rank) * (H // 2) + c * (H // 16) + ws[lo : lo + H // 16] = words[t, c * (H // 16) : (c + 1) * (H // 16)] + return ws + + +@pytest.mark.parametrize("parity", [0, 1]) +@pytest.mark.parametrize("cell", CELLS, ids=[f"{k}@{m}{kind}" for k, m, kind, _ in CELLS]) +def test_push_words_and_output(cell, parity) -> None: + k, m, kind, tactic = cell + world, rank = 2, 1 + g = torch.Generator(device="cuda").manual_seed(k + m + parity) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + x[0, :8] = ( + 0.0 # zero / signed-zero outputs exercise the -0 canonicalization through the weights below + ) + w = (torch.randn(H, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + w[:4] = 0.0 + w[4:8] = -0.0 + flags = ( + torch.randint(0, 7, (MAXT * 8,), generator=g, device="cuda", dtype=torch.int32) * 2 + parity + ) + ref = torch.empty(m, H, dtype=torch.bfloat16, device="cuda") + if kind == "d": + fi_direct(x, w.t(), ref, True, FiDirect(*tactic)) + else: + fi_splitk(x, w.t(), None, ref, True, FiSplitK(*tactic)) + + def run(out, ws): + ar = (ws.data_ptr(), flags.data_ptr(), rank, world, MAXT) + if kind == "d": + m3_lowm_direct.run_direct_dense_push( + x, w.t(), out, True, m3_lowm_direct.DirectTactic(*tactic), ar + ) + else: + m3_lowm_splitk.run_splitk_dense_push( + x, w.t(), out, True, m3_lowm_splitk.SplitKTactic(*tactic), ar + ) + + want = _expected_words(ref, flags, world, rank) + for graph in (False, True): + out = torch.empty(m, H, dtype=torch.bfloat16, device="cuda") + ws = torch.full_like(want, EMPTY) + if graph: + run(out, ws) # compile outside capture + torch.cuda.synchronize() + ws.fill_(EMPTY) + gr = torch.cuda.CUDAGraph() + with torch.cuda.graph(gr): + run(out, ws) + ws.fill_(EMPTY) + gr.replay() + else: + run(out, ws) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), f"output (graph={graph})" + assert torch.equal(ws, want), f"exchange words (graph={graph})" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..964e62f86558 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_input entry (embedding + both RMSNorms + concat in one launch), against its reference and +against the in-tree drafter's modules.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_input import ( # noqa: E501 + m3_draft_input as op, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + VOCAB, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, h = spec["m"], spec["h"] + table = randn((VOCAB, h), g) + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + if spec["invalid_id"]: + ids[0] = -1 + hidden = randn((4 * m if spec["gather"] else m, h), g) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 2 if spec["gather"] else None + w_e = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + w_h = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + out = torch.full((m, 2 * h), 7.0, dtype=torch.bfloat16, device="cuda") + residual = torch.full((m, h), 7.0, dtype=torch.bfloat16, device="cuda") + return (ids, table, hidden, rows, w_e, w_h, EPS, EPS, out, residual, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Normalized halves within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic in a + different summation order, rsqrtf); the gathered residual and the zero embedding of an invalid id exact.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out, residual = op(*args) + ref_out, ref_res = op.reference(*args) + torch.cuda.synchronize() + assert_bf16_close(out, ref_out) + assert torch.equal(residual, ref_res) + if cell.spec["invalid_id"]: + assert not bool(out[0, : cell.spec["h"]].float().any()) + + +def test_guard_refuses_rows_past_hidden() -> None: + gather = next(c for c in op.CELLS if c.spec["gather"]) + args = list(_case(gather.spec, 1)) + args[3] = args[3] + 100 + with pytest.raises(AssertionError, match="rows must index hidden"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("m,gather", [(4, False), (32, False), (8, True)], ids=str) +def test_the_in_tree_draft_input(m, gather) -> None: + """Against the in-tree drafter's own modules: the Embedding lookup, then the layer's input_layernorm and + hidden_norm (in-tree RMSNorm) and the concat of Eagle3DecoderLayer.forward. Within one bf16 ulp on at most 0.1% of + the elements (the RMS statistic summed in another order); the residual (the carried hidden rows) exact.""" + g = gen(seed_of(f"in-tree draft input {m} {gather}")) + h = HIDDEN + embed = Embedding(VOCAB, h, dtype=torch.bfloat16).cuda() + embed.weight.data = randn((VOCAB, h), g) + norms = [] + for _ in range(2): + norm = RMSNorm(hidden_size=h, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + norms.append(norm) + input_layernorm, hidden_norm = norms + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + hidden = randn((4 * m if gather else m, h), g, 2.0) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 3 if gather else None + carried = hidden[rows] if gather else hidden + stock = torch.cat([input_layernorm(embed(ids)), hidden_norm(carried)], dim=-1) + out, residual = m3_spec_tail.draft_input( + ids, embed.weight, hidden, rows, input_layernorm.weight, EPS, hidden_norm.weight, EPS + ) + torch.cuda.synchronize() + assert_bf16_close(out, stock) + assert torch.equal(residual, carried) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_push_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_push_take.py new file mode 100644 index 000000000000..523385b5dd06 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_push_take.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter boundary push bookkeeping (``M3Eagle3Draft.arm_push`` / ``take_pushed``, ``DRAFT_AR_PUSH``): the +boundary all-reduce runs prepushed only when the armed call pushed, and only on that call's own output; any other +input raises instead of reducing rows the exchange does not hold. Host state only (no GPU).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) + + +@pytest.fixture +def draft(): + """A drafter with only its push bookkeeping (the module and its weights are not built).""" + d = object.__new__(modeling.M3Eagle3Draft) + d._push_arm, d._pushed = None, 0 + return d + + +def test_not_pushed_is_false_and_disarms(draft) -> None: + draft._push_arm = (1, 2, 3, 0, 8, 8) + assert draft.take_pushed(torch.zeros(4)) is False + assert draft._push_arm is None + + +def test_pushed_output_is_true_once(draft) -> None: + out = torch.zeros(4, 8) + draft._pushed = out.data_ptr() + assert draft.take_pushed(out) is True + assert draft._pushed == 0 + assert draft.take_pushed(out) is False + + +def test_other_input_raises(draft) -> None: + out = torch.zeros(4, 8) + other = out.clone() + draft._pushed = out.data_ptr() + with pytest.raises(RuntimeError, match="not the boundary all-reduce's input"): + draft.take_pushed(other) + assert draft._pushed == 0 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py new file mode 100644 index 000000000000..a817e8ce00dc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_silu entry (the drafter MLP's plain SwiGLU as one C++ launch): bit for bit the Triton +silu_and_mul at its defaults, over every bf16 gate and every bf16 up pattern.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_silu import ( # noqa: E501 + m3_draft_silu as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" + +ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) + + +def _rows(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """[rows, 2d] gate | up from equally long flat gate and up values, d = 9216 (the TP2 drafter shard).""" + d = 9216 + n = gate.numel() + rows = -(-n // d) + pad = rows * d - n + g = torch.cat([gate, gate[:pad]]).view(rows, d) + u = torch.cat([up, up[:pad]]).view(rows, d) + return torch.cat([g, u], dim=1).contiguous() + + +def _assert_bitwise(x: torch.Tensor, what: str) -> None: + got = op(x) + assert got is not None, what + ref = op.reference(x) + torch.cuda.synchronize() + diff = got.view(torch.int16) != ref.view(torch.int16) + assert not bool(diff.any()), f"{what}: {int(diff.sum())} of {diff.numel()} differ" + + +@pytest.mark.parametrize("up", [0.0, 1.0, -1.0, 3.5, -0.0, 1e-3], ids=lambda v: f"up={v}") +def test_every_gate_pattern_is_bitwise_the_triton_kernel(up) -> None: + gate = ALL.cuda() + _assert_bitwise(_rows(gate, torch.full_like(gate, up)), f"gate sweep, up {up}") + + +@pytest.mark.parametrize("gate", [1.0, -1.0, 0.5, -3.0, 20.0, -0.0], ids=lambda v: f"gate={v}") +def test_every_up_pattern_is_bitwise_the_triton_kernel(gate) -> None: + up = ALL.cuda() + _assert_bitwise(_rows(torch.full_like(up, gate), up), f"up sweep, gate {gate}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed) -> None: + g = gen(100 + seed) + n = 1 << 22 + gate = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + up = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + _assert_bitwise(_rows(gate.view(torch.bfloat16), up.view(torch.bfloat16)), f"seed {seed}") + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The drafter's gate_up output range (and past it), in a CUDA graph as the drafter runs it.""" + g = gen(seed_of(cell.why)) + rows, d = cell.spec["rows"], cell.spec["d"] + x = (4.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :6] = torch.tensor([60.0, -60.0, 0.0, -0.0, 88.0, -88.0], dtype=torch.bfloat16) + certifying(op) + ref = op.reference(x) + eager = op(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(x) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(eager.view(torch.int16), ref.view(torch.int16)) + assert torch.equal(captured.view(torch.int16), ref.view(torch.int16)) + + +def test_declines_what_it_does_not_cover() -> None: + x = torch.zeros(2, 2 * 9216, dtype=torch.bfloat16, device="cuda") + assert op(x[:, 8:]) is None # strided rows + assert ( + op(torch.zeros(2, 2 * 12, dtype=torch.bfloat16, device="cuda")) is None + ) # d not a multiple of 8 + assert op(x.float()) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..0528b9f1abb3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,814 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_step entry (one fused greedy draft pass: the BF16 draft head, streamed from its lossless +12-bit coding or as BF16 fragments, the cross-rank argmax and the draft loop's bookkeeping), and its tokens against the +in-tree drafter's head and greedy draft sampling. + +The in-tree path: each rank's LM head shard GEMM (BF16 out), ``LogitsProcessor``'s float cast, then either the argmax +of the gathered logits (``torch.argmax``) or ``greedy_sample_draft_with_tp_gather``'s per-rank ``torch.max`` and the +pick among the ranks' (index, value) pairs. The edge cases use integer data, whose fp32 sums are exact in any order, +so the kernel and the in-tree path see the same BF16 logits and must agree bit for bit; every in-tree comparison runs +with the coded head (the default) and with the BF16 fragments, and on integer data the coded head's winning keys are +the BF16 fragments' bit for bit. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_draft_step as draft_step_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_LEN, + DRAFT_SHARD, + HIDDEN, + MAX_TOKENS, + bf16_fragments, + from_fragments, + head_weight, + shard_logits, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_step import ( # noqa: E501 + m3_draft_step as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + WORLD, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" + +T = DRAFT_LEN + 1 +_LOW32 = 0xFFFFFFFF + +# ------------------------------------------------------------------ the per-row rule against the in-tree token +# +# The kernel and the in-tree head (a cuBLAS BF16 GEMM) both round an FP32 sum of the same exact BF16 products to BF16, +# in different summation orders. A row's token can differ from the in-tree one only where some candidate's FP32 sum +# can land on either side of a BF16 rounding midpoint. Each row is held to: +# +# kernel token == in-tree token, or the row is provably undecidable: with the fp64 reference logit L_i (exact BF16 +# products, fp64 sum) and the FP32-order bound E_i = gamma_K * sum_k |w_ik * x_k|, the BF16 value of any FP32 sum of +# candidate i lies in [lo_i, hi_i] = [bf16(L_i - E_i), bf16(L_i + E_i)]; the kernel's token must win for some +# choice of values within those intervals (larger value, then the lower global id), and a BF16 rounding midpoint +# must lie within E of the kernel's or the in-tree token's L. +# +# gamma_K = K * 2^-22 is la-xh2's bound (4x the round-to-nearest unit roundoff 2^-24 per addition: any summation +# order, and a truncating tensor-core accumulator); K counts the row's non-zero products (adding an exact zero does +# not round). lo / hi take one more FP32 step outward before the BF16 rounding, which covers rounding L -+ E to FP32. +# Each undecidable row is reported as a warning with its midpoint distances, and each test reports its row count; any +# other mismatch fails. +GAMMA_UNIT = 2.0**-22 + + +def _bf16_outward(v: torch.Tensor, direction: float) -> torch.Tensor: + f = v.float() + f = torch.nextafter(f, torch.full_like(f, direction)) + b = f.to(torch.bfloat16).float() + return torch.where(b == 0, torch.zeros_like(b), b) + + +def bounds(h: torch.Tensor, w: torch.Tensor): + """(fp64 logits L, bound E, lo, hi) [rows, N] of the head input rows h against the head rows w.""" + cols = (h != 0).any(dim=0) # zero columns add exact zeros + hd, wd = h[:, cols].double(), w[:, cols].double() + logits = hd @ wd.t() + bound = (h != 0).sum(dim=1, keepdim=True).double() * GAMMA_UNIT * (hd.abs() @ wd.abs().t()) + del wd + inf = float("inf") + return logits, bound, _bf16_outward(logits - bound, -inf), _bf16_outward(logits + bound, inf) + + +def midpoint_distance(v: float) -> float: + """Distance of v to the nearest BF16 rounding midpoint.""" + b = torch.tensor(v, dtype=torch.float64).float().to(torch.bfloat16).float() + if float(b) == 0.0: + return abs(v) + ulp = 2.0 ** (int(torch.frexp(b).exponent) - 8) + return min(abs(v - (float(b) + ulp / 2)), abs(v - (float(b) - ulp / 2))) + + +def _winnable(tok: int, ids: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Whether candidate ``tok`` wins the row for some values within the candidates' [lo, hi].""" + k = int((ids == tok).nonzero()[0, 0]) + others = ids != tok + beat = (hi[k] > lo[others]) | ((hi[k] == lo[others]) & (tok < ids[others])) + return bool(beat.all()) + + +def check_row(name: str, kernel_tok: int, stock_tok: int, ids, logits, bound, lo, hi) -> bool: + """The per-row rule for one row over the candidates ``ids`` [C] (logits / bound / lo / hi [C]); True when the + row is undecidable (reported as a warning), False when the tokens match; a mismatch outside the bound fails.""" + if kernel_tok == stock_tok: + return False + assert _winnable(kernel_tok, ids, lo, hi), ( + f"{name}: token {kernel_tok} differs from the in-tree {stock_tok} outside the FP32-order bound" + ) + ik = int((ids == kernel_tok).nonzero()[0, 0]) + i_stock = int((ids == stock_tok).nonzero()[0, 0]) + lk, ls, ek, es = ( + float(logits[ik]), + float(logits[i_stock]), + float(bound[ik]), + float(bound[i_stock]), + ) + dk, ds = midpoint_distance(lk), midpoint_distance(ls) + assert dk <= ek or ds <= es, f"{name}: neither logit is within its bound of a BF16 midpoint" + warnings.warn( + f"[draft-step undecidable] {name}: kernel {kernel_tok} (L {lk:.9g}, midpoint distance {dk:.3e}, bound " + f"{ek:.3e}) vs in-tree {stock_tok} (L {ls:.9g}, midpoint distance {ds:.3e}, bound {es:.3e})" + ) + return True + + +def report(name: str, rows: int, undecidable: int) -> None: + warnings.warn(f"[draft-step rows] {name}: {rows} rows, {undecidable} undecidable") + + +def key_value(keys: torch.Tensor) -> torch.Tensor: + """The fp32 value a key orders by (the inverse of the kernel's order-preserving image).""" + u = (keys >> 32) & _LOW32 + bits = torch.where(u >= 0x80000000, u - 0x80000000, _LOW32 - u) + bits = torch.where(bits >= 2**31, bits - 2**32, bits) # as a signed 32-bit pattern + return bits.to(torch.int32).view(torch.float32) + + +def key_index(keys: torch.Tensor) -> torch.Tensor: + return _LOW32 - (keys & _LOW32) + + +def _bookkeeping(nreq: int, step: int, n_x_rows: int, g: torch.Generator): + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if step == 0 + else None + ) + seq_lens = torch.full((nreq,), T, dtype=torch.int32, device="cuda") if step == 0 else None + next_new = torch.full((nreq, T), -5, dtype=torch.int32, device="cuda") + pos_src = torch.randint(100, 3000, (n_x_rows,), generator=g, device="cuda", dtype=torch.int32) + kv_lens = torch.randint(200, 3000, (nreq,), generator=g, device="cuda", dtype=torch.int32) + ids_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + pos_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + return num_acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next + + +def _head(w: torch.Tensor, coded: bool): + """(wf, (exc_off, exc_pos, exc_w, exp_base)) of a BF16 shard: its 12-bit coding (checked to round-trip) or its + BF16 fragments.""" + if not coded: + return bf16_fragments(w), (None, None, None, 0) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + return packed, (exc_off, exc_pos, exc_w.view(-1), base) + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=(None, None, None, 0), +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + coded head's (exc_off, exc_pos, exc_w, exp_base).""" + acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next = _bookkeeping( + nreq, step, x.shape[0], g + ) + if num_acc is None: + num_acc = acc + if uc is None: + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + out_keys = torch.zeros(nreq, dtype=torch.int64, device="cuda") + tiles = wf.shape[0] if coded[0] is None else coded[0].numel() - 1 + return (x, x_rows, nreq, wf, rank * tiles * 16, keys, counters, uc, uc.data_ptr(), rank, world, + next_new[:, 1:], step, ids_next, pos_src, pos_next, kv_lens, num_acc, seq_lens, DRAFT_LEN, num_sms(), + True, out_keys, *coded) # fmt: skip + + +def _case(spec, seed): + g = gen(seed) + nreq, step, rank, world, vs, k = ( + spec["nreq"], + spec["step"], + spec["rank"], + spec["world"], + spec["vs"], + spec["k"], + ) + pass0 = step == 0 + x = randn((nreq * T if pass0 else nreq, k), g) + wf, coded = _head(randn((vs, k), g, 0.02), spec["coded"]) + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if pass0 + else None + ) + x_rows = (torch.arange(nreq, device="cuda") * T + num_acc.long() - 1) if pass0 else None + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + h = x[x_rows] if pass0 else x + own = shard_logits(h, wf, *coded).max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=MAX_TOKENS, ties=spec.get("ties", False) + ) + return _args(x, x_rows, nreq, wf, rank, world, step, g, uc, keys, counters, num_acc, coded) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The bookkeeping is exact on every row (next positions, KV and sequence lengths). Each row's token is the + in-tree one (this rank's BF16 shard logits from the head's GEMM, their first maximum, against the peers' keys) or + the row is provably undecidable (check_row); the winning key's index is the token, and its value lies in the + token's [lo, hi] (or is the peer's pushed key when a peer wins).""" + args = _case(cell.spec, seed_of(cell.why)) + x, x_rows, nreq, wf, offset, _, counters, uc = args[:8] + rank, world = args[9], args[10] + half = int(counters[1].item()) & 1 + pushed = ( + uc[half * MAX_TOKENS * world : (half * MAX_TOKENS + nreq) * world].view(nreq, world).clone() + ) + h = x[x_rows] if x_rows is not None else x[:nreq] + w = head_weight(wf, x.shape[1], *args[23:27]) + best, idx = local_argmax(torch.nn.functional.linear(h, w).float(), offset) + stock = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, MAX_TOKENS) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert args[6].tolist() == [0, 1] and not bool(args[5].any()), "counters / keys not re-armed" + step = args[12] + others = [c for c in range(args[11].shape[1]) if c != step] + assert torch.equal(args[11][:, others], ref["draft_out"][:, others]) + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + assert torch.equal(args[15], ref["pos_next"]) + assert torch.equal(args[16], ref["kv_lens"]) + if args[18] is not None: + assert torch.equal(args[18], ref["seq_lens"]) + # Candidates per row: this rank's head rows (fp64 bounds) and each peer's pushed key (an exact value). + logits, bound, lo, hi = bounds(h, w) + own_ids = offset + torch.arange(w.shape[0], device="cuda") + peers = [p for p in range(world) if p != rank] + undecidable = 0 + for r in range(nreq): + pv = key_value(pushed[r, peers]).double() + ids = torch.cat([own_ids, key_index(pushed[r, peers])]) + undecidable += check_row( + f"{cell.why[:40]} row {r}", + int(args[13][r]), + int(stock[r]), + ids, + torch.cat([logits[r], pv]), + torch.cat([bound[r], torch.zeros_like(pv)]), + torch.cat([lo[r], pv.float()]), + torch.cat([hi[r], pv.float()]), + ) + report(cell.why[:40], nreq, undecidable) + out_keys = args[22] + assert torch.equal(key_index(out_keys), args[13].long()), "the winning key's index is the token" + got = key_value(out_keys) + for r in range(nreq): + c = int(args[13][r]) - offset + if 0 <= c < w.shape[0]: + assert float(lo[r, c]) <= float(got[r]) <= float(hi[r, c]), ( + f"row {r}: winning value outside [lo, hi]" + ) + else: + assert int(out_keys[r]) in pushed[r].tolist(), ( + f"row {r}: a peer won with a key it did not push" + ) + + +def test_fragments_round_trip() -> None: + g = gen(seed_of("fragments")) + w = randn((64, 96), g) + wf = bf16_fragments(w) + assert wf.shape == (4, 6, 256) + assert torch.equal(from_fragments(wf), w) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 3) # K 96: no coding, the BF16 fragments + assert not prepared.coded + assert torch.equal(prepared.wf, wf) and prepared.vocab_offset == 3 * 64 and prepared.weight is w + + +# ------------------------------------------------------------------ against the in-tree drafter, bit for bit + + +def _stock_tokens(shards: list, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The in-tree greedy draft tokens of the rows x over the vocab-sharded BF16 head: (argmax of the gathered float + logits, greedy_sample_draft_with_tp_gather's per-rank (index, value) pick).""" + logits = [torch.nn.functional.linear(x, w).float() for w in shards] + gathered = torch.argmax(torch.cat(logits, dim=-1), dim=-1).to(torch.int32) + combined = [] + for rank, shard_logit in enumerate(logits): + worker = SimpleNamespace(mapping=SimpleNamespace(tp_rank=rank)) + combined.append(Eagle3OneModelWorker._get_local_max_and_combined(worker, shard_logit)) + pick = Eagle3OneModelWorker._get_draft_tokens_from_gathered + pick = getattr(pick, "_torchdynamo_orig_callable", pick) + tp_gather = pick(None, torch.cat(combined, dim=-1)).to(torch.int32) + return gathered, tp_gather + + +def _fused_tokens(shards: list, x: torch.Tensor, coded: bool, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's four-rank pass with + its peers' keys in its slots, as their multicast pushes would put them.""" + g = gen(seed_of("fused tokens")) + nreq = x.shape[0] + heads = [_head(w, coded) for w in shards] + own_keys = [] + for rank, (wf, cd) in enumerate(heads): + args = list(_args(x, None, nreq, wf, rank, 1, step, g, coded=cd)) + args[9] = 0 # one-rank run: this rank's slot only, global ids of `rank` + op(*args) + own_keys.append(args[22].clone()) + tokens = [] + for rank, (wf, cd) in enumerate(heads): + uc, keys, counters = head_workspace(WORLD, slot_rows=MAX_TOKENS) + for peer in range(WORLD): + if peer != rank: + for t in range(nreq): + uc[t * WORLD + peer] = own_keys[peer][t] + args = _args(x, None, nreq, wf, rank, WORLD, step, g, uc, keys, counters, coded=cd) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + tokens.append(args[13].clone()) + return tokens + + +def _integer_head(g: torch.Generator, rows: int = WORLD * DRAFT_SHARD) -> torch.Tensor: + """A [rows, 6144] head of small integers over the first 4 columns (zero elsewhere): its logits against integer + rows are exact integers in any summation order, many of which share a BF16 value.""" + w = torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, :4] = torch.randint(-60, 61, (rows, 4), generator=g, device="cuda").to(torch.bfloat16) + return w + + +def _shards(w: torch.Tensor) -> list: + return [w[r * DRAFT_SHARD : (r + 1) * DRAFT_SHARD].contiguous() for r in range(WORLD)] + + +def _assert_stock(w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None) -> None: + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + if expect is not None: + assert gathered.tolist() == expect, (gathered.tolist(), expect) + for rank, tokens in enumerate(_fused_tokens(shards, x, coded)): + assert torch.equal(tokens, gathered), ( + f"rank {rank}: {tokens.tolist()} vs stock {gathered.tolist()}" + ) + + +def _ones_rows(nreq: int) -> torch.Tensor: + x = torch.zeros((nreq, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[:, :4] = 1.0 + return x + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_forced_bf16_ties_take_the_lowest_global_id(coded) -> None: + """A head row copied onto every rank that wins every row of ones (logit 244, every other at most 240): the lowest + global id must win, inside a shard and across ranks. Then integer rows in [-2, 2], whose logits (up to 480 in + magnitude, BF16 spacing 2 above 256) collide often, exactly or after BF16 rounding.""" + g = gen(seed_of("forced ties")) + w = _integer_head(g) + win = torch.tensor([61.0, 61.0, 61.0, 61.0], device="cuda").to(torch.bfloat16) + for gid in (DRAFT_SHARD + 5000, DRAFT_SHARD + 77, DRAFT_SHARD + 9000, 95000): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), coded, expect=[95000] * 5) + # Random integer rows: ties after BF16 rounding happen naturally. + x = torch.randint(-2, 3, (8, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + x[:, 4:] = 0 + _assert_stock(w, x, coded) + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_signed_zero_maxima_take_the_lowest_global_id(coded) -> None: + """Rows whose maximum is zero: an all-zero head input (every logit +-0) and logits <= 0 with zeros by cancellation + on several ranks. -0 and +0 compare equal, so the lowest global id reaching zero wins.""" + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 0] = -1.0 + w[:, 1] = -1.0 + for gid in (DRAFT_SHARD + 5, DRAFT_SHARD + 40000, DRAFT_SHARD + 70007, 49999): + w[gid, 0] = 1.0 + w[gid, 1] = -1.0 + x = torch.zeros((3, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[1, :2] = 1.0 + x[2, :2] = 2.0 + x[2, 2] = -1.0 # w[:, 2] is zero: a -0 product on every row + _assert_stock(w, x, coded, expect=[0, 49999, 49999]) + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + (1000, 1001, "lower"), # 1001 rounds to 1000 in BF16: a tie, the lower id wins + (1001, 1002, "lower"), # 1001 and the tie-to-even midpoint 1002 both round to 1000 + (1000, 1004, "higher"), # adjacent BF16 values: the larger wins + (1003, 1000, "lower"), # 1003 rounds up to 1004 > 1000 + ], + ids=["one_fp32_step", "midpoint", "adjacent_bf16", "rounds_up"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where, coded) -> None: + """Top-2 logits one to four fp32 steps apart: the winner is the one the BF16 logits decide (the in-tree logits' + dtype), not the fp32 sum. The pair sits inside rank 1's shard or on ranks 0 and 1; every other logit is at most + 240 in magnitude.""" + g = gen(seed_of(f"near top2 {lower} {higher} {where}")) + w = _integer_head(g) + if where == "in_shard": + lo_id, hi_id = DRAFT_SHARD + 100, DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD - 1000, DRAFT_SHARD + 3 + for gid, value in ((lo_id, lower), (hi_id, higher)): + q, r = divmod(value, 4) # four integers <= 251, exact in BF16 + w[gid, :4] = torch.tensor([q + (i < r) for i in range(4)], device="cuda").to(torch.bfloat16) + expect = lo_id if winner == "lower" else hi_id + _assert_stock(w, _ones_rows(2), coded, expect=[expect, expect]) + + +# ------------------------------------------------------------------ non-integer rows at a BF16 rounding midpoint + +#: A BF16 rounding midpoint (between 6.0 and 6.03125; a tie there rounds to the even 6.0) and the FP32 spacing there. +MIDPOINT = 6.0 + 2.0**-6 +FP32_ULP = 2.0**-21 + + +def _near_midpoint_head(g: torch.Generator, targets: dict): + """A full-vocabulary head [2 * 100032, 6144] and one head-input row x with 65 non-zero entries: 1, 2^-12, 2^-20 + and 2^-28 on columns 0..3 and 61 random non-integer BF16 values. Every head row carries random non-integer BF16 + values on those 61 columns (logits about 0.2 in magnitude); each row of ``targets`` (global id -> fp64 value) also + gets columns 0..3 set so its fp64 logit is the target to within 2^-30 (a coarse value, then three finer + corrections, each BF16-rounded).""" + x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device="cuda") + scales = (1.0, 2.0**-12, 2.0**-20, 2.0**-28) + for col, scale in enumerate(scales): + x[0, col] = scale + x[0, 4:65] = randn((61,), g, 0.5) + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 4:65] = randn((WORLD * DRAFT_SHARD, 61), g, 0.05) + xd = x[0].double() + for gid, target in targets.items(): + rest = float(target - (w[gid].double() * xd).sum()) + for col, scale in enumerate(scales): + v = torch.tensor(rest / scale, dtype=torch.float64).to(torch.bfloat16) + w[gid, col] = v + rest -= float(v) * scale + assert abs(rest) <= 2.0**-30, (gid, rest) + return w, x + + +def _check_fused_against_stock( + name: str, w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None +) -> int: + """Each rank's fused tokens against the in-tree tokens under the per-row rule (full-vocabulary candidates); + ``expect`` (a token per row) additionally pins rows the bound decides. Returns the undecidable rows.""" + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + fused = _fused_tokens(shards, x, coded) + for tokens in fused[1:]: + assert torch.equal(tokens, fused[0]), "the ranks disagree on the global token" + logits, bound, lo, hi = bounds(x, w) + ids = torch.arange(w.shape[0], device="cuda") + undecidable = 0 + for r in range(x.shape[0]): + undecidable += check_row( + f"{name} row {r}", + int(fused[0][r]), + int(gathered[r]), + ids, + logits[r], + bound[r], + lo[r], + hi[r], + ) + if expect is not None: + assert int(gathered[r]) == expect[r] and int(fused[0][r]) == expect[r], ( + name, + int(gathered[r]), + int(fused[0][r]), + expect[r], + ) + others = ids != expect[r] + assert _winnable(expect[r], ids, lo[r], hi[r]), ( + f"{name}: the bound does not let {expect[r]} win" + ) + rivals = [int(i) for i in ids[others][hi[r][others] >= lo[r][expect[r]]]] + assert not any(_winnable(i, ids, lo[r], hi[r]) for i in rivals), ( + f"{name}: the row is not decided" + ) + report(name, x.shape[0], undecidable) + return undecidable + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + ( + -2, + 2, + None, + ), # straddling the midpoint by 2 FP32 ulps: either side is in-tree up to the order + (2, -2, None), + (0, -2, None), # exactly on the midpoint (ties to the even 6.0) against just below it + ( + -2048, + 2048, + "higher", + ), # ~1e-3 either side, well outside the bound: decided by BF16 rounding + ( + -4096, + -2048, + "lower", + ), # both below the midpoint, ~1e-3 apart: a BF16 tie the lower id wins + ], + ids=["straddle", "straddle_reversed", "on_midpoint", "decided_up", "decided_tie"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where, coded) -> None: + """Non-integer rows whose top-2 fp64 logits sit the given number of FP32 ulps from a BF16 rounding midpoint, inside + rank 1's shard or on ranks 0 and 1, every other logit far below. The per-row rule holds on every row; outside the + bound the token is the one the BF16 rounding decides, bit for bit.""" + g = gen(seed_of(f"midpoint {lower} {higher} {where}")) + if where == "in_shard": + lo_id, hi_id = DRAFT_SHARD + 100, DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD - 1000, DRAFT_SHARD + 3 + w, x = _near_midpoint_head( + g, {lo_id: MIDPOINT + lower * FP32_ULP, hi_id: MIDPOINT + higher * FP32_ULP} + ) + expect = None if winner is None else [lo_id if winner == "lower" else hi_id] + _check_fused_against_stock(f"midpoint {where} {lower:+d}/{higher:+d} ulp", w, x, coded, expect) + + +# ------------------------------------------------------------------ the coded head against the BF16 fragments + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_coded_head_is_the_bf16_head_bit_for_bit(nreq) -> None: + """Integer data (fp32 sums exact in any order) with exceptions on both sides of the exponent window: the coded + head's winning keys (value and global id) and its bookkeeping are the BF16 fragments' bit for bit, on one rank and + across the ranks.""" + g = gen(seed_of(f"coded vs bf16 {nreq}")) + sign = torch.randint(0, 2, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * 2 - 1 + w = (torch.randint(1, 41, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * sign).to( + torch.bfloat16 + ) + # Exceptions on both sides of the window of 1..40: exact zeros, and 2^15 (fp32 sums stay below 2^24). + w[::5, ::512] = 0.0 + w[::13, 7::600] = 32768.0 + x = torch.randint(-2, 3, (nreq, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + wf_c, cd = _head(w, True) + assert cd[1].numel() > 0, "the shard must carry exceptions" + wf_b, cb = _head(w, False) + outs = [] + for wf, c in ((wf_c, cd), (wf_b, cb)): + args = _args(x, None, nreq, wf, 1, 1, 1, gen(seed_of("bookkeeping")), coded=c) + args = list(args) + args[9] = 0 + with validating(op): + op(*args) + torch.cuda.synchronize() + outs.append(args) + coded_args, bf16_args = outs + for i in (11, 13, 15, 16, 22): # draft_out, ids_next, pos_next, kv_lens, out_keys + assert torch.equal(coded_args[i], bf16_args[i]), i + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), DRAFT_SHARD) + assert torch.equal(coded_args[22], make_keys(best, idx)), ( + "the winning keys are not the in-tree ones" + ) + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_tiny_exception_products_are_exact(nreq) -> None: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1, and a one-hot + activation on column 0: the exceptions' exact products are the only non-zero logits, so the coded head must pick + row 1 with the BF16 head's key bit for bit (a correction against a non-zero placeholder cancels them to 0).""" + n, k = 256, 1024 + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + x = torch.zeros((nreq, k), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + wf_c, cd = _head(w, True) + assert cd[3] == 113 + keys = [] + for wf, c in ((wf_c, cd), _head(w, False)): + args = _args(x, None, nreq, wf, 0, 1, 1, gen(seed_of("tiny exceptions")), coded=c) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert args[13].tolist() == [1] * nreq + keys.append(args[22]) + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 0) + assert torch.equal(keys[0], keys[1]) and torch.equal(keys[0], make_keys(best, idx)) + + +def test_prepare_head_codes_the_shard_and_fails_closed(monkeypatch) -> None: + """With the coding on, prepare_head keeps the coded shard when it decodes to the weight bit for bit; a coding + that does not is dropped for the BF16 fragments. (The tp2 target ships HEAD_CODED = False.)""" + monkeypatch.setattr(draft_step_kernels, "HEAD_CODED", True) + g = gen(seed_of("prepare coded")) + w = randn((4096, 1024), g, 0.02) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 2) + assert prepared.coded and prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(head_weight(prepared.wf, 1024, prepared.exc_off, prepared.exc_pos, prepared.exc_w, + prepared.exp_base), w) # fmt: skip + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[11] ^= 0x10 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(draft_step_kernels, "encode_chunked", corrupt) + dropped = draft_step_kernels.prepare_head(head, 2) + assert not dropped.coded and dropped.exc_off is None + assert torch.equal(dropped.wf, bf16_fragments(w)) + + +def test_prepare_head_keeps_bf16_when_coding_is_off(monkeypatch) -> None: + """HEAD_CODED = False keeps the BF16 fragments even for a shard whose coding would decode bit for bit.""" + g = gen(seed_of("prepare bf16")) + w = randn((4096, 1024), g, 0.02) + monkeypatch.setattr(draft_step_kernels, "HEAD_CODED", False) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2) + assert not prepared.coded and prepared.exc_off is None + assert torch.equal(prepared.wf, bf16_fragments(w)) + + +def test_guard_refuses_head_rows_past_x() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[1] = args[1] + 1000 + with pytest.raises(AssertionError, match="x_rows must index x"): + with validating(op): + op(*args) + + +def test_guard_refuses_too_few_ctas() -> None: + args = list(_case(op.CELLS[1].spec, 2)) + args[20] = 8 + with pytest.raises(AssertionError, match="tiles per CTA"): + with validating(op): + op(*args) + + +# ------------------------------------------------- one pass as two launches (uGPU, ops/m3_draft_step.UGPU) + + +@pytest.mark.parametrize( + "cell", + [c for c in op.CELLS if not c.spec["coded"]], + ids=[c.why[:44] for c in op.CELLS if not c.spec["coded"]], +) +def test_two_launches_per_domain_are_the_one_launch_pass(cell) -> None: + """A pass split into two launches (tiles [0, split) and [split, n), one per locality domain stream, the exchange + and bookkeeping by the last of both launches' CTAs) writes exactly what the one-launch pass writes: tokens, + winning keys, draft columns, next positions, KV and sequence lengths, and re-arms the counters and keys.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_locality import ( # noqa: E501 + LocalityMap, + ) + + one = _case(cell.spec, seed_of(cell.why)) + two = _case(cell.spec, seed_of(cell.why)) + op(*one) + lm = LocalityMap.get(torch.cuda.current_device()) + streams = lm.domain_streams() + n_tiles = two[3].shape[0] + split = max(1, n_tiles // 2) + per = min(lm.sms_per_domain, split, n_tiles - split) + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + for st, lo, hi in ((streams[0], 0, split), (streams[1], split, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + op(*two[:20], per, False, *two[22:], tile_lo=lo, tile_hi=hi, total_ctas=2 * per) + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) + torch.cuda.synchronize() + for i, name in ( + (11, "draft_out"), + (13, "ids_next"), + (15, "pos_next"), + (16, "kv_lens"), + (22, "out_keys"), + ): + assert torch.equal(one[i], two[i]), name + if one[18] is not None: + assert torch.equal(one[18], two[18]), "seq_lens" + assert two[6].tolist() == [0, 1] and not bool(two[5].any()), "counters / keys not re-armed" + + +# ----------------------------------------------- equal-chunk split (ops/m3_draft_step.EQ_CHUNKS: bnd / bnd_cnt set) + + +def _eq_buffers(n_tiles: int): + return ( + torch.zeros(n_tiles, 32, 4, dtype=torch.float32, device="cuda"), + torch.zeros(n_tiles, dtype=torch.int32, device="cuda"), + ) + + +@pytest.mark.parametrize( + "cell", + [c for c in op.CELLS if not c.spec["coded"]], + ids=[c.why[:44] for c in op.CELLS if not c.spec["coded"]], +) +@pytest.mark.parametrize("launches", [1, 2]) +def test_equal_chunk_split_is_the_whole_tile_pass(cell, launches) -> None: + """Each CTA streaming an equal share of the (tile, chunk) space, the CTA-boundary tiles summed from two partials in + bnd: the same tokens, winning global ids, draft columns, positions and lengths as whole tiles per CTA (the logits + differ at most by the order of fp32 additions: their BF16 keys by at most one ulp), and bnd / bnd_cnt are left + zero for the next call, over one launch and over the two per-domain launches.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_locality import ( # noqa: E501 + LocalityMap, + ) + + whole = _case(cell.spec, seed_of(cell.why)) + eq = _case(cell.spec, seed_of(cell.why)) + n_tiles = eq[3].shape[0] + bnd, bnd_cnt = _eq_buffers(n_tiles) + op(*whole) + if launches == 1: + op(*eq, tile_lo=0, tile_hi=0, total_ctas=0, bnd=bnd, bnd_cnt=bnd_cnt) + else: + lm = LocalityMap.get(torch.cuda.current_device()) + streams = lm.domain_streams() + split = max(1, n_tiles // 2) + per = min(lm.sms_per_domain, split, n_tiles - split) + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + for st, lo, hi in ((streams[0], 0, split), (streams[1], split, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + op(*eq[:20], per, False, *eq[22:], tile_lo=lo, tile_hi=hi, total_ctas=2 * per, bnd=bnd, + bnd_cnt=bnd_cnt) # fmt: skip + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) + torch.cuda.synchronize() + for i, name in ((11, "draft_out"), (13, "ids_next"), (15, "pos_next"), (16, "kv_lens")): + assert torch.equal(whole[i], eq[i]), name + assert torch.equal(key_index(whole[22]), key_index(eq[22])), "winning global ids" + if whole[18] is not None: + assert torch.equal(whole[18], eq[18]), "seq_lens" + assert not bool(bnd.any()) and not bool(bnd_cnt.any()), "bnd / bnd_cnt not restored to zero" + assert eq[6].tolist()[0] == 0 and not bool(eq[5].any()), "counters / keys not re-armed" + + +def test_guard_refuses_an_equal_chunk_split_below_a_tile_per_cta() -> None: + """With fewer tiles than CTAs a tile could gather three partials; the op refuses the split.""" + spec = dict(op.CELLS[0].spec) + args = _case(spec, 7) + n_tiles = args[3].shape[0] + bnd, bnd_cnt = _eq_buffers(n_tiles) + with pytest.raises(RuntimeError, match="a tile per CTA"): + op( + *args[:20], + n_tiles + 1, + *args[21:], + tile_lo=0, + tile_hi=0, + total_ctas=0, + bnd=bnd, + bnd_cnt=bnd_cnt, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py new file mode 100644 index 000000000000..7c9a826c1d3d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the forward-tail spec fold (ops/m3_fold_tail and M3SpecStepFold.in_forward). + +The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers +(step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing +when the step was not armed. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( + m3_fold_tail, + m3_spec_tail, +) + +assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" + +VOCAB = 200064 +NUM_SLOTS = 9 +MAX_ROWS = 8 +STALE = -4 + + +def _store(widths, g): + wn, wx, wd = widths + i32 = dict(dtype=torch.int32, device="cuda") + return types.SimpleNamespace( + new_tokens=torch.randint( + 0, VOCAB, (wn, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_new_tokens=torch.randint( + 0, VOCAB, (wx, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_draft_tokens=torch.randint( + 0, VOCAB, (NUM_SLOTS, wd), generator=g, device="cuda", dtype=torch.int32 + ), + new_tokens_lens=torch.randint(1, 5, (NUM_SLOTS,), generator=g, **i32), + ) + + +def _clone(ns): + return types.SimpleNamespace(**{k: v.clone() for k, v in vars(ns).items()}) + + +def _outputs(rows, t, g): + d = t - 1 + return { + "new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_draft_tokens": torch.randint( + 0, VOCAB, (rows, d), generator=g, device="cuda", dtype=torch.int32 + ), + "new_tokens_lens": torch.randint( + 1, t + 1, (rows,), generator=g, device="cuda", dtype=torch.int32 + ), + } + + +def _engine(t): + i32 = dict(dtype=torch.int32, device="cuda") + return ( + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS * (t - 1),), STALE, **i32), + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS,), STALE, **i32), + ) + + +def _reference(store, engine, slots, out, t): + """slot_scatter into the store, then step_gather of the next step's inputs from it.""" + assert m3_spec_tail.slot_scatter( + slots, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["new_tokens_lens"], + store.new_tokens_lens, + out["next_draft_tokens"], + store.next_draft_tokens, + ) + p = slots.numel() + ids, draft, pos_off, kv_off = engine + slots32 = slots.to(torch.int32) + assert m3_spec_tail.step_gather( + store.next_new_tokens, + store.next_draft_tokens, + store.new_tokens_lens, + slots32, + slots32.repeat_interleave(t), + p, + t, + t - 1, + ids, + 0, + draft, + 0, + pos_off, + kv_off, + 0, + ) + + +def _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=False): + ids, draft, pos_off, kv_off = engine + assert m3_fold_tail.fold_tail( + ctrl, + slots_buf, + rows, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["next_draft_tokens"], + store.next_draft_tokens, + out["new_tokens_lens"], + store.new_tokens_lens, + ids, + draft, + pos_off, + kv_off, + t, + cpp=cpp, + ) + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("widths", [(4, 4, 3), (8, 8, 7)], ids=["store=step", "store wider"]) +def test_tail_matches_scatter_then_gather(rows, widths, cpp) -> None: + """Store bit for bit; engine rows [0, P*T) / [0, P) bit for bit (the forward reads only those rows).""" + t = 4 + g = torch.Generator(device="cuda").manual_seed(1000 * rows + widths[0]) + store = _store(widths, g) + out = _outputs(rows, t, g) + slots = torch.randperm(NUM_SLOTS, generator=g, device="cuda")[:rows] + ref_store, tail_store = _clone(store), _clone(store) + ref_engine, tail_engine = _engine(t), _engine(t) + _reference(ref_store, ref_engine, slots, out, t) + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + slots_buf = torch.zeros(MAX_ROWS, dtype=torch.long, device="cuda") + slots_buf[:rows] = slots + _tail(tail_store, tail_engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(ref_store, name), getattr(tail_store, name)), name + spans = (rows * t, rows * (t - 1), rows * t, rows) + for k, (ref, got, n) in enumerate(zip(ref_engine, tail_engine, spans)): + assert torch.equal(ref[:n], got[:n]), k + assert torch.all(got[n:] == STALE), k + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +def test_unarmed_tail_writes_nothing(cpp) -> None: + """ctrl[0] = 0 (the step was not armed): the store and the engine buffers are untouched.""" + t, rows = 4, 2 + g = torch.Generator(device="cuda").manual_seed(7) + store = _store((4, 4, 3), g) + out = _outputs(rows, t, g) + before, engine = _clone(store), _engine(t) + ctrl = torch.zeros(1, dtype=torch.int32, device="cuda") + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(before, name), getattr(store, name)), name + for buf in engine: + assert torch.all(buf == STALE) + + +def _fold(): + fold = modeling.M3SpecStepFold() + t = 4 + engine = _engine(t) + fold.bind_engine_buffers(*engine, max_batch_size=MAX_ROWS) + g = torch.Generator(device="cuda").manual_seed(11) + fold.bind_sampler_store(_store((4, 4, 3), g)) + return fold, g, t + + +def test_fold_handshake_stages_only_confirmed_rows() -> None: + """begin_step arms; the sampler confirms only the armed rows from the tail's outputs; the next step is staged only + for exactly those rows in that order.""" + fold, g, t = _fold() + assert fold.begin_step([3, 5], 2) is False # nothing staged before the first forward + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [3, 5]) is True + assert fold.begin_step([3, 5], 2) is True + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert ( + fold.sampler_rows_stored(out, 0, [5, 3]) is False + ) # other order: the sampler stores itself + assert fold.begin_step([5, 3], 2) is False + + +def test_fold_handshake_needs_the_tail_outputs() -> None: + """Outputs the tail did not run on (another dict, a mixed step, an unarmed step) are never reported stored.""" + fold, g, t = _fold() + fold.begin_step([1], 1) + fold.store_in_forward(_outputs(1, t, g)) + assert fold.sampler_rows_stored(_outputs(1, t, g), 0, [1]) is False + fold.begin_step(None, 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [1]) is False + fold.begin_step([1], 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 1, [1]) is False # rows were skipped (chunked context) + + +def test_a_padded_step_after_an_armed_one_writes_nothing() -> None: + """K3 #150 check #5: CUDA-graph padding rows carry no seq slot, so the engine passes rows=None for a padded step. + Even right after an armed step, that step uploads ctrl 0 and its forward tail leaves the store and the engine + buffers untouched (a dummy row never writes a live slot).""" + fold, g, t = _fold() + fold.begin_step([2, 4], 2) + fold.store_in_forward(_outputs(2, t, g)) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 2 + + fold.begin_step(None, 4) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 0 + store_before = _clone(fold.store) + engine_before = [buf.clone() for buf in fold.step_inputs] + fold.store_in_forward(_outputs(4, t, g)) + torch.cuda.synchronize() + for name in vars(store_before): + assert torch.equal(getattr(store_before, name), getattr(fold.store, name)), name + for before, buf in zip(engine_before, fold.step_inputs): + assert torch.equal(before, buf) + + +def test_the_fold_stores_in_the_forward_once_bound() -> None: + """The fold always stores and stages at the forward's tail; before the engine binds its buffers nothing is armed.""" + fold = modeling.M3SpecStepFold() + assert fold.in_forward + assert fold.begin_step([0], 1) is False + + +def test_the_store_copy_event_is_recorded_at_binding(): + """An external event the sampler records after its side-stream store copies and the forward's tail waits on, + recorded once at binding so a captured wait has a completed record to wait on.""" + fold = modeling.M3SpecStepFold() + device = torch.device("cuda") + fold.bind_engine_buffers( + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + 4, + ) + assert fold.store_copies_done is not None + fold.store_copies_done.synchronize() + + +def test_the_store_wait_precedes_only_a_plain_launch_under_capture(): + """The store wait is an external-event wait node under graph capture, which drops the next kernel's PDL edge: + it may precede the fold tail (a plain launch); any other launch fails the capture (eager steps are not checked).""" + fold = modeling.M3SpecStepFold() + device = torch.device("cuda") + fold.bind_engine_buffers( + *(torch.zeros(8, dtype=torch.int32, device=device) for _ in range(4)), 4 + ) + assert fold.STORE_WAIT_BEFORE == frozenset({"fold_tail"}) + fold._wait_store_copies(before="draft_step") # eager: allowed + x = torch.zeros(4, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + fold._wait_store_copies(before="fold_tail") + x.add_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(x, torch.ones(4, device=device)) + bad = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="PDL edge"): + with torch.cuda.graph(bad, stream=stream): + fold._wait_store_copies(before="draft_step") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py new file mode 100644 index 000000000000..c102a3374bc1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the per-request forced-acceptance draw (``FORCE_AL_PER_REQUEST``, the ``m3_force_vec_req`` entry, +synth only): each request's forced count comes from its own key (its KV length at the verify), so it does not depend +on the call history, on the request's slot in the batch, or on the other requests.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec_req, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SEED, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _pool(): + g = torch.Generator(device="cpu") + g.manual_seed(_FORCE_ACCEPT_RNG_SEED) + return torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, dtype=torch.float32, generator=g).cuda() + + +def _draw(pool, keys, base_total=2, max_total=4, frac=0.6): + out = torch.empty(keys.numel(), dtype=torch.int32, device="cuda") + m3_force_vec_req( + pool, + keys, + out, + keys.numel(), + base_total, + max_total, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (keys.numel() - 1).bit_length()), + ) + return out + + +def test_a_request_draws_the_same_whatever_ran_before_or_beside_it() -> None: + pool = _pool() + keys = torch.tensor([165_521, 40_003, 323_007, 7], dtype=torch.int32, device="cuda") + first = _draw(pool, keys) + for _ in range(5): # no shared counter: repeated calls draw the same + assert torch.equal(_draw(pool, keys), first) + perm = torch.tensor([2, 0, 3, 1], device="cuda") + assert torch.equal(_draw(pool, keys[perm]), first[perm]) # slot-independent + alone = torch.cat([_draw(pool, keys[i : i + 1]) for i in range(4)]) + assert torch.equal(alone, first) # independent of the other requests + + +def test_the_counts_follow_the_forced_fraction() -> None: + pool = _pool() + keys = torch.arange(100_000, 100_000 + 64 * 400, dtype=torch.int32, device="cuda") + counts = torch.cat([_draw(pool, keys[i : i + 64]) for i in range(0, keys.numel(), 64)]) + assert set(counts.tolist()) <= {2, 3} + extra = (counts == 3).float().mean().item() + assert abs(extra - 0.6) < 0.02, extra diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..3bad6ccc5d2a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec entry of the MiniMax-M3 tp2 target: the fractional forced-acceptance draw.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(5) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + counter = torch.tensor([41], dtype=torch.int64, device="cuda") + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + counter, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert int(counter.item()) == 42, "the draw bumps the step counter once" + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py new file mode 100644 index 000000000000..53a58df60bef --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec_req entry of the MiniMax-M3 tp2 target: the per-request fractional forced-acceptance +draw (one draw per request at its own key, no shared counter).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec_req as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(7) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + # Each request's key is its KV length at the verify: long, distinct, arbitrary. + keys = torch.randint( + 1, 1 << 20, (spec["batch"],), generator=g, device="cuda", dtype=torch.int64 + ).to(torch.int32) + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + keys, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py new file mode 100644 index 000000000000..6d585d8afa32 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (ops/m3_lowm_splitk, +modeling.DRAFT_LOWM_OWN_SPLITK): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's +linears, eager and in a CUDA graph, with and without PDL.""" + +import pytest +import torch +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, +) + +assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" + +# Every certified (N, K, rows) of the TP2 drafter table that a split-K tactic serves: (N, K, rows, tactic). +CELLS = [ + (nk[0], nk[1], m, tactic_for(nk[0], nk[1], m)[1]) + for nk, rows in sorted(ROWS.items()) + for m in rows + if tactic_for(nk[0], nk[1], m)[0] == "s" +] + + +def test_the_table_has_the_certified_split_k_cells() -> None: + assert sorted((n, k, m) for n, k, m, _ in CELLS) == [ + (6144, 4096, 4), + (6144, 4096, 8), + (6144, 4096, 16), + (6144, 9216, 4), + (6144, 9216, 16), + (6144, 9216, 32), + (6144, 18432, 16), + ] + + +@pytest.mark.parametrize("use_pdl", [False, True], ids=["plain", "pdl"]) +@pytest.mark.parametrize("cell", CELLS, ids=[f"{n}x{k}@{m}" for n, k, m, _ in CELLS]) +def test_bitwise_flashinfer(cell, use_pdl) -> None: + n, k, m, tactic = cell + g = torch.Generator(device="cuda").manual_seed(n + 7 * k + m) + for trial in range(3): + x = (torch.randn(m, k, generator=g, device="cuda") * (1.0 + trial)).to(torch.bfloat16) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + ref = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + got = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + fi_run(x, w.t(), None, ref, use_pdl, FiTactic(*tactic)) + run_splitk_dense(x, w.t(), None, got, use_pdl, SplitKTactic(*tactic)) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), f"trial {trial}" + # In a CUDA graph, as the drafter runs it. + cap = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_splitk_dense(x, w.t(), None, cap, use_pdl, SplitKTactic(*tactic)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(cap.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..7a59742e1302 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp2 m3_sampled_accept entry (fused temperature + top-p verify over the 100032-column TP2 shard) +and its workspace size.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_sampled_accept import ( # noqa: E501 + m3_sampled_accept as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_sampled_accept_ws_words import ( # noqa: E501 + m3_sampled_accept_ws_words as ws_words, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" + +DRAFT_LEN = 3 +#: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). +HEAD = 12 +#: Draws per distribution cell, and the total-variation band against the reference distribution. With 12 +#: categories and >= 4000 draws the expected TV of an exact sampler is ~0.02 (sd ~0.006); 0.06 is ~6 sd above it +#: while a wrong nucleus cut (one category too many or too few, >= 1 / 12 of the mass moved) lands at >= 0.08. +DRAWS = 4000 +TV_BAND = 0.06 +#: The cells of one kernel launch (<= 48 verify rows; every certified cell): the tests below call the kernel itself. +SINGLE = [c for c in op.CELLS if c.spec["nreq"] * (DRAFT_LEN + 1) <= 48] + + +def _logits(spec, g, rows): + vs = spec["vs"] + if spec["check"] == "flat": + return (torch.randn((rows, vs), generator=g, device="cuda") * 0.3).to(torch.bfloat16) + if spec["check"] == "argmax": + return (torch.randn((rows, vs), generator=g, device="cuda") * 3).to(torch.bfloat16) + row = torch.full((vs,), -20.0, device="cuda") + hot = torch.randperm(vs, generator=g, device="cuda")[:HEAD] + row[hot] = torch.linspace(2.0, 0.2, HEAD, device="cuda") + return row.to(torch.bfloat16).expand(rows, vs).contiguous() + + +def _case(spec, seed): + g = gen(seed) + nreq = spec["nreq"] + rows = nreq * (DRAFT_LEN + 1) + logits = _logits(spec, g, rows) + words = ws_words(1) + uc = torch.zeros(words, dtype=torch.int64, device="cuda") + counters = torch.zeros(2, dtype=torch.int32, device="cuda") + tokens = torch.zeros(max(48, rows), dtype=torch.int32, device="cuda") + extra = torch.zeros(1, dtype=torch.int32, device="cuda") + temperatures = torch.full((rows,), spec["temperature"], dtype=torch.float32, device="cuda") + top_ps = torch.full((rows,), spec["top_p"], dtype=torch.float32, device="cuda") + seeds = torch.randint(0, 2**62, (rows,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.zeros(rows, dtype=torch.int64, device="cuda") + # Drafts drawn from the likeliest tokens, so some prefixes match. + top = logits[0].float().topk(3).indices.to(torch.int32) + spec["vocab_offset"] + draft = top[torch.randint(0, 3, (nreq, DRAFT_LEN), generator=g, device="cuda")].contiguous() + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + force = spec.get("force", "none") + force_total = 3 if force == "total" else 0 + force_vec = (torch.randint(1, DRAFT_LEN + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if force == "vec" else None) # fmt: skip + return [logits, spec["vocab_offset"], temperatures, top_ps, seeds, offsets, uc, uc.data_ptr(), counters, tokens, + 0, 1, draft, batch_idx, acc, num, gather, force_total, next_new, force_vec, extra, True] # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Statistical for the draw, exact for everything derived from it. + + Every drawn token has nonzero reference probability (lies in the nucleus); the acceptance outputs equal the + reference's strict acceptance of the drawn tokens, every call; the call counter advances once per call; and, + for the distribution cells, the draws over many calls (fresh Philox offsets) are within TV_BAND of the + reference distribution (see DRAWS / TV_BAND). + """ + spec = cell.spec + args = _case(spec, seed_of(cell.why)) + logits, offset, rows = args[0], spec["vocab_offset"], args[0].shape[0] + certifying(op) + ref = op.reference(*args) + probs = ref["probs"] + calls = 1 if spec["check"] == "argmax" else -(-DRAWS // rows) + counts = torch.zeros(logits.shape[1], dtype=torch.float64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + drawn = args[9][:rows].clone() + local = drawn.long() - offset + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + assert bool((probs.gather(1, local[:, None]) > 0).all()), ( + "a drawn token outside the nucleus" + ) + acc = ref["accept"](drawn) + assert torch.equal(args[14], acc["accepted"]) + assert torch.equal(args[15], acc["num_accepted"]) + assert torch.equal(args[16], acc["gather_ids"]) + assert torch.equal(args[18], acc["next_new"]) + # One launch per call (every cell is <= 48 rows). + assert args[8].tolist() == [0, (c + 1) * -(-rows // 48)] + counts.index_add_(0, local, torch.ones(rows, dtype=torch.float64, device="cuda")) + if spec["check"] == "distribution": + emp = counts / counts.sum() + tv = 0.5 * (emp - probs[0].double()).abs().sum().item() + assert tv <= TV_BAND, f"total variation {tv:.3f} against the reference nucleus distribution" + + +@pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) +def test_prewait_parameter_loads_change_nothing(cell) -> None: + """ops/m3_sampled_accept.PREWAIT_PARAMS moves only where the sampling parameters load (before or after the PDL + wait): the same draws and acceptance outputs, bit for bit, call after call.""" + outs = [] + for prewait in (True, False): + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + got = [] + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept(*args, prewait) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + outs.append(got) + for call, (a, b) in enumerate(zip(*outs)): + for x, y in zip(a, b): + assert torch.equal(x, y), f"call {call}: prewait changes {x.tolist()} -> {y.tolist()}" + + +@pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) +def test_the_early_trigger_changes_nothing(cell) -> None: + """ops/m3_sampled_accept.EARLY_TRIGGER moves only where the grid releases its dependents: the same draws and + acceptance outputs, bit for bit, call after call, PDL launch on.""" + outs = [] + for early in (True, False): + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + got = [] + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept(*args, True, early) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + outs.append(got) + for call, (a, b) in enumerate(zip(*outs)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the early trigger changes {x.tolist()} -> {y.tolist()}" + ) + + +def _walk_outputs(args, ilp, calls=3, offset_step=1000): + """The draws and acceptance outputs of `calls` calls with the owner walk one token (ilp False) or 8 tokens at a + time (ilp True), each call at its own Philox offsets.""" + rows = args[0].shape[0] + got = [] + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * offset_step) + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept(*args, True, True, ilp) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + return got + + +@pytest.mark.parametrize("cell", SINGLE, ids=[c.why[:44] for c in SINGLE]) +def test_the_ilp_walk_changes_nothing(cell) -> None: + """ops/m3_sampled_accept.ILP_WALK (SA-PW) only computes the owner's token weights 8 at a time; the walk adds + them in the same order with the same early exit, so the draws and acceptance outputs are the one-token walk's, + bit for bit, call after call.""" + one = _walk_outputs(_case(cell.spec, seed_of(cell.why)), False) + ilp = _walk_outputs(_case(cell.spec, seed_of(cell.why)), True) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.95]) +def test_the_ilp_walk_changes_nothing_over_many_rounds(top_p) -> None: + """A flat tail (many tokens of equal weight) under a tiny and a wide nucleus: rows take many exchange rounds and + each round's walk exits at a different token; the 8-token walk draws the same tokens as the one-token walk.""" + spec = dict(nreq=1, vs=100032, vocab_offset=100032, check="flat", top_p=top_p, temperature=1.0) + one = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), False, calls=24) + ilp = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), True, calls=24) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +def test_the_ilp_walk_is_the_same_under_graph_replay() -> None: + """The 8-token walk captured in a CUDA graph and replayed draws what the one-token walk draws eagerly.""" + cell = op.CELLS[0] + eager = _walk_outputs(_case(cell.spec, seed_of(cell.why)), False) + _walk_outputs( + _case(cell.spec, seed_of("ilp walk warm-up")), True, calls=1 + ) # first launch outside the capture + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + torch.ops.trtllm.minimax_m3_sm107tp2_sampled_accept(*args, True, True, True) + torch.cuda.current_stream().wait_stream(stream) + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + graph.replay() + torch.cuda.synchronize() + for x, y in zip(eager[c], [args[i] for i in (9, 14, 15, 16, 18)]): + assert torch.equal(x, y), ( + f"replay {c}: the graphed 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.0002]) +def test_long_rejection_chains_stay_in_the_nucleus(top_p) -> None: + """A tiny nucleus over a flat tail: rows that need more exchange rounds than the workspace has slots. + + Each round's candidate is drawn over the tokens heavier than the last one, so a row takes ~ln(1 / top_p) + 1 + rounds on average, with a tail past the 16 exchange slots (about 1% of the rows at top_p 0.001). Rounds reuse the + slots cyclically, told apart by a lap-parity bit in the wire tags, and no round accepts unconditionally: every + draw, however many rounds it took, lies in the reference nucleus. World 1 (one rank's 8 CTAs per row exchange + through the same slot ring); the cross-rank path needs the multi-device job. + """ + spec = dict(nreq=1, vs=100032, vocab_offset=100032, check="flat", top_p=top_p, temperature=1.0) + args = _case(spec, seed_of(f"flat tail top_p {top_p}")) + logits, rows = args[0], args[0].shape[0] + probs = op.reference(*args)["probs"] + calls = -(-DRAWS // rows) + outside = torch.zeros((), dtype=torch.int64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + local = args[9][:rows].long() - spec["vocab_offset"] + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + outside += (probs.gather(1, local[:, None]) == 0).sum() + assert args[8].tolist() == [0, c + 1] + draws = calls * rows + rounds = 2.0 + args[20].item() / draws + kept = int((probs > 0).sum(1).min().item()) + assert outside.item() == 0, ( + f"{outside.item()} of {draws} draws outside the nucleus (kept >= {kept} of {logits.shape[1]}, " + f"{rounds:.2f} rounds per row on average)" + ) + assert rounds > 4.0, ( + f"only {rounds:.2f} rounds per row: the rows no longer stress the slot ring" + ) + + +@pytest.mark.parametrize("cell", ws_words.CELLS, ids=[c.why[:44] for c in ws_words.CELLS]) +def test_workspace_words(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(ws_words) + world = cell.spec["world"] + assert ws_words(world) == ws_words.reference(world) + + +def test_guard_refuses_top_k_rows() -> None: + args = _case(op.CELLS[0].spec, 1) + top_ks = torch.full((args[0].shape[0],), 20, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="top-k rows"): + with validating(op): + op(*args, top_ks=top_ks) + + +def test_guard_refuses_fp64_temperatures() -> None: + args = _case(op.CELLS[0].spec, 2) + args[2] = args[2].double() + with pytest.raises(AssertionError, match="fp32 sampling parameters"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..5c6a99bb495a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_shard_accept entry (greedy verify from this rank's logits shard).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + craft_drafts, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" + + +def _case(spec, seed, edit_logits=None): + g = gen(seed) + nreq, d, vs, world, rank = ( + spec["nreq"], + spec["draft_len"], + spec["vs"], + spec["world"], + spec["rank"], + ) + m = nreq * (d + 1) + if spec["strided"]: + logits = randn((m, vs + 72), g, 3.0)[:, :vs] + else: + logits = randn((m, vs), g, 3.0) + if edit_logits is not None: + edit_logits(logits) + uc, keys, counters = head_workspace(world, slot_rows=SHARD_ROWS) + own = logits.float().max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=SHARD_ROWS, ties=spec.get("ties", False) + ) + offset = rank * vs + best, idx = local_argmax(logits.float(), offset) + if world == 1: + tokens = idx.to( + torch.int32 + ) # no peers: this rank's own pick (any row count, incl. W2's chunked calls) + else: + tokens = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, SHARD_ROWS) + draft = craft_drafts(tokens, nreq, d) + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + force_total = 3 if spec["force"] == "total" else 0 + force_vec = (torch.randint(1, d + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if spec["force"] == "vec" else None) # fmt: skip + return (logits, offset, keys, counters, uc, uc.data_ptr(), rank, world, draft, batch_idx, acc, num, gather, + force_total, next_new, num_sms(), force_vec) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bit-exact: the op reduces BF16 logits without arithmetic, ties (within the shard and across ranks) go to + the lowest global index on both sides.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) + # One call per launch: a chunked (W2, > 48-row) call launches once per 12-request chunk. + launches = -(-args[0].shape[0] // SHARD_ROWS) + assert args[3].tolist() == [0, launches] and not bool(args[2].any()) + + +def test_guard_refuses_nan_logits() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0][1, 7] = float("nan") + with pytest.raises(AssertionError, match="NaN"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_batch_index_past_the_requests() -> None: + args = list(_case(op.CELLS[0].spec, 2)) + args[9] = args[9] + 1 + with pytest.raises(AssertionError, match="batch_idx"): + with validating(op): + op(*args) + + +def _signed_zero_maxima(logits: torch.Tensor) -> None: + """Every row's maximum is zero, reached as -0.0 at a lower index and +0.0 at a higher one.""" + logits.fill_(-1.0) + logits[:, 3] = -0.0 + logits[:, 9] = 0.0 + + +@pytest.mark.parametrize("ties", [False, True], ids=["peer_rows", "peer_ties"]) +def test_signed_zero_maxima_take_the_lower_index(ties) -> None: + """-0.0 and +0.0 compare equal, so the lower index wins as with torch.argmax (also across ranks).""" + spec = dict(op.CELLS[1].spec, ties=ties) + args = _case(spec, seed_of(f"signed-zero-{ties}"), edit_logits=_signed_zero_maxima) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..f0417a37a805 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_slot_scatter entry (the spec sampler's per-slot stores in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_slot_scatter import ( # noqa: E501 + m3_slot_scatter as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" + +STALE = -4 + + +def _case(spec, seed): + """Worker outputs and the store's transposed [slots, width] views of its [width, slots, 1] buffers.""" + g = gen(seed) + n, nslots = spec["n"], spec["slots"] + + def rnd(shape): + return torch.randint(0, VOCAB, shape, generator=g, device="cuda", dtype=torch.int32) + + store_new = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_next = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_lens = torch.full((nslots,), STALE, dtype=torch.int32, device="cuda") + store_draft = torch.full((nslots, spec["store_d"]), STALE, dtype=torch.int32, device="cuda") + slots = torch.randperm(nslots, generator=g, device="cuda")[:n] + lens = torch.randint(1, 5, (n,), generator=g, device="cuda", dtype=torch.int32) + return (slots, rnd((n, spec["src_w"])), store_new.squeeze(-1).T, rnd((n, spec["src_w"])), + store_next.squeeze(-1).T, lens, store_lens, rnd((n, spec["src_d"])), store_draft) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, zero padding included, other slots untouched.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in (("da", args[2]), ("db", args[4]), ("dl", args[6]), ("dd", args[8])): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_repeated_slots() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0] = args[0].clone() + args[0][1] = args[0][0] + with pytest.raises(AssertionError, match="distinct"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py new file mode 100644 index 000000000000..b023cbd47add --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for m3_spec_offsets (RuntimeHooks.spec_offsets_op): the overlap scheduler's two speculative step corrections +in one plain launch, bit for bit the in-place ATen adds, eager and in a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_spec_offsets import ( # noqa: E501 + m3_spec_offsets as op, +) + +assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" + +I32 = torch.iinfo(torch.int32) + + +def _inputs(tokens, seqs, seed, wrap=False): + g = torch.Generator(device="cuda").manual_seed(seed) + lo, hi = (I32.min, I32.max) if wrap else (0, 1 << 20) + pos = torch.randint(lo, hi, (tokens,), generator=g, device="cuda", dtype=torch.int32) + kv = torch.randint(lo, hi, (seqs,), generator=g, device="cuda", dtype=torch.int32) + pos_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (tokens,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + kv_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (seqs,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + if wrap: + pos[:2], pos_off[:2] = I32.max, 1 # overflow wraps, as the ATen add does + kv[:1], kv_off[:1] = I32.min, -1 + return pos, pos_off, kv, kv_off + + +CELLS = [c.spec for c in op.CELLS] + [ + dict(tokens=1, seqs=1), + dict(tokens=300, seqs=75), + dict(tokens=4, seqs=0), +] + + +@pytest.mark.parametrize("spec", CELLS, ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items())) +def test_bitwise_aten_adds(spec) -> None: + for trial in range(3): + pos, pos_off, kv, kv_off = _inputs( + spec["tokens"], spec["seqs"], 97 * trial + spec["tokens"], spec.get("wrap", False) + ) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + assert op(pos, pos_off, kv, kv_off) is True + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"trial {trial}" + + +def test_graph_replays_read_the_live_offsets() -> None: + pos, pos_off, kv, kv_off = _inputs(32, 8, 5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + for step in range(3): + fresh = _inputs(32, 8, 11 + step) + for dst, src in zip((pos, pos_off, kv, kv_off), fresh): + dst.copy_(src) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"step {step}" + + +def test_runner_views_are_accepted() -> None: + """The runner passes position_ids[0, ctx:] and kv_lens_cuda[ctx_reqs:seqs]: 1-D slices of contiguous buffers.""" + position_ids = torch.arange(40, dtype=torch.int32, device="cuda").view(1, 40) + kv_lens = torch.arange(12, dtype=torch.int32, device="cuda") + pos_off = torch.full((40,), 2, dtype=torch.int32, device="cuda") + kv_off = torch.full((12,), 3, dtype=torch.int32, device="cuda") + assert op(position_ids[0, 8:], pos_off[:32], kv_lens[2:10], kv_off[:8]) is True + torch.cuda.synchronize() + assert torch.equal( + position_ids[0], torch.cat([torch.arange(8), torch.arange(8, 40) + 2]).int().cuda() + ) + assert torch.equal( + kv_lens, + torch.cat([torch.arange(2), torch.arange(2, 10) + 3, torch.arange(10, 12)]).int().cuda(), + ) + + +@pytest.mark.parametrize( + "bad", + ["int64", "strided", "length", "2d"], +) +def test_declines_other_layouts_without_writing(bad) -> None: + pos, pos_off, kv, kv_off = _inputs(8, 2, 3) + if bad == "int64": + kv = kv.long() + elif bad == "strided": + pos = torch.zeros(16, dtype=torch.int32, device="cuda")[::2] + elif bad == "length": + pos_off = pos_off[:7] + else: + pos = pos.view(2, 4) + before = (pos.clone(), kv.clone()) + assert op(pos, pos_off, kv, kv_off) is False + torch.cuda.synchronize() + assert torch.equal(pos, before[0]) and torch.equal(kv, before[1]) + + +def test_launches_plainly() -> None: + """It writes step inputs that pre-wait readers take before their grid wait (MSA contract, dispatch + PRODUCER_PREWAIT_STEP_INPUTS): its kernel node carries no programmatic (PDL) serialization.""" + from cuda.bindings import runtime as rt + + pos, pos_off, kv, kv_off = _inputs(32, 8, 7) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + err, _, n = rt.cudaGraphGetNodes(raw, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + flags = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind == rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + flags.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + assert flags == [0], flags diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..e6e732174a8b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_gather entry (the overlap scheduler's step-input gathers in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_step_gather import ( # noqa: E501 + m3_step_gather as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" + +#: Rows of the engine's flat buffers (the max batch), and their stale fill. +MAX_ROWS = 8 +STALE = -4 + + +def _case(spec, seed): + g = gen(seed) + p, t, d, nslots, b0 = spec["p"], spec["t"], spec["d"], spec["slots"], spec["b0"] + i32 = dict(dtype=torch.int32, device="cuda") + new_tok = torch.randint(0, VOCAB, (t, nslots, 1), generator=g, device="cuda", dtype=torch.int32) + if spec.get("two_d"): + new_tok = new_tok.squeeze(-1) + next_draft = torch.randint(0, VOCAB, (nslots, d), generator=g, device="cuda", dtype=torch.int32) + lens = torch.randint(1, t + 1, (nslots,), generator=g, device="cuda", dtype=torch.int32) + slots = torch.randperm(nslots, generator=g, device="cuda")[:p].to(torch.int32) + ppi = slots.repeat_interleave(t) + ids = torch.full((MAX_ROWS * t,), STALE, **i32) + draft = torch.full((MAX_ROWS * d,), STALE, **i32) + pos_off = torch.full((MAX_ROWS * t,), STALE, **i32) + kv_off = torch.full((MAX_ROWS,), STALE, **i32) + return ( + new_tok, + next_draft, + lens, + slots, + ppi, + p, + t, + d, + ids, + spec["ids0"], + draft, + spec["d0"], + pos_off, + kv_off, + b0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in ( + ("ids", args[8]), + ("draft", args[10]), + ("pos_off", args[12]), + ("kv_off", args[13]), + ): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_a_slot_past_the_store() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[3] = args[3].clone() + args[3][0] = 8 + with pytest.raises(AssertionError, match="slots must index the store"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py new file mode 100644 index 000000000000..13fc115c0afa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_prologue entry (the decode graph's step-input prologue in one launch).""" + +import random +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_step_prologue import ( # noqa: E501 + m3_step_prologue as op, +) +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" + +MAX_DESC = StepInputPack.MAX_DESC +WINDOW = StepInputPack.GRAPH_WINDOW +DATA_START = StepInputPack.HEADER_BYTES + StepInputPack.DESC_BYTES * MAX_DESC +ARENA = 1 << 20 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _layout(spec, seed): + """Copies as (dst, src, bytes) arena offsets (a host copy's src: a window offset) and the host copies' bytes.""" + rng = random.Random(seed) + h2d, d2d, data = [], [], [] + used = DATA_START + pos = 0 + for _ in range(spec["h2d"]): + n = rng.choice([1, 3, 4, 7, 16, 24, 32, 48, 64, 5, 12]) + pos += rng.randint(0, 40) # gaps and every alignment + h2d.append((pos, used, n)) + data.append(bytes(rng.getrandbits(8) for _ in range(n))) + pos += n + used += -(-n // 16) * 16 + assert pos < ARENA // 2 and used <= WINDOW + # Device copies write the last quarter; unchained ones read the third. + dpos, spos = 3 * ARENA // 4, ARENA // 2 + for i in range(spec["d2d"]): + if spec.get("chained") and h2d: + dst_h, _, n_h = h2d[i % len(h2d)] + off = rng.randint(0, n_h - 1) + src, n = dst_h + off, n_h - off + else: + n = rng.choice([4, 8, 16, 32, 6, 64]) + spos += rng.randint(0, 24) + src = spos + spos += n + dpos += rng.randint(0, 24) + d2d.append((dpos, src, n)) + dpos += n + assert dpos < ARENA and spos < 3 * ARENA // 4 + return h2d, d2d, data + + +def _window(layout, base: int) -> torch.Tensor: + """The window of ``layout`` for an arena at device address ``base``.""" + h2d, d2d, data = layout + w = torch.zeros(WINDOW, dtype=torch.uint8) + words = [len(h2d), len(d2d)] + for dst, src, n in h2d: + words += [base + dst, src, n] + for dst, src, n in d2d: + words += [base + dst, base + src, n] + w[: 8 * len(words)] = torch.tensor(words, dtype=torch.int64).view(torch.uint8) + for (_, off, n), b in zip(h2d, data): + w[off : off + n] = torch.frombuffer(bytearray(b), dtype=torch.uint8) + return w.cuda() + + +def _arenas(seed): + a = torch.randint(0, 256, (ARENA,), generator=gen(seed), device="cuda", dtype=torch.int32) + a = a.to(torch.uint8) + return a, a.clone() + + +@pytest.mark.parametrize("with_epoch", [True, False], ids=["epoch", "no_epoch"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, with_epoch, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Byte for byte the default prologue's destinations, header clear and epoch.""" + seed = seed_of(cell.why) + layout = _layout(cell.spec, seed) + arena, ref_arena = _arenas(seed) + window, ref_window = _window(layout, arena.data_ptr()), _window(layout, ref_arena.data_ptr()) + epoch = torch.full((1,), 41, dtype=torch.int64, device="cuda") if with_epoch else None + ref_epoch = epoch.clone() if with_epoch else None + certifying(op) + op(window, epoch, MAX_DESC) + op.reference(ref_window, ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena) + # Each window holds its own arena's addresses: past the consumed header, both are left as staged. + for w, a in ((window, arena), (ref_window, ref_arena)): + assert not w[:16].any(), "the header is consumed" + assert torch.equal(w[16:], _window(layout, a.data_ptr())[16:]) + if with_epoch: + assert int(epoch.item()) == int(ref_epoch.item()) == 42 + + +def test_graph_replays_scatter_each_new_window() -> None: + """Captured once on a static window: every replay writes that replay's window and bumps the epoch once.""" + arena, ref_arena = _arenas(7) + window = torch.zeros(WINDOW, dtype=torch.uint8, device="cuda") + epoch = torch.zeros(1, dtype=torch.int64, device="cuda") + ref_epoch = epoch.clone() + op(window, epoch, MAX_DESC) # the empty window: writes nothing + op.reference(window.clone(), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(window, epoch, MAX_DESC) + for step, spec in enumerate( + [dict(h2d=24, d2d=8, chained=True), dict(h2d=0, d2d=0), dict(h2d=40, d2d=4)] + ): + layout = _layout(spec, 100 + step) + window.copy_(_window(layout, arena.data_ptr())) + graph.replay() + op.reference(_window(layout, ref_arena.data_ptr()), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena), step + assert not window[:16].any(), step + assert int(epoch.item()) == int(ref_epoch.item()) == step + 2, step + + +def _pack_case(prologue_op): + pack = StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=True, + in_graph=True, + early_copy=True, + step_epoch=True, + prologue_op=prologue_op, + ) + a, b, c = (torch.zeros(16, dtype=torch.int32, device="cuda") for _ in range(3)) + out = torch.zeros(48, dtype=torch.int32, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _pack_step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + + def host(v): + return torch.tensor(v, dtype=torch.int32).pin_memory() + + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, host(values[0])) + step_inputs.copy_(b, host(values[1])) + if overlap: + step_inputs.copy_(a[:4], host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if not flush: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_pack_with_the_one_launch_prologue_matches_the_default() -> None: + """The same pack steps through the default prologue and the one-launch one: the same outputs and epochs, and the + one-launch prologue is engaged (the default's warm-up compile is not what bumped its epoch).""" + packs = [_pack_case(None), _pack_case(op)] + assert packs[0][0].prologue_op is None and packs[1][0].prologue_op is op + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + outs = [] + for pack, graph, dsts, out in packs: + # A replay with a new pack, one with an eagerly issued pack, and one with nothing to scatter. + _pack_step(pack, graph, dsts, (va, vb), overlap=step == 2, flush=step != 3) + torch.cuda.synchronize() + outs.append((out.tolist(), int(pack.epoch.item()))) + assert outs[0] == outs[1], step + assert outs[1][1] == step + 2, step # one bump at warm-up, then one per replay diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py new file mode 100644 index 000000000000..f89a2ddde2e6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_wide_draft_logits.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""W2b (SpecTail, wide target): the in-tree draft loop drafts from the drafter head's local vocab shard.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.modeling import ( # noqa: E501 + M3Eagle3Worker, + _M3ShardDraftLogits, +) +from tensorrt_llm._torch.modules.linear import TensorParallelMode + +pytestmark = pytest.mark.cpu_only + + +class _Head(nn.Module): + tp_mode = TensorParallelMode.COLUMN + padding_size = 0 + tp_size = 2 + bias = None + + def apply_linear(self, x, bias): + assert bias is None + return x + 1 # stands for this rank's shard GEMM + + +class _Original(nn.Module): + def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits=False): + return hidden_states * 0 - 5 + + +def test_the_shard_processor_returns_the_local_shard_for_selected_rows() -> None: + proc = _M3ShardDraftLogits(_Original()) + x = torch.zeros(3, 4, dtype=torch.bfloat16) + assert torch.equal(proc(x, _Head(), None, True), x + 1) + # Rows not selected by the caller take the original processor (it selects them itself). + assert torch.equal(proc(x, _Head(), None, False), x * 0 - 5) + # The shared original is not a submodule of the swap (the target model owns it). + assert list(proc.children()) == [] + + +def _worker(**over): + w = M3Eagle3Worker.__new__(M3Eagle3Worker) + if isinstance(w, nn.Module): + nn.Module.__init__( + w + ) # attribute storage only; the worker's own init needs a full model config + w.is_mtp_eagle = over.get("mtp", False) + w.guided_decoder = over.get("guided", None) + w.mapping = SimpleNamespace(enable_attention_dp=over.get("adp", False), tp_size=2, tp_rank=0) + return w + + +def _draft(head=None): + return SimpleNamespace( + lm_head=head if head is not None else _Head(), logits_processor=_Original() + ) + + +def test_the_wide_draft_logits_engage_only_where_the_shard_argmax_is_valid() -> None: + meta = SimpleNamespace(wants_advanced_draft_sampling=False) + ok = _worker()._m3_wide_draft_logits_ok(meta, _draft()) + assert ok == bool(dispatch.WIDE_SPEC_TAIL) + if not dispatch.WIDE_SPEC_TAIL: + return + assert not _worker()._m3_wide_draft_logits_ok( + SimpleNamespace(wants_advanced_draft_sampling=True), _draft() + ) + assert not _worker(guided=object())._m3_wide_draft_logits_ok(meta, _draft()) + assert not _worker(adp=True)._m3_wide_draft_logits_ok(meta, _draft()) + assert not _worker(mtp=True)._m3_wide_draft_logits_ok(meta, _draft()) + assert not _worker()._m3_wide_draft_logits_ok(meta, None) + padded = _Head() + padded.padding_size = 8 + assert not _worker()._m3_wide_draft_logits_ok(meta, _draft(padded)) + + +def test_the_base_draft_loop_routes_the_head_through_the_swapped_processor() -> None: + """The swap only matters if the in-tree loop calls draft_model.logits_processor and then sample_draft_tokens; if + the base loop stops doing so, W2b goes inert again (the frozen 2e0a73900b was: the head's processor all-gathered + and cast to fp32, so the sharded check never passed). Fail here instead of silently in-model.""" + import inspect + + from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + + # _forward_draft_loop dispatches to the linear loop, which runs the head and the pick. + dispatcher = inspect.getsource(Eagle3OneModelWorker._forward_draft_loop) + assert "self._forward_linear_draft_loop(" in dispatcher + src = inspect.getsource(Eagle3OneModelWorker._forward_linear_draft_loop) + assert "draft_model.logits_processor(" in src + assert "self.sample_draft_tokens(" in src + + +def test_a_wide_draft_pick_reaches_the_shard_argmax(monkeypatch) -> None: + """The path the in-tree loop takes on a wide step: the swapped processor returns the BF16 shard, and + sample_draft_tokens -> greedy_sample_draft_with_tp_gather hands it to the fused shard argmax, whose tokens come + back (engagement at wide rows, without a GPU).""" + if not dispatch.WIDE_SPEC_TAIL: + pytest.skip("WIDE_SPEC_TAIL is off") + calls = [] + + def fake_shard_accept(logits, vocab_offset, mapping, draft, batch_idx, force_total, *a, **k): + calls.append((tuple(logits.shape), logits.dtype, vocab_offset, tuple(draft.shape))) + m = logits.shape[0] + acc = (torch.arange(m, dtype=torch.int32) + 7).view(m, 1) + return acc, torch.ones(m, dtype=torch.int32), torch.zeros(m, dtype=torch.int64) + + monkeypatch.setattr(dispatch.m3_shard_accept_kernel, "shard_accept", fake_shard_accept) + w = _worker() + w._d2t = None + w.mapping.tp_rank = 1 + shard, rows = 100032, 46 + meta = SimpleNamespace(wants_advanced_draft_sampling=False, draft_vocab_size=2 * shard) + draft_model = _draft() + assert w._m3_wide_draft_logits_ok(meta, draft_model) + + class _ShardHead(_Head): + def apply_linear(self, x, bias): + return torch.zeros(x.shape[0], shard, dtype=torch.bfloat16) + + proc = _M3ShardDraftLogits(draft_model.logits_processor) + logits = proc(torch.zeros(rows, 8), _ShardHead(), None, True) + tokens = w.sample_draft_tokens(logits, meta, rows, draft_step=0) + assert calls == [((rows, shard), torch.bfloat16, shard, (rows, 0))] + assert torch.equal(tokens, torch.arange(rows, dtype=torch.int32) + 7) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..9e16b0a9fc2b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic layer's feed-forward exit on Eagle3 capture layers, in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``). + +Every layer's MoE / MLP output goes unreduced to the fused boundary (all-reduce + residual + the next layer's input +norm, the final norm after the last layer), and a capture layer captures the reduced residual sum that boundary +writes, with no residual of its own. CPU only: the attention, the feed-forward and the boundaries are stand-ins. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) + +LAYERS = 4 + + +class _Spec: + def __init__(self, layers): + self.layers = set(layers) + + def is_layer_capture(self, i): + return i in self.layers + + +@pytest.fixture(params=["prefill", "decode"]) +def layer(request, monkeypatch): + """(core, the target's generic layer bound to a target over ``core``).""" + calls = [] + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + c = SimpleNamespace( + w=w, + num_layers=LAYERS, + moe_layers=frozenset({2, 3}), + rms_eps=1e-6, + calls=calls, + # The 9..48-row MC finalize boundary declines on these stand-ins (the stand-in boundary is under test). + _decode=SimpleNamespace(_generic_mc_finalize=lambda i, h, r, norm_w: None), + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, + "_native_allreduce_norm", + lambda core, x, r, norm: (("normed", x, norm), ("updated", x, r)), + ) + + def boundary(self, x, x2, residual, norm_w): + calls.append((x, x2, norm_w)) + return ("next", norm_w), ("sum", x, x2) + + if request.param == "prefill": + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i))) + monkeypatch.setattr( + cls, "_swiglu_mlp", lambda self, h, prefix, low_m, dense=False: ("mlp", prefix) + ) + else: + cls = modeling.DecodeTarget + monkeypatch.setattr( + cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i), False, False) + ) + monkeypatch.setattr( + cls, + "_swiglu_mlp", + lambda self, h, prefix, low_m, push=None, dense=False: (("mlp", prefix), False), + ) + # The 33..48-row MC finalize declines here: the stand-in boundary is the one under test. + monkeypatch.setattr(cls, "_generic_mc_finalize", lambda self, i, h, r, norm_w: None) + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + monkeypatch.setattr(cls, "_generic_boundary", boundary) + target = cls(c) + return c, target._layer if request.param == "prefill" else target._generic_layer + + +@pytest.mark.parametrize("i", range(LAYERS)) +def test_a_capture_layer_captures_the_boundary_residual_sum(layer, i) -> None: + core, run = layer + captured = [] + out = run( + i, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([i]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + norm = core.w[f"l{i + 1}_in_norm"] if i + 1 < LAYERS else "final" + assert core.calls[-1][2] == norm, ( + "the boundary applies the next layer's norm (the final norm after the last)" + ) + assert captured == [(i, out[1], None)] + if i in core.moe_layers: + assert core.calls[-1][1] == ("routed", i), "the routed part folds into the boundary" + + +def test_a_layer_without_capture_captures_nothing(layer) -> None: + _, run = layer + captured = [] + run( + 1, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([]), + lambda *a: captured.append(a), + ) + assert not captured + + +def test_a_prefill_moe_layer_takes_the_mc_finalize_boundary_when_it_folds(monkeypatch) -> None: + """PrefillTarget at a MoE layer: the decode target's generic MC finalize (9..48 rows) ends the layer when it folds, + its updated residual is the Eagle3 capture, and the generic boundary does not run.""" + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + seen = [] + + def mc_finalize(i, h, r, norm_w): + seen.append((i, h, r, norm_w)) + return ("normed", i), ("updated", i) + + c = SimpleNamespace( + w=w, + num_layers=LAYERS, + moe_layers=frozenset({2, 3}), + rms_eps=1e-6, + _decode=SimpleNamespace(_generic_mc_finalize=mc_finalize), + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, "_native_allreduce_norm", lambda core, x, r, norm: (("normed", x), ("upd", x, r)) + ) + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + + def no_boundary(*a, **k): + raise AssertionError("the generic boundary runs after a folded MC finalize") + + monkeypatch.setattr(cls, "_generic_boundary", no_boundary) + monkeypatch.setattr( + cls, "_moe", lambda self, *a, **k: pytest.fail("the MoE runs outside the fold") + ) + target = cls(c) + captured = [] + out = target._layer( + 2, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([2]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + assert out == (("normed", 2), ("updated", 2)) + assert seen and seen[0][0] == 2 and seen[0][3] == "in_norm3" + assert captured == [(2, ("updated", 2), None)] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py new file mode 100644 index 000000000000..1394b8bd5924 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_decode_dense_route.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A dense layer's MLP runs its SwiGLU-OAI activation through ``_dense_swiglu``, the seam of the dense swigluoai C++ +op (``m3_dense_swiglu``), in the decode-graph layer (``DecodeTarget._layer``) and in the prefill layer, eager or in a +decode graph (r4.10.1 called the module directly, so the op never ran in a decode graph). The shared expert keeps the +plain ``m3_silu_and_mul``. + +CPU only: the GEMMs, the activations and the boundaries are stubbed; the test reads which seam each call takes.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, +) + + +@pytest.fixture +def seams(monkeypatch): + """Records which activation seam each MLP call takes; the GEMMs pass their input through.""" + calls = [] + monkeypatch.setattr( + modeling, + "_linear", + lambda core, x, w, s, low_m_gemm, push=None: (x, push is not None), + ) + monkeypatch.setattr( + modeling, "_dense_swiglu", lambda core, gu: calls.append(("dense", gu.shape[0])) or gu + 1 + ) + monkeypatch.setattr( + modeling, + "m3_silu_and_mul", + lambda gu, *args: calls.append(("plain", gu.shape[0])) or gu + 2, + ) + return calls + + +def _core(): + return SimpleNamespace( + swiglu_limit=7.0, + swiglu_alpha=1.702, + swiglu_beta=1.0, + num_layers=2, + moe_layers=set(), + w={"l0_post_norm": None, "l1_in_norm": None}, + ) + + +@pytest.mark.parametrize("decode_graph", [True, False]) +@pytest.mark.parametrize("rows", [1, 4, 8]) +def test_decode_dense_layer_runs_through_dense_swiglu(monkeypatch, seams, rows, decode_graph): + monkeypatch.setattr( + modeling, "flashinfer_mxfp8_decode_graph_capture_active", lambda: decode_graph + ) + target = SimpleNamespace(core=_core()) + target._attention = lambda i, position_ids, hidden, step, armed: (hidden, False) + target._attn_push_target = lambda num_tokens: None + target._boundary = lambda mc_ar, h, residual, norm, **kwargs: (h, residual) + target._swiglu_mlp = lambda *args, **kwargs: modeling.DecodeTarget._swiglu_mlp( + target, *args, **kwargs + ) + x = torch.zeros(rows, 16) + plan = SimpleNamespace(attn_mc_ar=False, moe_mc_ar=False, qkv_prequant=False) + out, _ = modeling.DecodeTarget._layer(target, 0, plan, None, x, x.clone(), {}) + assert seams == [("dense", rows)] + assert torch.equal(out, x + 1) + + +@pytest.mark.parametrize("rows", [1, 64]) +def test_prefill_dense_mlp_runs_through_dense_swiglu(seams, rows): + target = SimpleNamespace(core=_core()) + x = torch.zeros(rows, 16) + out = modeling.PrefillTarget._swiglu_mlp(target, x, "l0_mlp", True, dense=True) + assert seams == [("dense", rows)] + assert torch.equal(out, x + 1) + + +def test_shared_expert_keeps_the_plain_activation(seams): + target = SimpleNamespace(core=_core()) + x = torch.zeros(4, 16) + out, pushed = modeling.DecodeTarget._swiglu_mlp(target, x, "l3_sh", True) + assert seams == [("plain", 4)] and not pushed + assert torch.equal(out, x + 2) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..49e21a3a471b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,771 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp2 (Rubin) decode target's fast-path table (``modeling.py``: the row ranges and the selection functions): its +row ranges against the catalog cells that certify them, and the selection for every decode row count. + +Every range's covered shapes are listed in ``ENABLED``; a range not listed covers nothing. + +CPU only: nothing here launches a kernel. +""" + +import inspect +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_draft_step as draft_step_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + triton_sparse_decode as tsd, +) + +ROWS = range(1, dispatch.TINY_MAX_TOKENS + 2) # every decode row count and the first one past it + + +@pytest.fixture(autouse=True) +def _policy(): + """Every test starts outside require mode with nothing logged, and leaves it so.""" + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_every_covered_shape_is_a_certified_cell(fast_path) -> None: + uncertified = fast_path.keys - fast_path.certified() + assert not uncertified, ( + f"{fast_path.path} covers {sorted(uncertified, key=str)} but " + f"{type(fast_path.entry).__name__}.CELLS certifies none of them: add a cell or narrow the range" + ) + + +def test_the_fast_path_names_are_unique() -> None: + names = [r.path for r in dispatch.FAST_PATHS] + assert len(names) == len(set(names)) + + +def test_the_decode_target_takes_the_gen_workers_batch() -> None: + """Up to eight requests verifying 1 + 3 drafts (32 rows) take the decode target; a wider step is the prefill + target's. Every fast path still keeps its own rows (the batch-1/2 target's TINY_MAX_TOKENS unless raised).""" + assert dispatch.TINY_MAX_TOKENS == 8 + rows = range(1, dispatch.DECODE_MAX_TOKENS + 2) + if dispatch.B8_DECODE: + assert dispatch.DECODE_MAX_TOKENS == 32 + elif dispatch.B4_DECODE: + assert dispatch.DECODE_MAX_TOKENS == 16 + else: + assert dispatch.DECODE_MAX_TOKENS == dispatch.TINY_MAX_TOKENS + expected = [True] * dispatch.DECODE_MAX_TOKENS + [False] + assert [dispatch.is_decode_step(n) for n in rows] == expected + + +# ------------------------------------------------------------------ the cells turned on + + +def test_the_moe_engine_takes_every_decode_step() -> None: + """The MoE lane's m3_moe engine (TP2 width 1536, shared expert folded in) at 1..8 rows, unless killed; its sum + pushes into the multicast boundary all-reduce at the engine's rows while that all-reduce is on.""" + expected = set() if dispatch.KILL_MOE_ENGINE else set(range(1, 9)) + assert set(dispatch.MOE_ENGINE.keys) == expected + for rows in ROWS: + if rows in expected: + want = dispatch.MOE_ENGINE.path + elif dispatch.MC_FINALIZE.covers(rows): # B3: past the engine, the multicast finalize + want = dispatch.MC_FINALIZE.path + else: + want = None + assert dispatch.moe_path(rows) == want, rows + push_rows = dispatch.MOE_ENGINE_MAX_TOKENS if dispatch.MC_AR_TP2 else 0 + assert dispatch.MOE_PUSH_MAX_TOKENS == push_rows + assert set(dispatch.MOE_PUSH.keys) == set(range(1, push_rows + 1)) + + +# ------------------------------------------------------------------ on: the sparse layers' block selection + + +def test_the_index_scorer_takes_one_to_forty_eight_requests_with_their_schedules() -> None: + # W-IDX1 (9..48 requests) is parked: the default covers 1..8; its schedules below stay certified for the knob. + assert dispatch.INDEX_SCORE.keys == frozenset(range(1, 9)) + assert dispatch.index_score_schedule(1) == (1536, True) + assert dispatch.index_score_schedule(2) == (768, True) + assert dispatch.index_score_schedule(3) == (424, True) + assert dispatch.index_score_schedule(4) == (318, True) + for requests in range(5, 9): + assert dispatch.index_score_schedule(requests) == (1272 // requests, True) + # W-IDX1: 9..48 requests in 7 split buckets, about 1272 CTAs in total, early issue + assert len({dispatch.index_score_schedule(r)[0] for r in range(9, 49)}) == 7 + for requests in range(9, 49): + split, early = dispatch.index_score_schedule(requests) + assert early and requests * split <= 1272 + split + # every covered request count has a split (the graph captures it) + assert set(dispatch.INDEX_SCORE.keys) <= set(dispatch.INDEX_SCORE_SPLIT_K) + + +def test_the_long_select_takes_every_key_width_up_to_its_maximum() -> None: + assert dispatch.LONG_SELECT.keys == frozenset(dispatch.LONG_SELECT_KEY_WIDTHS) + assert (dispatch.LONG_SELECT_MIN_BLOCKS, dispatch.LONG_SELECT_MAX_BLOCKS) == (129, 9216) + + +def test_the_384_thread_select_takes_every_key_width_with_its_launch_key() -> None: + assert dispatch.SELECT_LONG_DSL is True + assert dispatch.LONG_SELECT_DSL.keys == frozenset(dispatch.LONG_SELECT_DSL_KEY_WIDTHS) + keys = dispatch.long_select_dsl_keys() + assert [key[1:] for key in keys] == [(False, dispatch.SELECT_LONG_TRIGGER_PHASE)] + + +@pytest.mark.parametrize("num_blocks", [129, 1294, 5888, 9216]) +def test_a_long_select_decline_inside_its_range_is_a_miss(num_blocks) -> None: + """The kernel declining operands its cells certify must not fall back silently under require.""" + assert dispatch.long_select_covers(num_blocks) + # Outside require: logged once, the in-tree selector runs. + dispatch.long_select_fallback(num_blocks, 8) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="long_select"): + dispatch.long_select_fallback(num_blocks, 8) + + +@pytest.mark.parametrize("num_blocks", [1, 16, 128, 9217, 20000]) +def test_rows_outside_the_long_select_range_are_uncovered_even_under_require(num_blocks) -> None: + """Short rows and rows past LONG_SELECT_MAX_BLOCKS take the in-tree selector in every mode (long contexts must not + fail under require).""" + assert not dispatch.long_select_covers(num_blocks) + dispatch.set_require(True) + dispatch.long_select_fallback(num_blocks, 8) + + +def test_taken_logs_without_raising_under_require() -> None: + dispatch.set_require(True) + dispatch.taken(dispatch.INDEX_SCORE.path, 1) + dispatch.taken(dispatch.LONG_SELECT.path, 8) + + +def test_an_index_score_cell_of_another_schedule_certifies_nothing() -> None: + spec = dict(seq_lens=(165521,), qlen=4, heads=2, split_k=512, early=True) + assert tuple(dispatch.INDEX_SCORE.certifies(spec)) == () + assert tuple(dispatch.INDEX_SCORE.certifies(dict(spec, split_k=1536, heads=1))) == () + assert tuple(dispatch.INDEX_SCORE.certifies(dict(spec, split_k=1536))) == (1,) + + +#: Every TP2 MXFP8 projection (N, K): sparse-layer QKV + index, dense QKV, o_proj, shared gate_up / down, dense +#: gate_up / down. +TP2_PROJECTIONS = ( + (4992, 6144), + (4608, 6144), + (6144, 4096), + (3072, 6144), + (6144, 1536), + (12288, 6144), + (6144, 6144), +) +_DECODE_ROWS = range(1, dispatch.TINY_MAX_TOKENS + 1) +# Requests of a fused spec-tail step (B4 spec tail: up to the target's steps). +_SPEC_REQS = range(1, dispatch.SPEC_TAIL_MAX_REQUESTS + 1) +# W2: the fused verify's requests (up to WIDE_MAX_TOKENS rows, unless WIDE_SPEC_TAIL is killed). +_VERIFY_REQS = range(1, dispatch.WIDE_VERIFY_MAX_REQUESTS + 1) +#: B2: the GEMVs' rows at batch 4 (1..16 unless MX_GEMV_B4 is killed). The bf16 GEMV declines dense gate_up / down at +#: 9..16 rows (24 k-tiles per split at MMA width 16, past its in-kernel quantize limit of 16). +_GEMV_HI = ( + (dispatch.DECODE_MAX_TOKENS if dispatch.B8_MX_GEMV_32 else dispatch.B4_MAX_TOKENS) + if dispatch.MX_GEMV_B4 + else dispatch.TINY_MAX_TOKENS +) +_BF16_WIDE_DECLINES = frozenset({(12288, 6144), (6144, 6144)}) + + +#: Batch 8: at 17..32 rows (MMA width 32) the bf16 GEMV's in-kernel quantize serves only shared down (6 k-tiles per +#: split); the other projections keep their 16-row reach. +_BF16_B8_SERVES = frozenset({(6144, 1536)}) + + +#: Step B (GEMM-AR 794929): dense gate_up / down and o_proj keep the 16-row GEMV even with B8_MX_GEMV_32. +_GEMV_B4_ROWS_ONLY = frozenset({(12288, 6144), (6144, 6144), (6144, 4096)}) + + +def _mx_hi(nk) -> int: + if nk in _GEMV_B4_ROWS_ONLY: + return min(_GEMV_HI, 16) + # W-G1: the four projections it raises reach their served rows up to 48 (dispatch._GEMV_W48_ROWS). + return max(dispatch._GEMV_W48_ROWS.get(nk, ()), default=_GEMV_HI) + + +def _bf16_hi(nk) -> int: + if nk in _BF16_WIDE_DECLINES: + return dispatch.TINY_MAX_TOKENS + if _GEMV_HI > 16 and nk not in _BF16_B8_SERVES: + return 16 + return _GEMV_HI + + +#: Fast paths a lane has turned on, each with the shapes it covers (the rest of the skeleton stays off). +ENABLED = { + # GEMM-AR (agg G-d): the tensor-core router of an unfused MoE at 9..48 rows, as the agg canon runs it. + "m3_router_tc": frozenset( + range(dispatch.ROUTER_TC_GEN_MIN_ROWS, dispatch.ROUTER_TC_GEN_MAX_ROWS + 1) + ), + # SpecTail: sampled (temperature + top-p) verify of 1-2 requests at draft length 3 (TP2 shard kernel). + "m3_sampled_accept": frozenset( + (3, n) for n in range(1, dispatch.SAMPLED_VERIFY_MAX_REQUESTS + 1) + ), + # SpecTail: greedy verify (shard argmax + accept) and the fused draft loop of 1-2 requests at draft length 3, + # and that loop's draft-input rows (1 / 2 later-pass rows, 4 / 8 first-pass rows). + "m3_shard_accept": frozenset((3, n) for n in _VERIFY_REQS), + "m3_draft_step": frozenset((3, n) for n in _SPEC_REQS), + "m3_draft_input": frozenset(r for n in _SPEC_REQS for r in (n, 4 * n)), + # SpecTail: the drafter's BF16 linears on FlashInfer's low-M kernels (TP8 port) at the certified b1 rows. + "m3_draft_lowm": frozenset( + { + ((4608, 12288), 1), + ((4608, 12288), 4), + ((6144, 4096), 1), + ((6144, 4096), 4), + ((18432, 6144), 1), + ((6144, 9216), 1), + ((6144, 9216), 4), + ((6144, 18432), 1), + } + # Batch 4: the 16-row draft pass 0 (o, down, fc), unless DRAFT_LOWM_MAX_ROWS is set back to 4. + | ( + {((6144, 4096), 16), ((6144, 9216), 16), ((6144, 18432), 16)} + if dispatch.DRAFT_LOWM_MAX_ROWS >= 16 + else set() + ) + # Batch 8: the 8-row passes' o and the 32-row pass 0's down, unless DRAFT_LOWM_MAX_ROWS is set lower. + | ({((6144, 4096), 8)} if dispatch.DRAFT_LOWM_MAX_ROWS >= 8 else set()) + | ({((6144, 9216), 32)} if dispatch.DRAFT_LOWM_MAX_ROWS >= 32 else set()) + ), + # SpecTail: the fused first drafter inputs of 1-2 fused-verify requests, the capture norm and the chunk norm at + # their 4 / 8 verify rows (own ranges and kills since the TP2 fast-path audit). + "m3_fused_first_inputs": frozenset(_SPEC_REQS), + "m3_capture_norm": frozenset(4 * n for n in _SPEC_REQS), + "m3_chunk_norm": frozenset(4 * n for n in _SPEC_REQS), + # MSA: the long-context sparse decode with the MXFP8 o_proj hand-off, every decode step, and the sparse layers' + # NVFP4 QKV + indexer producer at the TP2 heads. + "m3_sparse_decode_lc": frozenset({"mxfp8"}), + "m3_producer": frozenset({"nvfp4_tp2"}), + "sparse_decode_producer": frozenset({"per_row"}), + # MoE: the m3_moe engine (TP2 width 1536, shared expert folded in) at 1..8 rows, unless killed. + "m3_moe_engine": frozenset() if dispatch.KILL_MOE_ENGINE else frozenset(range(1, 9)), + # Indexer: the index scorer at 1-8 requests, the long-row select at every key width (129..9216 blocks). + "m3_index_score": frozenset(range(1, dispatch.INDEX_SCORE_MAX_REQUESTS + 1)), + "long_select": frozenset({2, 4, 9}), + "long_select_dsl": frozenset({4, 11, 24}), + # LongAttn: the dense layers' full-context decode of 1-2 requests, 1 + 3 verify rows or one draft-pass row. + "m3_fullctx_decode": frozenset({(4, 1), (1, 1), (4, 2), (1, 2)}), + # MoE (uGPU): the die-placed target head at 4 / 8 / 12 / 16 verify rows, unless killed. + "m3_ugpu_head": frozenset() if dispatch.KILL_UGPU_HEAD else frozenset({4, 8, 12, 16}), + # GEMM/AR (+ B2): the decode MXFP8 GEMVs at every TP2 projection and decode row, on an MXFP8 and a bf16 + # activation, raised to 16 rows where the kernel serves them. + "m3_mx_gemv": frozenset((nk, m) for nk in TP2_PROJECTIONS for m in range(1, _mx_hi(nk) + 1)) + if dispatch.MX_GEMV_TP2 + else frozenset(), + "m3_mx_gemv_bf16": frozenset( + (nk, m) for nk in TP2_PROJECTIONS for m in range(1, _bf16_hi(nk) + 1) + ) + if dispatch.MX_GEMV_TP2 + else frozenset(), + # GEMM/AR: the multicast boundary all-reduce, the o_proj's push into it, the fused embedding + first RMSNorm. + # (+ B3: 9..16 rows on the 16-token bucket, unless B4_MC_AR is killed.) + "mc_allreduce_norm": frozenset(range(1, dispatch.PLAN_MC_AR_MAX_TOKENS + 1)) + if dispatch.MC_AR_TP2 + else frozenset(), + # (+ B6: 9..16 rows into the 16-token bucket, unless B4_ATTN_PUSH is killed; the o_proj GEMV stops at 16.) + "attn_push": frozenset( + range(1, min(dispatch.PLAN_MC_AR_MAX_TOKENS, dispatch.B4_MAX_TOKENS) + 1) + if dispatch.B4_ATTN_PUSH + else _DECODE_ROWS + ) + if (dispatch.MC_AR_TP2 and dispatch.ATTN_PUSH_TP2) + else frozenset(), + "dense_push": frozenset(_DECODE_ROWS) + if (dispatch.MC_AR_TP2 and dispatch.DENSE_PUSH_TP2) + else frozenset(), + # B3: the MoE's multicast finalize + all-reduce + norm at 9..16 rows (past the engine), unless killed. + "mc_finalize_allreduce_norm": frozenset( + range(dispatch.MC_FINALIZE_MIN_TOKENS, dispatch.PLAN_MC_AR_MAX_TOKENS + 1) + ) + if dispatch.B4_MC_FINALIZE + else frozenset(), + # MoE: the engine's push into the multicast boundary all-reduce, at the engine's rows. + "moe_push": frozenset(range(1, dispatch.MOE_PUSH_MAX_TOKENS + 1)), + # (+ B5: 9..16 rows, unless B4_EMBED_FIRSTNORM is killed; 17..32 keep the stock layer 0.) + "embed_firstnorm": frozenset( + range(1, dispatch.B4_MAX_TOKENS + 1) if dispatch.B4_EMBED_FIRSTNORM else _DECODE_ROWS + ) + if dispatch.EMBED_FIRSTNORM_TP2 + else frozenset(), + # SpecTail: the fused draft loop's all-reduce boundaries (1..8 rows) on the multicast all-reduce (opts 16). + "m3_draft_mc_allreduce_norm": frozenset(range(1, dispatch.DRAFT_MC_AR_MAX_ROWS + 1)), + # LongAttn: the Eagle3 drafter layer's generation passes of 1-2 requests on the same recipes. + "m3_drafter_fullctx": frozenset({(4, 1), (1, 1), (4, 2), (1, 2)}), +} + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_skeleton_no_fast_path_covers_a_shape(fast_path) -> None: + expected = ENABLED.get(fast_path.path, frozenset()) + assert fast_path.keys == expected, f"{fast_path.path} covers {sorted(fast_path.keys, key=str)}" + + +def test_the_sparse_decode_writes_the_mxfp8_handoff() -> None: + """The sparse decode (LC2, the Triton merge) writes the MXFP8 o_proj hand-off; the range keys the target's + SPARSE_DECODE_LC / SPARSE_DECODE_PRODUCER cover are the kernel module's.""" + assert tsd.MERGE_MXFP8_HANDOFF is True + assert dispatch.SPARSE_DECODE_LC.keys is tsd.LC_KEYS + assert dispatch.SPARSE_DECODE_PRODUCER.keys is tsd.PRODUCER_KEYS + + +def test_sparse_decode_lc_covers_the_decode_steps_with_the_handoff() -> None: + """Every decode step (at most DECODE_MAX_TOKENS rows) writes the MXFP8 o_proj hand-off, so the long-context + sparse decode needs only its "mxfp8" variant, on the stock log-domain recurrence and without pre-wait staging.""" + assert dispatch.SPARSE_DECODE_LC.keys == frozenset({"mxfp8"}) + assert tsd.MERGE_MXFP8_MIN_TOKENS <= 1 and dispatch.DECODE_MAX_TOKENS <= tsd.MergeMxfp8.MAX_ROWS + assert tsd.LC_PREWAIT is False + + +def test_the_dense_layers_take_the_fullctx_decode_at_the_gen_workers_shapes() -> None: + """One or two requests, 1 + 3 verify rows or one draft-pass row, each at its own pinned recipe.""" + assert dispatch.FULLCTX_DECODE_ENABLED + assert dispatch.FULLCTX_DECODE.keys == {(4, 1), (1, 1), (4, 2), (1, 2)} + assert set(dispatch.FULLCTX_DECODE_RECIPES) == set(dispatch.FULLCTX_DECODE.keys) + for recipe in dispatch.FULLCTX_DECODE_RECIPES.values(): + assert 2 <= recipe["splits"] <= 128 # the standalone reducer's range + + +def test_the_fullctx_entry_cells_run_at_the_forward_plan_capacity() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention import ( # noqa: E501 + m3_fullctx_decode as entry, + ) + + assert entry.MAX_KV_LEN == dispatch.FULLCTX_DECODE_MAX_KV_LEN + + +def test_the_fullctx_capacity_covers_the_agentx_cache_bound() -> None: + """The forward compares the cache's length bound (max_seq_len plus the speculative rows it reserves) with the plan + capacity; at the TP2 AgentX max_seq_len of 1 << 20 with up to 8 draft rows the dense layers must still be taken.""" + assert dispatch.FULLCTX_DECODE_MAX_KV_LEN >= (1 << 20) + 8 + + +def test_the_drafter_takes_the_fullctx_decode_at_the_same_shapes() -> None: + assert dispatch.DRAFTER_FULLCTX_ENABLED + assert dispatch.DRAFTER_FULLCTX.keys == dispatch.FULLCTX_DECODE.keys + + +@pytest.mark.parametrize("rows", range(1, 50)) +def test_the_tensor_core_router_takes_9_to_48_rows(rows) -> None: + """agg G-d: the tensor-core router of an unfused MoE at 9..48 rows, as the agg canon runs it; 1..8 rows stay on + the MoE engine's fused router (the batch-1/2 target's) and 49+ on the fp32 router.""" + assert dispatch.ROUTER_TC_GEN.covers(rows) == (9 <= rows <= 48) + + +# ------------------------------------------------------------------ landed: the decode MXFP8 GEMVs (GEMM/AR lane) + + +def test_the_gemv_rows_follow_the_kill_constant() -> None: + expected = set(TP2_PROJECTIONS) if dispatch.MX_GEMV_TP2 else set() + assert set(dispatch.MX_GEMV_ROWS) == expected + assert set(dispatch.MX_GEMV_BF16_MAX_TOKENS) == expected + + +@pytest.mark.parametrize("rows", range(1, dispatch.TINY_MAX_TOKENS + 1)) +@pytest.mark.parametrize("nk", TP2_PROJECTIONS, ids=str) +def test_every_projection_takes_the_gemv_at_every_decode_row(nk, rows) -> None: + """MXFP8 activation (a producer quantized it) -> the GEMV; bf16 -> the GEMV quantizing in its prologue (staged: + the stock decode-graph GEMMs while MX_GEMV_TP2 is off).""" + on = dispatch.MX_GEMV_TP2 + assert dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) == ( + dispatch.MX_GEMV_PATH if on else dispatch.PREQUANT_GEMM_PATH + ) + assert dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) == ( + dispatch.MX_GEMV_BF16_PATH if on else dispatch.GENERIC_GEMM_PATH + ) + + +@pytest.mark.parametrize("nk", TP2_PROJECTIONS, ids=str) +def test_wider_than_the_decode_step_is_the_generic_gemm(nk) -> None: + # Past the decode step and past every GEMV row it serves (W-G1 raises four projections to 48). + rows = max(dispatch.DECODE_MAX_TOKENS, _mx_hi(nk)) + 1 + assert ( + dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) + == dispatch.PREQUANT_GEMM_PATH + ) + assert ( + dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) + == dispatch.GENERIC_GEMM_PATH + ) + + +def test_an_unknown_projection_is_the_generic_gemm() -> None: + assert ( + dispatch.linear_path(256, 512, 1, prequant=True, low_m_gemm=True) + == dispatch.PREQUANT_GEMM_PATH + ) + assert ( + dispatch.linear_path(256, 512, 1, prequant=False, low_m_gemm=True) + == dispatch.GENERIC_GEMM_PATH + ) + + +def test_taken_logs_once_per_path_and_rows(monkeypatch) -> None: + lines = [] + monkeypatch.setattr(dispatch.logger, "info", lines.append) + for _ in range(3): + dispatch.taken(dispatch.MX_GEMV_PATH, 4, "[6144, 4096] mxfp8 in") + dispatch.taken(dispatch.MX_GEMV_PATH, 8, "[6144, 4096] mxfp8 in") + assert len(lines) == 2 and all("m3_mx_gemv taken at" in line for line in lines) + + +def test_skeleton_the_dense_decode_split_is_the_built_in_one() -> None: + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {} + + +def test_skeleton_speculation_runs_the_in_tree_loop() -> None: + # SpecTail: fused greedy / sampled verify and the fused draft loop up to the target's steps (B4 spec tail; the + # batch-1/2 rows with B4_SPEC_TAIL off; BF16 head; no coded heads). + reqs = dispatch.SPEC_TAIL_MAX_TOKENS // (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert dispatch.SPEC_TAIL_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if dispatch.B4_SPEC_TAIL else dispatch.TINY_MAX_TOKENS + ) + # W2: the fused verify alone widens to WIDE_MAX_TOKENS rows (WIDE_SPEC_TAIL off keeps SPEC_TAIL_MAX_TOKENS). + verify = ( + dispatch.WIDE_MAX_TOKENS if dispatch.WIDE_SPEC_TAIL else dispatch.SPEC_TAIL_MAX_TOKENS + ) // (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert dispatch.WIDE_VERIFY_MAX_REQUESTS == verify + assert dispatch.SHARD_ACCEPT_MAX_REQUESTS == verify + # The sampled verify stays at SPEC_TAIL_MAX_TOKENS rows (past it, the in-tree sampler). + assert dispatch.SAMPLED_VERIFY_MAX_REQUESTS == reqs + assert dispatch.DRAFT_STEP_MAX_REQUESTS == min(reqs, 8) + assert dispatch.FC_EARLY_MAX_ROWS == 0 + assert not draft_step_kernel.HEAD_CODED + + +def test_a_miss_raises_only_under_require() -> None: + dispatch.missed("some_path", 1, "a precondition") + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="some_path"): + dispatch.missed("some_path", 1, "a precondition") + + +def test_skeleton_first_drafter_inputs_are_in_tree() -> None: + # SpecTail turned the sampled verify on (SAMPLED_VERIFY_MAX_REQUESTS 2, test_skeleton_speculation_...). + assert dispatch.FIRST_DRAFTER_INPUTS_MAX_REQUESTS == 0 + + +@pytest.mark.parametrize( + "rows, capture_normed, taken", + [ + (4, False, True), + (8, False, True), + (8, True, True), + # 3 / 4 requests: the B4 spec tail's rows (in-tree with B4_SPEC_TAIL off). + (12, False, dispatch.B4_SPEC_TAIL), + (12, True, dispatch.B4_SPEC_TAIL), + (16, True, dispatch.B4_SPEC_TAIL), + ], +) +def test_the_fused_first_inputs_take_only_their_covered_rows(rows, capture_normed, taken) -> None: + """The fused first drafter inputs (and the fc fork) take a fused-verify step only at the requests + FUSED_FIRST_INPUTS covers and, unless the captures applied the fc_norm, the rows CHUNK_NORM covers; anything else + runs the in-tree first drafter inputs.""" + + worker = object.__new__(dispatch.M3Eagle3Worker) + worker._m3_capture_norm = (None, 1e-6) if capture_normed else None + meta = SimpleNamespace(runtime_draft_len=dispatch.SHARD_ACCEPT_DRAFT_LEN) + assert worker._m3_first_inputs_covered(rows, meta) is taken + + +# ------------------------------------------------------- landed: the multicast boundary all-reduce (GEMM/AR lane) + + +def test_both_boundaries_take_the_multicast_all_reduce() -> None: + """One 8-token bucket for the whole decode step; the MoE boundary publishes the next QKV's MXFP8 rows (killed: + every boundary on the built-in all-reduce while MC_AR_TP2 is off).""" + on = dispatch.MC_AR_TP2 + tiny = dispatch.DECODE_PLANS[0] + assert (tiny.name, tiny.max_tokens) == ("tiny", dispatch.TINY_MAX_TOKENS) + assert (tiny.attn_mc_ar, tiny.moe_mc_ar, tiny.qkv_prequant) == (on, on, on), tiny + # B6: the o_proj push follows the boundary's rows up to the o_proj GEMV's 16, unless B4_ATTN_PUSH is killed. + push_rows = ( + min(dispatch.PLAN_MC_AR_MAX_TOKENS, dispatch.B4_MAX_TOKENS) + if dispatch.B4_ATTN_PUSH + else dispatch.TINY_MAX_TOKENS + ) + assert dispatch.ATTN_PUSH_MAX_TOKENS == (push_rows if on and dispatch.ATTN_PUSH_TP2 else 0) + assert not dispatch.B4_ATTN_PUSH or dispatch.B4_MC_AR + assert dispatch.ATTN_PUSH_MAX_TOKENS <= max( + r for (nk, r) in dispatch.MX_GEMV.keys if nk == dispatch.O_PROJ + ) + assert dispatch.DENSE_PUSH_MAX_TOKENS == ( + dispatch.TINY_MAX_TOKENS if on and dispatch.DENSE_PUSH_TP2 else 0 + ) + for rows in _DECODE_ROWS: + assert dispatch.plan_for(rows) is dispatch.DECODE_PLANS[0] + # The MoE pushes only from its engine, at the engine's rows. + assert dispatch.MOE_PUSH_MAX_TOKENS == (dispatch.MOE_ENGINE_MAX_TOKENS if on else 0) + + +def test_the_generic_boundaries_take_the_multicast_all_reduce() -> None: + """Prefill, mixed steps and the Eagle3 capture layers: the multicast all-reduce's 48-token bucket, alone and folded, + and the MoE finalize folded into it, as the agg canon runs them.""" + rows = 48 if dispatch.MC_AR_TP2 else 0 + assert dispatch.GENERIC_MC_AR_MAX_TOKENS == rows + assert dispatch.GENERIC_MC_AR_FOLDED_MAX_TOKENS == rows + assert dispatch.GENERIC_MC_FINALIZE_MAX_TOKENS == rows + + +def test_the_tensor_core_router_reaches_48_rows() -> None: + assert dispatch.ROUTER_TC_GEN_MAX_ROWS == 48 + + +def test_the_sampled_verify_owner_walk_is_one_token_by_default() -> None: + assert dispatch.m3_sampled_accept_kernel.ILP_WALK is False + + +def test_the_fused_embedding_rows_follow_the_kill_constant() -> None: + # B5: 9..16 rows too (not 17..32), unless B4_EMBED_FIRSTNORM is killed. + rows = dispatch.B4_MAX_TOKENS if dispatch.B4_EMBED_FIRSTNORM else dispatch.TINY_MAX_TOKENS + assert dispatch.EMBED_FIRSTNORM_MAX_TOKENS == (rows if dispatch.EMBED_FIRSTNORM_TP2 else 0) + + +# ------------------------------------------------------- landed: MX-TAC (pinned decode-graph MXFP8 GEMM tactic) + + +def test_the_decode_mxfp8_decision_is_pinned_to_cute_dsl() -> None: + from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, + ) + + assert dispatch.MXFP8_DECODE_QUANT_TACTIC == MXFP8QuantizeRunner.CUTE_DSL + assert dispatch.MXFP8_DECODE_GEMM_TACTIC == FlashInferMXFP8GemmRunner.CUTE_DSL + + +@pytest.mark.parametrize("target", ["DecodeTarget", "PrefillTarget"]) +def test_the_moe_forward_reports_the_engine_taken(target) -> None: + """The engagement evidence comes from the forward (not the engine build's warm call).""" + + src = inspect.getsource(getattr(dispatch, target)._moe) + assert src.index("taken(MOE_ENGINE.path") < src.index("m3_moe_engine(h.contiguous()") + + +# ------------------------------------------------------------------ the drafter's low-M linears (TP8 port) + + +def test_draft_lowm_covers_exactly_the_certified_rows_of_its_tactic_table() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm import ( # noqa: E501 + m3_draft_lowm as entry, + ) + + assert set(entry.TACTICS) == set(entry.ROWS) == set(dispatch.DRAFT_LOWM_ROWS) + assert dispatch.DRAFT_LOWM.keys == frozenset( + (nk, m) + for nk, rows in entry.ROWS.items() + for m in rows + if m <= dispatch.DRAFT_LOWM_MAX_ROWS + ) + for nk, table in entry.TACTICS.items(): + rows = [row[0] for row in table] + assert rows == sorted(rows) and len(set(rows)) == len(rows), nk + for m in entry.ROWS[nk]: + assert entry.tactic_for(*nk, m) is not None, (nk, m) + + +def test_prebuild_draft_head_is_idempotent_and_keeps_replaced_copies(monkeypatch) -> None: + """A repeated cache_derived_state with an unchanged head weight keeps the built (placed) copy; a changed weight + rebuilds it and keeps the old copy alive, so a CUDA graph captured earlier never reads unmapped memory.""" + + import torch + + built = [] + + def fake_prepare(lm_head, tp_rank): + head = SimpleNamespace(weight=lm_head.weight) + built.append(head) + return head + + monkeypatch.setattr(draft_step_kernel, "prepare_head", fake_prepare) + worker = object.__new__(dispatch.M3Eagle3Worker) + worker._m3_draft_head = None + worker._m3_draft_head_version = None + worker._m3_draft_head_graveyard = [] + worker.mapping = SimpleNamespace(tp_rank=0) + monkeypatch.setattr(worker, "_m3_drafter_supported", lambda dm: True, raising=False) + model = SimpleNamespace(lm_head=SimpleNamespace(weight=torch.zeros(4, 4))) + + assert worker.prebuild_draft_head(model) and len(built) == 1 + first = worker._m3_draft_head + assert worker.prebuild_draft_head(model) and len(built) == 1 and worker._m3_draft_head is first + model.lm_head.weight.add_(1.0) # new values: version bump + assert worker.prebuild_draft_head(model) and len(built) == 2 + assert worker._m3_draft_head is built[1] and worker._m3_draft_head_graveyard == [first] + + +# ----------------------------------------------------------------------- batch 4 (this target): the b16 plan + + +def test_the_batch_8_steps_take_the_b16_and_b32_plans() -> None: + """9..16-row steps run the "b16" plan (16-token bucket) and 17..32-row steps the "b32" plan (48-token bucket), both + on the multicast all-reduce (B4_MC_AR) with no MXFP8 copy for the next QKV; 1..8 rows keep the "tiny" plan.""" + if not dispatch.B4_DECODE: + assert len(dispatch.DECODE_PLANS) == 1 + return + names = [p.name for p in dispatch.DECODE_PLANS] + assert names == (["tiny", "b16", "b32"] if dispatch.B8_DECODE else ["tiny", "b16"]) + b16 = dispatch.DECODE_PLANS[1] + assert b16.max_tokens == dispatch.B4_MAX_TOKENS == 16 + assert dispatch.DECODE_MAX_TOKENS == (32 if dispatch.B8_DECODE else 16) + b3 = dispatch.B4_MC_AR + for plan in dispatch.DECODE_PLANS[1:]: + quant = dispatch.B8_QKV_PREQUANT if plan.name == "b32" else dispatch.B4_QKV_PREQUANT + assert (plan.attn_mc_ar, plan.moe_mc_ar, plan.qkv_prequant) == (b3, b3, quant) + assert dispatch.PLAN_MC_AR_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if b3 else dispatch.TINY_MAX_TOKENS + ) + for rows in range(1, dispatch.DECODE_MAX_TOKENS + 1): + want = "tiny" if rows <= dispatch.TINY_MAX_TOKENS else "b16" if rows <= 16 else "b32" + assert dispatch.plan_for(rows).name == want, rows + + +def test_the_batch_4_gemv_rows() -> None: + """B2: the decode GEMVs serve 1..16 rows (kill MX_GEMV_B4: 1..8), the bf16 one except dense gate_up / down.""" + if not dispatch.MX_GEMV_TP2: + return + for nk, rows in dispatch.MX_GEMV_ROWS.items(): + assert rows == frozenset(range(1, _mx_hi(nk) + 1)), nk + assert dispatch.MX_GEMV_BF16_MAX_TOKENS[nk] == _bf16_hi(nk), nk + + +@pytest.mark.parametrize("nk", sorted(_BF16_WIDE_DECLINES), ids=str) +def test_the_bf16_gemv_declines_take_the_generic_gemm_at_batch_4(nk) -> None: + """Dense gate_up / down on a bf16 activation: the GEMV at 1..8 rows, the generic GEMM at 9..16 (the MXFP8 GEMV + still serves them on a prequantized activation).""" + if not dispatch.MX_GEMV_TP2: + return + for rows in range(1, dispatch.DECODE_MAX_TOKENS + 1): + bf16 = dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=False) + mx = dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=False) + assert bf16 == ( + dispatch.MX_GEMV_BF16_PATH + if rows <= dispatch.TINY_MAX_TOKENS + else dispatch.GENERIC_GEMM_PATH + ), rows + assert mx == ( + dispatch.MX_GEMV_PATH if rows <= _mx_hi(nk) else dispatch.PREQUANT_GEMM_PATH + ), rows + + +def test_no_other_fast_path_is_raised_past_the_batch_1_2_rows() -> None: + """Only the paths this target raises for 16 rows (the GEMVs, B3's multicast all-reduce) cover more than + TINY_MAX_TOKENS rows; every other row-keyed path declines a 9..16-row step to stock inside the decode target.""" + raised = { + dispatch.MX_GEMV.path, + dispatch.MX_GEMV_BF16.path, + dispatch.MC_ALLREDUCE.path, + dispatch.MC_FINALIZE.path, + dispatch.ROUTER_TC_GEN.path, # agg G-d: 9..48 rows + } + # SpecTail TR-5b: B5's fused embedding and B6's o_proj push past 8 rows, unless killed. + if dispatch.B4_EMBED_FIRSTNORM: + raised.add(dispatch.EMBED_FIRSTNORM.path) + if dispatch.B4_ATTN_PUSH: + raised.add(dispatch.ATTN_PUSH.path) + if dispatch.B4_SPEC_TAIL: + # The B4 spec tail: the fused verify's and draft loop's row-keyed paths up to SPEC_TAIL_MAX_TOKENS. + raised |= { + dispatch.DRAFT_INPUT.path, + dispatch.CAPTURE_NORM.path, + dispatch.CHUNK_NORM.path, + } + # Keys that are not row counts (both selects' key widths; the index scorer's request counts, W-IDX1), or rows + # already certified past 8 in the batch-1/2 target (the die-split head's 12 / 16-row cells, reached only with the + # fused verify). + not_rows = { + dispatch.LONG_SELECT.path, + dispatch.LONG_SELECT_DSL.path, + dispatch.UGPU_HEAD.path, + dispatch.INDEX_SCORE.path, + } + for path in dispatch.FAST_PATHS: + if path.path in raised or path.path in not_rows: + continue + for key in path.keys: + if isinstance(key, int): + assert key <= dispatch.TINY_MAX_TOKENS, (path.path, key) + + +def test_the_batch_4_moe_finalizes_on_the_multicast_kernel() -> None: + """1..8 rows keep the MoE engine; 9..16 rows (past it) take the multicast finalize + all-reduce + norm on the + 16-token bucket, or the generic MoE when B4_MC_FINALIZE is killed; nothing past the decode cap.""" + for rows in range(1, dispatch.TINY_MAX_TOKENS + 1): + assert dispatch.moe_path(rows) == ( + dispatch.MOE_ENGINE.path if not dispatch.KILL_MOE_ENGINE else None + ), rows + for rows in range(dispatch.TINY_MAX_TOKENS + 1, dispatch.DECODE_MAX_TOKENS + 1): + expected = dispatch.MC_FINALIZE.path if dispatch.B4_MC_FINALIZE else None + assert dispatch.moe_path(rows) == expected, rows + assert dispatch.moe_path(dispatch.DECODE_MAX_TOKENS + 1) is None + + +# ------------------------------------------------------------------ W-G1: the MXFP8 GEMV at 33..48 rows (agg wide +# decode) + +_W48 = (dispatch.QKV_SPARSE, dispatch.QKV, dispatch.SHARED_GATE_UP, dispatch.SHARED_DOWN) + + +def test_w48_rows_follow_the_kill_constant() -> None: + """WIDE_MX_GEMV_48 raises the four projections it wins on to their served rows past 32 (at most 48); o_proj and the + dense MLP keep theirs.""" + on = ( + dispatch.MX_GEMV_TP2 + and dispatch.WIDE_MX_GEMV_48 + and dispatch.GEMV_MAX_TOKENS == dispatch.DECODE_MAX_TOKENS + ) + for nk in _W48: + rows = dispatch._GEMV_W48_ROWS.get(nk, frozenset()) + if not on: + assert not rows, nk + continue + assert rows and min(rows) == dispatch.GEMV_MAX_TOKENS + 1, nk + assert max(rows) <= dispatch.WIDE_GEMV_MAX_TOKENS, nk + assert rows == frozenset(range(min(rows), max(rows) + 1)), nk + assert rows <= dispatch.MX_GEMV_ROWS[nk], nk + for nk in (dispatch.O_PROJ, dispatch.DENSE_GATE_UP, dispatch.DENSE_DOWN): + assert nk not in dispatch._GEMV_W48_ROWS and nk not in dispatch.MX_GEMV_QUANT_ROWS, nk + + +@pytest.mark.parametrize("nk", _W48, ids=str) +def test_w48_bf16_activation_takes_quantize_plus_gemv(nk) -> None: + """A bf16 activation on rows the bf16 GEMV does not serve but the MXFP8 GEMV does: linear_path returns the GEMV + (M3MXFP8LinearMethod.apply quantizes with the decode-graph quantize first); rows the bf16 GEMV serves keep it; rows + past the MXFP8 GEMV's reach keep the generic GEMM.""" + if not (dispatch.MX_GEMV_TP2 and dispatch.WIDE_MX_GEMV_48): + return + bf16_hi = dispatch.MX_GEMV_BF16_MAX_TOKENS.get(nk, 0) + for rows in range(1, 65): + path = dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=False) + if rows <= bf16_hi: + assert path == dispatch.MX_GEMV_BF16_PATH, rows + elif rows in dispatch.MX_GEMV_ROWS[nk]: + assert rows in dispatch.MX_GEMV_QUANT_ROWS[nk], rows + assert path == dispatch.MX_GEMV_PATH, rows + else: + assert path == dispatch.GENERIC_GEMM_PATH, rows + if rows in dispatch.MX_GEMV_ROWS[nk]: + assert ( + dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=False) + == dispatch.MX_GEMV_PATH + ), rows + + +def test_w48_rows_are_certified_by_the_gemv_cells() -> None: + """Every (projection, rows) the dispatch can send to the MXFP8 GEMV is one of m3_mx_gemv's CELLS (served_cells to + 48).""" + cells = {(c.spec["nk"], c.spec["m"]) for c in dispatch.m3_mx_gemv.CELLS} + for nk, rows in dispatch.MX_GEMV_ROWS.items(): + for m in rows: + assert (nk, m) in cells, (nk, m) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py new file mode 100644 index 000000000000..a7dba3ce84f3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_draft_attention_bcg_route.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Eagle3 drafter layer's attention under a breakable prefill CUDA graph (``M3Eagle3Draft.attention``): it runs +in the eager draft attention boundary on the step's metadata, read at the call (as the in-tree drafter's attention +boundary reads it), and writes the live rows of the padded output; eagerly it calls the op as before. A metadata that +is not the step's is refused. CPU only: the attention op, the linears and the reduction are stand-ins.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) + +pytestmark = pytest.mark.cpu_only + +Q_SIZE = 8 + + +@pytest.fixture +def draft(monkeypatch): + """(drafter, calls of the attention op stand-in, the step's metadata).""" + calls = [] + step_md = SimpleNamespace(num_tokens=3) + + def op(attn, qkv, md): + calls.append((attn, qkv, md)) + return torch.full((int(md.num_tokens), Q_SIZE), 7.0, dtype=torch.bfloat16) + + monkeypatch.setattr(modeling, "m3_draft_attention", op) + monkeypatch.setattr(modeling, "_record", lambda *a: None) + monkeypatch.setattr( + modeling, "get_model_extra_attrs", lambda: {"attention_metadata": lambda: step_md} + ) + d = object.__new__(modeling.M3Eagle3Draft) + d.__dict__.update(attn="attn", q_size=Q_SIZE, _bcg_key="draft-under-test") + d.linear = lambda x, name: x + monkeypatch.setitem(modeling._BCG_DRAFTS, d._bcg_key, d) + return d, calls, step_md + + +def test_breakable_graph_runs_the_eager_boundary_on_the_steps_metadata(draft, monkeypatch) -> None: + d, calls, step_md = draft + monkeypatch.setattr(modeling, "is_in_breakable_cuda_graph", lambda: True) + x = torch.zeros(5, 12, dtype=torch.bfloat16) # 3 live rows + 2 pad rows of the captured bucket + out = modeling.M3Eagle3Draft.attention(d, x, step_md, False) + assert len(calls) == 1 and calls[0][0] == "attn" and calls[0][2] is step_md + assert out.shape == (5, Q_SIZE) + assert torch.equal(out[:3], torch.full((3, Q_SIZE), 7.0, dtype=torch.bfloat16)) + + +def test_breakable_graph_refuses_a_metadata_that_is_not_the_steps(draft, monkeypatch) -> None: + d, calls, _ = draft + monkeypatch.setattr(modeling, "is_in_breakable_cuda_graph", lambda: True) + with pytest.raises(RuntimeError, match="metadata is not the step's"): + modeling.M3Eagle3Draft.attention( + d, torch.zeros(3, 12), SimpleNamespace(num_tokens=3), False + ) + assert not calls + + +def test_eager_calls_the_op_directly(draft, monkeypatch) -> None: + d, calls, _ = draft + monkeypatch.setattr(modeling, "is_in_breakable_cuda_graph", lambda: False) + md = SimpleNamespace(num_tokens=2) + out = modeling.M3Eagle3Draft.attention(d, torch.zeros(2, 12), md, False) + assert len(calls) == 1 and calls[0][2] is md and out.shape == (2, Q_SIZE) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py new file mode 100644 index 000000000000..7da5d1206a57 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter's full-context decode (M3DrafterFullCtxFmha) is selected through FmhaManager, which caches a library's +support per cell of the generation request count (a ceiling bucket of _FMHA_CACHE_BATCH_SIZE_GRID, e.g. 41..47 -> 47). +The runner supports exact (query rows, requests) recipes, so every recipe's request count must be alone in its cell: +otherwise the first count seen in a cell decides for the others (a 46-request recipe would share 47's cell, and an +eager 41-request pass would reach the runner without a plan). CPU only.""" + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling, +) +from tensorrt_llm._torch.attention.backends.fmha import manager + +pytestmark = pytest.mark.cpu_only + + +def _cell(value: int, grid) -> int: + return manager._normalize_fmha_cache_grid_value(value, grid) + + +def test_every_drafter_recipe_is_alone_in_its_cache_cells(): + keys = set(modeling.FULLCTX_DECODE_RECIPES) + batch, qlen = manager._FMHA_CACHE_BATCH_SIZE_GRID, manager._FMHA_CACHE_SEQ_LEN_Q_GRID + for q, n in keys: + same_batch = {m for m in range(1, 4 * batch[-1]) if _cell(m, batch) == _cell(n, batch)} + same_qlen = {r for r in range(1, 4 * qlen[-1]) if _cell(r, qlen) == _cell(q, qlen)} + assert same_batch == {n}, ( + f"({q}, {n}): requests {sorted(same_batch)} share its FmhaManager cell" + ) + assert same_qlen == {q}, ( + f"({q}, {n}): query rows {sorted(same_qlen)} share its FmhaManager cell" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..37b04013094a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode robustness without a GPU kernel: the multicast all-reduce epilogue options keep the cluster-barrier bit, +an engine build that fails on one TP rank drops that engine on every rank (two ranks simulated with threads), and +the Eagle3 fc fork's output is recorded on the stream that reads it.""" + +import threading +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = modeling.require_mode() + yield + modeling.set_require(require) + modeling.reset_logged() + + +@pytest.mark.parametrize("finalize", [False, True]) +@pytest.mark.parametrize("bucket", [8, 16, 32, 48]) +def test_ar_opts_keep_cluster_barrier_bit(monkeypatch, bucket, finalize): + # Tables without bit 2 and with buckets missing: every lookup, fallback included, keeps the barrier. + monkeypatch.setattr(mc, "AR_OPTS_BY_M", (0, 1, 4, 5, 8, 9, 12, 13)) + monkeypatch.setattr(mc, "AR_OPTS_BY_BUCKET", {}) + monkeypatch.setattr(mc, "AR_FIN_OPTS_BY_BUCKET", {}) + for m in range(1, min(bucket, 8) + 1): + opts = mc._ar_opts(m, bucket, finalize=finalize) + assert opts & mc._OPT_CLUSTER_BARRIER and 0 <= opts <= 15, (m, bucket, finalize, opts) + + +def test_ar_opts_shipped_tables_unchanged(): + for m in range(1, 9): + assert mc._ar_opts(m, 8) == mc.AR_OPTS_BY_M[m - 1] + for b, o in mc.AR_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b) == o + for b, o in mc.AR_FIN_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b, finalize=True) == o + + +class _FakeComm: + """An MPI-like TP communicator for ``world`` threads: allreduce sums the ranks' ints.""" + + def __init__(self, world): + self.world = world + self.barrier = threading.Barrier(world) + self.lock = threading.Lock() + self.vals = [] + self.calls = 0 + + def allreduce(self, v): + with self.lock: + self.vals.append(v) + self.calls += 1 + self.barrier.wait() + total = sum(self.vals) + self.barrier.wait() + with self.lock: + self.vals.clear() + self.barrier.wait() + return total + + def Get_size(self): # noqa: N802 (MPI name) + return self.world + + +def _mapping(tp_size=2, adp=False): + return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) + + +def _run_ranks(fns): + out, errs = [None] * len(fns), [] + + def run(i): + try: + out[i] = fns[i]() + except BaseException as e: # noqa: BLE001 + errs.append(e) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errs, errs + return out + + +@pytest.fixture +def two_rank_comm(monkeypatch): + from tensorrt_llm._torch.distributed import ops + + comm = _FakeComm(2) + monkeypatch.setattr(mc, "_tp_comm", lambda mapping: comm) + monkeypatch.setattr(ops, "mpi_disabled", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return comm + + +def test_tp_agree_single_rank_and_adp_take_no_collective(monkeypatch): + monkeypatch.setattr(mc, "_tp_comm", mock.Mock(side_effect=AssertionError("no collective"))) + assert mc.tp_agree(_mapping(1), True) is True + assert mc.tp_agree(_mapping(1), False) is False + assert mc.tp_agree(_mapping(4, adp=True), False) is False + + +def test_tp_agree_refuses_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError): + mc.tp_agree(_mapping(2), True) + + +@pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) +def test_tp_agree_built_two_ranks(two_rank_comm, oks): + built = [object() if ok else None for ok in oks] + out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + assert out == (built if all(oks) else [None, None]) + + +def _core(fail: bool) -> SimpleNamespace: + """The engine-building half of the core (``MiniMaxM3Tp2ModelingV2``) on one of two simulated TP ranks, its + engine builds stubbed (and failing on one rank).""" + core = SimpleNamespace(mapping=_mapping(2), moe_engines={}) + cls = modeling.MiniMaxM3Tp2ModelingV2 + + def build(i, num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + core._build_engine = build + for name in ("_prebuild_engines", "_try_build", "engine", "_lookup"): + setattr(core, name, getattr(cls, name).__get__(core)) + return core + + +@pytest.mark.parametrize("fail", [(False, False), (False, True)]) +def test_engine_build_agreed_across_ranks(two_rank_comm, fail): + cores = [_core(f) for f in fail] + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + keys = sorted(cores[0].moe_engines[3]) + assert keys and two_rank_comm.calls == 2 * len(keys) + a, b = cores[0].moe_engines[3], cores[1].moe_engines[3] + assert a.keys() == b.keys() + for key in a: + assert (a[key] is None) == (b[key] is None) == any(fail), key + # Both ranks then take the same path in the forward: the engine, or (after a miss) the generic MoE. + modeling.set_require(False) + for n in sorted(modeling.MOE_ENGINE.keys): + got = [c.engine(3, n, shared=True) for c in cores] + assert (got[0] is None) == (got[1] is None) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + M3Eagle3Worker = modeling.M3Eagle3Worker + main, fork = mock.Mock(name="main"), mock.Mock(name="fork") + hidden = mock.Mock(name="hidden", is_cuda=True) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *a: main) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + worker = SimpleNamespace(_m3_fc_hidden=hidden) + M3Eagle3Worker._m3_fc_join(worker, fork) + main.wait_stream.assert_called_once_with(fork) + if capturing: + hidden.record_stream.assert_not_called() + else: + hidden.record_stream.assert_called_once_with(main) + worker._m3_fc_hidden = None + M3Eagle3Worker._m3_fc_join(worker, fork) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py new file mode 100644 index 000000000000..828bdba6d48a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_wide_lc_cap.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The wide target's MXFP8 o_proj hand-off, and so its long-context sparse decode, stops at one 128-row scale tile. + +LC2 is certified as the MXFP8 hand-off step, so the hand-off's row cap is LC2's row cap. A wide decode step past it +(33..48 requests x 4 verify rows) takes the Triton producer + merge and the o_proj quantizes its BF16 input itself +(attention/test_modeling_v2_m3_sparse_decode.py: test_plain_route_is_bitwise_the_stock_decode). +""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( + triton_sparse_decode as tsd, +) + + +def test_the_hand_off_cap_is_one_scale_tile(): + """The MXFP8 hand-off (and so the long-context sparse decode) serves up to one 128-row scale tile; the target's + 129..192-row steps take the Triton producer + merge and the o_proj quantizes its BF16 input itself.""" + assert tsd.MergeMxfp8.MAX_ROWS == 128 + assert dispatch.WIDE_MAX_TOKENS > tsd.MergeMxfp8.MAX_ROWS + + +def test_lc2_certifies_only_the_hand_off_step(): + assert dispatch.SPARSE_DECODE_LC.covers("mxfp8") + assert not dispatch.SPARSE_DECODE_LC.covers("plain") + + +@pytest.mark.parametrize("rows", [1, 4, 127, 128]) +def test_the_hand_off_is_the_steps_own_scale_tiles(rows): + """The published scales are the stock MXFP8 quantization's size for the step (the o_proj GEMM checks the swizzled + scales against the step's padded rows): one 128-row tile for every step the hand-off serves.""" + width = 4096 + buffers = tsd.MERGE_MXFP8.buffers(width, torch.device("cuda")) + q8, sf = tsd.MergeMxfp8.handoff(buffers, rows) + _, stock_sf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(rows, width, dtype=torch.bfloat16, device="cuda"), True + ) + assert q8.shape == (rows, width) + assert sf.numel() == stock_sf.numel() + assert sf.data_ptr() == buffers[1].data_ptr() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..b58bcfacd648 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_silu_and_mul entry of the MiniMax-M3 tp4 target: SwiGLU-OAI and plain SwiGLU.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.activation.m3_silu_and_mul import ( # noqa: E501 + m3_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + # Gate and up spread past the clamp (|x| up to ~12) so both clamps engage. + x = (torch.randn(spec["m"], 2 * spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + return (x, *OAI) if spec["oai"] else (x, None, None, None) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few + fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); + measured on GB300: worst element 0.5 unit, relative RMS 0.01.""" + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(*args) + assert out.shape == (cell.spec["m"], cell.spec["inter"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=1.0, rms_ulp=0.1) + + +def test_the_gate_sees_a_missing_clamp() -> None: + """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" + spec = op.CELLS[0].spec + x, *oai = _build(dict(spec, m=64), 3) + out = op(x, *oai) + with pytest.raises(AssertionError): + assert_within_ulp(out, op.reference(x, None, oai[1], oai[2]), element_ulp=1.0, rms_ulp=0.1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_attn_runtime_tp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_attn_runtime_tp4.py new file mode 100644 index 000000000000..7bfd8de384f9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_attn_runtime_tp4.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real MiniMax-M3 attention runtime state at the TP4 per-rank geometry, for the backend-step entries' GPU tests. + +One hybrid NVFP4 KV cache manager (``MiniMaxM3KVCacheManagerV2``) laid out as the target's: a dense layer in the FP8 +half, a sparse MSA layer in the NVFP4 half with its FP8 index-K cache, and (``draft=True``) the one-model Eagle3 +drafter's shared layer, read through the manager's FP8 draft view as the runtime's draft swap reads it. Metadata is +the runtime's own type, filled by its own ``prepare()``; the backends are built the way the target builds them. + +The geometry is one rank's: 16 query heads over one KV head and one index head of dim 128. The mapping is a single +rank holding exactly those heads, so nothing is sharded again. + +History rows are written by the in-tree scatters (``fused_write_layer_caches`` for the FP8 layers, +``fused_write_layer_caches_nvfp4`` for the NVFP4 layer and its index-K), at the slots the manager assigned, which are +not in request order. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + MiniMaxM3MsaSparseAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, + fused_write_layer_caches_nvfp4, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HEADS = 16 +KV_HEADS = 1 +INDEX_HEADS = 1 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +DENSE, SPARSE, DRAFT = 0, 1, 2 +MAX_BATCH = 8 +MAX_SEQ_LEN = 8192 +MAX_NUM_TOKENS = 8192 +#: The Eagle3 drafter checkpoint's RoPE: NeoX over the whole head, theta 5e6. +DRAFT_ROPE_THETA = 5.0e6 + +#: E4M3-exact planted index levels, 1/8 apart in relative terms (0.25 .. 15): 48 distinct block scores. +LEVELS = tuple(2.0**e * (1 + m / 8) for e in range(-2, 4) for m in range(8)) +#: The index value of a block's non-peak rows; its score (8) is below every level's (>= 16). +LOW = 0.125 + + +def mapping() -> Mapping: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + +def sparse_attention_config() -> MiniMaxM3SparseAttentionConfig: + """The target's runtime: MSA, the FP8 indexer, the fused five-way projection (M3's checkpoint selection rule: + no forced initial blocks, the query's own block forced).""" + return MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_num_index_heads=INDEX_HEADS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + ) + + +def make_manager(*, draft: bool) -> MiniMaxM3KVCacheManagerV2: + torch.cuda.init() + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_tokens=4 * MAX_SEQ_LEN, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + event_buffer_max_size=0, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=[KV_HEADS] * (3 if draft else 2), + head_dim=HEAD_DIM, + tokens_per_block=PAGE, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + mapping=mapping(), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") if draft else None + ), + sparse_layer_ids=[SPARSE], + sparse_attention_config=sparse_attention_config(), + ) + + +def make_msa_backend(layer: int) -> MiniMaxM3MsaSparseAttention: + """One layer's MSA backend as ``MiniMaxM3Tp4ModelingV2`` builds it (NVFP4 KV quant config on every layer).""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + backend = create_attention( + "TRTLLM", + layer, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=None, + quant_config=quant, + q_scaling=1.0, + sparse_params=sparse_attention_config().to_sparse_params(layer_idx=layer), + ) + backend.update_quant_config(quant) + assert isinstance(backend, MiniMaxM3MsaSparseAttention), type(backend).__name__ + return backend + + +def make_draft_backend() -> TrtllmAttention: + """The drafter layer's attention as ``M3Eagle3Draft`` builds it: fused NeoX RoPE, an FP8 KV cache (the target + overrides the drafter's KV dtype to FP8), no other quantization.""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + backend = create_attention( + "TRTLLM", + DRAFT, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=HEAD_DIM, theta=DRAFT_ROPE_THETA, max_positions=MAX_SEQ_LEN), + ), + quant_config=quant, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + backend.update_quant_config(quant) + assert type(backend) is TrtllmAttention, type(backend).__name__ + return backend + + +def make_msa_metadata(manager) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = MiniMaxM3MsaSparseAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + max_num_sequences=MAX_BATCH, + kv_cache_manager=manager, + mapping=mapping(), + sparse_metadata_params=sparse_attention_config().to_sparse_metadata_params(), + ) + # The target's MSA preparation choice that changes what prepare() stages (the identity tables of short steps). + metadata.msa_policy = MsaRuntimePolicy(short_sequence_identity=True) + return metadata + + +def make_draft_metadata(manager) -> TrtllmAttentionMetadata: + """Drafter metadata over the shared manager's FP8 draft view, the manager the runtime's draft swap installs.""" + view = manager.get_draft_kv_cache_view() + assert view is not None and view.dtype == DataType.FP8 + return TrtllmAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=mapping(), + ) + + +def allocate(manager, request_ids: Sequence[int], qlens: Sequence[int], cached: Sequence[int]): + manager.add_dummy_requests(list(request_ids), token_nums=[c + n for c, n in zip(cached, qlens)]) + + +def prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + c + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +def flat_slots(manager, layer: int, request_id: int, positions: torch.Tensor) -> torch.Tensor: + """The cache slots (page slot id * PAGE + offset) of a request's positions in ``layer``.""" + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0], + dtype=torch.int64, + device="cuda", + ) + positions = positions.to(device="cuda", dtype=torch.int64) + return (pages[positions // PAGE] * PAGE + positions % PAGE).to(torch.int32) + + +def write_fp8_rows(manager, layer, request_id, positions, k, v) -> None: + """K / V rows [n, 128] (E4M3) into an FP8 layer's cache at ``positions``.""" + view = manager.get_buffers(layer, "HND") + assert view.dtype == torch.float8_e4m3fn + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches(view[:, 0], view[:, 1], None, slots, k, v, None) + + +def write_nvfp4_rows(manager, metadata, layer, request_id, positions, k, v, idx_k, quant) -> None: + """BF16 K / V rows quantized into the NVFP4 layer's cache, and index-K rows into its FP8 index cache.""" + data = manager.get_buffers(layer, "HND") + scales = manager.get_block_scale_buffers(layer, "HND") + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches_nvfp4( + data[:, 0], + data[:, 1], + scales[:, 0], + scales[:, 1], + metadata.msa_idx_k_cache(layer), + slots, + k, + v, + idx_k, + quant, + ) + + +def index_directions(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """Two orthogonal ±1 index directions (the halves of the 128 dims), E4M3-exact.""" + signs = torch.randint(0, 2, (HEAD_DIM,), generator=gen, device="cuda").float() * 2 - 1 + u0 = torch.zeros(HEAD_DIM, device="cuda") + u1 = torch.zeros(HEAD_DIM, device="cuda") + u0[: HEAD_DIM // 2] = signs[: HEAD_DIM // 2] + u1[HEAD_DIM // 2 :] = signs[HEAD_DIM // 2 :] + return u0, u1 + + +def planted_index_keys(length: int, u0, u1, gen: torch.Generator) -> torch.Tensor: + """Index-K rows [length, 128] (bf16, E4M3-exact) whose per-block maximum score is unambiguous. + + Every block has one peak row at a random position, ``a_b u0 + c_b u1`` with ``a`` and ``c`` two independent + permutations of ``LEVELS``; its other rows are ``LOW (u0 + u1)``. A query along ``u0`` then ranks the blocks by + ``a``, one along ``u1`` by ``c``, every score 1/8 or more apart from the next.""" + blocks = (length + PAGE - 1) // PAGE + assert blocks <= len(LEVELS) + levels = torch.tensor(LEVELS, device="cuda") + a = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + c = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + rows = (LOW * (u0 + u1)).expand(length, HEAD_DIM).clone() + for b in range(blocks): + span = min(PAGE, length - b * PAGE) + peak = b * PAGE + int(torch.randint(0, span, (1,), generator=gen, device="cuda")) + rows[peak] = a[b] * u0 + c[b] * u1 + return rows.to(torch.bfloat16) + + +def nvfp4_quant(*values: torch.Tensor) -> torch.Tensor: + """The (Q, K, V) quantization scales ``fused_write_layer_caches_nvfp4`` takes: each role's amax at the top of the + E2M1 x E4M3 range.""" + quant = torch.ones(3, dtype=torch.float32, device="cuda") + for role, x in enumerate(values, start=1): + quant[role] = (448.0 * 6.0) / x.abs().max().float() + return quant diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_nvfp4_cache.py new file mode 100644 index 000000000000..2254f6c0b68c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/_m3_nvfp4_cache.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A MiniMax-M3 TP4 sparse-layer decode step over a packed NVFP4 cache, for the attention entries' tests. + +The cache is written by the in-tree scatter (``fused_write_layer_caches_nvfp4``, +which quantizes with ``torch.ops.trtllm.fp4_quantize``), into one +``[pages, 2 (K, V), 1, 128, .]`` pool per role exactly as ``msa_paged_kv`` and +``get_block_scale_buffers`` hand the K and V views over. Pages are shuffled so +a kernel that ignored the block table would read the wrong ones. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +NUM_HEADS = 16 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 + + +def token_kv_len(seq_len: int, decode_query_len: int, offset: int) -> int: + """The KV extent query token ``offset`` of a request attends, as every decode kernel derives it.""" + return max(seq_len - decode_query_len + offset + 1, 0) + + +def build_step( + seq_lens: Sequence[int], decode_query_len: int, *, identity: bool, seed: int +) -> SimpleNamespace: + """One decode step: E4M3 q [T, 16, 128], the packed NVFP4 K/V views and their scales, a [1, T, 16] + head-major top-k table (the identity table, or a random ascending selection with a -1 tail), the block + table and per-request lengths.""" + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = max(1, max((s + PAGE - 1) // PAGE for s in seq_lens)) + num_pages = batch * max_blocks + num_slots = num_pages * PAGE + block_table = ( + torch.randperm(num_pages, device="cuda", generator=g) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + data_pool = torch.zeros(num_pages, 2, 1, PAGE, HEAD_DIM // 2, dtype=torch.uint8, device="cuda") + scale_pool = torch.zeros( + num_pages, 2, 1, PAGE, HEAD_DIM // 16, dtype=torch.uint8, device="cuda" + ) + quant = torch.ones(3, dtype=torch.float32, device="cuda") + values = [] + for role in (0, 1): + x = torch.randn(num_slots, HEAD_DIM, device="cuda", generator=g, dtype=torch.float32).to( + torch.bfloat16 + ) + quant[role + 1] = (448.0 * 6.0) / x.abs().max().float().item() + values.append(x) + assert fused_write_layer_caches_nvfp4( + data_pool[:, 0], + data_pool[:, 1], + scale_pool[:, 0], + scale_pool[:, 1], + None, + torch.arange(num_slots, device="cuda", dtype=torch.int32), + values[0], + values[1], + None, + quant, + ) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) + topk = torch.full((1, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + nblocks = (token_kv_len(seq_lens[req], decode_query_len, offset) + PAGE - 1) // PAGE + if identity: + assert nblocks <= TOPK, "an identity step has every request within 16 blocks" + chosen = torch.arange(nblocks, device="cuda") + else: + chosen = torch.randperm(nblocks, device="cuda", generator=g)[:TOPK].sort().values + topk[0, tok, : chosen.numel()] = chosen.to(torch.int32) + return SimpleNamespace( + q=q, + k_paged=data_pool[:, 0], + v_paged=data_pool[:, 1], + k_block_scale=scale_pool[:, 0], + v_block_scale=scale_pool[:, 1], + k_global_scale=(1.0 / quant[1:2]).contiguous(), + v_global_scale=(1.0 / quant[2:3]).contiguous(), + topk=topk, + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device="cuda"), + decode_query_len=decode_query_len, + sm_scale=HEAD_DIM**-0.5, + source=values, + quant=quant, + num_slots=num_slots, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py new file mode 100644 index 000000000000..824cb6488e58 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_dense_decode entry: the dense layers' trtllm-gen decode at the KV split this target +sizes per request count. + +Every cell runs against the fp32 reference, and a capped cell also against the same call at the model maximum (the +split the in-tree step hands it): the two differ only by the order the CTAs' partial softmax sums are merged in, held +here to a few BF16 ulps. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 + CAPPED, + HEAD_DIM, + HEADS, + KV_HEADS, + MODEL_MAX, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 + m3_dense_decode as op, +) +from tensorrt_llm._torch.attention.backends.fmha.msa_decode import dense_decode_max_seq_len + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_dense_decode requires a CUDA device" + +PAGE = 32 +SUBPAGES = 2 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the kernel asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor): + self.pool = pool + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, SUBPAGES + + +def _inputs(kv_lens, qlen: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * SUBPAGES, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + # Slots handed out non-contiguously, as the block manager does. + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1: a third of each row's softmax mass sits there, far + # outside the band, so a split that missed the tail of a context longer than its max_seq_len fails. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + first_full = (start + PAGE - 1) // PAGE + last = (n + PAGE - 1) // PAGE + if first_full < last: + v_idx = block_table[b, first_full:last].long() * SUBPAGES + 1 + pool[v_idx] += 1.0 + if start % PAGE: + v0 = int(block_table[b, start // PAGE]) * SUBPAGES + 1 + pool[v0, :, start % PAGE :] += 1.0 + q = torch.randn((batch * qlen, HEADS, HEAD_DIM), device="cuda", generator=g).to(torch.bfloat16) + return ( + q, + _Pool(pool.to(torch.float8_e4m3fn)), + block_table, + torch.tensor(kv_lens, device="cuda", dtype=torch.int32), + ) + + +def _args(q, manager, block_table, seq_lens, qlen, max_seq_len): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + return ( + q, + manager, + 0, + block_table, + seq_lens, + SM_SCALE, + out, + qlen, + max_seq_len, + seq_lens.shape[0], + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], spec["qlen"], zlib.crc32(cell.why.encode()) % 2**31 + ) + args = _args(q, manager, block_table, seq_lens, spec["qlen"], spec["max_seq_len"]) + certifying(op) + ref = op.reference(*args) + out = op(*args).clone() + torch.testing.assert_close(out.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + if spec["capped"]: + uncapped = op(*_args(q, manager, block_table, seq_lens, spec["qlen"], MODEL_MAX)).clone() + # The split reorders only the merge of the CTAs' fp32 partial sums before the one BF16 rounding: every element + # is within 4 of its own BF16 spacing (2^(e - 8) for |x| = m * 2^e, m in [0.5, 1)). + ref_bf16 = uncapped.float() + _, exp = torch.frexp(ref_bf16.abs()) + spacing = torch.where( + ref_bf16 == 0, + torch.full_like(ref_bf16, 2.0**-133), + torch.ldexp(torch.ones_like(ref_bf16), exp - 8), + ) + assert bool(((out.float() - ref_bf16).abs() <= 4 * spacing).all()) + assert_within_ulp(out, uncapped, element_ulp=4.0, rms_ulp=1.0) + + +@pytest.mark.skipif( + not dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH, + reason="the TP4 SM 10.7 target has no dense-decode cap until one is measured and certified on Rubin", +) +def test_the_cap_is_the_dispatch_constant() -> None: + """The entry's capped cells run at the split the target hands trtllm-gen at one request, and the in-tree step + resolves the cap from the attention module only at the request counts it names.""" + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {1: CAPPED} + attn = type( + "Attn", + (), + {"dense_decode_max_seq_len_by_batch": dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH}, + )() + assert dense_decode_max_seq_len(attn, MODEL_MAX, 1) == CAPPED + assert dense_decode_max_seq_len(attn, MODEL_MAX, 2) == MODEL_MAX + assert dense_decode_max_seq_len(attn, 4096, 1) == 4096 + assert dense_decode_max_seq_len(object(), MODEL_MAX, 1) == MODEL_MAX + + +def test_guard_refuses_more_requests_than_the_counters() -> None: + q, manager, block_table, seq_lens = _inputs((3000, 3000), 1, 3) + args = list(_args(q, manager, block_table, seq_lens, 1, CAPPED)) + args[9] = 1 + with pytest.raises(AssertionError, match="counters cover"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..98c6a27fc039 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7): the TP4 dense layers' FP8 q/k/v (minimax_m3_sm107tp4_dense_qk_norm_rope_fp8, head-count +generic, at the TP4 per-rank 16 query heads over 1 KV head) carries the bits of the stock path, the BF16 +fused_qk_norm_rope output cast to E4M3, for every norm / RoPE variant: Q and K are rounded to BF16 before E4M3 (K is +stored in the FP8 cache), V is cast.""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers the trtllm ops) + +assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) + +HQ, HK, HV, D, ROT = 16, 1, 1, 128, 64 + + +def _entry(): + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_dense_qk_norm_rope_fp8 import ( # noqa: E501 + m3_dense_qk_norm_rope_fp8, + ) + + return m3_dense_qk_norm_rope_fp8 + + +def test_certified_cells() -> None: + """The entry (ops/attention/m3_dense_qk_norm_rope_fp8) at every cell: its output is its reference's bit for bit; + with ``kv_write`` the fused-insert form also returns the plain form's bytes and the compact Q.""" + op = _entry() + for cell in op.CELLS: + spec = cell.spec + m, hq, hk = spec["m"], spec["q"], spec["kv"] + g = torch.Generator(device="cuda").manual_seed(m * 7 + hq) + qkv = (torch.randn(m, (hq + 2 * hk) * D, device="cuda", generator=g) * 2).to(torch.bfloat16) + pos = torch.randint(0, 1 << 20, (m,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + args = (qkv, hq, hk, hk, ROT, 1e-6, qw, kw, 5e6, True, pos, True) + ref = op.reference(*args) + if spec.get("kv_write"): + pages, tpb = 4, 128 + cache = torch.zeros(pages, 2, hk, tpb, D, device="cuda", dtype=torch.float8_e4m3fn) + slots = torch.randperm(pages * tpb, device="cuda", generator=g)[:m].to(torch.int32) + q_compact = torch.empty(m, hq, D, device="cuda", dtype=torch.float8_e4m3fn) + out = op(*args, cache, slots, m, q_compact) + assert torch.equal( + q_compact.view(m, hq * D).view(torch.uint8), out[:, : hq * D].view(torch.uint8) + ), cell.why + else: + out = op(*args) + assert torch.equal(out.view(torch.uint8), ref.view(torch.uint8)), cell.why + + +@pytest.mark.parametrize("tokens", [1, 4, 8, 4096]) +@pytest.mark.parametrize("use_gemma", [False, True], ids=["rmsnorm", "gemma"]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_dense_fp8_qkv_is_the_bf16_kernel_cast(tokens, use_gemma, is_neox) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8"): + pytest.skip("this build does not carry minimax_m3_sm107tp4_dense_qk_norm_rope_fp8") + g = torch.Generator(device="cuda").manual_seed(tokens * 4 + int(use_gemma) * 2 + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + ours = torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8( + qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, use_gemma + ) + stock = qkv.clone() + torch.ops.trtllm.fused_qk_norm_rope( + stock, HQ, HK, HV, D, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, + 1.0, 0.0, 0.0, 1.0, True, use_gemma, False, 0, 0, + ) # fmt: skip + assert torch.equal(ours.view(torch.uint8), stock.to(torch.float8_e4m3fn).view(torch.uint8)) + + +@pytest.mark.parametrize("tokens,live", [(1, 1), (4, 4), (8, 5), (64, 64)]) +@pytest.mark.parametrize("tokens_per_block", [32, 128]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_fused_kv_insert_stores_the_bytes_of_write_kv_slots_and_the_compact_q( + tokens, live, tokens_per_block, is_neox +) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + g = torch.Generator(device="cuda").manual_seed(tokens * 1000 + tokens_per_block + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + pages, layers, layer = 16, 3, 1 + # The layer's HND view of a multi-layer pool, strided as get_buffers returns it. + raw = torch.randint( + 0, + 256, + (pages, layers, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.randperm(pages * tokens_per_block, device="cuda", generator=g)[:tokens].to( + torch.int32 + ) + slots[live:] = -1 # padded rows own no slot + + def view(t): + return t.view(torch.float8_e4m3fn)[:, layer].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, False) + q_compact = torch.empty(tokens, HQ, D, device="cuda", dtype=torch.float8_e4m3fn) + ours = torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, live, q_compact + ) + plain = torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8(*args) + assert torch.equal(ours.view(torch.uint8), plain.view(torch.uint8)) + assert torch.equal( + q_compact.view(tokens, HQ * D).view(torch.uint8), plain[:, : HQ * D].view(torch.uint8) + ) + + ref_raw = raw.clone() + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D) + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D) + ref = view(ref_raw) + write_kv_slots(ref[:, 0], slots, k, live, layout="HND") + write_kv_slots(ref[:, 1], slots, v, live, layout="HND") + assert torch.equal(ours_raw, ref_raw) + + +def test_the_fused_kv_insert_skips_a_negative_slot() -> None: + """A live row whose slot is the -1 sentinel stores nothing (no write before the pool row).""" + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + tokens, tokens_per_block, pages = 4, 32, 8 + g = torch.Generator(device="cuda").manual_seed(7) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + raw = torch.randint( + 0, + 256, + (pages, 3, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.tensor([5, -1, 40, 77], device="cuda", dtype=torch.int32) + + def view(t): + return t.view(torch.float8_e4m3fn)[:, 1].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, True, pos, False) + torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, tokens + ) + plain = torch.ops.trtllm.minimax_m3_sm107tp4_dense_qk_norm_rope_fp8(*args) + keep = torch.tensor([0, 2, 3], device="cuda") + ref_raw = raw.clone() + ref = view(ref_raw) + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D)[keep] + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D)[keep] + write_kv_slots(ref[:, 0], slots[keep], k, 3, layout="HND") + write_kv_slots(ref[:, 1], slots[keep], v, 3, layout="HND") + assert torch.equal(ours_raw, ref_raw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..a594c099c2ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_attention entry: the Eagle3 drafter layer's attention against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp4``): the target's hybrid NVFP4 +``MiniMaxM3KVCacheManagerV2`` with its one-model Eagle3 draft layer, the requests allocated in it, drafter +``TrtllmAttentionMetadata`` over the manager's FP8 draft view (the manager the runtime's draft swap installs) filled +by its ``prepare()``, and the drafter's ``TrtllmAttention`` built as ``M3Eagle3Draft`` builds it (fused NeoX RoPE at +the drafter checkpoint's theta, FP8 KV). Earlier rows go into the draft layer's cache through the in-tree scatter at +the slots the manager assigned; the call rotates and appends the step's own. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp4 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. +# The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: +# it takes the rotated Q in E4M3 (rounding the reference's Q to E4M3 takes the RMS gap from ~9 to ~7 ulps) and rounds +# the softmax probabilities before the PV product. Measured on GB300 over these cells (max element / RMS): 13.6 / 8.6 +# (4 x 4 rows), 9.9 / 7.2 (8 x 1), 20.4 / 7.7 (512 context rows). The band leaves about 1.5x; a reference shifted +# by one row is far outside it. +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 + +# Cached tokens per request of each cell; the new rows per request come from the cell. +_CACHED = { + (4, 4): (200, 129, 1000, 64), + (8, 1): (5, 127, 128, 300, 1000, 2047, 64, 900), + (1, None): (0,), +} +#: Rows past the live ones in the fused QKV, as a padded step hands them over. +_PAD = 3 + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + +# Non-paged context FMHA on SM107 does not apply RoPE (this fixture's path; the served drafter runs the paged context +# FMHA, which rotates): the context cell is a strict expected failure until the runtime rotates there too. +_UNROTATED_CONTEXT = pytest.mark.xfail( + strict=True, + reason="non-paged context FMHA on SM107 does not apply RoPE (shared runtime)", +) +_CELLS = [ + pytest.param(c, marks=_UNROTATED_CONTEXT) if c.spec.get("context") else c for c in op.CELLS +] + + +@pytest.mark.parametrize("cell", _CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + context = spec.get("context") + cached = _CACHED[spec["reqs"], spec.get("qlen")] + qlens = (context,) if context else (spec["qlen"],) * spec["reqs"] + assert len(cached) == spec["reqs"] + contexts = 1 if context else 0 + rids = list(range(spec["reqs"])) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = rt.make_draft_metadata(manager) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + assert int(metadata.num_tokens) == total + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + certifying(op) + out = op(backend, qkv, metadata) + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{cell.why}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + # The call's append: each live row's rotated K and its V at the row's position, E4M3 (unit KV scale), to + # within one E4M3 step of the float64 rotation (the kernel rotates in FP32). + pool = manager.get_buffers(rt.DRAFT, "HND") + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + want_k = _rope(rows[:, q_size : q_size + rt.HEAD_DIM].float(), positions) + want_v = rows[:, q_size + rt.HEAD_DIM :].to(torch.float8_e4m3fn) + got_k = pool[slots // rt.PAGE, 0, 0, slots % rt.PAGE].to(torch.float64) + got_v = pool[slots // rt.PAGE, 1, 0, slots % rt.PAGE] + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), "K is appended rotated, in E4M3" + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose rows are shifted by one token is outside it. + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py new file mode 100644 index 000000000000..e5cec6259f9a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for ops/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V +append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row +first draft pass and the 1-row later passes of one request. + +Lengths sweep the KV tile edges each query row in turn ends on (the rows of a Q CTA with no key in its trailing +splits); every call is held to the fp32 reference over the cache after the append (E4M3 query, unit KV scales, no +RoPE) within the FP8-KV band, finite, for both split-KV reducers. +""" + +import math +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_tokmap_guard as tokmap_guard, +) + +assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +FP8_TOL = 6e-2 +# Persistent like the FMHA manager's workspace and the spec-decoding metadata of a real step: the in-tree preprocess +# kernels must not see these move between calls. +_PERSISTENT = {} +# Row counts of the one-request token maps the guard wrote; the model installs it in prebuild_kernels. +_TOKMAP_TAKEN = [] + + +@pytest.fixture(scope="module", autouse=True) +def _tokmap_guard_installed(): + # A one-request spec pass places its rows by the workspace token map, which the op builds only for several + # requests; without the guard the call reads whatever the workspace holds. + tokmap_guard.install(torch.device("cuda"), dispatch.DECODE_MAX_TOKENS, _TOKMAP_TAKEN.append) + yield + + +def _persistent(name: str, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _edge_lens(plan, qlen: int): + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +def _library(qlen: int, nreq: int, own_reducer: bool): + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = own_reducer + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = torch.cuda.get_device_properties(0).multi_processor_count + lib._on_taken = lambda rows: None + # The in-tree preprocess pass (its RoPE-free cache append); the one-launch preprocess has its own bitwise test. + lib._fused_rope_mode = None + lib._pdl = False + lib._on_rope_taken = None + return lib + + +def _attn(qlen: int): + return SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=0, theta=10000.0, scale_type=0, scale=1.0, max_positions=1024 + ), + ) + + +def _run(lib, qlen: int, kv_lens, seed: int, spec: bool = True): + """One drafter generation pass over a fresh FP8 cache; returns (output, fp32 reference). ``spec`` passes the + spec-decoding lengths/positions (a multi-token generation step); without them a multi-row pass is the drafter's + plain generation step of input_seq_length rows, as it runs in-model.""" + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g).clamp_(-4, 4) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + for r, n in enumerate( + kv_lens + ): # shift the tail third of V, so a missed trailing split fails the band + start = 2 * n // 3 + for pg in range(start // PAGE, math.ceil(n / PAGE)): + pool[int(perm[r, pg]), 1, :, (start % PAGE if pg == start // PAGE else 0) :] += 1.0 + pool = pool.to(torch.float8_e4m3fn) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + pointers = torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64) + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=dev, generator=g) + seq = torch.tensor(kv_lens, dtype=torch.int32, device=dev) + one = torch.ones(1, dtype=torch.float32, device=dev) + multi = qlen > 1 and spec + out = torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=_attn(qlen), + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=pointers, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=seq, + spec_decoding_generation_lengths=( + _persistent( + f"gen_lengths{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=( + _persistent( + f"position_offsets{qlen}x{nreq}", + lambda: ( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous() + ), + ) + if multi + else None + ), + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=out, + ) + taken = len(_TOKMAP_TAKEN) + lib.run_generation(params) + torch.cuda.synchronize() + if multi and nreq == 1: + assert _TOKMAP_TAKEN[taken:] == [ntok], ( + "the one-request spec pass ran without its token map" + ) + # fp32 reference over the cache after the append: E4M3 query, unit scales, causal over the step's rows. + q = qkv[:, : H * D].float().view(nreq, qlen, H, D).to(torch.float8_e4m3fn).float() + flat = pool.float() + ref = torch.empty(nreq, qlen, H, D, device=dev) + for r, n in enumerate(kv_lens): + blocks = perm[r, : math.ceil(n / PAGE)] + k = flat[blocks, 0].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + v = flat[blocks, 1].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + k = k.repeat_interleave(H // KVH, dim=0) + v = v.repeat_interleave(H // KVH, dim=0) + s = torch.einsum("qhd,hkd->hqk", q[r], k) * D**-0.5 + visible = ( + torch.arange(n, device=dev)[None, :] + < (n - qlen + 1 + torch.arange(qlen, device=dev))[:, None] + ) + s = s.masked_fill(~visible[None], float("-inf")) + ref[r] = torch.einsum("hqk,hkd->qhd", torch.softmax(s, dim=-1), v) + return out.float().view(nreq, qlen, H, D), ref + + +@pytest.mark.parametrize("qlen,nreq", sorted(dispatch.FULLCTX_DECODE_RECIPES)) +def test_the_drafter_pass_at_tile_edges(qlen, nreq) -> None: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if fullctx.has_own_reducer() else (False,) + for own in reducers: + fullctx.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_LEN, own, get_env_enable_pdl() + ) + lens = _edge_lens(fullctx.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + lens = [n for n in lens if n > qlen] + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for i, kv_lens in enumerate(batches): + # The first draft pass runs in-model as a plain multi-row generation step; every other batch also takes the + # spec-decoding form. + spec = qlen > 1 and i % 2 == 1 + for own in reducers: + got, ref = _run(_library(qlen, nreq, own), qlen, kv_lens, sum(kv_lens), spec=spec) + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +def _drafter_model(scale: float): + """A drafter whose one attention module holds a TrtllmAttention backend as ``attn`` (the backend is no module).""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + backend = TrtllmAttention.__new__(TrtllmAttention) + backend.kv_cache_scaling_factor = torch.full((1,), scale, device="cuda") + backend.is_mla_enable, backend.kv_lora_rank, backend.v_head_dim, backend.head_dim = ( + False, + None, + None, + D, + ) + backend._fmha_manager = SimpleNamespace(fmha_libs=[]) + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.zeros(1, device="cuda")) + model.self_attn = torch.nn.Module() + model.self_attn.attn = backend + return model, backend + + +def test_install_declines_a_scaled_kv_cache() -> None: + """At unit KV scales the arithmetic is the stock decode's; a scaled cache keeps the stock path.""" + reasons = [] + model, layer = _drafter_model(0.5) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 0 + assert layer._fmha_manager.fmha_libs == [] and reasons == ["a scaled FP8 KV cache"] + model, layer = _drafter_model(1.0) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 1 + assert isinstance(layer._fmha_manager.fmha_libs[0], drafter.M3DrafterFullCtxFmha) + + +def test_a_cache_bound_past_the_plan_capacity_keeps_the_stock_path() -> None: + """The plans are compiled for ``max_kv_len`` of K/V: a step whose cache length bound is longer falls through.""" + + from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase + + class _Backend: + num_heads, num_kv_heads, head_dim, quant_mode = H, KVH, D, FP8_KV_QUANT_MODE + kv_scale_orig_quant = kv_scale_quant_orig = None + + backend = _Backend() + backend.kv_scale_orig_quant, backend.kv_scale_quant_orig = object(), object() + reasons = [] + lib = _library(1, 1, False) + lib._on_uncovered = reasons.append + lib._attn_ref = weakref.ref(backend) + fwd = SimpleNamespace( + attention_input_type=None, + kv_scale_orig_quant=backend.kv_scale_orig_quant, + kv_scale_quant_orig=backend.kv_scale_quant_orig, + ) + meta = SimpleNamespace(num_contexts=0, num_generations=1, max_seq_len=MAX_LEN + 8) + q = torch.empty(1, H * D, device="cuda") + assert not lib._is_supported(q, None, None, meta, fwd, phase=FmhaPhase.GENERATION) + assert reasons == ["the cache length bound is beyond the plans' K/V capacity"] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py new file mode 100644 index 000000000000..aefbea361b74 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -0,0 +1,371 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for minimax_m3_sm107tp4_drafter_rope_fp8_kv, the TP4 drafter layer's one-launch generation +preprocess: bitwise against the in-tree ``thop.trtllm_gen_generation_preprocess`` followed by the torch E4M3 cast +of its BF16 query, on the query bytes and on every byte of the paged cache, with GPT-NeoX RoPE from a cos/sin table +(full and partial rotary width), multi-token (spec-decoding positions) and plain passes, one and two requests, and +query values past the E4M3 range. The library's fused pass is also held bitwise to its preprocess pass end to end. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_drafter_rope_fp8_kv"): + pytest.skip("minimax_m3_sm107tp4_drafter_rope_fp8_kv is not built", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +ROPE_GPT_NEOX = 2 +_PERSISTENT = {} + + +def _persistent(name, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _case(qlen, kv_lens, rotary_dim, multi, seed): + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = (torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g) * 2).to( + torch.float8_e4m3fn + ) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, device=dev, generator=g) * 4 + # A few query / key values near and past the E4M3 range (both casts saturate at 448). + hot = torch.rand(qkv.shape, device=dev, generator=g) < 0.01 + qkv = torch.where( + hot, qkv.sign() * (430 + 80 * torch.rand(qkv.shape, device=dev, generator=g)), qkv + ) + qkv = qkv.to(torch.bfloat16).contiguous() + max_pos = max(kv_lens) + 64 + cos_sin = torch.randn(max_pos, rotary_dim // 2, 2, device=dev, generator=g).float().contiguous() + attn = SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=ROPE_GPT_NEOX, + rotary_inv_freq=None, + rotary_cos_sin=cos_sin, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=rotary_dim, theta=10000.0, scale_type=0, scale=1.0, max_positions=max_pos + ), + ) + one = torch.ones(1, dtype=torch.float32, device=dev) + spec_offsets = None + if multi: + spec_offsets = _persistent( + f"pos{qlen}x{nreq}", + lambda: torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous(), + ) + # Spec positions are not the plain row index: shift them per request. + spec_offsets.copy_( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1) + + torch.arange(nreq, dtype=torch.int32, device=dev)[:, None] + ) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=attn, + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=torch.tensor(kv_lens, dtype=torch.int32, device=dev), + spec_decoding_generation_lengths=( + _persistent( + f"len{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=spec_offsets, + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev), + ) + return params, pool + + +def _preprocess(params, mpc): + """run_generation's in-tree preprocess call and the E4M3 cast of its query.""" + attn, meta, fwd, rope = params.attn, params.meta, params.fwd, params.attn.rope_params + q_processed = thop.trtllm_gen_generation_preprocess( + params.qkv_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope.dim, + rope.theta, + int(rope.scale_type), + rope.scale, + rope.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + False, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + mpc, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + )[0] + rows = params.num_tokens * attn.num_heads * attn.head_dim + return ( + q_processed.flatten()[:rows] + .to(torch.float8_e4m3fn) + .view(params.num_tokens, attn.num_heads, attn.head_dim) + ) + + +def _op(params, pool, mode): + attn = params.attn + nreq, qlen = params.batch_size, params.input_seq_length + kv_pages = pool.view(-1, KVH, PAGE, D) + tables = params.meta.kv_cache_block_offsets[0, :nreq] + q = torch.empty( + params.num_tokens, + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=pool.device, + ) + multi = params.spec_decoding_generation_lengths is not None and qlen > 1 + torch.ops.trtllm.minimax_m3_sm107tp4_drafter_rope_fp8_kv( + params.qkv_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + attn.rope_params.dim, + PAGE, + mode, + True, + ) + return q + + +# One-request multi-token passes keep the in-tree preprocess (it reads an unbuilt token map), so the spec cases hold +# two requests. +CASES = [ + (qlen, kv_lens, rotary, multi) + for qlen, multi in ((1, False), (4, True), (4, False)) + for kv_lens in (((4, 130) if multi else (4,)) if qlen > 1 else (130,), (8191, 131073)) + for rotary in (128, 64) +] + + +def _bytes(t): + return t.contiguous().view(torch.uint8) + + +def _matches(case, seed): + """Per rope mode: (query bytes equal, cache bytes equal) against the in-tree preprocess + cast.""" + qlen, kv_lens, rotary, multi = case + kv_lens = [max(n, qlen) for n in kv_lens] + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(qlen, kv_lens, rotary, multi, seed) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + pool_ref = pool.clone() + out = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + out[mode] = ( + bool(torch.equal(_bytes(q), _bytes(q_ref))), + bool(torch.equal(_bytes(pool), _bytes(pool_ref))), + ) + return out + + +@pytest.mark.parametrize( + "case", + CASES, + ids=lambda c: f"q{c[0]}-{'x'.join(map(str, c[1]))}-r{c[2]}-{'spec' if c[3] else 'plain'}", +) +def test_the_pinned_rope_mode_is_bitwise(case) -> None: + result = _matches(case, seed=7) + print(f"[rope modes] {case}: {result}", flush=True) + assert result[dispatch.DRAFTER_ROPE_MODE] == (True, True), result + + +def test_the_fused_pass_matches_the_preprocess_pass() -> None: + mpc = torch.cuda.get_device_properties(0).multi_processor_count + for qlen, nreq in sorted(dispatch.FULLCTX_DECODE_RECIPES): + for multi in (False, True) if qlen > 1 else (False,): + kv_lens = [70001 + 517 * r for r in range(nreq)] + outs = [] + for fused in (False, True): + params, pool = _case(qlen, kv_lens, 128, multi, seed=11) + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = True + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = mpc + lib._fused_rope_mode = dispatch.DRAFTER_ROPE_MODE if fused else None + lib._pdl = True + taken = [] + lib._on_rope_taken = lambda rows, what: taken.append(what) + lib._on_taken = lambda rows: None + fullctx.prebuild(torch.device("cuda"), lib._recipes, MAX_LEN, True, True) + lib.run_generation(params) + torch.cuda.synchronize() + assert bool(taken) == (fused and not (multi and nreq == 1)) + outs.append((params.output.clone(), pool.clone())) + assert torch.equal(_bytes(outs[0][1]), _bytes(outs[1][1])), (qlen, nreq, multi, "cache") + assert torch.equal(_bytes(outs[0][0]), _bytes(outs[1][0])), ( + qlen, + nreq, + multi, + "output", + ) + + +def _round_bf16(x): + u = x.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000 + return u.to(torch.int64).to(torch.int32).view(torch.float32) + + +def _discriminating(n, seed): + """(cos, sin, x, pair) quadruples whose rotated first-half output's E4M3 byte differs between the rope modes + (fp32 CPU emulation; the GPU comparison is the verdict).""" + g = torch.Generator().manual_seed(seed) + keep = [] + while sum(k.shape[0] for k in keep) < n: + m = 1 << 20 + c, s = torch.randn(m, generator=g), torch.randn(m, generator=g) + x = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + xp = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + sn = -s + r0 = (c * x) + (sn * xp) + r1 = ((c.double() * x.double()) + (sn * xp).double()).float() + r2 = ((sn.double() * xp.double()) + (c * x).double()).float() + b = [ + _round_bf16(r.contiguous()).to(torch.float8_e4m3fn).view(torch.uint8) + for r in (r0, r1, r2) + ] + diff = (b[0] != b[1]) | (b[0] != b[2]) | (b[1] != b[2]) + diff &= (r0.abs() < 400) & (r1.abs() < 400) + keep.append(torch.stack([c, s, x, xp], 1)[diff]) + return torch.cat(keep)[:n] + + +def test_crafted_inputs_pin_the_rope_mode() -> None: + """Rotations whose E4M3 byte depends on how the two products are summed: the in-tree preprocess and exactly the + pinned mode agree on every one.""" + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(1, [130], 128, False, seed=3) + half = 64 + quads = _discriminating(H * half, seed=5).cuda() + pos = 129 # the one row's rope position: cache length 130, one query row + table = params.attn.rotary_cos_sin.view(-1, half, 2) + # One (cos, sin) per dim: every head shares the position's table row, so the heads take the same dims' quads with + # their own (x, pair). + table[pos, :, 0] = quads[:half, 0] + table[pos, :, 1] = quads[:half, 1] + qkv = params.qkv_input.view(-1, H + 2 * KVH, D) + for h in range(H): + qkv[0, h, :half] = quads[:half, 2].to(torch.bfloat16) + qkv[0, h, half:] = quads[:half, 3].to(torch.bfloat16) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + agree = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + agree[mode] = int((_bytes(q)[0, :, :half] == _bytes(q_ref)[0, :, :half]).sum()) + print( + f"[rope modes crafted] agreeing first-half query bytes of {H * half}: {agree}", flush=True + ) + assert agree[dispatch.DRAFTER_ROPE_MODE] == H * half, agree diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..529a556ea4c3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp8_indexer_qk_norm_rope entry of the MiniMax-M3 tp4 target: index-Q out and index-K in the +cache against the in-place fused norm + RoPE followed by the E4M3 cast, byte for byte up to the 0.1% budget the +in-tree test of the kernel allows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 + m3_fp8_indexer_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" + +DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 + + +def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: + assert a.shape == b.shape and a.dtype == b.dtype == torch.float8_e4m3fn + differ = int((a.view(torch.uint8) != b.view(torch.uint8)).sum()) + assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + m, heads = cell.spec["m"], cell.spec["heads"] + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + qk = torch.randn(m, (heads + 1) * DIM, generator=g, device="cuda").to(torch.bfloat16) + qw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (m,), generator=g, device="cuda", dtype=torch.int32) + pages = (m + PAGE - 1) // PAGE + 1 + cache = torch.zeros(pages, 1, PAGE, DIM, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(pages * PAGE, generator=g, device="cuda")[:m].to(torch.int32) + before = qk.clone() + args = (qk, cache, slots, heads, DIM, ROTARY, EPS, qw, kw, THETA, pos) + with validating(op): + q = op(*args) + assert torch.equal(qk, before), "the BF16 input is left untouched" + q_ref, k_ref = op.reference(*args) + _bytes_close(q.reshape(m, heads, DIM), q_ref) + flat = cache.view(-1, DIM) + _bytes_close(flat[slots.long()], k_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_saturate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_saturate.py new file mode 100644 index 000000000000..a60894191c22 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_saturate.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""paged_cache.fp8_e4m3_saturate (FP8_SATURATE): the E4M3 cache writes and the transient FP8 query casts saturate to ++-448 like the CUDA writers (__NV_SATFINITE); every finite in-range value keeps the plain cast's bits, values already +in E4M3 pass through, and the kill switch is exactly the plain cast. Where torch's cast already saturates (this image) +the conversion is the cast alone, one kernel; otherwise it clamps first, the same bytes.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import paged_cache + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fp8_e4m3_saturate_matches_satfinite(dtype, monkeypatch) -> None: + g = torch.Generator(device="cuda").manual_seed(0) + inside = (torch.randn(4096, generator=g, device="cuda") * 100).clamp(-448, 448).to(dtype) + plain = inside.to(torch.float8_e4m3fn) + assert torch.equal( + paged_cache.fp8_e4m3_saturate(inside).view(torch.uint8), plain.view(torch.uint8) + ) + edge = torch.tensor( + [449.0, 463.0, 464.0, 1e4, float("inf"), -500.0, -float("inf")], dtype=dtype, device="cuda" + ) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0], device="cuda") + assert torch.equal(paged_cache.fp8_e4m3_saturate(edge).float(), want) + nan = paged_cache.fp8_e4m3_saturate(torch.tensor([float("nan")], dtype=dtype, device="cuda")) + assert torch.isnan(nan.float()).all() + already = edge.clamp(-448, 448).to(torch.float8_e4m3fn) + assert paged_cache.fp8_e4m3_saturate(already) is already + monkeypatch.setattr(paged_cache, "FP8_SATURATE", False) + killed = paged_cache.fp8_e4m3_saturate(edge) + assert torch.equal(killed.view(torch.uint8), edge.to(torch.float8_e4m3fn).view(torch.uint8)) + + +def test_cache_write_saturates() -> None: + """The cache-dtype conversion every write_kv_slots branch applies (compared as bytes: torch has no float8 + index_copy_ kernel), and an E4M3 source passing through untouched.""" + values = torch.tensor([[1.0, 500.0, -1e5, float("inf")]], dtype=torch.bfloat16, device="cuda") + out = paged_cache._to_cache_dtype(values, torch.float8_e4m3fn) + assert out.dtype == torch.float8_e4m3fn + want = torch.tensor([[1.0, 448.0, -448.0, 448.0]], device="cuda").to(torch.float8_e4m3fn) + assert torch.equal(out.view(torch.uint8), want.view(torch.uint8)) + assert paged_cache._to_cache_dtype(out, torch.float8_e4m3fn) is out + bf = paged_cache._to_cache_dtype(values, torch.bfloat16) + assert bf.dtype == torch.bfloat16 and torch.equal(bf, values) + + +_EDGES = [ + 448.0, + 449.0, + 463.0, + 464.0, + 1e4, + float("inf"), + -448.0, + -449.0, + -464.0, + -1e4, + -float("inf"), + float("nan"), + 0.0, +] + + +def _cuda_kernels(fn) -> int: + from torch.profiler import ProfilerActivity, profile + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + return sum(1 for e in prof.events() if e.device_type.name == "CUDA") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("cast_saturates", [True, False], ids=["cast_only", "clamp_fallback"]) +def test_fp8_e4m3_saturate_bytes_equal_clamp_then_cast(dtype, cast_saturates, monkeypatch) -> None: + """Both paths give the bytes of clamp(+-448) followed by the cast, edge values included.""" + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", cast_saturates) + x = torch.tensor(_EDGES, dtype=dtype, device="cuda") + want = x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + assert torch.equal(paged_cache.fp8_e4m3_saturate(x).view(torch.uint8), want.view(torch.uint8)) + + +def test_fp8_e4m3_saturate_takes_the_cast_only_path_here(monkeypatch) -> None: + """This image's torch cast saturates, on the CPU (the import-time probe) and on CUDA, so the conversion is one + kernel; with the probe failing it falls back to clamp + cast, two kernels.""" + assert paged_cache._CAST_SATURATES + x = torch.tensor(_EDGES, dtype=torch.bfloat16, device="cuda") + assert torch.equal( + x.to(torch.float8_e4m3fn).view(torch.uint8), + x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).view(torch.uint8), + ) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 1 + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", False) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 2 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py new file mode 100644 index 000000000000..0014a17a7b01 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for the attention/m3_fullctx_decode entry: the dense layers' full-context FP8 decode on PrimTS +split-KV decode at the recipe the target pins per (query rows per request, requests). + +Every cell runs against the fp32 reference (q rounded to E4M3 first, as the kernel runs it), eager and replayed from a +CUDA graph (the same bits), and twice (the same bits). The K/V pages are read through the strided views of a flat +sub-page pool, as the KV cache manager hands them out. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_fullctx_decode import ( # noqa: E501 + HEAD_DIM, + HEADS, + KV_HEADS, + MAX_KV_LEN, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_fullctx_decode import ( # noqa: E501 + m3_fullctx_decode as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) + +PAGE = 128 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the entry asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor, subpages: int): + self.pool = pool + self.subpages = subpages + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, self.subpages + + +def _inputs(kv_lens, qlen: int, seed: int, subpages: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * subpages, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1, so a split that missed the tail fails the band. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + for page in range(start // PAGE, (n + PAGE - 1) // PAGE): + v = int(block_table[b, page]) * subpages + 1 + pool[v, :, (start % PAGE if page == start // PAGE else 0) :] += 1.0 + # q as the layer hands it over on decode steps: the FP8 query columns of the fused norm/RoPE kernel's + # [q | k | v] output (row stride wider than q), so no cast makes it compact. + qkv = torch.randn( + (batch * qlen, (HEADS + 2 * KV_HEADS) * HEAD_DIM), device="cuda", generator=g + ).to(torch.float8_e4m3fn) + q = qkv[:, : HEADS * HEAD_DIM].view(batch * qlen, HEADS, HEAD_DIM) + seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32) + return q, _Pool(pool.to(torch.float8_e4m3fn), subpages), block_table, seq_lens + + +@pytest.mark.parametrize("subpages", [2, 3], ids=["kv", "kv+index"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, subpages, certifying) -> None: # noqa: F811 + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], qlen, zlib.crc32(cell.why.encode()) % 2**31, subpages + ) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + certifying(op) + ref = op.reference(*args) + eager = op(*args).clone() + torch.testing.assert_close(eager.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + again = op(*args).clone() + assert torch.equal(again, eager), "not deterministic" + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args) + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager), "the CUDA-graph replay differs from the eager call" + + +# Plans compiled at other K/V stage depths stay alive for the whole session: a compiled decode is cached per config +# and is only ever run with the buffers of a live plan. +_OTHER_DEPTH_PLANS = [] + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_kv_stage_depth_keeps_the_bits(cell) -> None: + """kv_stages only sets how many K/V tiles the load pipeline keeps in flight: the recipe's plan and plans compiled + at 12 and at 4 stages give the same bits.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + if not kernel.has_own_reducer(): + pytest.skip("this build does not carry minimax_m3_sm107tp4_fullctx_reduce") + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs(spec["kv_lens"], qlen, 13, 2) + k_pages, v_pages = kernel.kv_page_views(manager.pool, manager.subpages) + kernel.prebuild(q.device, {(qlen, nreq): recipe}, MAX_KV_LEN, True, get_env_enable_pdl()) + plans = [kernel.plan_for(q.device, qlen, nreq, True)] + for stages in (12, 4): + config = {**recipe["config"], "kv_stages": stages} + plans.append( + kernel.FullCtxPlan( + q.device, + qlen, + nreq, + MAX_KV_LEN, + {**recipe, "config": config}, + True, + get_env_enable_pdl(), + ) + ) + _OTHER_DEPTH_PLANS.append(plans[-1]) + outs = [] + for plan in plans: + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + plan.run(q, k_pages, v_pages, block_table, seq_lens, SM_SCALE, out) + outs.append(out) + assert torch.equal(outs[0], outs[1]) and torch.equal(outs[0], outs[2]) + + +def test_the_cells_cover_the_dispatch_rows() -> None: + certified = {(c.spec["qlen"], len(c.spec["kv_lens"])) for c in op.CELLS} + assert set(dispatch.FULLCTX_DECODE.keys) <= certified + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: + """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once + in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + + if not kernel.has_own_reducer(): + pytest.skip("this build does not carry minimax_m3_sm107tp4_fullctx_reduce") + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs(spec["kv_lens"], qlen, 11, 2) + outs = [] + for own in (True, False): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + op(q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe, own_reducer=own) + outs.append(out.float()) + ours, primts = outs + _, exp = torch.frexp(primts.abs()) + spacing = torch.where( + primts == 0, + torch.full_like(primts, 2.0**-133), + torch.ldexp(torch.ones_like(primts), exp - 8), + ) + assert bool(((ours - primts).abs() <= spacing).all()) + + +def _tile_edge_kv_lens(plan, qlen: int): + """KV lengths around every kind of boundary the split ranges have, so each query row of the Q CTA in turn ends + its keys on a KV tile edge (k * tile + d, d in -1 .. qlen - 1): the first tiles, the reducer's 8-group stride, the + lengths where a split's range grows (splits * insts tiles) and a split edge past it.""" + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +@pytest.mark.parametrize("qlen,nreq", sorted(dispatch.FULLCTX_DECODE_RECIPES)) +def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: + """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys + (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode + step first hit (NaN rows from exp2(-inf - -inf)).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, False, get_env_enable_pdl() + ) + lens = _tile_edge_kv_lens(kernel.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, sum(kv_lens), 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +@pytest.mark.parametrize("qlen,nreq", sorted(dispatch.FULLCTX_DECODE_RECIPES)) +def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: + """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just + past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = dispatch.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + lengths = [(1 << 20) + 1, (1 << 20) + 8] + ([1_500_001] if nreq == 1 else []) + batches = [[n] * nreq if nreq == 1 else [n, (1 << 20) - 3] for n in lengths] + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, 13, 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, own, get_env_enable_pdl() + ) + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +def _swizzled_scale_offsets(rows: int, width: int) -> torch.Tensor: + """Offsets of rows [0, rows)'s UE8M0 scales in the 128x4-swizzled layout of mxfp8_quantize(x, True).""" + k_sf = width // 32 + k_tiles = (k_sf + 3) // 4 + r = torch.arange(rows, device="cuda").view(-1, 1) + k = torch.arange(k_sf, device="cuda").view(1, -1) + offsets = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + return offsets.reshape(-1) + + +@pytest.mark.parametrize("splits,rows,seq_len_q", [(108, 8, 1), (64, 8, 1), (96, 32, 4)]) +def test_the_merge_mxfp8_copy_is_stock_mxfp8_quantize(splits, rows, seq_len_q) -> None: + """The TP4 merge's MXFP8 copy (the dense o_proj hand-off) equals mxfp8_quantize(out, True) bit for bit: E4M3 rows + and the real rows' swizzled UE8M0 scales, with an all-zero head row (amax 0) and saturating outliers; the BF16 + output is the merge without the copy.""" + ops = torch.ops.trtllm + if not hasattr(ops, "minimax_m3_sm107tp4_fullctx_reduce") or "q8" not in str( + ops.minimax_m3_sm107tp4_fullctx_reduce.default._schema + ): + pytest.skip("needs the TP4 merge op with the MXFP8 copy") + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(splits * 7 + rows) + nreq, kvh = 2, 1 + tokens = nreq * seq_len_q + width = kvh * rows // seq_len_q * 128 + for dtype in (torch.float16, torch.bfloat16): + po = torch.randn(nreq, kvh, splits, rows, 128, device=dev, generator=g) * 3 + po[0, 0, :, 1] = 0.0 + po[1, 0, :, 0, 5] = 6.0e4 + po = po.to(dtype).contiguous() + ps = (torch.randn(nreq, kvh, splits, rows, device=dev, generator=g) * 4).contiguous() + seq = torch.tensor([165521, 120011], dtype=torch.int32, device=dev) + args = (po, ps, seq) + tail = (seq_len_q, seq_len_q, 128, 1, seq_len_q > 1, False) + plain = torch.zeros(tokens, width, dtype=torch.bfloat16, device=dev) + ops.minimax_m3_sm107tp4_fullctx_reduce(*args, plain, *tail) + out = torch.zeros(tokens, width, dtype=torch.bfloat16, device=dev) + q8 = torch.empty(128, width, dtype=torch.float8_e4m3fn, device=dev) + sf = torch.zeros(128 * (width // 32), dtype=torch.uint8, device=dev) + ops.minimax_m3_sm107tp4_fullctx_reduce(*args, out, *tail, q8, sf) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), plain.view(torch.int16)) + stock_q, stock_sf = ops.mxfp8_quantize(out, True) + assert torch.equal( + q8[:tokens].view(torch.uint8), stock_q.view(torch.uint8).reshape(tokens, -1) + ), dtype + at = _swizzled_scale_offsets(tokens, width) + assert torch.equal(sf[at], stock_sf.view(torch.uint8).reshape(-1)[at]), dtype diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..1d4353c466cf --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score entry: the MSA index scorer with its pre-wait schedule. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score``) bit for bit, which the pre-wait prefetch and the +early tile issue must not change. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score import ( # noqa: E501 + PAGE, + SPLIT_K, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 1 # index heads per rank at TP4 +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + args = (idx_q, index_k, block_table, seq, score, qlen, SPLIT_K, spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen, SPLIT_K + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +# The 2-request decode step also runs the early-issue kernel on this target (dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS). +# The op file is the sm_103 copy (test_modeling_v2_m3_tp4_drift), so these cells live here. +TWO_REQUEST_EARLY = ( + ((20000, 3000), "two requests verifying 1 + 3 rows each over 20000 and 3000 tokens"), + ( + (70000, 129), + "two requests: one past 512 blocks (its CTAs loop), one whose early tiles are drained", + ), +) + + +def test_the_two_request_step_takes_the_early_kernel() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, + ) + + assert dispatch.INDEX_SCORE_EARLY_MAX_REQUESTS == 2 + + +@pytest.mark.parametrize( + "seq_lens,why", TWO_REQUEST_EARLY, ids=[w[:44] for _, w in TWO_REQUEST_EARLY] +) +def test_two_request_early_issue_is_the_in_tree_scorer(seq_lens, why) -> None: + idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, 4, zlib.crc32(why.encode())) + score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) + args = (idx_q, index_k, block_table, seq, score, 4, SPLIT_K, True) + ref = op.reference(*args) + op(*args) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4, SPLIT_K + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, SPLIT_K, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4, SPLIT_K + ) + assert torch.equal(score, stock) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, SPLIT_K, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py new file mode 100644 index 000000000000..be30baa9bd5b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score_t2sc entry (T2SC, dispatch.INDEX_SCORE_T2SC): the TP2 / TP8 MSA index +scorer with its pre-wait schedule, as the TP4 target's copy at its 1 index head per rank. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score`` at its own 256-CTA split) bit for bit, which the +pre-wait prefetch, the early tile issue and the cell's split must not change. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score_t2sc import ( # noqa: E501 + PAGE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score_t2sc import ( # noqa: E501 + m3_index_score_t2sc as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 1 # index heads per rank at TP4 (4 index heads over 4 KV groups) +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + assert spec["heads"] == HEADS + args = (idx_q, index_k, block_table, seq, score, qlen, spec["split_k"], spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, 1536, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4 + ) + assert torch.equal(score, stock) + + +@pytest.mark.parametrize("qlen", [1, 2, 3, 4]) +def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> None: + """Lengths around the edge between pre-staged (all-old) and post-wait first blocks, 1 and 2 requests.""" + for k in (1, 2, 1294): + for length in (k * PAGE - qlen - 1, k * PAGE - qlen, k * PAGE - qlen + 1, k * PAGE + 1): + for seq_lens, split in ( + ((length,), 1536), + ((length,), 1272), + ((length, max(qlen, length // 3)), 768), + ): + idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) + score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) + op(idx_q, index_k, block_table, seq, score, 4, split, True) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + assert torch.equal(score, stock), (qlen, seq_lens) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, 1536, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..70cbd7b25768 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_msa_attention entry: one layer's MSA backend step against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp4``): a hybrid NVFP4 ``MiniMaxM3KVCacheManagerV2`` +with the requests allocated in it, ``MiniMaxM3MsaSparseAttentionMetadata`` filled by its ``prepare()``, and the layer's +``MiniMaxM3MsaSparseAttention`` built as the target builds it. Earlier rows (and, on a sparse layer, the step's own +rows, which the target's producer writes before the call) go into the cache through the in-tree scatters, at the +slots the manager assigned. The sparse layers' index-K rows are planted so every block's score is 1/8 or more apart +from every other's: the top-k selection the reference recomputes is then unambiguous. + +The cells are the target's calls (``_attention_core`` in modeling.py): a dense layer with live bf16 K / V in a +prefill chunk and live FP8 K / V in a decode step (the dense layers sit in the hybrid cache's FP8 half), a sparse +layer whose producer already wrote the caches (``k = v = None``) in a prefill chunk beyond 16 blocks and in a mixed +step, both with the NVFP4 scales the target passes. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_msa_attention import ( # noqa: E501 + m3_msa_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp4 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" + +# Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, +# per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 +# FMHA does), so what is left is each FMHA's own arithmetic, measured on GB300 over these cells (max element / RMS): +# * fmha_sm100 over the FP8 pool (dense context rows): 10.1 / 6.1. It runs E4M3 x E4M3 and its softmax +# probabilities are rounded before the PV product. +# * trtllm-gen over the FP8 pool (dense generation rows): 7.0 / 6.1. It rounds the probabilities to E4M3 (scaled by +# 448 against the running maximum): emulating that brings a single-tile request to 0.6 RMS. +# * the NVFP4 CSR prefill (sparse context rows): 34.4 / 11.8 at 4096 rows, 26.8 / 9.1 at 512. It stages the +# reconstructed K / V in E4M3 (the in-tree A/B test, test_minimax_m3_sparse_attn_decode, notes it and holds it +# to 3e-2). +# * the Triton NVFP4 sparse decode (sparse generation rows): 1.0 / 0.25 -- the band of its own entry +# (m3_sparse_decode): exact operands, FP16 probabilities, BF16 rounding. +# Each band leaves about 1.5x the measured maximum; a shifted or wrongly selected reference is far outside them. +_FP8_CONTEXT = dict(element_ulp=16.0, rms_ulp=9.0) +_FP8_DECODE = dict(element_ulp=12.0, rms_ulp=9.0) +_NVFP4_CONTEXT = dict(element_ulp=48.0, rms_ulp=16.0) +_NVFP4_DECODE = dict(element_ulp=2.0, rms_ulp=0.5) +_BANDS = { + ("dense", "context"): (_FP8_CONTEXT, None), + ("dense", "generation"): (None, _FP8_DECODE), + ("sparse", "context"): (_NVFP4_CONTEXT, None), + ("sparse", "mixed"): (_NVFP4_CONTEXT, _NVFP4_DECODE), +} + + +def _assert_in_bands(output, ref, bands, context_rows: int) -> None: + """Gate the context rows and the generation rows each on their own kernel's band.""" + context_band, decode_band = bands + if context_rows: + assert_within_ulp(output[:context_rows], ref[:context_rows], **context_band) + if context_rows < output.shape[0]: + assert_within_ulp(output[context_rows:], ref[context_rows:], **decode_band) + + +# (request rows, cached tokens per request, new tokens per request, context rows) per cell. +_SHAPES = { + ("dense", "context"): ((384,), (1024,), 1), + ("dense", "generation"): ((127, 128, 1000, 2500), (1, 1, 1, 1), 0), + ("sparse", "context"): ((0,), (4096,), 1), + ("sparse", "mixed"): ((0, 2999, 4500), (512, 1, 1), 1), +} + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """The two quantities ``assert_within_ulp`` gates, in BF16 ulps.""" + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _shift(x: torch.Tensor) -> torch.Tensor: + """``x``'s rows moved down by one (byte-wise, so E4M3 rows move too).""" + return x.view(torch.uint8).roll(1, 0).view(x.dtype) if x.element_size() == 1 else x.roll(1, 0) + + +def _shape(cell): + spec = cell.spec + cached, qlens, contexts = _SHAPES[spec["layer"], spec["phase"]] + if spec["phase"] == "context": + assert sum(qlens) == spec["tokens"] + elif spec["phase"] == "generation": + assert len(qlens) == spec["reqs"] + else: + assert sum(qlens[:contexts]) == spec["tokens"] and len(qlens) - contexts == spec["reqs"] + return cached, qlens, contexts + + +def _run_dense(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + rt.write_fp8_rows(manager, rt.DENSE, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + # A prefill chunk arrives in BF16 (the hybrid cache's dense layers take FP8 Q / K / V only on decode steps). + dtype = torch.bfloat16 if contexts else torch.float8_e4m3fn + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + k = torch.randn(total, rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + v = torch.randn(total, rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + # The NVFP4 target passes its layer scales on every layer; an FP8 layer's write and FMHA take none. + inv = torch.tensor([1.0, 0.5, 2.0], device="cuda") + fwd = 1.0 / inv + return (q, k, v, None, inv, fwd), dict(rids=rids, cached=cached, qlens=qlens, contexts=contexts) + + +def _run_sparse(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + u0, u1 = rt.index_directions(gen) + lengths = [c + n for c, n in zip(cached, qlens)] + keys = [torch.randn(n, rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths] + values = [torch.randn(n, rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths] + quant = rt.nvfp4_quant(torch.cat(keys), torch.cat(values)) + for rid, n, kk, vv in zip(rids, lengths, keys, values): + # Every row, the step's own included: the producer writes a sparse layer's caches before the call. + idx_k = rt.planted_index_keys(n, u0, u1, gen) + rt.write_nvfp4_rows( + manager, + metadata, + rt.SPARSE, + rid, + torch.arange(n), + kk.to(torch.bfloat16), + vv.to(torch.bfloat16), + idx_k, + quant, + ) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + # Query tokens alternate between the two planted directions, so neighbouring tokens select different blocks. + idx_q = torch.stack([u0, u1])[torch.arange(total, device="cuda") % 2].to(torch.float8_e4m3fn) + return (q, None, None, idx_q, quant, 1.0 / quant), dict( + rids=rids, cached=cached, qlens=qlens, contexts=contexts + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + sparse = spec["layer"] == "sparse" + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=False) + try: + backend = rt.make_msa_backend(rt.SPARSE if sparse else rt.DENSE) + metadata = rt.make_msa_metadata(manager) + run = _run_sparse if sparse else _run_dense + (q, k, v, idx_q, inv, fwd), step = run(cell, gen, manager, backend, metadata) + output = torch.full( + (q.shape[0], rt.HEADS * rt.HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + certifying(op) + result = op(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert result is output + ref = op.reference(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert torch.isfinite(output.float()).all(), "every live row is written" + bands = _BANDS[spec["layer"], spec["phase"]] + context_rows = sum(step["qlens"][: step["contexts"]]) + for name, rows in ( + ("context", slice(0, context_rows)), + ("generation", slice(context_rows, None)), + ): + if output[rows].numel(): + element, rms = _ulp_stats(output[rows], ref[rows]) + print( + f"\n[{cell.why}] {name} rows: max element {element:.2f} ulp, rms {rms:.3f} ulp" + ) + _assert_in_bands(output, ref, bands, context_rows) + + if k is not None: + # The call's own cache write: the live K / V rows land, as a plain E4M3 cast, at their positions. + pool = manager.get_buffers(rt.DENSE, "HND") + first = 0 + for rid, c, n in zip(step["rids"], step["cached"], step["qlens"]): + slots = rt.flat_slots(manager, rt.DENSE, rid, torch.arange(c, c + n)).long() + for role, src in ((0, k), (1, v)): + got = pool[slots // rt.PAGE, role, 0, slots % rt.PAGE] + want = src[first : first + n].to(torch.float8_e4m3fn) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose query rows are shifted by one token is outside it. + shifted = op.reference(backend, _shift(q), k, v, idx_q, metadata, output, sparse, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, shifted, bands, context_rows) + if sparse: + # And the selection is load-bearing: attending every causal block instead is outside it too. + dense = op.reference(backend, q, k, v, idx_q, metadata, output, False, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, dense, bands, context_rows) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py new file mode 100644 index 000000000000..be10d54aca2a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_n_valid_contract.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-token valid-block count is written only before the step's first layer. + +The long-row select (``ops/m3_select_long_dsl.py``, ``N_VALID_PREWAIT``) reads its row's entry of +``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the +buffer happens before the step's first layer: + +- the allocation (``_alloc_msa_proxy_scratch``); +- the host-to-device copy in ``prepare()`` (``_build_step_plans``); +- ``on_update_kv_lens``, which the runner calls before the forward. It writes through the plain + ``<<<>>>`` launch of ``msa_kv_update_kernel`` or through torch copies, so any PDL grid of the step starts only + after it completes. + +This test scans ``tensorrt_llm/_torch`` and fails on any other write site (including writes through a local alias of +the buffer), so a new writer has to revisit the pre-wait read. A buffer handed to an opaque kernel launch is not +tracked; such a launch inside a layer must be reviewed by hand. + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +_NAME = "msa_n_valid_blocks" +# In-place tensor methods that write their receiver. +_MUTATORS = { + "copy_", + "fill_", + "zero_", + "index_put_", + "index_copy_", + "scatter_", + "masked_fill_", + "add_", + "sub_", +} +# Calls whose first argument is the destination they write: ``step_inputs.copy_(dst, src)`` and the H2D helpers. +_COPY_HELPERS = {"copy_", "copy_h2d_if_changed", "copy_h2d_pool_shared"} +_BACKEND = "attention/backends/sparse/minimax_m3/msa_backend.py" +# (file relative to tensorrt_llm/_torch, enclosing function) -> writes allowed there +_ALLOWED = { + (_BACKEND, "_alloc_msa_proxy_scratch"), # the allocation + (_BACKEND, "_build_step_plans"), # the H2D copy in prepare() + (_BACKEND, "on_update_kv_lens"), # the device-side correction before the forward +} + + +def _base(node): + """The innermost receiver of an attribute / subscript / call chain.""" + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute) and node.attr == _NAME: + return node + node = node.func if isinstance(node, ast.Call) else node.value + return node + + +def _mentions(node, aliases) -> bool: + """Whether the expression's receiver chain is ``*.msa_n_valid_blocks`` or a local alias of it.""" + b = _base(node) + if isinstance(b, ast.Attribute): + return b.attr == _NAME + return isinstance(b, ast.Name) and (b.id == _NAME or b.id in aliases) + + +def _refers(node, aliases) -> bool: + """Whether any sub-expression (e.g. a branch of a conditional) names the buffer or an alias.""" + return any( + (isinstance(n, ast.Attribute) and n.attr == _NAME) + or (isinstance(n, ast.Name) and (n.id == _NAME or n.id in aliases)) + for n in ast.walk(node) + ) + + +def _write_sites(tree: ast.AST): + sites = [] + + def visit(node, func, aliases): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + func, aliases = node.name, set() + if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + value = node.value + for t in targets: + if isinstance(t, ast.Name) and value is not None and _refers(value, aliases): + aliases.add(t.id) # a local view of the buffer + elif ( + not isinstance(t, ast.Name) + and _mentions(t, aliases) + and not isinstance(node, ast.AnnAssign) + ): + sites.append((func, node.lineno)) + elif ( + isinstance(t, ast.Attribute) + and t.attr == _NAME + and isinstance(node, ast.AnnAssign) + ): + pass # a dataclass field declaration + if isinstance(node, ast.Call): + f = node.func + if isinstance(f, ast.Attribute) and f.attr in _MUTATORS and _mentions(f.value, aliases): + sites.append((func, node.lineno)) + name = ( + f.id + if isinstance(f, ast.Name) + else (f.attr if isinstance(f, ast.Attribute) else "") + ) + if name in _COPY_HELPERS and node.args and _mentions(node.args[0], aliases): + sites.append((func, node.lineno)) + for child in ast.iter_child_nodes(node): + visit(child, func, aliases) + + visit(tree, None, set()) + return sites + + +def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: + found = set() + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if _NAME not in src: + continue + rel = path.relative_to(_ROOT).as_posix() + for func, line in _write_sites(ast.parse(src)): + assert (rel, func) in _ALLOWED, ( + f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " + "wait (ops/m3_select_long_dsl.py, N_VALID_PREWAIT), so a writer inside the " + "step must be ordered before the select's grid or the read moved after the wait" + ) + found.add((rel, func)) + assert found == _ALLOWED, f"expected write sites {_ALLOWED}, found {found}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..f759d0d20d3a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_producer entry: the pre-wait-staged NVFP4 QKV + indexer producer. + +Every cell is held to the native producer (``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert``) bit for bit on q, +index-q and every K / V / scale / index-K cache byte, and to the fp32 torch reference: the FP8 outputs within one E4M3 +step, the decoded NVFP4 K / V within one E2M1 step of their block scale. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_producer import ( # noqa: E501 + FP4_BLOCK, + HEADS_Q, + PAGE, + decode_nvfp4_token, + scale_index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_producer import ( # noqa: E501 + m3_producer as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_producer requires a CUDA device" + +PAGES = 1024 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _caches(): + return ( + torch.zeros(PAGES, 2, 1, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, 1, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _inputs(tokens: int, scale: float, bad_slots: bool, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + weights = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + packed = ( + torch.randn(tokens, (HEADS_Q + 4) * 128, device="cuda", generator=g) * scale + ).bfloat16() + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return packed, slots, quant, weights, rotary, positions + + +def _args(inputs, caches, early: bool = True): + packed, slots, quant, weights, rotary, positions = inputs + return (packed, *caches, slots, quant, EPS, *weights, rotary, positions, HEADS_Q, True, early) + + +def _native(inputs, caches): + packed, slots, quant, weights, rotary, positions = inputs + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, *caches, slots, quant, HEADS_Q, 1, 1, 128, 64, EPS, *weights, rotary, positions + ) + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + early = spec.get("early", True) + inputs = _inputs( + spec["tokens"], spec["scale"], spec.get("bad_slots", False), zlib.crc32(cell.why.encode()) + ) + stock_caches, caches = _caches(), _caches() + stock = _native(inputs, stock_caches) + args = _args(inputs, caches, early) + certifying(op) + ref = op.reference(*args) + out = op(*args) + torch.cuda.synchronize() + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, s, o in zip(names, (*stock, *stock_caches), (*out, *caches)): + assert torch.equal(s.view(torch.uint8), o.view(torch.uint8)), ( + f"{name} differs from the native producer" + ) + + _within_fp8(out[0], ref["q"], "q") + _within_fp8(out[1][:, 0], ref["index_q"], "index_q") + kv_data, kv_scale, index_k = caches + slots = inputs[1] + checked = 0 + for t in range(spec["tokens"]): + slot = int(slots[t]) + if slot < 0 or slot >= PAGES * PAGE: + continue + page, within = divmod(slot, PAGE) + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") + for plane, name in ((0, "k"), (1, "v")): + scales = kv_scale[page, plane, 0].reshape(-1)[scale_index(within, plane).cuda()] + dq = decode_nvfp4_token(kv_data[page, plane, 0, within], scales, KV_QUANT[1 + plane]) + r = ref[name][t].bfloat16().float() + # Round to nearest E2M1 under the block's E4M3-rounded scale s (in value units s / quant_scale): at most + # half the widest E2M1 step (one unit), plus the clipping at 6 units when s rounded down; a scale that + # underflows E4M3 zeroes the block. A slack of one E4M3 step covers the kernel's approximate reciprocals. + qs = KV_QUANT[1 + plane] + amax = r.view(-1, FP4_BLOCK).abs().amax(dim=1) + unit = (qs * amax / 6.0).to(torch.float8_e4m3fn).float() / qs + bound = torch.where( + unit > 0, 1.15 * unit + (amax - 6.0 * unit / 1.15).clamp_min(0.0), amax + ) + bound = bound.repeat_interleave(FP4_BLOCK) + 1e-12 + assert bool(((dq - r).abs() <= bound).all()), (name, t) + checked += 1 + assert checked >= spec["tokens"] - 2 + + +def test_early_trigger_does_not_change_the_bits() -> None: + inputs = _inputs(48, 1.0, False, 7) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=True)) + out_b = op(*_args(inputs, b, early=False)) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +@pytest.mark.parametrize( + "tokens, scale, bad_slots, early", + [ + (1, 1.0, False, True), + (4, 3.0, False, True), + (16, 1e-3, False, True), + (32, 40.0, True, True), + (48, 1.0, False, True), + (1000, 1.0, False, False), + ], +) +def test_prewait_step_inputs_is_bitwise_the_shared_op(tokens, scale, bad_slots, early) -> None: + """dispatch.PRODUCER_PREWAIT_STEP_INPUTS: this target's copy (minimax_m3_sm107_tp4_producer_pw, the position ids, + slots and RoPE rows read before its grid wait) writes exactly the shared op's q, index-q and cache bytes.""" + assert hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp4_producer_pw"), ( + "the target's producer op is not loaded (its TP4 extension)" + ) + inputs = _inputs(tokens, scale, bad_slots, 11) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=early)) + out_b = op(*_args(inputs, b, early=early), prewait_step_inputs=True) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +def test_guard_refuses_a_repeated_slot() -> None: + inputs = list(_inputs(4, 1.0, False, 3)) + inputs[1] = inputs[1].clone() + inputs[1][1] = inputs[1][0] + with pytest.raises(AssertionError, match="must not repeat a slot"): + with validating(op): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_operand_off_the_device() -> None: + inputs = list(_inputs(4, 1.0, False, 9)) + inputs[5] = inputs[5].cpu() # position ids on the host + with pytest.raises(RuntimeError, match="packed's CUDA device"): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_unaligned_index_k_page_stride() -> None: + """The index-K rows are written as 32-bit words: a page stride that is not a multiple of 4 is refused.""" + kv_data, kv_scale, _ = _caches() + backing = torch.zeros(PAGES * (128 * 128 + 1), dtype=torch.float8_e4m3fn, device="cuda") + index_k = backing.as_strided((PAGES, 1, PAGE, 128), (128 * 128 + 1, 128 * 128, 128, 1)) + with pytest.raises(RuntimeError, match="index_k"): + op(*_args(_inputs(4, 1.0, False, 13), (kv_data, kv_scale, index_k))) + + +def test_op_refuses_unaligned_weights_or_rope_table() -> None: + """The norm weights are read as 8-byte words and the RoPE rows as float4: offset views are refused.""" + packed, slots, quant, weights, rotary, positions = _inputs(4, 1.0, False, 17) + offset_weight = torch.zeros(129, dtype=torch.bfloat16, device="cuda")[1:] + with pytest.raises(RuntimeError, match="norm weights"): + op( + *_args( + (packed, slots, quant, [offset_weight, *weights[1:]], rotary, positions), _caches() + ) + ) + flat = torch.zeros(rotary.numel() + 1, device="cuda") + offset_rope = flat[1:].view(rotary.shape) + with pytest.raises(RuntimeError, match="rotary_cos_sin"): + op(*_args((packed, slots, quant, weights, offset_rope, positions), _caches())) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..071795f8c3fa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_qk_norm_rope_fp8 entry of the MiniMax-M3 tp4 target: the out-of-place FP8 form of the fused +per-head Gemma RMSNorm + partial NeoX RoPE against the in-place kernel followed by the E4M3 cast. V is a cast, so it +matches bit for bit; Q and K round once from fp32 where the composition rounds through bf16 first, so each element is +within one E4M3 step of it.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 + m3_qk_norm_rope_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" + +HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 + + +def _args(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q, kv = spec["q"], spec["kv"] + qkv = torch.randn(spec["m"], (q + 2 * kv) * HEAD_DIM, generator=g, device="cuda").to( + torch.bfloat16 + ) + qw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (spec["m"],), generator=g, device="cuda", dtype=torch.int32) + return ( + qkv, + q, + kv, + kv, + HEAD_DIM, + ROTARY_DIM, + EPS, + qw, + kw, + THETA, + True, + pos, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = args[0].clone() + with validating(op): + out = op(*args) + assert torch.equal(args[0], before), "the FP8 form must leave qkv untouched" + ref = op.reference(*args) + assert out.dtype == torch.float8_e4m3fn and out.shape == ref.shape + v_start = (cell.spec["q"] + cell.spec["kv"]) * HEAD_DIM + assert torch.equal(out[:, v_start:].view(torch.uint8), ref[:, v_start:].view(torch.uint8)) + assert _e4m3_steps(out[:, :v_start], ref[:, :v_start]) <= 1 + + +def _e4m3_steps(a, b) -> int: + """The largest distance between two E4M3 tensors in representable values (E4M3 codes order by magnitude).""" + + def signed(t): + c = t.view(torch.uint8).int() + # Sign-magnitude to a line: -0 and +0 both map to 0. + return torch.where(c >= 128, 128 - c, c) + + return int((signed(a) - signed(b)).abs().max()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..a7ee9f877707 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_qkv_insert entry: the generic horizontal producer on an FP8 main cache. + +Q, index-Q and every written K / V / index-K cache row are held within one E4M3 step of the fp32 reference. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_qkv_insert import ( # noqa: E501 + m3_qkv_insert as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" + +PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 16, 1, 1, 128, 64 +MAX_POSITIONS = 262144 +EPS = 1e-6 + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +def _args(tokens: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + tokens, (HEADS + 2 * KV + INDEX + 1) * D, device="cuda", generator=g + ).bfloat16() + kv = torch.zeros(PAGES, 2, KV, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + index_k = torch.zeros(PAGES, 1, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + weights = [(torch.randn(D, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + rotary = torch.randn(MAX_POSITIONS, 2, ROTARY // 2, device="cuda", generator=g) + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + return ( + packed, + kv, + None, + index_k, + slots, + None, + HEADS, + KV, + INDEX, + D, + ROTARY, + EPS, + *weights, + rotary, + positions, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _args(cell.spec["m"], zlib.crc32(cell.why.encode())) + certifying(op) + ref = op.reference(*args) + q, index_q = op(*args) + torch.cuda.synchronize() + _within_fp8(q.view(ref["q"].shape), ref["q"], "q") + _within_fp8(index_q.view(ref["index_q"].shape), ref["index_q"], "index_q") + kv, index_k, slots = args[1], args[3], args[4] + for t in range(cell.spec["m"]): + page, within = divmod(int(slots[t]), PAGE) + _within_fp8(kv[page, 0, :, within], ref["k"][t], f"k token {t}") + _within_fp8(kv[page, 1, :, within], ref["v"][t], f"v token {t}") + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..6f23d21ce10f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_blocks entry of the MiniMax-M3 tp4 target: MSA top-16 block selection with the forced +init / local blocks and per-row valid-block masks.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_blocks import ( # noqa: E501 + m3_select_blocks as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" + +TOPK, INIT, LOCAL = 16, 0, 1 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + n_blocks, total_q = spec["n_blocks"], spec["total_q"] + # Distinct scores (a permutation) so the selection has no ties to break. + scores = ( + torch.randperm(n_blocks * total_q, generator=g, device="cuda") + .float() + .view(1, n_blocks, total_q) + ) + if spec["mixed"]: + n_valid = torch.randint( + 17, n_blocks + 1, (total_q,), generator=g, device="cuda", dtype=torch.int32 + ) + else: + n_valid = torch.full((total_q,), n_blocks, dtype=torch.int32, device="cuda") + with validating(op): + out = op(scores, TOPK, n_valid, INIT, LOCAL, True) + ref = op.reference(scores, TOPK, n_valid, INIT, LOCAL, True) + assert out.shape == ref.shape + assert torch.equal(out, ref) + assert out.permute(1, 0, 2).is_contiguous(), "head-major backing" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..6f4d512ba75e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_long entry (top-16 MSA block selection for long decode rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long import ( # noqa: E501 + m3_select_long as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key +# widths, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 2048, 2049, 4096, 4097, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # re-read path). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + # Odd tokens: almost everything -inf, so picks run past the finite scores. + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +def _mass_tie_rows(scores, nvb, init_blocks, local_blocks): + """[total_q, num_kv_heads] mask of rows where more than 1024 valid blocks share the 16th-largest score.""" + num_kv_heads, num_blocks, total_q = scores.shape + blocks = torch.arange(num_blocks, device=scores.device) + raw = nvb[:total_q].long() + valid = raw.clamp(0, num_blocks) + forced = scores.permute(2, 0, 1).clone() # [total_q, heads, blocks] + forced = torch.where(blocks < init_blocks, 1.0e30, forced) + forced = torch.where(blocks >= (raw - local_blocks).clamp_min(0)[:, None, None], 1.0e29, forced) + in_row = (blocks[None, :] < valid[:, None])[:, None, :] + ranked = torch.where(in_row, forced, float("-inf")).sort(dim=-1, descending=True).values + kth = ranked[..., min(TOPK, num_blocks) - 1 : min(TOPK, num_blocks)] + return ((forced == kth) & in_row).sum(dim=-1) > 1024 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + + # Selection is an exact claim: the same blocks or a different attention. + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # And it is the in-tree selector's answer, bit for bit (the op the target falls back to) -- except on rows + # where more than 1024 blocks tie at the 16th-largest score: there the in-tree op does not break the tie to + # the smallest blocks, while this op and the reference do. + scores, nvb, _, init_blocks, local_blocks, _ = args + in_tree = select_blocks_from_maxscore( + scores, + topk=TOPK, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + mass_tie = _mass_tie_rows(scores, nvb, init_blocks, local_blocks) + assert torch.equal(out[~mass_tie], in_tree[~mass_tie]), ( + "selected blocks differ from the in-tree minimax_m3_select_blocks" + ) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + """Releasing the dependent grid right after the select's own wait changes only when the next kernel may launch.""" + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + """Releasing the dependent grid once the row's scores are loaded changes only when the next kernel may launch.""" + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +@pytest.mark.parametrize("cell", range(len(op.CELLS))) +def test_warp_tail_selects_the_same(cell: int) -> None: + """The warp-0 tail (rows of at most 64 candidates) and the all-thread rank pick the same blocks + (dispatch.SELECT_LONG_SL3 and its kill value).""" + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=7 + cell) + rank_out = torch.empty_like(out) + for phase in (0, 1): + op(scores, nvb, out, init_blocks, local_blocks, True, False, phase, True) + op(scores, nvb, rank_out, init_blocks, local_blocks, True, False, phase, False) + assert torch.equal(out, rank_out) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) + + +def test_guard_refuses_out_aliasing_scores() -> None: + scores, nvb, _, init_blocks, local_blocks, _ = _case(op.CELLS[0].spec, seed=3) + aliased = ( + scores.view(torch.int32) + .reshape(-1)[: scores.shape[2] * TOPK] + .view(scores.shape[2], 1, TOPK) + ) + with pytest.raises(AssertionError, match="alias"): + with validating(op): + op(scores, nvb, aliased, init_blocks, local_blocks, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py new file mode 100644 index 000000000000..816351089a99 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp4 (Rubin) m3_select_long_dsl entry: the long-row top-16 select from one 384-thread CTA per row, +against the reference and bit for bit against the 1024-thread m3_select_long on every cell.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( + m3_select_long_dsl as select_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, + m3_select_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_long_dsl import ( # noqa: E501 + m3_select_long_dsl as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key +# widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 1536, 1537, 4224, 4225, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # whole-row re-rank). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # Every row, mass ties included, is the 1024-thread kernel's answer. + scores, nvb, _, init_blocks, local_blocks, _ = args + long_out = torch.empty_like(out) + m3_select_long(scores, nvb, long_out, init_blocks, local_blocks, False) + assert torch.equal(out, long_out), "selected blocks differ from m3_select_long" + + +@pytest.mark.parametrize("layout", ["kv_blocks_q", "kv_q_blocks"]) +@pytest.mark.parametrize("valid", [17, 1294, 2524, 3907, 9216]) +def test_layouts_and_widths_match_the_1024_thread_select(layout, valid) -> None: + """Both score layouts the scorer can hand over, at a valid count in each key width.""" + g = torch.Generator(device="cuda").manual_seed(valid) + nb = max(valid, 129) + s = torch.randn(1, nb, 8, generator=g, device="cuda") + if layout == "kv_q_blocks": + s = s.permute(0, 2, 1).contiguous().permute(0, 2, 1) + nvb = torch.tensor([valid - (t % 2) for t in range(8)], dtype=torch.int32, device="cuda") + out = torch.empty(1, 8, TOPK, dtype=torch.int32, device="cuda").permute(1, 0, 2) + ref = torch.empty_like(out) + op(s, nvb, out, 0, 1, False) + m3_select_long(s, nvb, ref, 0, 1, False) + assert torch.equal(out, ref) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +@pytest.mark.parametrize("early_trigger", [False, True]) +@pytest.mark.parametrize("cell", range(len(op.CELLS))) +def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: + """ops/m3_select_long_dsl.N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the grid + wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and tied + rows).""" + phase = 0 if early_trigger else 1 + spec = op.CELLS[cell].spec + scores, nvb, before, init_blocks, local_blocks, _ = _case(spec, seed=11 + cell) + after = torch.empty_like(before) + select_kernels._compile((True, early_trigger, phase, True))( + scores, nvb, before, init_blocks, local_blocks + ) + select_kernels._compile((True, early_trigger, phase, False))( + scores, nvb, after, init_blocks, local_blocks + ) + assert torch.equal(before, after) + + +def test_graph_replay_selects_the_same() -> None: + spec = op.CELLS[1].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=9) + ref = op.reference(scores, nvb, out, init_blocks, local_blocks) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + for _ in range(3): + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..1c72256dac47 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,442 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode entry (the target's Triton block-sparse decode and its routes).""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_report, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode as stock_decode, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import ( # noqa: E402 — needs the path declared above + HEAD_DIM, + NUM_HEADS, + build_step, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_sparse_decode requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each [16, 128] head row's largest magnitude, and of the relative RMS. Every route +# computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference +# is the BF16 rounding on both sides (one ulp of the row at most), the FP16 probabilities before PV, the split-K +# partials kept in FP32 and merged by their log2-sum-exp, and approximate exp2/log2: 2 ulps per element, the two +# roundings (each at most half an ulp) dominating the RMS, so half an ulp. Measured on GB300 over the cells: at most +# 1.80 ulps per element and 0.20 ulps RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, output, identity): + return op( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + identity, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=spec["identity"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + certifying(op) + _call(case, output, spec["identity"]) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + spec["identity"], + ) + assert torch.isfinite(output.float()).all(), "padding rows must be zeros, not NaN" + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_flag_does_not_change_the_answer() -> None: + """On an identity table the identity routes (the producer, or the multi-query producer from 17 rows) and the + selected-table route (the long-context kernel) compute the same attention.""" + case = build_step((2048, 130, 5, 0, 1777, 900, 2000, 0), 4, identity=True, seed=17) + outs = [] + for identity in (True, False): + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, identity) + outs.append(output) + assert_within_ulp(outs[0], outs[1], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_guard_refuses_a_padded_selection_slot() -> None: + """The Triton producers index the block table with every one of a token's first min(16, blocks) ids, + -1 included.""" + case = build_step((1500,), 1, identity=True, seed=23) + case.topk[0, 0, 3] = -1 + output = torch.empty(1, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="blocks of the request"): + with validating(op): + _call(case, output, True) + + +def test_guard_refuses_a_false_identity_promise() -> None: + """identity_blocks lets the multi-query producer read one table for the whole request.""" + case = build_step((1500, 1600, 1700, 1800, 1500, 1600, 1700, 1800), 4, identity=True, seed=29) + case.topk[0, 5, :2] = case.topk[0, 5, :2].flip(0) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="identity_blocks promises"): + with validating(op): + _call(case, output, True) + + +def _stock(case, output): + stock_decode( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + sm_scale=case.sm_scale, + output=output, + decode_query_len=case.decode_query_len, + k_block_scale=case.k_block_scale, + v_block_scale=case.v_block_scale, + k_global_scale=case.k_global_scale, + v_global_scale=case.v_global_scale, + ) + + +def _identity_table( + seq_lens: torch.Tensor, decode_query_len: int, heads: int, lag: int = 0 +) -> torch.Tensor: + """The identity selection of every verify token, staged from lengths `lag` tokens behind seq_lens.""" + seq = (seq_lens.cpu().to(torch.int64) - lag).clamp(min=0).repeat_interleave(decode_query_len) + intra = torch.arange(seq.shape[0]) % decode_query_len + n_valid = ((seq - decode_query_len + intra + 1).clamp(min=0) + 127) // 128 + ar = torch.arange(16) + table = torch.where(ar[None, :] < n_valid[:, None], ar[None, :], -1).to(torch.int32) + return table.unsqueeze(0).expand(heads, -1, -1).contiguous().cuda() + + +# Each crosses a 128-token edge within its four verify tokens or between the stale and the true lengths. +_EDGE_LENS = (1284, 1287, 1281, 129, 132, 2048, 1028, 5, 643, 1920, 256, 1) + + +@pytest.mark.parametrize("fault", ["stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_identity_route_reads_its_own_lengths_not_the_table(requests, fault, monkeypatch) -> None: + """The Triton identity route (single / pair producer, up to 16 rows) derives block i for slot i from the kv lengths + the kernel reads: a staged table that lags the lengths by a verify step (a 128-token edge crossed), or is garbage, + gives the stock decode's rows over the true identity selection bit for bit. (The LC identity route has its own test + below.)""" + _triton_route(monkeypatch) + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=400 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + assert torch.equal(case.topk, true_table), ( + "build_step's identity table is the identity selection" + ) + if fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _stock(case, stock) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + _call(case_staged, target, True) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("requests", [1, 2, 4]) +def test_identity_route_is_bitwise_the_stock_decode(requests, monkeypatch) -> None: + """The target's Triton producer + split-K merge takes the stock in-tree decode's split-K at every row count + (4 / 8 / 16 identity rows here) and so writes its rows bit for bit.""" + _triton_route(monkeypatch) + case = build_step((1500, 2048, 900, 130)[:requests], 4, identity=True, seed=60 + requests) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _call(case, target, True) + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), case.seq_lens.tolist() + + +def _triton_route(monkeypatch) -> None: + """Route selected steps to the target's Triton producer + merge (this target takes the LC whenever it applies).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( + m3_attn_lc, + ) + + monkeypatch.setattr(m3_attn_lc, "lc_applies", lambda *args, **kwargs: False) + + +def _short_lens(requests: int, g: torch.Generator) -> tuple: + """Short contexts of 1..16 blocks of 128 tokens (the identity-route regime), a few exactly on a block edge.""" + lens = torch.randint(4, 16 * 128 + 1, (requests,), generator=g).tolist() + if requests > 1: + lens[0] = 128 * int(torch.randint(1, 17, (1,), generator=g)) + return tuple(lens) + + +@pytest.mark.parametrize("seed", range(12)) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_short_table_stress_is_bitwise_the_stock_decode(requests, seed, monkeypatch) -> None: + """The fail-closed id / page guard leaves valid inputs alone: short contexts (1..16 blocks) at 1..4 requests x 4 + verify rows, the identity route and a random selected route with -1 tails past each token's blocks, write the + stock decode's rows bit for bit.""" + _triton_route(monkeypatch) + g = torch.Generator().manual_seed(1000 * requests + seed) + for identity in (True, False): + case = build_step(_short_lens(requests, g), 4, identity=identity, seed=7 * seed + requests) + target = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + stock = torch.empty_like(target) + _call(case, target, identity) + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + identity, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("fault", ["neg_in_window", "id_past_row", "page_out_of_pool"]) +def test_invalid_slot_fails_closed(fault, monkeypatch) -> None: + """A corrupt decode input (a -1 inside a token's block window, an id past the block-table row, or a block-table + page outside the pool) is read as an absent block instead of a wild KV read: the rows equal the reference over the + remaining blocks.""" + _triton_route(monkeypatch) + case = build_step((1500, 2048, 700), 4, identity=False, seed=303) + tok, slot = 5, 2 + if fault == "neg_in_window": + case.topk[0, tok, slot] = -1 + elif fault == "id_past_row": + case.topk[0, tok, slot] = case.block_table.shape[1] + 7 + else: + req = tok // 4 + blk = int(case.topk[0, tok, slot]) + case.block_table = case.block_table.clone() + case.block_table[req, blk] = case.k_paged.shape[0] + 1000 + # Every token of the request that selected that block loses it. + hits = (case.topk[:, req * 4 : req * 4 + 4] == blk).nonzero().tolist() + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(case, output, False) + torch.cuda.synchronize() + # The reference walks each token's first min(16, blocks) entries and drops -1: the corrupt slots as -1. + ref_topk = case.topk.clone() + if fault == "page_out_of_pool": + for h, t, k in hits: + ref_topk[h, req * 4 + t, k] = -1 + else: + ref_topk[0, tok, slot] = -1 + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + ref_topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert torch.isfinite(output.float()).all() + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +@pytest.mark.parametrize("fault", ["none", "stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_lc_identity_route_reads_its_own_lengths_not_the_table( + requests, fault, monkeypatch +) -> None: + """An identity step runs the long-context kernel, which takes block i for CTA i from the kv lengths it reads. + Whatever the staged table holds (true, lagging a verify step, all -1, garbage), the rows are the LC's on the true + identity table (the selected-table route) bit for bit, and in the band of the float64 reference.""" + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=700 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + if fault == "none": + staged = true_table.clone() + elif fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + selected = torch.empty_like(target) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + case_true = SimpleNamespace(**{**vars(case), "topk": true_table}) + taken = [] + monkeypatch.setattr(m3_report, "taken", lambda path, rows, what="": taken.append(path)) + _call(case_staged, target, True) + monkeypatch.undo() + # The LC engagement, plus its one-division merge's. + assert sorted(taken) == sorted(["m3_sparse_decode_lc", "sparse_decode_lc_one_div_merge"]), taken + _call(case_true, selected, False) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), selected.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + true_table, + case.block_table, + case.seq_lens, + case.sm_scale, + target, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert_within_ulp(target, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +@pytest.mark.parametrize("requests", [1, 2]) +def test_lc_identity_step_reads_q_after_the_wait(requests, monkeypatch) -> None: + """An identity step has no select, so the LC must read q after its grid wait: the op gets identity_step (and the + device identity selection) on every identity step.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_attn_lc, + m3_prequant, + ) + + lens = tuple(_EDGE_LENS[(requests * 3 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=900 + requests) + true_table = _identity_table(case.seq_lens, 4, case.topk.shape[0]) + case_true = SimpleNamespace(**{**vars(case), "topk": true_table}) + rows = case.q.shape[0] + seen = [] + real = m3_attn_lc.m3_sparse_decode_lc + + def spy(*args, **kwargs): + seen.append((kwargs.get("identity", False), kwargs.get("identity_step", False))) + return real(*args, **kwargs) + + monkeypatch.setattr(m3_attn_lc, "m3_sparse_decode_lc", spy) + out = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(case_true, out, True) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert seen[0] == (True, True), seen + assert bool(torch.isfinite(out.float()).all()) + + +@pytest.mark.parametrize("requests", [2, 3]) +def test_lc_fails_closed_on_ids_past_the_row(requests) -> None: + """A selected id past the token's row (a stale table, e.g. a padded row keeping a 500K request's n_valid) is + skipped, as the Triton producer skips it: the rows equal those with that slot at -1, bit for bit, and stay finite + even though the block-table tail it would reach points at a page whose block scales are NaN (0xFF).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_prequant, + ) + + lens = (165521, 700, 3000)[:requests] + case = build_step(lens, 4, identity=False, seed=1300 + requests) + dql = case.decode_query_len + good = case.topk.clone() + bad = case.topk.clone() + for tok in range(case.q.shape[0]): + req, qoff = divmod(tok, dql) + kv_len = max(int(case.seq_lens[req]) - dql + qoff + 1, 0) + nblk = (kv_len + 127) // 128 + past = nblk + 1 + if 2 <= nblk < 16 and past < case.block_table.shape[1]: + bad[:, tok, nblk - 1] = past + good[:, tok, nblk - 1] = -1 + page = int(case.block_table[req, past]) + case.k_block_scale[page].fill_(0xFF) + case.v_block_scale[page].fill_(0xFF) + assert not torch.equal(bad, good) + outs = [] + for table in (bad, good): + out = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(SimpleNamespace(**{**vars(case), "topk": table}), out, False) + outs.append(out) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert torch.isfinite(outs[0].float()).all() + assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py new file mode 100644 index 000000000000..0c0becea2433 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py @@ -0,0 +1,526 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode_lc entry (one-launch long-context sparse decode).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + WIDTH, + dequant_nvfp4_pages, + mxfp8_handoff_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode_lc import ( # noqa: E501 + m3_sparse_decode_lc as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import build_step # noqa: E402 — needs the path declared above +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_sparse_decode_lc requires a CUDA device" + +#: Hand-off buffers are sized as MERGE_MXFP8 sizes them. +_HANDOFF_ROWS = 128 + +# Band, in BF16 ulps (2^-8) of each output row's largest magnitude, and of the relative RMS. +# +# The reference is float64 attention over the exactly dequantized blocks, rounded once to BF16. The kernel's +# operands are exact too (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so what separates the two is: the final +# BF16 rounding on both sides (up to one ulp of the element, at most one of the row's scale), the softmax +# probabilities rounded to FP16 before PV (2^-11 relative each, averaging down over the block), ex2.approx / +# lg2.approx (~2^-22) and FP32 accumulation order. Hence 2 ulps per element; the RMS is dominated by the two +# roundings (each at most half an ulp, about 0.3 in aggregate), so half an ulp. Measured on GB300 over the cells: +# at most 1.91 ulps per element and 0.25 ulps RMS. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, spec, out, q8, sf): + return op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + spec["linear_softmax"], + False, + True, + ) + + +def _buffers(total_q, mxfp8): + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + if not mxfp8: + return out, None, None + q8 = torch.zeros(_HANDOFF_ROWS, WIDTH, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_HANDOFF_ROWS * (WIDTH // 32), dtype=torch.uint8, device="cuda") + return out, q8, sf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=False, + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + total_q = case.q.shape[0] + out, q8, sf = _buffers(total_q, spec["mxfp8"]) + certifying(op) + _call(case, spec, out, q8, sf) + ref = op.reference( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + spec["linear_softmax"], + False, + True, + ) + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + # Padding rows attend nothing and must be exact zeros, not NaN. + assert torch.isfinite(out.float()).all() + if spec["mxfp8"]: + # The hand-off is a deterministic function of the op's own BF16 rows: exact. + ref_q8, ref_sf = mxfp8_handoff_reference(out, _HANDOFF_ROWS) + assert torch.equal(q8[:total_q], ref_q8), ( + "MXFP8 rows differ from the recipe applied to the BF16 rows" + ) + assert torch.equal(sf, ref_sf), "UE8M0 scales differ from the recipe / swizzle" + + +def test_reference_dequant_matches_the_quantizer() -> None: + """The reference reads the paged bytes through its own layout claims (K scales token-major, V scales in + 4x4 token quads). Hold them to the quantizer's own flat output: every slot was written in order, so slot + ``s`` is page ``s // 128``, token ``s % 128``, and the flat scales are token-major by construction.""" + case = build_step((3000,), 1, identity=False, seed=5) + pages = torch.arange(case.k_paged.shape[0], device="cuda") + views = ((case.k_paged, case.k_block_scale, False), (case.v_paged, case.v_block_scale, True)) + for role, (data, scale, quad) in enumerate(views): + flat_data, flat_scale = torch.ops.trtllm.fp4_quantize( + case.source[role].reshape(case.num_slots, 1, -1).contiguous(), + case.quant[role + 1 : role + 2], + 16, + False, + False, + ) + expected = dequant_nvfp4_pages( + flat_data.view(torch.uint8).reshape(-1, 1, 128, 64), + flat_scale.view(torch.uint8).reshape(-1, 1, 128, 8), + pages, + quad_scales=False, + ) + got = dequant_nvfp4_pages(data, scale, pages, quad_scales=quad) + assert torch.equal(got, expected), ( + f"role {role}: paged dequantization disagrees with the quantizer" + ) + + +def test_the_variants_agree() -> None: + """The two softmax recurrences and the hand-off switch must not change the BF16 rows beyond rounding.""" + spec = dict(seq_lens=(5000, 12345), decode_query_len=4) + case = build_step(spec["seq_lens"], 4, identity=False, seed=11) + outs = [] + for linear in (False, True): + for mxfp8 in (False, True): + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + _call(case, dict(linear_softmax=linear), out, q8, sf) + outs.append(out) + for other in outs[1:]: + assert_within_ulp(other, outs[0], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + torch.testing.assert_close(outs[1], outs[0], rtol=0, atol=0) # hand-off on/off: same rows + + +def test_op_refuses_the_uninstantiated_late_trigger() -> None: + case = build_step((3000,), 1, identity=False, seed=3) + out, q8, sf = _buffers(1, True) + with pytest.raises(RuntimeError, match="early_trigger"): + op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + 1, + case.sm_scale, + True, + False, + False, + ) + + +def test_guard_refuses_a_block_past_the_request() -> None: + """A selected id at or past the request's block count reads a page the token cannot attend; with every + position masked the CTA's softmax is NaN.""" + case = build_step((3000,), 1, identity=False, seed=4) + case.topk[0, 0, 0] = 40 + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="block count"): + with validating(op): + _call(case, dict(linear_softmax=False), out, q8, sf) + + +def test_guard_refuses_strided_page_rows() -> None: + """The kernel walks a page's token rows at a fixed 64-byte pitch; a padded row stride is misread.""" + case = build_step((3000,), 1, identity=False, seed=6) + wide = torch.zeros(*case.k_paged.shape[:3], 80, dtype=torch.uint8, device="cuda")[..., :64] + wide.copy_(case.k_paged) + case.k_paged = wide + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="64-byte rows"): + with validating(op): + _call(case, dict(linear_softmax=False), out, q8, sf) + + +def test_guard_refuses_a_narrow_handoff() -> None: + case = build_step((3000,), 1, identity=False, seed=7) + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="num_sf_k_tiles"): + with validating(op): + op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + 8, + 1, + case.sm_scale, + False, + False, + True, + ) + + +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + [ + ((165521,), 4), + ((5000, 12345), 4), + ((700,), 1), + ((1500, 300, 90000), 2), + ((200000, 3000, 64, 70000), 4), + ], + ids=["1x165k", "2req", "short-1row", "3req-mixed", "4req-incl-64"], +) +def test_merge_once_is_bitwise_the_image_op(seq_lens, decode_query_len) -> None: + """dispatch.SPARSE_DECODE_LC_MERGE_ONCE: the target's op (merge weights formed once per CTA) writes exactly the + image op's rows and hand-off, at both softmax recurrences, including requests with fewer than 16 selected blocks + (empty chunks, -inf LSE). With SPARSE_DECODE_LC_PREFETCH_OPROJ it also pulls two ranges the size of the TP4 + o_proj weight and scales (one not a multiple of 16 bytes) into L2: the same bits again. With + SPARSE_DECODE_LC_PREFETCH_NEXT_KV it records each row's selection (the first min(16, blocks) ids, -1 after) and + prefetches a next layer's recorded blocks, stale ids past a request's blocks included: the same bits, and the + record.""" + assert hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp4_sparse_decode_lc"), ( + "the target's sparse decode op is not loaded (its TP4 extension)" + ) + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + ) + o_proj_weight = torch.empty(6144, 2048, dtype=torch.uint8, device="cuda") + o_proj_scale = torch.empty(6144 * 64 + 9, dtype=torch.uint8, device="cuda") + rows = case.q.shape[0] + n1_self = torch.full((128, 16), -7, dtype=torch.int32, device="cuda") + n1_next = torch.randint( + -1, 2 * (max(seq_lens) // 128 + 1), (128, 16), dtype=torch.int32, device="cuda" + ) + next_kv = (n1_self, n1_next, case.k_paged, case.v_paged, case.k_block_scale, case.v_block_scale) + variants = ( + (False, None, None), + (True, None, None), + (True, (o_proj_weight, o_proj_scale), None), + (True, (o_proj_weight, o_proj_scale), next_kv), + ) + for linear in (False, True): + for mxfp8 in (False, True): + got = [] + for merge_once, prefetch, n1 in variants: + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + linear, + False, + True, + merge_once=merge_once, + prefetch=prefetch, + next_kv=n1, + ) + got.append((out, q8, sf)) + torch.cuda.synchronize() + a, aq, asf = got[0] + for i, (b, bq, bsf) in enumerate(got[1:], 1): + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), (linear, mxfp8, i) + if mxfp8: + assert torch.equal(aq, bq) and torch.equal(asf, bsf), (linear, mxfp8, i) + # The record: each row's first min(16, blocks) selected ids, -1 past them; rows past the step untouched. + want = case.topk[0].clone() + for tok in range(rows): + req, qoff = divmod(tok, case.decode_query_len) + kv_len = max(int(case.seq_lens[req]) - case.decode_query_len + qoff + 1, 0) + want[tok, min(16, (kv_len + 127) // 128) :] = -1 + assert torch.equal(n1_self[:rows], want), "the recorded selection" + assert bool((n1_self[rows:] == -7).all()), "rows past the step" + + +def _tp4_op(case, out, q8, sf, linear, use_pdl, one_div): + torch.ops.trtllm.minimax_m3_sm107_tp4_sparse_decode_lc( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + linear, + use_pdl, + True, + one_div=one_div, + ) + + +_ONE_DIV_STEPS = [ + ((165521,), 4), + ((5000, 12345), 4), + ((700,), 1), + ((1500, 300, 90000), 2), + ((200000, 3000, 64, 70000), 4), +] + + +@pytest.mark.parametrize("linear", [False, True]) +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + _ONE_DIV_STEPS, + ids=["1x165k", "2req", "short-1row", "3req-mixed", "4req-incl-64"], +) +def test_one_div_merge_is_deterministic_and_in_band(seq_lens, decode_query_len, linear) -> None: + """The one-division merge (on in this target): it writes the same rows on every run + (5 eager runs, PDL on and off) and under CUDA-graph replay (3 replays == eager), MXFP8 hand-off included; the rows + are within the float64 reference's band, and within summation-order distance of the per-chunk merge.""" + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + 1 + ) + total_q = case.q.shape[0] + runs = [] + for i in range(6): + out, q8, sf = _buffers(total_q, True) + _tp4_op(case, out, q8, sf, linear, i == 5, True) + runs.append((out, q8, sf)) + torch.cuda.synchronize() + a, aq, asf = runs[0] + for b, bq, bsf in runs[1:]: + assert torch.equal(a.view(torch.int16), b.view(torch.int16)) + assert torch.equal(aq, bq) and torch.equal(asf, bsf) + out = torch.empty_like(a) + _tp4_op(case, out, None, None, linear, True, True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _tp4_op(case, out, None, None, linear, True, True) + for _ in range(3): + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), a.view(torch.int16)) + per, _, _ = _buffers(total_q, False) + _tp4_op(case, per, None, None, linear, False, False) + ref = op.reference( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + a, + None, + None, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + linear, + False, + True, + ) + assert torch.isfinite(a.float()).all() + assert_within_ulp(a, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + assert_within_ulp(a, per, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +_BT_PREWAIT_STEPS = [ + ((165521,), 4), + ((5000, 12345), 4), + ((700,), 1), + ((1500, 300, 90000), 2), + ((200000, 3000, 64, 70000), 4), + ((524288,), 4), # exactly the 4096-entry staging cap + ((524289, 2000), 4), # one block past the cap: the global lookup, beside a staged request + ((600000,), 1), +] + + +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + _BT_PREWAIT_STEPS, + ids=[ + "1x165k", + "2req", + "short-1row", + "3req-mixed", + "4req-incl-64", + "at-cap", + "past-cap+short", + "600k-1row", + ], +) +def test_bt_prewait_is_bitwise(seq_lens, decode_query_len) -> None: + """The op's ``bt_prewait`` option (not used by this target): staging the request's block-table row in shared memory + before the grid wait writes exactly the rows and MXFP8 hand-off of the global lookup, at both softmax recurrences, + with and without the one-division merge, PDL on and off, eager and under CUDA-graph replay; rows at the 4096-entry + cap are staged, longer ones keep the global lookup; the o_proj / next-KV prefetch hints alongside change nothing.""" + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + 7 + ) + total_q = case.q.shape[0] + n1_self = torch.full((128, 16), -7, dtype=torch.int32, device="cuda") + n1_next = torch.randint( + -1, 2 * (max(seq_lens) // 128 + 1), (128, 16), dtype=torch.int32, device="cuda" + ) + next_kv = (n1_self, n1_next, case.k_paged, case.v_paged, case.k_block_scale, case.v_block_scale) + + def run(out, q8, sf, linear, one_div, bt, use_pdl, hints=False): + torch.ops.trtllm.minimax_m3_sm107_tp4_sparse_decode_lc( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + WIDTH // 128, + case.decode_query_len, + case.sm_scale, + linear, + use_pdl, + True, + *((None, None) + next_kv if hints else ()), + one_div=one_div, + bt_prewait=bt, + ) + + for linear in (False, True): + for one_div in (False, True): + base = _buffers(total_q, True) + run(*base, linear, one_div, False, False) + arms = [] + for bt, use_pdl, hints in ( + (True, False, False), + (True, True, False), + (True, True, True), + ): + bufs = _buffers(total_q, True) + run(*bufs, linear, one_div, bt, use_pdl, hints) + arms.append(bufs) + torch.cuda.synchronize() + a, aq, asf = base + for i, (b, bq, bsf) in enumerate(arms): + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), (linear, one_div, i) + assert torch.equal(aq, bq) and torch.equal(asf, bsf), (linear, one_div, i) + out = torch.empty_like(a) + run(out, None, None, linear, one_div, True, True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(out, None, None, linear, one_div, True, True) + for _ in range(3): + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), a.view(torch.int16)), ( + linear, + one_div, + "graph", + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py new file mode 100644 index 000000000000..013e12cd3703 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for ops/m3_tokmap_guard: a one-request multi-token call of +thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a +stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the +appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for +every length. Without the guard the stale map moves rows (the negative control), and a several-request call passes +through with its own map untouched. + +The fresh-map reference forces the op's own build (a scaled RoPE type with no inv-freq cache, at rotary dim 0, so no +row is rotated either way). +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_tokmap_guard as guard, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = 16, 1, 128, 128 +BUCKET = 4 +_PEER_EXTRA = 777 # the second request's extra length behind the first +FP8_KV_QUANT_MODE = 128 +_DEV = torch.device("cuda") # unindexed on purpose: install() and the lookup normalize it +_TAKEN = [] + + +@pytest.fixture(scope="module", autouse=True) +def _installed(): + guard.install(_DEV, 48, _TAKEN.append) + yield + + +class _Step: + """One request (or ``nreq``) of BUCKET padded rows over a fresh FP8 paged cache, persistent across calls like a + real step's buffers.""" + + def __init__(self, kv_len: int, nreq: int = 1, seed: int = 0): + g = torch.Generator(device=_DEV).manual_seed(seed + kv_len) + self.nreq = nreq + # Per-request inputs (page tables, generation lengths, position offsets) for at least two requests, the page + # tables at the peer's longer length (_seq_with_peer): the stale map sends rows to request 1, which must have + # real entries behind it, also on the unguarded calls. + tables = max(nreq, 2) + pages = math.ceil((kv_len + _PEER_EXTRA + BUCKET) / PAGE) + self.nblocks = tables * pages + 2 + self.pool0 = ( + torch.randn(self.nblocks, 2, KVH, PAGE, D, device=_DEV, generator=g) + .clamp_(-4, 4) + .to(torch.float8_e4m3fn) + ) + self.pool = self.pool0.clone() + perm = torch.randperm(self.nblocks, device=_DEV, generator=g)[: tables * pages].view( + tables, pages + ) + self.offsets = torch.zeros(1, tables, 2, pages, dtype=torch.int32, device=_DEV) + self.offsets[0, :, 0] = 2 * perm + self.offsets[0, :, 1] = 2 * perm + 1 + self.pointers = torch.tensor([[self.pool.data_ptr(), 0]], dtype=torch.int64) + self.mapping = torch.zeros(1, 2, dtype=torch.int32) + self.qkv = torch.randn( + nreq * BUCKET, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=_DEV, generator=g + ) + # Request lengths after this step's append; generation lengths (accepted rows) set per case. + self.seq = torch.full((nreq,), kv_len, dtype=torch.int32, device=_DEV) + self.gen = torch.full((tables,), BUCKET, dtype=torch.int32, device=_DEV) + self.pos = ( + torch.arange(BUCKET, dtype=torch.int32, device=_DEV).repeat(tables, 1).contiguous() + ) + self.one = torch.ones(1, dtype=torch.float32, device=_DEV) + self.ws = torch.empty(64 << 20, dtype=torch.uint8, device=_DEV) + self.kv_len = kv_len + + def call(self, op, fresh_map: bool = False): + """One generation preprocess; returns a copy of the processed query.""" + out = op( + self.qkv, + self.ws, + self.seq, + self.gen, + self.pos, + self.offsets, + self.pointers, + self.mapping, + self.one, + self.one, + None, + None, + None, + None, + 0, + 0, + H, + KVH, + D, + PAGE, + FP8_KV_QUANT_MODE, + 1 << 20, + 1 << 20, + self.nreq * BUCKET, + self.nreq, + BUCKET, + self.kv_len, + 0, + 10000.0, + 1 if fresh_map else 0, + 1.0, + 1 << 20, + 0, + 1.0, + 1.0, + False, + BUCKET, + 0, + torch.cuda.get_device_properties(0).multi_processor_count, + self.nblocks, + 2, + True, + False, + ) + return out[0].clone() + + def poison_map(self) -> None: + """A stale multi-request map in the op's one-request map slice: rows 1.. on request 1.""" + layout = thop.get_trtllm_gen_generation_workspace_layout( + self.qkv.dtype, 1, self.nreq * BUCKET, H, D, 0, KVH, 0, False, False + ) + off = int(layout["tokens_info_offset"]) + stale = torch.tensor( + [[0, 0]] + [[1, t] for t in range(BUCKET - 1)], dtype=torch.int32, device=_DEV + ) + self.ws[off : off + 8 * BUCKET].view(torch.int32).view(BUCKET, 2).copy_(stale) + + def reset(self) -> None: + self.pool.copy_(self.pool0) + + +def _seq_with_peer(step: _Step) -> None: + # A second request's length behind the first, so a row the stale map moves to request 1 reads a real length. + step.seq = torch.tensor( + [step.kv_len, step.kv_len + _PEER_EXTRA], dtype=torch.int32, device=_DEV + ) + + +@pytest.mark.parametrize("kv_len", [130, 8191, 34117]) +@pytest.mark.parametrize("n", [1, 2, 3, 4]) +def test_guard_equals_a_fresh_map(kv_len, n) -> None: + step = _Step(kv_len) + _seq_with_peer(step) + step.gen.fill_(n) + step.poison_map() + step.reset() + q_ref = step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + + step.poison_map() + step.reset() + q = step.call(thop.trtllm_gen_generation_preprocess) + # The processed query's n live rows (rows past the request's length are padding the op leaves unwritten). + assert torch.equal(q[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)) + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + assert BUCKET in _TAKEN + + +def test_stale_map_moves_rows_without_the_guard() -> None: + step = _Step(34117) + _seq_with_peer(step) + step.gen.fill_(BUCKET) + step.poison_map() + step.reset() + step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + step.poison_map() + step.reset() + step.call(guard._original) + assert not torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + + +def test_one_captured_graph_serves_every_length() -> None: + step = _Step(34000) + _seq_with_peer(step) + refs = {} + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + q = step.call(guard._original, fresh_map=True) + refs[n] = (q, step.pool.clone()) + step.gen.fill_(BUCKET) + step.call(thop.trtllm_gen_generation_preprocess) # eager first, as the warmup runs it + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = thop.trtllm_gen_generation_preprocess + q_static = step.call(out) + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + g.replay() + torch.cuda.synchronize() + q_ref, pool_ref = refs[n] + assert torch.equal(q_static[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)), n + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)), n + + +def test_several_requests_pass_through() -> None: + step = _Step(8191, nreq=2) + layout = thop.get_trtllm_gen_generation_workspace_layout( + step.qkv.dtype, 2, 2 * BUCKET, H, D, 0, KVH, 0, False, False + ) + step.reset() + step.call(thop.trtllm_gen_generation_preprocess) + pool_guarded = step.pool.clone() + step.reset() + step.call(guard._original) + assert torch.equal(step.pool.view(torch.uint8), pool_guarded.view(torch.uint8)) + assert int(layout["tokens_info_size"]) >= 8 * 2 * BUCKET diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..a1682f829082 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the msa_kv_update entry (the one-launch MSA KV-length correction).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.msa_kv_update import ( # noqa: E501 + TOKENS_PER_BLOCK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" + +#: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. +_SPARE = 3 + + +def _case(spec, seed): + """The metadata's buffers for one step, with stale values in every output so a skipped write shows.""" + g = torch.Generator(device="cpu").manual_seed(seed) + batch, qlen = spec["batch"], spec["query_len"] + total_q = batch * qlen + staged = torch.randint(qlen, 5000, (batch,), generator=g, dtype=torch.int32) + accept = spec["accept"] + if accept == "full": + kv = staged.clone() + elif accept == "page_edge": + staged = torch.tensor([130, 257, 2 * TOKENS_PER_BLOCK + qlen][:batch], dtype=torch.int32) + kv = staged - torch.tensor([0, 1, qlen - 1][:batch], dtype=torch.int32) + else: + kv = staged - torch.randint(0, qlen, (batch,), generator=g, dtype=torch.int32) + if accept == "padded": + kv[batch // 2 :] = 0 + staged[batch // 2 :] = 0 + ncols = (5000 + qlen) // TOKENS_PER_BLOCK + 2 + table = ( + torch.randperm(batch * ncols + 7, generator=g)[: batch * ncols] + .to(torch.int32) + .reshape(batch, ncols) + ) + dev = dict(device="cuda", dtype=torch.int32) + stale = -12345 + return ( + torch.cat([kv, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.cat([staged, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.full((batch + _SPARE,), qlen, **dev), + torch.arange(batch, **dev).repeat_interleave(qlen), + torch.arange(qlen, **dev).repeat(batch), + table.cuda(), + TOKENS_PER_BLOCK, + batch, + total_q, + torch.full((batch + _SPARE,), stale, **dev), + torch.full((batch + 1 + _SPARE,), stale, **dev) if spec["cu_kv"] else None, + torch.full((total_q + _SPARE,), stale, **dev), + torch.full((total_q + _SPARE,), stale, **dev) if spec["n_valid"] else None, + ) + + +def _check(args, ref): + seq, cu_kv, loc, n_valid = args[9], args[10], args[11], args[12] + batch, total_q = args[7], args[8] + ref_seq, ref_cu, ref_loc, ref_nv = ref + # Integer bookkeeping: exact, and nothing past the live rows is touched. + assert torch.equal(seq[:batch], ref_seq) + assert torch.equal(loc[:total_q], ref_loc) + assert bool((seq[batch:] == -12345).all()) and bool((loc[total_q:] == -12345).all()) + assert (cu_kv is None) == (ref_cu is None) and (n_valid is None) == (ref_nv is None) + if cu_kv is not None: + assert torch.equal(cu_kv[: batch + 1], ref_cu) + assert bool((cu_kv[batch + 1 :] == -12345).all()) + if n_valid is not None: + assert torch.equal(n_valid[:total_q], ref_nv) + assert bool((n_valid[total_q:] == -12345).all()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + op(*args) + _check(args, op.reference(*args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_hook_takes_every_cell(cell) -> None: + """The target's hook (what the metadata calls) takes each certified step and writes what the op does.""" + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + 1) + assert hook(*args) is True + _check(args, op.reference(*args)) + + +def test_the_hook_declines_what_the_op_would_misread() -> None: + """A strided operand is read wrongly by the op, so the hook declines it (the metadata then runs its own + torch update) and writes nothing.""" + args = list(_case(op.CELLS[2].spec, seed=5)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + before = args[11].clone() + assert hook(*args) is False + assert torch.equal(args[11], before) + + +def test_guard_refuses_a_strided_operand() -> None: + args = list(_case(op.CELLS[2].spec, seed=6)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_index_past_the_batch() -> None: + args = list(_case(op.CELLS[2].spec, seed=7)) + args[3] = args[3].clone() + args[3][-1] = args[7] + with pytest.raises(AssertionError, match="q_batch_row"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py new file mode 100644 index 000000000000..127859c5cb2a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the ``m3_allreduce`` entry, 4 ranks: its cells against the entry's reference, and the +fused draft loop's layer boundaries built on it (``M3Eagle3Draft.ar_norm``: the plain all-reduce, then FlashInfer's +fused add + RMSNorm) against stock Eagle3's, bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy the single-node drafter can run (AUTO, +ONESHOT, TWOSHOT, NCCL; MNNVL is multi-node only), at tp4 shapes (H = 6144, eps = 1e-6, 1 to 32 rows). The inputs are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own 4-rank launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_allreduce.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py +""" + +import math +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 4, 8, 32) +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +RANK = 0 +WORLD = 1 + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on rank 0, x/4 on ranks 1 and 2, 0 on rank 3. + parts = [x / 2, x / 4, x / 4] + [torch.zeros_like(x)] * (WORLD - 3) + return ( + [p.bfloat16() for p in parts[:WORLD]], + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A, -A, +B, -B + d) onto a unit residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a, -a, b, (-b + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return [p.bfloat16() for p in parts[:WORLD]], residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +def _ulps_at(out: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the magnitude ``scale`` (the operands'), not at the result's.""" + return ((out.double() - ref.double()).abs() / _ulp_at(scale.abs())).max().item() + + +def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary (``M3Eagle3Draft.ar_norm``): ``m3_allreduce`` plain, then the catalog's + ``flashinfer_fused_add_rmsnorm``.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + + return flashinfer_fused_add_rmsnorm(m3_allreduce(x, ar, None), residual.clone(), w, EPS) + + +def _check_cells(mapping: Any) -> None: + """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum + to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude + u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final + roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). + The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm.functional import AllReduceStrategy + + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating + + ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + for i, cell in enumerate(m3_allreduce.CELLS): + spec = cell.spec + xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) + params = ( + None + if spec.get("plain") + else AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual.clone(), + norm_weight=w, + eps=EPS, + trigger_completion_at_end=False, + ) + ) + # The reference first: the fused path may write the updated residual into ``params.residual``. + ref = m3_allreduce.reference(xs, ar, params) + with validating(m3_allreduce): + out = m3_allreduce(xs[RANK], ar, params) + torch.cuda.synchronize() + total = sum(x.double() for x in xs) + if params is None: + err = _ulps_at(out, ref, total) + assert err <= 1.0, (cell.why, err) + else: + err = _ulps_at(out[1], ref[1], torch.maximum(total.abs(), residual.double().abs())) + assert err <= 2.5, (cell.why, "residual", err) + h = out[1].double() + own = h * torch.rsqrt(h.square().mean(-1, keepdim=True) + EPS) * w.double() + assert _ulps(out[0], own) <= 2.0, (cell.why, "normed", _ulps(out[0], own)) + dist.barrier() + if RANK == 0: + print(f"{RESULT_TAG} cell {cell.why} passed", flush=True) + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + _check_cells(mapping) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + ulp, res_ulp = 0.0, 0.0 + for m in TOKENS: + xs, residual, w = build(m) + target_out, target_res = _target(ar, xs[RANK], residual, w) + stock_out, stock_res = _stock(ar, xs[RANK], residual, w) + torch.cuda.synchronize() + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {name} m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {name:34s} | target = stock vs fp64: normed {ulp:.3f}" + f" | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the drafter runs at tp4, got {WORLD} ranks" + try: + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + 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", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py new file mode 100644 index 000000000000..db9dba58f108 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CAPTURE_NORM_IN_BOUNDARY on the tp4 multicast all-reduce + RMSNorm op, at 4 ranks. + +With ``cap_out`` the boundary also writes the Eagle3 capture of its updated residual. Asserted bit for bit against +``m3_spec_tail.capture_norm`` of the op's own updated residual, in a strided [M, H] block of a [rows, 3H] capture +buffer (the other bytes untouched), and the op's normed / updated / MXFP8 outputs bitwise the same call's without +``cap_out``. Every kernel a boundary can run: + +- the 8-CTA cluster kernel in the stock oneshot order (options 39 / 47, plain calls up to 8 rows, x and x + x2, + prepushed 2); +- the row kernel (plain options 7 / 15 and the carveout instance 71 / 79, 1 to 8 rows on the 8-token bucket, 12 rows + on the 16-token bucket, 32 rows on the 48-token bucket, prepushed 2); + +(The tp4 MoE finalize runs the shared GB300 op, which has no capture epilogue.) + +at residual scales 1e-3 to 1e3 and the fc_norm eps 1e-6 / 1e-5. + +Its own 4-rank launcher, as ``_m3_draft_ar_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_capture_in_boundary.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_capture_in_boundary_op_matrix.py +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +ar_module: Any = None +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + handle = _make_mnnvl_mcast_buffer(dist.group.WORLD, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=dist.group.WORLD, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor) -> None: + """Store this rank's rows into every rank's exchange (set A), as a producer does before a prepushed call.""" + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = parity.view(m, CLUSTER, 1) * half + (t * WORLD + RANK) * (HIDDEN // 2) + c * chunk + w + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, (where, out.dtype, ref.dtype) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +def _capture_reference(updated, weight, eps, out) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_spec_tail, + ) + + m3_spec_tail.capture_norm(updated, None, weight, eps, out) + + +def _one_case( + kind: str, m: int, maxt: int, opts: int, scale: float, quant: bool, seed: int +) -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written, + ) + + cap_eps = (1e-6, 1e-5)[seed % 2] + ws = WORKSPACES[maxt] + residual = (_rows(seed + 7, -1, m).float() * 4.0 * scale).to(torch.bfloat16) + norm_w, cap_w = _norm_w(seed), _norm_w(seed + 101) + results = [] + for with_cap in (False, True): + buf = torch.full((m + 3, 3 * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + cap = buf[:m, HIDDEN : 2 * HIDDEN] if with_cap else None + cw = cap_w if with_cap else None + flags = ws.flags_for(CLUSTER) + x = _rows(seed, RANK, m, scale) + x2 = _rows(seed + 3, RANK, m, scale) if kind == "x2" else None + prepushed = 2 if kind == "prepushed" else 0 + if prepushed: + _push(ws, x) + outs = allreduce_entry.m3_mc_allreduce_norm( + x, x2, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), flags, RANK, WORLD, CLUSTER, False, maxt, + quant, prepushed, opts, cap, cw, cap_eps, + ) # fmt: skip + torch.cuda.synchronize() + results.append((outs, buf)) + dist.barrier() + (plain, _), (capped, buf) = results + where = f"{kind} opts {opts} bucket {maxt} rows {m} scale {scale} quant {quant}" + assert len(plain) == len(capped), where + for i, (a, b) in enumerate(zip(plain, capped)): + if i == 3: + written = sf_written(m, HIDDEN // 32, a.device) + a, b = a[written], b[written] + _assert_bitwise(b, a, f"{where}: output {i} with vs without cap_out") + ref = torch.empty((m, HIDDEN), dtype=torch.bfloat16, device="cuda") + _capture_reference(capped[1], cap_w, cap_eps, ref) + torch.cuda.synchronize() + _assert_bitwise(buf[:m, HIDDEN : 2 * HIDDEN], ref, f"{where}: capture vs capture_norm") + untouched = buf.clone() + untouched[:m, HIDDEN : 2 * HIDDEN] = 7.0 + assert bool((untouched == 7.0).all()), f"{where}: bytes outside the capture block changed" + dist.barrier() + + +def _cases(): + """(kind, rows, bucket, opts, scale, quant).""" + cases = [] + scales = (1.0, 1e3, 1e-3) + # The cluster kernel in the oneshot order (plain calls up to 8 rows): options 39 and 47. + for i, m in enumerate((1, 4, 8)): + for opts in (39, 47): + cases.append(("x" if i != 1 else "x2", m, 8, opts, scales[i], opts == 47)) + cases.append(("prepushed", 4, 8, 39, 1.0, False)) + # The row kernel: plain options 7 / 15 and the carveout instance (+ 64). + for i, (m, opts) in enumerate(((1, 7), (4, 15), (8, 7 | 64), (3, 15 | 64))): + cases.append(("x2" if i % 2 else "x", m, 8, opts, scales[i % 3], i == 1)) + cases.append(("prepushed", 8, 8, 7, 1e-3, False)) + cases.append(("x", 12, 16, 15, 1.0, True)) + cases.append(("x2", 12, 16, 15 | 64, 1e3, False)) + cases.append(("x", 32, 48, 15, 1e-3, False)) + return cases + + +def _run_matrix() -> None: + for n, (kind, m, maxt, opts, scale, quant) in enumerate(_cases()): + _one_case(kind, m, maxt, opts, scale, quant, 7000 + 37 * n) + if RANK == 0: + print( + f"[case] {kind} opts {opts} bucket {maxt} rows {m} scale {scale}: bitwise", + flush=True, + ) + + +def _run_one_rank() -> int: + global RANK, WORLD, allreduce_entry, ar_module + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the tp4 boundary runs at 4 ranks, got {WORLD}" + try: + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + + allreduce_entry, ar_module = _plain, _ar + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_draft_ar_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + 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", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py new file mode 100644 index 000000000000..13f7835c7d66 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py @@ -0,0 +1,533 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The fused draft loop's layer boundaries (``modeling.M3Eagle3Worker._m3_draft_ar_norm``) against stock Eagle3's, +bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy the single-node drafter can run (AUTO, +ONESHOT, TWOSHOT, NCCL; MNNVL is multi-node only), at tp4 shapes (H = 6144, eps = 1e-6, 1 to 32 rows). The inputs are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +The worker method runs with stand-ins for what it reads (the worker's ``mapping``; the drafter's ``w[norm_key]``, +``rms_eps``, ``allreduce`` and its own ``M3Eagle3Draft.ar_norm``), over the multicast workspace the target's +``m3_mc_allreduce.get_workspace`` returns (built here over the torch.distributed group, as the sibling multicast +matrices build it, and registered where ``get_workspace`` finds it). Two paths per strategy, case and row count: + +- ``boundary``: the method as the target runs it. At the rows ``DRAFT_MC_AR`` covers (1..``DRAFT_MC_AR_MAX_ROWS``) + the multicast all-reduce + residual + RMSNorm in its drafter form (row kernel, opts bit 16; the strategy is not + used); above them the built-in path, ``M3Eagle3Draft.ar_norm`` on the strategy's all-reduce. +- ``declined``: the method with ``get_workspace`` returning None (no multicast workspace), so the covered rows take + its fallback (``missed``, then ``M3Eagle3Draft.ar_norm`` on the strategy's all-reduce) as well. + +Which of the two ran is checked per call (the drafter-form multicast call returned rows exactly at the covered rows +of ``boundary``, never in ``declined``). Not driven: ``pushed=True`` (``DRAFT_AR_PUSH``, prepushed 2; it needs the +projection's split-K push) and ``DRAFT_MC_AR_CLUSTER`` (off; exact-class, not bitwise). + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own 4-rank launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the collected entry point is +``test_modeling_v2_m3_draft_ar_norm.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_draft_ar_norm_op_matrix.py +""" + +import contextlib +import math +import os +import signal +import socket +import subprocess +import sys +import time +import types +from pathlib import Path +from typing import Any, Callable, Iterator, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 2, 3, 4, 5, 6, 7, 8, 9, 32) +"""Every row count of the drafter-form multicast boundary (1..8: its option table differs by row), the first one +past it and a wide one (the built-in path).""" +NORM_KEYS = ("post_norm", "final_norm") +"""The drafter's two boundaries per pass (o_proj -> post-attention norm, down -> final norm), alternated by row +count: the method reads only ``w[norm_key]``, so they differ in the weight's key alone.""" +CLUSTER = 8 +EMPTY_WORD = -2147483648 +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +RANK = 0 +WORLD = 1 + +# Bound inside the rank body: importing the target pulls in tensorrt_llm. +modeling: Any = None +ar_module: Any = None + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on rank 0, x/4 on ranks 1 and 2, 0 on rank 3. + parts = [x / 2, x / 4, x / 4] + [torch.zeros_like(x)] * (WORLD - 3) + return ( + [p.bfloat16() for p in parts[:WORLD]], + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A, -A, +B, -B + d) onto a unit residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a, -a, b, (-b + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return [p.bfloat16() for p in parts[:WORLD]], residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +class _McCalls: + """Wraps ``m3_mc_allreduce.m3_mc_allreduce_norm`` (the method reads it off the module at each call): counts the + drafter-form calls that returned rows, i.e. the boundaries the multicast kernel ran.""" + + def __init__(self, fn: Callable) -> None: + self.fn = fn + self.ran = 0 + + def __call__(self, *args, **kwargs): + out = self.fn(*args, **kwargs) + if out is not None and kwargs.get("drafter_norm"): + self.ran += 1 + return out + + +@contextlib.contextmanager +def _no_workspace() -> Iterator[None]: + """``get_workspace`` returning None, as the target sees it with no multicast workspace: the method's decline.""" + real = ar_module.get_workspace + ar_module.get_workspace = lambda mapping, hidden: None + try: + yield + finally: + ar_module.get_workspace = real + + +def _stand_ins(mapping: Any, ar: Any, w: torch.Tensor) -> Tuple[Any, Any]: + """(worker, drafter) with what ``M3Eagle3Worker._m3_draft_ar_norm`` reads: the worker's ``mapping``; the + drafter's norm weights, ``rms_eps``, its all-reduce and its own ``ar_norm`` (the target's method, bound).""" + worker = types.SimpleNamespace(mapping=mapping) + draft = types.SimpleNamespace( + w={key: w for key in NORM_KEYS}, rms_eps=EPS, allreduce=ar, dtype=torch.bfloat16 + ) + draft.ar_norm = types.MethodType(modeling.M3Eagle3Draft.ar_norm, draft) + return worker, draft + + +def _target( + mapping: Any, + ar: Any, + x: torch.Tensor, + residual: torch.Tensor, + w: torch.Tensor, + norm_key: str, + mc: _McCalls, + declined: bool, +) -> Tuple[torch.Tensor, torch.Tensor, bool]: + """The draft loop's boundary as the target runs it (``M3Eagle3Worker._m3_draft_ar_norm``, ``pushed`` False); + with ``declined`` the multicast workspace is unavailable. Returns (normed, residual sum, whether the multicast + drafter form ran).""" + worker, draft = _stand_ins(mapping, ar, w) + before = mc.ran + with _no_workspace() if declined else contextlib.nullcontext(): + out, res = modeling.M3Eagle3Worker._m3_draft_ar_norm( + worker, draft, x, residual.clone(), norm_key + ) + return out, res, mc.ran > before + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _register_workspace(mapping: Any) -> Any: + """The multicast workspace ``m3_mc_allreduce.get_workspace(mapping, HIDDEN)`` returns: one bucket per + ``BUCKETS`` entry, each built as ``McAllReduceWorkspace.create`` builds it but over the torch.distributed group + (these ranks have no MPI), with the 8-CTA call counters ``get_workspace`` creates; registered in its cache.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + + wss = [] + for maxt in ar_module.BUCKETS: + words = 4 * maxt * WORLD * (HIDDEN // 2) + handle = _make_mnnvl_mcast_buffer(dist.group.WORLD, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=dist.group.WORLD, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + wss.append(ws) + buckets = ar_module.McBuckets(wss) + ar_module._WORKSPACE[HIDDEN] = buckets + assert ar_module.get_workspace(mapping, HIDDEN) is buckets + return buckets + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + # The covered rows are 1..DRAFT_MC_AR_MAX_ROWS, every one of them driven; the row kernel's drafter form, the one + # the bitwise claim is about (the cluster form is exact-class). + covered = set(modeling.DRAFT_MC_AR.keys) + assert covered == set(range(1, modeling.DRAFT_MC_AR_MAX_ROWS + 1)), sorted(covered) + assert covered <= set(TOKENS), (sorted(covered), TOKENS) + assert any(m not in covered for m in TOKENS), ( + "drive at least one row count past the multicast boundary" + ) + assert not modeling.DRAFT_MC_AR_CLUSTER, ( + "the cluster drafter form is exact-class, not bitwise stock" + ) + # A declined multicast boundary logs (``missed``) and falls back instead of raising. + modeling.set_require(False) + _register_workspace(mapping) + mc = _McCalls(ar_module.m3_mc_allreduce_norm) + ar_module.m3_mc_allreduce_norm = mc + if RANK == 0: + print( + f"{RESULT_TAG} multicast drafter form at rows {sorted(covered)} (DRAFT_MC_AR_MAX_ROWS" + f" {modeling.DRAFT_MC_AR_MAX_ROWS}); built-in path at the others and when declined", + flush=True, + ) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + # The multicast drafter kernel takes no strategy: it sums the ranks in FP32 in rank order, which is stock's ONESHOT + # (the strategy a served drafter resolves AUTO to). Under NCCL its rows are therefore compared with stock ONESHOT; + # NCCL itself (which can round cancelling partials differently) stays the reference of the built-in rows. + oneshot = AllReduce(mapping=mapping, strategy=AllReduceStrategy.ONESHOT, dtype=torch.bfloat16) + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + for declined in (False, True): + path = "declined" if declined else "boundary" + ulp, res_ulp = 0.0, 0.0 + for i, m in enumerate(TOKENS): + xs, residual, w = build(m) + target_out, target_res, ran_mc = _target( + mapping, ar, xs[RANK], residual, w, NORM_KEYS[i % 2], mc, declined + ) + expect_mc = not declined and m in covered + stock_ar = oneshot if expect_mc and strategy == AllReduceStrategy.NCCL else ar + stock_out, stock_res = _stock(stock_ar, xs[RANK], residual, w) + torch.cuda.synchronize() + assert ran_mc == expect_mc, ( + f"{strategy.name} {path} {name} m={m}: the multicast drafter form " + f"{'ran' if ran_mc else 'did not run'}" + ) + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {path} ({'multicast' if ran_mc else 'built-in'}) {name}" + f" m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, path, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {path:8s} | {name:34s} | target = stock vs fp64:" + f" normed {ulp:.3f} | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | both | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD, modeling, ar_module + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the drafter runs at tp4, got {WORLD} ranks" + try: + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.modeling as _modeling # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_mc_allreduce as _ar # noqa: E501 + + modeling, ar_module = _modeling, _ar + assert modeling.m3_mc_allreduce is ar_module, ( + "the method calls the kernels module this patches" + ) + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + 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", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..852149f12cc7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,665 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``, 4 ranks. + +Its own launcher, as the catalog's collective matrices are: run plainly it spawns +one rank process per device in CUDA_VISIBLE_DEVICES and kills them all on the +first failure or on the deadline. A rank whose peers never push polls forever, +so the deadline is what keeps a broken exchange from wedging the calling run. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_mc_allreduce_norm_op_matrix.py + +Each rank builds the three multicast workspaces (8 / 16 / 48 tokens) the target +builds, over the torch.distributed group (the non-MPI path of +``_make_mnnvl_mcast_buffer``). Every rank can regenerate every rank's inputs from +the shared seed convention, which is what keeps each reference arithmetic rather +than a second collective. Producer pushes (``prepushed`` 2 / 3) are emulated by +storing the canonicalized rows through the multicast mapping at the call parity +the flags select, which is exactly what the producing kernels store. + +Not a pytest module, despite the ``check_*`` bodies: they are one fixed sequence +inside a single 4-rank job whose workspace state (call parities) carries from +one to the next. The collected entry point is +``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +finalize_entry: Any = None +ar_module: Any = None +validating: Any = None +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +# ---------------------------------------------------------------- inputs + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + """One rank's [m, H] BF16 rows; ``rank`` -1 for tensors every rank holds alike.""" + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return _rows(seed + 7, -1, m, 4.0) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _finalize_inputs(seed: int, rank: int, m: int, topk: int, dropped: int): + """One rank's unfinalized MoE outputs: fc2 [m * topk, H], perm [m, topk], fp32 scales, shared [m, H].""" + g = _gen(seed + 23, rank) + rows = m * topk + fc2 = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + perm = torch.randperm(rows, generator=g, device="cuda").to(torch.int32) + if dropped: + drop = torch.randperm(rows, generator=g, device="cuda")[:dropped] + perm[drop] = -1 + perm = perm.view(m, topk).contiguous() + scales = torch.rand(m, topk, generator=g, device="cuda", dtype=torch.float32) + shared = _rows(seed + 29, rank, m, 0.5) + return fc2, perm, scales, shared + + +# ---------------------------------------------------------------- workspace / producer emulation + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor, exchange_set: int) -> None: + """Store this rank's rows into every rank's exchange, as a producer kernel does before a prepushed call. + + Slot of (token t, cluster CTA c): set * 2 * half + parity(t, c) * half + (t * world + rank) * H / 2 + c * H / 16 + words, parity read from the CTA's flag; -0.0 halves are stored as +0.0 (0x80000000 is the empty word). + """ + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = ( + exchange_set * 2 * half + + parity.view(m, CLUSTER, 1) * half + + (t * WORLD + RANK) * (HIDDEN // 2) + + c * chunk + + w + ) + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +# ---------------------------------------------------------------- gates + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +#: Share of normed elements one BF16 step off the float64 reference: measured at most 1.1e-4 per call over +#: every call of this matrix on GB300 (see ``_assert_normed``). The gate sits ~18x above that and far below what a +#: wrong statistic produces (most elements move). +NORMED_STEP_FRACTION = 0.002 +NORMED_STEP_SEEN: List[float] = [] + + +def _assert_normed(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """At most one BF16 rounding step from the float64 reference, and rarely that. + + The kernel's RMS statistic is a shuffle / warp / cluster tree of FP32 partial sums and its ``rsqrtf`` is + approximate (~2^-22 relative): the scaled value moves by ~1e-6 relative, which changes the BF16 rounding only + for values within that distance of a rounding midpoint -- one step, never two. The fraction bound catches a + statistic that is systematically off (a missing rank, a wrong H), which moves most elements. + """ + assert out.dtype == ref.dtype == torch.bfloat16 and out.shape == ref.shape, where + o, r = out.float(), ref.float() + _, exponent = torch.frexp(torch.maximum(o.abs(), r.abs())) + step = torch.ldexp(torch.ones_like(o), exponent - 8) + diff = (o - r).abs() + assert bool((diff <= step).all()), f"{where}: max diff {diff.max().item()} beyond one BF16 step" + frac = (diff != 0).float().mean().item() + NORMED_STEP_SEEN.append(frac) + assert frac <= NORMED_STEP_FRACTION, f"{where}: {frac:.4f} of normed elements off by one step" + + +def _assert_quant(outs, ref, m: int, where: str) -> None: + """The MXFP8 copy: bit-exact against the recipe applied to the op's own normed rows (a pure function of + them), and equal to the reference's wherever the reference's normed 32-column block is the op's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + mxfp8_rows, + sf_written, + ) + + q, sf = outs[2], outs[3] + own_q, own_sf = mxfp8_rows(outs[0]) + written = sf_written(m, HIDDEN // 32, sf.device) + _assert_bitwise(q, own_q, f"{where} e4m3 rows") + _assert_bitwise(sf[written], own_sf[written], f"{where} UE8M0 scales") + same_block = (outs[0].view(m, -1, 32) == ref[0].view(m, -1, 32)).all(-1) + qb = q.view(torch.uint8).view(m, -1, 32) + rqb = ref[2].view(torch.uint8).view(m, -1, 32) + assert bool((qb == rqb).all(-1)[same_block].all()), ( + f"{where}: e4m3 differs on a block whose normed rows agree" + ) + + +def _check_outputs(outs, ref, m: int, quant: bool, where: str) -> None: + assert len(outs) == (4 if quant else 2), (where, len(outs)) + _assert_bitwise(outs[1], ref[1], f"{where} updated residual") + _assert_normed(outs[0], ref[0], f"{where} normed") + if quant: + _assert_quant(outs, ref, m, where) + + +# ---------------------------------------------------------------- calls + + +def _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, opts): + return allreduce_entry.m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + ws.max_tokens, + quant, + prepushed, + opts, + ) + + +def _plain_case(spec: Dict[str, Any], seed: int, where: str) -> None: + m, maxt, quant, prepushed = spec["m"], spec["maxt"], spec["quant"], spec["prepushed"] + ws = WORKSPACES[maxt] + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] if spec["x2"] else None + residual, norm_w = _residual(seed, m), _norm_w(seed) + x, x2 = xs[RANK], (None if x2s is None else x2s[RANK]) + if prepushed >= 2: + _push(ws, x, 0) + if prepushed == 3: + _push(ws, x2, 1) + with validating(allreduce_entry.m3_mc_allreduce_norm): + outs = _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, spec["opts"]) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, residual, norm_w, EPS, quant, spec.get("drafter", False) + ) + _check_outputs(outs, ref, m, quant, where) + dist.barrier() + + +# ---------------------------------------------------------------- checks + + +def check_plain_cells() -> None: + """Every cell of ``m3_mc_allreduce_norm.CELLS`` through ``__call__`` under ``validating``, gated on its + ``reference``.""" + for i, cell in enumerate(allreduce_entry.m3_mc_allreduce_norm.CELLS): + _plain_case(cell.spec, 1000 + 17 * i, cell.why) + + +def check_finalize_cells() -> None: + """Every cell of ``m3_mc_finalize_allreduce_norm.CELLS``, gated on its ``reference``.""" + op = finalize_entry.m3_mc_finalize_allreduce_norm + for i, cell in enumerate(op.CELLS): + spec = cell.spec + m, maxt, quant, topk = spec["m"], spec["maxt"], spec["quant"], spec["topk"] + seed = 2000 + 17 * i + ws = WORKSPACES[maxt] + per_rank = [_finalize_inputs(seed, r, m, topk, spec["dropped"]) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + with validating(op): + outs = op( + fc2, + perm, + scales, + shared, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + maxt, + quant, + spec["opts"], + ) + torch.cuda.synchronize() + ref = op.reference( + [p[0] for p in per_rank], + [p[1] for p in per_rank], + [p[2] for p in per_rank], + [p[3] for p in per_rank], + residual, + norm_w, + EPS, + quant, + ) + _check_outputs(outs, ref, m, quant, cell.why) + dist.barrier() + + +def check_the_target_wrappers_select_bucket_and_options() -> None: + """The target's own wrappers (bucket by token count, option table) on top of the entries, every token count of + the 8-token bucket plus one of each larger bucket, and the finalize wrapper; the MXFP8 publication is the + prequant slot's business and is left off.""" + buckets = ar_module.McBuckets([WORKSPACES[b] for b in BUCKETS]) + for m in (*range(1, 9), 11, 40): + seed = 3000 + m + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + normed, updated = ar_module.m3_mc_allreduce_norm( + xs[RANK], residual, norm_w, EPS, buckets, x2=x2s[RANK] + ) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, False) + _check_outputs([normed, updated], ref, m, False, f"wrapper m={m}") + dist.barrier() + for m in (9, 33): + seed = 3100 + m + per_rank = [_finalize_inputs(seed, r, m, 4, 2) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + normed, updated = ar_module.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, buckets + ) + torch.cuda.synchronize() + ref = finalize_entry.m3_mc_finalize_allreduce_norm.reference( + *[[p[j] for p in per_rank] for j in range(4)], residual, norm_w, EPS, False + ) + _check_outputs([normed, updated], ref, m, False, f"finalize wrapper m={m}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """Two boundaries captured in one graph (8- and 16-token buckets), replayed with fresh inputs. + + The decode graph holds every boundary call; each replay rotates the same per-CTA parities the eager calls do, + so a replay that reused a stale half or skipped a re-arm reads the previous replay's rows and fails the + bitwise residual gate (every replay gets new inputs). + """ + shapes = ((3, 8, 15), (12, 16, 15)) + norm_w = _norm_w(4000) + bufs = [] + for m, maxt, opts in shapes: + bufs.append( + dict( + m=m, + maxt=maxt, + opts=opts, + x=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + x2=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + residual=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + ) + + def body(): + return [ + _call_plain( + WORKSPACES[b["maxt"]], b["x"], b["x2"], b["residual"], norm_w, False, 0, b["opts"] + ) + for b in bufs + ] + + def fill(seed: int): + refs = [] + for j, b in enumerate(bufs): + s = seed + 31 * j + xs = [_rows(s, r, b["m"]) for r in range(WORLD)] + x2s = [_rows(s + 3, r, b["m"]) for r in range(WORLD)] + b["x"].copy_(xs[RANK]) + b["x2"].copy_(x2s[RANK]) + b["residual"].copy_(_residual(s, b["m"])) + refs.append( + allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, b["residual"].clone(), norm_w, EPS, False + ) + ) + torch.cuda.synchronize() + dist.barrier() + return refs + + refs = fill(4100) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm = body() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(warm[j], refs[j], b["m"], False, f"graph warm-up m={b['m']}") + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = body() + for rep in range(4): + refs = fill(4200 + 100 * rep) + graph.replay() + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(captured[j], refs[j], b["m"], False, f"replay {rep} m={b['m']}") + dist.barrier() + del graph + dist.barrier() + + +def check_the_gate_sees_a_missing_rank() -> None: + """The residual gate is not blind to the reduction: dropping any one rank's partial from the reference + moves most elements of the updated residual (each rank's rows are distinct by construction).""" + m, seed = 4, 5000 + xs = [_rows(seed, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + full = allreduce_entry.m3_mc_allreduce_norm.reference(xs, None, residual, norm_w, EPS, False) + for r in range(WORLD): + partial = [x if i != r else torch.zeros_like(x) for i, x in enumerate(xs)] + wrong = allreduce_entry.m3_mc_allreduce_norm.reference( + partial, None, residual, norm_w, EPS, False + ) + frac = (wrong[1] != full[1]).float().mean().item() + assert frac > 0.9, f"dropping rank {r} changes only {frac:.3f} of the residual" + dist.barrier() + + +def _expect_rejected(fn, match: str, where: str) -> None: + try: + fn() + except AssertionError as exc: + assert match in str(exc), f"{where}: rejected for the wrong reason: {exc}" + else: + raise AssertionError(f"{where}: the guard admitted it") + + +def check_guards() -> None: + """``is_valid`` refuses what the op accepts and answers wrongly. Nothing here reaches the kernel: the guard + raises first, on every rank alike, so no rank is left polling.""" + m, seed = 4, 6000 + ws8, ws16 = WORKSPACES[8], WORKSPACES[16] + x, residual, norm_w = _rows(seed, RANK, m), _residual(seed, m), _norm_w(seed) + op = allreduce_entry.m3_mc_allreduce_norm + + def call(ws_uc, flags, norm): + with validating(op): + op( + x, + None, + residual, + norm, + EPS, + ws_uc, + ws8.mc.data_ptr(), + flags, + RANK, + WORLD, + CLUSTER, + False, + 8, + False, + 0, + 15, + ) + + strided = torch.empty(2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[::2] + strided.copy_(norm_w) + _expect_rejected( + lambda: call(ws8.uc, ws8.flags_for(CLUSTER), strided), "contiguous", "strided norm_w" + ) + _expect_rejected( + lambda: call(ws8.uc[: ws8.uc.numel() // 2], ws8.flags_for(CLUSTER), norm_w), + "words", + "short workspace", + ) + _expect_rejected( + lambda: call(ws8.uc, ws16.flags_for(CLUSTER), norm_w), "counters", "another bucket's flags" + ) + + fin = finalize_entry.m3_mc_finalize_allreduce_norm + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, 9, 4, 0) + bad = perm.clone() + bad[0, 0] = fc2.shape[0] + + def call_fin(): + with validating(fin): + fin( + fc2, + bad, + scales, + shared, + _residual(seed, 9), + norm_w, + EPS, + ws16.uc, + ws16.mc.data_ptr(), + ws16.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + 16, + False, + 15, + ) + + _expect_rejected(call_fin, "perm spans", "perm past the FC2 rows") + dist.barrier() + + +CHECKS = ( + check_plain_cells, + check_finalize_cells, + check_the_target_wrappers_select_bucket_and_options, + check_cuda_graph_replays, + check_the_gate_sees_a_missing_rank, + check_guards, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, allreduce_entry, finalize_entry, ar_module, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD == 4, f"the entries are certified at 4 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as _finalize, + ) + + allreduce_entry, finalize_entry, ar_module = _plain, _finalize, _ar + # Spawned by file path, not collected by pytest: reach the sibling helper the plain way (confined to the + # rank processes). /comm -> tests/unittest/_torch/modeling_v2. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peers' polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print( + f"[rank {RANK}] {len(CHECKS)} checks passed; normed one-step fraction max " + f"{max(NORMED_STEP_SEEN):.5f} over {len(NORMED_STEP_SEEN)} calls", + flush=True, + ) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 4 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 4, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 4" + ) + 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", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + 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), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_rank_job.py new file mode 100644 index 000000000000..1115ff3b5976 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_rank_job.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run this target's collective rank job from pytest, and report what it did. + +The same contract as ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``: the +launcher beside this file owns the per-device rank processes and the deadline +that kills a wedged collective; this only selects the devices, starts it in a +fresh interpreter (by file path: this tree is not a package) and turns its exit +code into an assertion. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +WORLD_SIZE = 4 +"""tp4's world size -- the topology these entries are certified for.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the launcher's own deadline, so its message wins the race.""" + + +def _devices() -> str: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + devices = [d for d in visible.split(",") if d.strip()] + else: + import torch + + devices = [str(i) for i in range(torch.cuda.device_count())] + if len(devices) < WORLD_SIZE: + pytest.skip( + f"this entry is certified at world size {WORLD_SIZE}; only {len(devices)} device(s) are visible" + ) + return ",".join(devices[:WORLD_SIZE]) + + +def _load_launcher_constants(launcher: Path): + """Read the launcher's module-level budget without importing it by name (it imports only torch at + module scope).""" + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str) -> str: + """Run ``__op_matrix``'s launcher over ``WORLD_SIZE`` devices; returns its stdout.""" + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=_devices()) + + # The ranks import tensorrt_llm absolutely. From a source checkout that is not installed, put the repo root + # (tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root) on their path; from a tests + # tree run against an installed package, they import the installed one. + repo_root = Path(__file__).resolve().parents[8] + if (repo_root / "tensorrt_llm").is_dir(): + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p + ) + else: + assert importlib.util.find_spec("tensorrt_llm") is not None, ( + f"no tensorrt_llm/ at the repo root {repo_root} and no installed tensorrt_llm" + ) + + constants = _load_launcher_constants(launcher) + timeout = constants.DEADLINE_S + _LAUNCHER_GRACE_S + + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) + return completed.stdout diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_allreduce.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_allreduce.py new file mode 100644 index 000000000000..27d45f08edc4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_allreduce.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the ``m3_allreduce`` certification matrix. + +The matrix is ``_m3_allreduce_op_matrix.py`` beside this file, its own 4-rank launcher; see ``_rank_job``. It drives +every cell of the entry against its reference, then requires the fused draft loop's boundaries built on it (plain +all-reduce, then the fused add + RMSNorm) to equal stock Eagle3's bit for bit under every strategy, and reports both +against an fp64 reference as a warning, one line per strategy and case. +""" + +import warnings + +import _rank_job +import pytest +import torch + +pytestmark = [ + pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _rank_job.WORLD_SIZE, + reason=f"certified at world size {_rank_job.WORLD_SIZE}", + ), +] + +_RESULT_TAG = "[ulp]" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_allreduce_op_matrix() -> None: + out = _rank_job.run("m3_allreduce") + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + "draft boundary (= stock Eagle3) vs fp64, max BF16 ulps\n" + "\n".join(lines), stacklevel=1 + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py new file mode 100644 index 000000000000..ad7539d27829 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for CAPTURE_NORM_IN_BOUNDARY on the tp4 multicast all-reduce + RMSNorm op. + +The matrix is ``_m3_capture_in_boundary_op_matrix.py`` beside this file, its own 4-rank launcher; see ``_rank_job``. +It requires the boundary's capture output to equal ``m3_spec_tail.capture_norm`` of its updated residual bit for bit, +on both kernels the tp4 op runs (the oneshot cluster kernel, the row kernel), and the op's other outputs to be +unchanged by it. +""" + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "the multicast all-reduce requires CUDA devices" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_capture_norm_in_boundary_is_capture_norm() -> None: + out = _rank_job.run("m3_capture_in_boundary") + assert "[case]" in out, f"the launcher ran no case:\n{out}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py new file mode 100644 index 000000000000..3cc1592bc3c5 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the fast draft loop's layer boundaries against stock Eagle3's. + +The matrix is ``_m3_draft_ar_norm_op_matrix.py`` beside this file, its own 4-rank launcher; see ``_rank_job``. It +requires the target's drafter boundary (``modeling.M3Eagle3Worker._m3_draft_ar_norm``: the multicast all-reduce + +residual + RMSNorm in its drafter form at the rows ``DRAFT_MC_AR`` covers, else and when that declines +``M3Eagle3Draft.ar_norm`` on the strategy's all-reduce) to equal stock Eagle3's boundary (plain all-reduce, then the +``RMSNorm`` module with the residual) bit for bit under every strategy, and reports both against an fp64 reference as a +warning, one line per strategy, path and case. +""" + +import warnings + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_RESULT_TAG = "[ulp]" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_draft_ar_norm_is_stock_eagle3() -> None: + out = _rank_job.run("m3_draft_ar_norm") + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + "draft boundary (= stock Eagle3) vs fp64, max BF16 ulps\n" + "\n".join(lines), stacklevel=1 + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..31a054642e89 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the multicast all-reduce + RMSNorm entries' certification matrix. + +Covers both ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm`` (one +kernel, one exchange protocol). The matrix is ``_m3_mc_allreduce_norm_op_matrix.py`` +beside this file, its own 4-rank launcher; see ``_rank_job`` for why. +""" + +import _rank_job +import pytest +import torch + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + + +# Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. +@pytest.mark.no_xdist +def test_m3_mc_allreduce_norm_op_matrix() -> None: + _rank_job.run("m3_mc_allreduce_norm") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py new file mode 100644 index 000000000000..7cf1ffb5a141 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_finalize_allreduce_norm.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests of the ``m3_mc_finalize_allreduce_norm`` entry that need no collective. + +Its cells are certified, with ``m3_mc_allreduce_norm``'s (one kernel, one exchange protocol), by the 4-rank matrix +``_m3_mc_allreduce_norm_op_matrix.py`` that ``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py`` runs. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + finalize_partial, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as op, +) + + +def test_finalize_partial_rounds_the_routed_sum_before_the_shared_expert() -> None: + """As stock: the finalized routed output is BF16, and the shared expert is added to it (shared.add_(routed)). + Four routed rows of 1.0 at FP32 routing weight 0.2509765625 sum to 1.00390625, which rounds to BF16 1.0; plus the + shared -1 that is 0, where adding before the rounding would leave 0.00390625.""" + fc2 = torch.ones((4, 8), dtype=torch.bfloat16) + perm = torch.arange(4, dtype=torch.int32).view(1, 4) + scales = torch.full((1, 4), 0.2509765625, dtype=torch.float32) + shared = torch.full((1, 8), -1.0, dtype=torch.bfloat16) + out = finalize_partial(fc2, perm, scales, shared) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(out, shared.clone().add_(routed)) + assert torch.equal(out, torch.zeros(1, 8, dtype=torch.bfloat16)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_each_cells_partial_is_the_stock_finalize_then_shared_add(cell) -> None: + """Per cell shape (rows, top-k, dropped slots), every rank's partial in the entry's reference is the stock MoE's: + the finalized routed rows rounded to BF16, then ``shared.add_(routed)``. The 4-rank matrix drives the cells on the + GPU against this reference.""" + spec = cell.spec + g = torch.Generator().manual_seed(spec["m"] * 31 + spec["dropped"]) + m, topk, hidden = spec["m"], spec["topk"], 64 + fc2 = torch.randn(m * topk, hidden, generator=g).to(torch.bfloat16) + perm = torch.randperm(m * topk, generator=g).view(m, topk).to(torch.int32) + perm.view(-1)[: spec["dropped"]] = -1 + scales = torch.rand(m, topk, generator=g) + shared = torch.randn(m, hidden, generator=g).to(torch.bfloat16) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(finalize_partial(fc2, perm, scales, shared), shared.clone().add_(routed)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..4258edbc4638 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_embedding entry of the MiniMax-M3 tp4 target: an exact gather with zero rows outside the +table.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.embedding.m3_embedding import ( # noqa: E501 + m3_embedding as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" + +VOCAB, HIDDEN = 4096, 6144 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + ids = torch.randint( + 0, VOCAB, (spec["rows"],), generator=g, device="cuda", dtype=spec["id_dtype"] + ) + if spec["pad"]: + ids[1::3] = -1 + with validating(op): + out = op(ids, table) + assert torch.equal(out, op.reference(ids, table)) + if spec["pad"]: + assert not out[1::3].any(), "padded ids must embed to zero rows" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py new file mode 100644 index 000000000000..e44ad56d743d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" + +import torch + + +def swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): + """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" + v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) + return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..f17d7ec63189 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_bf16_linear entry of the MiniMax-M3 tp4 target (the drafter's projections).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_bf16_linear import ( # noqa: E501 + m3_bf16_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the + row, RMS well under one; measured on GB300: worst 0.5 unit.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w = (torch.randn(spec["n"], spec["k"], generator=g, device="cuda") * 0.02).to(torch.bfloat16) + with validating(op): + out = op(x, w) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(x, w), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py new file mode 100644 index 000000000000..c53155b7c02b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise oracle of the dense MLP's swigluoai C++ op (ops/gemm/m3_dense_swiglu, +trtllm::minimax_m3_sm107_tp4_swiglu_oai) against the Triton silu_and_mul kernel the dense MLP runs today +(modules/swiglu.py at limit 7, alpha 1.702, beta 1): every bf16 gate bit pattern against special and random up values, +every up bit pattern against special and random gates, random pairs, and the dense MLP's shapes at every decode row +count, eager and replayed from a CUDA graph, with and without the programmatic launch.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.silu_and_mul +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_dense_swiglu import ( # noqa: E501 + ALPHA, + BETA, + LIMIT, + m3_dense_swiglu, +) + +assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp4_swiglu_oai"): + pytest.skip( + "needs the TP4 op library with minimax_m3_sm107_tp4_swiglu_oai", allow_module_level=True + ) + +# The TP4 dense MLP's activation width (gate_up 3072 -> d 1536). +DENSE_D = dispatch.DENSE_DOWN[1] + + +def _triton(x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=LIMIT, swiglu_alpha=ALPHA, swiglu_beta=BETA + ) + + +def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" + return out + + +def _row(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + return torch.cat([gate.view(1, -1), up.view(1, -1)], dim=1).contiguous() + + +def _all_bf16() -> torch.Tensor: + return ( + torch.arange(1 << 16, dtype=torch.int32, device="cuda").to(torch.int16).view(torch.bfloat16) + ) + + +def _specials(n: int, seed: int) -> list: + g = torch.Generator(device="cuda").manual_seed(seed) + vals = [0.0, -0.0, 1.0, -1.0, 6.99, 7.0, 7.01, -7.0, -7.01, 1e-38, -1e-38, 1e-40, 88.0, -88.0, 3e38, -3e38, + float("inf"), float("-inf"), float("nan")] # fmt: skip + rows = [torch.full((n,), v, dtype=torch.bfloat16, device="cuda") for v in vals] + for scale in (0.05, 1.0, 4.0, 30.0): + rows.append((scale * torch.randn(n, generator=g, device="cuda")).to(torch.bfloat16)) + rows.append( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + return rows + + +def _assert_bitwise(got: torch.Tensor, ref: torch.Tensor, what: str) -> None: + bad = (got.view(torch.int16) != ref.view(torch.int16)).nonzero().flatten() + assert bad.numel() == 0, f"{what}: {bad.numel()} mismatches, first at {bad[:4].tolist()}" + + +def test_every_gate_pattern_is_bitwise_the_triton_kernel() -> None: + gate = _all_bf16() + for j, up in enumerate(_specials(gate.numel(), 1)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"up set {j}") + + +def test_every_up_pattern_is_bitwise_the_triton_kernel() -> None: + up = _all_bf16() + for j, gate in enumerate(_specials(up.numel(), 2)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"gate set {j}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: + g = torch.Generator(device="cuda").manual_seed(100 + seed) + n = 1 << 22 + + def pattern() -> torch.Tensor: + bits = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + return bits.to(torch.int16).view(torch.bfloat16) + + x = _row(pattern(), pattern()) + _assert_bitwise(_op(x), _triton(x), f"seed {seed}") + + +@pytest.mark.parametrize("pdl", [True, False], ids=["pdl", "plain"]) +@pytest.mark.parametrize("rows", range(1, dispatch.DECODE_MAX_TOKENS + 1)) +def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pdl: bool) -> None: + """The TP4 dense MLP's gate_up rows (a GEMV output's range plus values past the clamps), eager and replayed from a + CUDA graph after fresh rows are written in place, against the Triton kernel on the same rows.""" + g = torch.Generator(device="cuda").manual_seed(2000 + rows) + x = (3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, DENSE_D : DENSE_D + 8] = x[:, :8] + _assert_bitwise(_op(x, pdl), _triton(x), f"eager rows {rows}") + out = torch.empty(rows, DENSE_D, dtype=torch.bfloat16, device="cuda") + _op(x, pdl) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out.copy_(_op(x, pdl)) + x.copy_((3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(out, _triton(x), f"graph rows {rows}") + + +def test_the_op_declines_other_parameters_and_layouts() -> None: + x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + + +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: + """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton + kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" + rows, d = cell.spec["rows"], cell.spec["d"] + g = torch.Generator(device="cuda").manual_seed(3000 + rows) + x = (3.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, d : d + 8] = x[:, :8] + _assert_bitwise(_op(x), _triton(x), cell.why) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py new file mode 100644 index 000000000000..ebc36807e49b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_draft_lowm entry: the Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M +kernels (direct SIMT and split-K tensor-core), at the TP4 drafter shards, at the rows its table certifies. + +Every cell is held to the fp64 reference within the derived band. The drafter's projections (``M3Eagle3Draft._linear`` +over ``ops/m3_draft_lowm``) are held on the drafter's weight table: a covered row count is the entry's call bit for +bit, an uncovered one the plain GEMM's. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_draft_lowm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + TACTICS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 + m3_draft_lowm as op, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" + +# Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). +# +# Kernel and reference multiply the same BF16 operands (every product exact in fp32); the kernel sums in fp32 (the +# direct kernel one dot product in K order, the split-K kernel tensor-core blocks and then its split partials), the +# reference in fp64, and both round once to BF16. With w ~ N(0, 1 / K) and x ~ N(0, 1) the outputs are O(1) and the +# fp32 order difference is ~sqrt(K) * 2^-24 of the summed magnitudes, a few 1e-6 at K = 18432: far below half a BF16 +# ulp, so the two roundings disagree by at most one ulp of the element -- up to two units of the row scale for an +# element in its row maximum's binade -- and only where the sum sits that close to a rounding boundary, a few in ten +# thousand elements, hence 0.1 ulp RMS. The same derivation as the decode GEMV's band (test_modeling_v2_m3_mx_gemv), +# with no slack under it: a second ulp would mean a wrong sum. +ELEMENT_ULP = 2.0 +RMS_ULP = 0.1 + +# (N, K) -> the drafter's weight key (M3Eagle3Draft.w). +DRAFTER_LINEARS = { + (2304, 12288): "qkv", + (6144, 2048): "o", + (9216, 6144): "gu", + (6144, 4608): "dn", + (6144, 18432): "fc", +} + + +@pytest.fixture(autouse=True) +def _policy(): + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + return x, w + + +def test_the_table_is_the_tp4_drafter() -> None: + assert set(TACTICS) == set(DRAFTER_LINEARS) + assert {c.spec["nk"] for c in op.CELLS} == set(TACTICS) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + kind, tactic = tactic_for(n, k, m) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + args = ( + x, + w, + out, + kind, + tactic, + True, + ) # the target launches with PDL when the environment enables it + certifying(op) + assert op(*args) is out + ref = op.reference(*args) + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=RMS_ULP) + # PDL only changes the launch attribute. + plain = torch.empty_like(out) + op(x, w, plain, kind, tactic, False) + assert torch.equal(plain.view(torch.int16), out.view(torch.int16)) + + +def _draft(linears=DRAFTER_LINEARS): + """A drafter holding only its projection weights (the module is not built).""" + draft = object.__new__(dispatch.M3Eagle3Draft) + draft.w = {key: build(*nk, 1, seed=nk[0] + nk[1])[1] for nk, key in linears.items()} + draft._lowm_keys = frozenset() + return draft + + +@pytest.fixture(scope="module") +def prebuilt(): + draft = _draft() + draft.prebuild_linears() # 0 when an earlier test in this process compiled them + assert draft._lowm_keys == frozenset(DRAFTER_LINEARS.values()) + assert draft.prebuild_linears() == 0, "a second prebuild finds nothing left to compile" + return draft + + +@pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) +def test_the_projection_is_the_entry_and_falls_back_past_its_rows(prebuilt, nk) -> None: + n, k = nk + key = DRAFTER_LINEARS[nk] + w = prebuilt.w[key] + use_pdl = get_env_enable_pdl() + dispatch.set_require(True) + for m in ROWS[nk]: + assert dispatch.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt._linear(x, key) + kind, tactic = tactic_for(n, k, m) + want = op(x, w, torch.empty(m, n, dtype=x.dtype, device="cuda"), kind, tactic, use_pdl) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)), m + # Rows the table does not certify (past its rows): the plain GEMM, silently, even under require. + for m in sorted(set(range(1, dispatch.DECODE_MAX_TOKENS + 2)) - set(ROWS[nk])): + assert not dispatch.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt._linear(x, key) + assert torch.equal(got.view(torch.int16), F.linear(x, w).view(torch.int16)), m + + +def test_a_projection_outside_the_table_is_the_plain_gemm() -> None: + assert not m3_draft_lowm.installed(torch.empty(256, 512, dtype=torch.bfloat16)) + draft = _draft({(256, 512): "o"}) + draft._lowm_keys = frozenset(k for k in ("o",) if m3_draft_lowm.installed(draft.w[k])) + x, _ = build(256, 512, 1, seed=3) + got = draft._linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_a_covered_call_not_prebuilt_raises_under_require(monkeypatch) -> None: + monkeypatch.setattr(m3_draft_lowm, "_BUILT", set()) + monkeypatch.setattr(m3_draft_lowm, "_KEYS", frozenset(dispatch.DRAFT_LOWM.keys)) + w = build(6144, 2048, 1, seed=4)[1] + x, _ = build(6144, 2048, 1, seed=5) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not compiled before warmup"): + m3_draft_lowm.draft_linear(w, x, None) + # Outside require the drafter falls back to the plain GEMM. + dispatch.set_require(False) + assert m3_draft_lowm.draft_linear(w, x, None) is None + + +def test_guard_refuses_a_strided_activation() -> None: + x, w = build(6144, 2048, 2, seed=6) + wide = torch.zeros(2, 4096, dtype=torch.bfloat16, device="cuda") + wide[:, :2048] = x + kind, tactic = tactic_for(6144, 2048, 2) + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op( + wide[:, :2048], + w, + torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda"), + kind, + tactic, + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..fc71fc9f522a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp32_linear entry of the MiniMax-M3 tp4 target (the router beyond the decode routers' rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """torch's default fp32 matmul precision on this hardware may be TF32 (~2^-11 relative per product): the band is + stated in fp32 ulps of each row's largest logit and covers it (K = 6144 TF32 products: ~2^13 ulps worst element); + measured on GB300 with the engine's defaults: worst element ~1700 ulps, RMS ~300.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(spec["e"], spec["k"], generator=g, device="cuda") * 0.02 + with validating(op): + out = op(x, w) + assert out.dtype == torch.float32 and out.shape == (spec["m"], spec["e"]) + assert_within_ulp(out, op.reference(x, w), element_ulp=8192.0, rms_ulp=2048.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..0cd67d512573 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv entry: the decode MXFP8 GEMV on a producer-quantized activation. + +Operands as the target holds them: the weight through mxfp8_quantize with the swizzled scale layout (the MXFP8 linear +method's storage), the activation through mxfp8_quantize with the swizzled layout (what the boundary all-reduce and +the attention merge publish). +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" + + +# Element band, in ULP units of the row scale: see compare. +ELEMENT_ULP = 1.5 + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Element ELEMENT_ULP of the row scale, relative RMS 0.1 ulp. + + Kernel and reference multiply bit-identical dequantized operands; the kernel sums in fp32 (tensor-core blocks, + then the split-K partials), the reference in fp64. Over K <= 6144 that order difference is ~1e-6 relative, so the + two bf16 roundings disagree only where the sum sits within that of a rounding boundary: one bf16 ulp of that + element, which for an element in its row maximum's binade is up to two of ``assert_within_ulp``'s units of the row + scale, and a few in ten thousand elements in aggregate. Any fp32 sum shows it: torch's IEEE fp32 matmul of the + same dequantized operands measured at most 1.22 units against the fp64 reference over every row count 1..48 on + the o_proj and the QKV (6 seeds each, GB300), the kernel at most 1.25, both at most 0.009 ulp RMS. The band + leaves the margin between that and a full extra ulp. + """ + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=0.1) + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + return x, xq.view(m, k), xsf, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + _, xq, xsf, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(xq, xsf, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + compare(out, op.reference(xq, xsf, w, wsf)) + + +def test_an_unserved_shape_returns_none() -> None: + """Shapes outside the tuned table fall back at the call site; the op says so rather than computing.""" + _, xq, xsf, w, wsf = build(256, 512, 1, seed=1) + assert op(xq, xsf, w, wsf) is None + + +def test_guard_refuses_linear_activation_scales() -> None: + """A linear [M, K / 32] scale buffer is read as the swizzled one.""" + x, _, _, w, wsf = build(6144, 2048, 4, seed=2) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, False) + with pytest.raises(AssertionError, match="swizzled"): + with validating(op): + op(xq.view(4, 2048), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + _, xq, xsf, w, wsf = build(6144, 2048, 1, seed=3) + wide = torch.zeros(6144, 4096, dtype=w.dtype, device="cuda") + wide[:, :2048] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :2048], wsf) + + +@pytest.mark.parametrize("nk", [(2560, 6144), (2304, 6144)], ids=["qkv_sparse", "qkv_dense"]) +@pytest.mark.parametrize("m", range(1, 9)) +def test_qkv_split8_config(nk, m) -> None: + """The TP4 dispatch's (split, ring) override for the QKV at decode rows: within the band of the fp64 reference, + and the graph replay bit-identical to the eager call (one split-K order per config).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as dispatch, + ) + + cfg = dispatch.mx_gemv_cfg(nk[0], nk[1], m) + assert cfg == (8, 8) + _, xq, xsf, wq, wsf = build(nk[0], nk[1], m, zlib.crc32(f"qkv8-{nk}-{m}".encode())) + eager = op(xq, xsf, wq, wsf, cfg=cfg) + assert eager is not None, "the GEMV declined the split-8 config" + compare(eager, op.reference(xq, xsf, wq, wsf)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(xq, xsf, wq, wsf, cfg=cfg) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured.view(torch.int16), eager.view(torch.int16)) + + +def test_qkv_split8_config_is_bounded(monkeypatch) -> None: + """The overrides cover only the two QKV projections (split 8) and the o_proj (strided k-tiles, rule 1b) at decode + rows; every other call, and every call with its kill knob off, keeps the shared table.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as dispatch, + ) + + assert dispatch.mx_gemv_cfg(2560, 6144, 9) is None + assert dispatch.mx_gemv_cfg(6144, 2048, 4) == (2, 8, 0) + assert dispatch.mx_gemv_cfg(6144, 2048, 9) is None + assert dispatch.mx_gemv_cfg(6144, 3072, 4) is None + monkeypatch.setattr(dispatch, "GEMV_OPROJ_STRIDED_TP4", False) + assert dispatch.mx_gemv_cfg(6144, 2048, 4) is None + assert dispatch.mx_gemv_cfg(2560, 6144, 4) == (8, 8) + + +@pytest.mark.parametrize( + "nk", + [(2560, 6144), (2304, 6144), (6144, 2048)], + ids=["qkv_sparse", "qkv_dense", "o_proj_strided"], +) +def test_prebuild_covers_the_split8_config_while_frozen(nk) -> None: + """The core's prebuild (``prebuild_kernels``: ``_gemv_variants`` + ``_precompile_cfg``) compiles every (rows, cfg) + ``_linear`` can pick for this projection, the TP4 (split, ring) override included: with compiles frozen afterwards + (as the target freezes them before warmup), every decode-row MXFP8 call at the override is served, none + declines.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as dispatch, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, + ) + + n, k = nk + _, _, _, wq, wsf = build(n, k, 1, zlib.crc32(f"prebuild-{nk}".encode())) + mx_gemv.freeze(False) + mx_gemv.ensure_compiled(wq, wsf, True) + for rows, x_bf16, ar in dispatch._gemv_variants(n, k, None, 4): + cfg = None if x_bf16 else dispatch.mx_gemv_cfg(n, k, rows) + if cfg is not None: + dispatch._precompile_cfg(wq, wsf, rows, cfg, True, ar=ar) + else: + mx_gemv.precompile(wq, wsf, rows, x_bf16, use_pdl=True, ar=ar) + picked = [ + (rows, dispatch.mx_gemv_cfg(n, k, rows)) + for rows in sorted(dispatch.MX_GEMV_ROWS[nk]) + if dispatch.mx_gemv_cfg(n, k, rows) is not None + ] + assert picked, "the override covers decode rows of this projection" + for rows, cfg in picked: + resolved = mx_gemv._resolve(rows, n, k, False, True, cfg, True, None) + assert resolved is not None and mx_gemv.is_compiled(resolved[0]), (rows, cfg) + mx_gemv.freeze(True) + try: + for rows, cfg in picked: + _, xq, xsf, _, _ = build(n, k, rows, rows) + out = mx_gemv.m3_mx_gemv(xq, xsf, wq, wsf, cfg=cfg, use_pdl=True) + assert out is not None, (rows, cfg, mx_gemv.last_uncompiled) + finally: + mx_gemv.freeze(False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..2c3947b5cc81 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv_bf16 entry: the decode MXFP8 GEMV quantizing a bf16 activation in its prologue. + +The weight as the MXFP8 linear method stores it (mxfp8_quantize, swizzled scales). Besides the reference band, the +prologue's quantization is held to mxfp8_quantize itself: the GEMV on a bf16 x must equal, bit for bit, the +prequantized GEMV on mxfp8_quantize(x) -- the claim the target relies on when it mixes the two paths. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv_bf16 import ( # noqa: E501 + m3_mx_gemv_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" + +# The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. +ELEMENT_ULP = 1.5 + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return x, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(x, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + # The band of test_modeling_v2_m3_mx_gemv.compare: the quantized operands are the reference's exactly (checked + # below), the sums differ only in order. + assert_within_ulp(out, op.reference(x, w, wsf), element_ulp=ELEMENT_ULP, rms_ulp=0.1) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + pre = m3_mx_gemv(xq.view(m, k), xsf, w, wsf, use_pdl=True) + assert pre is not None and torch.equal(out.view(torch.int16), pre.view(torch.int16)), ( + "the prologue's MXFP8 quantization differs from mxfp8_quantize" + ) + + +def test_a_strided_activation_is_declined() -> None: + """The prologue reads x densely; the op declines a strided x (None) rather than computing on it.""" + x, w, wsf = build(6144, 2048, 2, seed=1) + wide = torch.zeros(2, 4096, dtype=torch.bfloat16, device="cuda") + wide[:, :2048] = x + assert op(wide[:, :2048], w, wsf) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_order.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_order.py new file mode 100644 index 000000000000..51cd471a9089 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_order.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The TP4 decode GEMV order overrides (modeling.mx_gemv_cfg: QKV split 8, o_proj strided k-tiles): +for every overridden (projection, rows) the output is the same on 3 eager calls and on a CUDA graph replay +(deterministic), and its max abs / max rel error vs the fp64 product of the dequantized operands stays within the +shipped order's band (the shared table's config on the same operands, over several seeds).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, + QKV, + QKV_SPARSE, + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 + mxfp8_dequant, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +CASES = [(nk, m) for nk in (QKV_SPARSE, QKV, O_PROJ) for m in (1, 2, 4, 8)] + + +def _err(y, ref): + d = (y.double() - ref).abs() + return float(d.max()), float((d / ref.abs().amax(dim=1, keepdim=True).clamp_min(1e-30)).max()) + + +@pytest.mark.parametrize("nk, m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +def test_order_override_is_deterministic_and_in_band(nk, m) -> None: + n, k = nk + cfg = dispatch.mx_gemv_cfg(n, k, m) + assert cfg is not None, "the override is off or does not cover this call" + worst_new, worst_ship = [0.0, 0.0], [0.0, 0.0] + for seed in range(4): + g = torch.Generator(device="cuda").manual_seed(1000 * seed + m) + w = (torch.randn(n, k, device="cuda", generator=g) * 2).to(torch.float8_e4m3fn) + ws = torch.randint(118, 124, (n * k // 32,), dtype=torch.uint8, device="cuda", generator=g) + x = (torch.randn(m, k, device="cuda", generator=g) * (0.5 + seed)).to(torch.bfloat16) + q, s = torch.ops.trtllm.mxfp8_quantize(x, True) + ref = ( + mxfp8_dequant(q.view(torch.uint8), s).double() + @ mxfp8_dequant(w.view(torch.uint8), ws).double().t() + ) + ys = [m3_mx_gemv(q, s, w, ws, True, cfg, True, None) for _ in range(3)] + assert all(y is not None for y in ys), "the GEMV declined the override" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + out = [] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out.append(m3_mx_gemv(q, s, w, ws, True, cfg, True, None)) + graph.replay() + torch.cuda.synchronize() + for y in ys[1:] + [out[0]]: + assert torch.equal(ys[0].view(torch.int16), y.view(torch.int16)), ( + "not deterministic (eager / graph)" + ) + ship = m3_mx_gemv(q, s, w, ws, True, None, True, None) + torch.cuda.synchronize() + for acc, y in ((worst_new, ys[0]), (worst_ship, ship)): + e = _err(y, ref) + acc[0], acc[1] = max(acc[0], e[0]), max(acc[1], e[1]) + print( + f"\nR1B1 tp4 {n}x{k} m={m} cfg {cfg}: max abs / rel vs fp64 {worst_new[0]:.3e} / {worst_new[1]:.3e} | shipped " + f"order {worst_ship[0]:.3e} / {worst_ship[1]:.3e}", + flush=True, + ) + # The band: the shipped order's own error (bf16 output rounding dominates), with 25 % headroom. + assert worst_new[0] <= 1.25 * worst_ship[0] and worst_new[1] <= 1.25 * worst_ship[1] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..54379ac1f3e2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_gemm entry of the MiniMax-M3 tp4 target (the decode-graph GEMM on a prequantized +activation).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most + 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, False), + element_ulp=2.0, + rms_ulp=0.25, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..b869ca3aef13 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_linear entry of the MiniMax-M3 tp4 target: MXFP8 quantize + GEMM on each backend the +engine selects (native CUTLASS, FlashInfer, the per-bucket tuned quantize + GEMM).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" + +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and + one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" + spec = cell.spec + if ( + spec["path"] == "flashinfer" + and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None + ): + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + args = (x, w, w_sf, alpha, torch.bfloat16, *_FLAGS[spec["path"]]) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py new file mode 100644 index 000000000000..a2ae95a27105 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_pinned_gemm entry (MX-TAC): the pinned decode-graph MXFP8 GEMM gives the same bits as +every candidate stock's AutoTuner can pick (CuTe DSL, CUTLASS), the same on every call, within the MXFP8 band of the +dequantized product.""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_pinned_gemm import ( # noqa: E501 + m3_mxfp8_pinned_gemm as op, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_pinned_gemm requires a CUDA device" + +TACTIC = modeling.MXFP8_DECODE_GEMM_TACTIC + +# The TP4 projections at rows the GEMVs do not serve (SHARED_DOWN 17..32, over 48) and the bf16 activations over the +# GEMVs' bf16 bounds. +CASES = [ + ((6144, 768), 17), + ((6144, 768), 32), + ((2560, 6144), 24), + ((2304, 6144), 64), + ((6144, 2048), 49), + ((1536, 6144), 20), + ((6144, 6144), 12), + ((6144, 3072), 56), +] + + +def test_the_pinned_tactic_is_cute_dsl() -> None: + assert TACTIC == FlashInferMXFP8GemmRunner.CUTE_DSL + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, TACTIC) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, TACTIC), + element_ulp=2.0, + rms_ulp=0.25, + ) + + +@pytest.mark.parametrize("nk,m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_pinned_decision_matches_every_stock_candidate(nk, m, scale) -> None: + n, k = nk + g = torch.Generator(device="cuda").manual_seed(n + k + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = (torch.randn(m, k, generator=g, device="cuda") * scale).to(torch.bfloat16) + x[0, ::32] *= 1e3 # one outlier per 32-element block on the first row + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + wq = wq.view(n, k) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + xq = xq.view(m, k) + out = op(xq, xsf, wq, wsf, torch.bfloat16, TACTIC) + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + for tactic in (FlashInferMXFP8GemmRunner.CUTE_DSL, FlashInferMXFP8GemmRunner.CUTLASS): + ref = runner.forward([xq, xsf, wq, wsf], tactic=tactic) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), tactic + again = op(xq, xsf, wq, wsf, torch.bfloat16, TACTIC) + assert torch.equal(out.view(torch.int16), again.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py new file mode 100644 index 000000000000..22e987e17979 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_pinned_quantize entry (MX-TAC): the pinned decode-graph MXFP8 activation quantize is the +stock quantize bit for bit, at every tactic stock's AutoTuner can pick.""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_pinned_quantize import ( # noqa: E501 + m3_mxfp8_pinned_quantize as op, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import MXFP8QuantizeRunner + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_pinned_quantize requires a CUDA device" + +TACTIC = modeling.MXFP8_DECODE_QUANT_TACTIC + + +def test_the_pinned_tactic_is_cute_dsl() -> None: + assert TACTIC == MXFP8QuantizeRunner.CUTE_DSL + + +def _check(x: torch.Tensor) -> None: + m, k = x.shape + with validating(op): + xq, xsf = op(x, TACTIC) + tq, tsf = op.reference(x, TACTIC) + assert torch.equal(xq.view(torch.uint8).view(m, k), tq.view(torch.uint8).view(m, k)) + assert torch.equal(xsf.view(torch.uint8).flatten(), tsf.view(torch.uint8).flatten()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_certified_cells(cell, scale) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * scale).to(torch.bfloat16) + x[0, ::32] *= 1e3 # one outlier per 32-element block on the first row + _check(x) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py new file mode 100644 index 000000000000..467cfd888e5e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py @@ -0,0 +1,253 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_target_head entry: the target lm_head shard GEMV over the lossless 12-bit coding of its +BF16 weight, against its reference, the in-tree BF16 head GEMM (``F.linear``) and fp64. + +* The coding round-trips bit for bit (code values and exception lists), and a corrupted coding fails the check. +* Integer-valued data (fp32 sums exact in any order): the kernel's logits equal the in-tree GEMM's bit for bit. +* Random data: the kernel's error against fp64 is the in-tree GEMM's (the same products, another fp32 order, one + BF16 rounding), its argmax is the in-tree one and its result is the same run to run. +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_target_head as target_head_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + CTA_WARPS, + HIDDEN, + TARGET_SHARD, + decode, + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + m3_target_head as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_target_head requires a CUDA device" + + +def _seed(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def _num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def _weight(n: int, k: int, seed: int, exceptions: bool = False) -> torch.Tensor: + """A [n, k] head of trained-weight scale; with ``exceptions``, values on both sides of the exponent window, exact + zeros and a large outlier in many tiles.""" + g = torch.Generator(device="cuda").manual_seed(seed) + w = torch.randn(n, k, device="cuda", generator=g) * 0.02 + if exceptions: + w[::97, ::131] = 1e-9 + w[::89, 7] = 0.0 + w[3, 5] = 7.5 + w[n - 1, k - 1] = -3.0 + return w.to(torch.bfloat16) + + +_CODED: dict = {} + + +def _coded(n: int, k: int, seed: int, exceptions: bool): + """(weight, encode_chunked(weight)), built once per shape.""" + key = (n, k, seed, exceptions) + if key not in _CODED: + _CODED.clear() + w = _weight(n, k, seed, exceptions) + _CODED[key] = (w, encode_chunked(w)) + return _CODED[key] + + +def _args(coded, x: torch.Tensor, n: int, grid: int): + packed, base, exc_off, exc_pos, exc_w = coded + out = torch.full((x.shape[0], n), float("nan"), dtype=torch.bfloat16, device="cuda") + part = torch.zeros(grid * CTA_WARPS * 256, dtype=torch.float32, device="cuda") + counter = torch.zeros((n + 15) // 16, dtype=torch.int32, device="cuda") + return (packed, x, out, part, counter, exc_off, exc_pos, exc_w.view(-1), base, n, grid, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The kernel against its reference and the in-tree BF16 GEMM: the error against fp64 is the in-tree GEMM's + (both are one BF16 rounding of an fp32 sum of the same exact products), the argmax is the in-tree one, nearly + every logit is the in-tree one bit for bit, and a second call gives the same bits.""" + spec = cell.spec + m, n, k = spec["m"], spec["n"], spec["k"] + exceptions = spec.get("exceptions", False) + w, coded = _coded(n, k, 1, exceptions) + if exceptions: + assert coded[3].numel() > 0, "the cell must exercise the exception lists" + g = torch.Generator(device="cuda").manual_seed(_seed(cell.why)) + x = torch.randn(m, k, device="cuda", generator=g).to(torch.bfloat16) + args = _args(coded, x, n, _num_sms()) + certifying(op) + ref = op.reference(*args) + out = op(*args).clone() + again = op(*_args(coded, x, n, _num_sms())).clone() + torch.cuda.synchronize() + assert not bool(args[4].any()), "the split-tile counters are not re-armed" + assert torch.equal(out, again), "the split-tile reduction is not deterministic" + stock = torch.nn.functional.linear(x, w) + exact = x.double() @ w.double().t() + # Each side: one BF16 rounding (<= half an ulp of the logit) of an fp32 sum whose order error is far below it. + err_out = (out.double() - exact).abs().max() + err_stock = (stock.double() - exact).abs().max() + err_ref = (ref.double() - exact).abs().max() + assert err_out <= 1.01 * max(err_stock, err_ref) + 1e-6, (err_out, err_stock, err_ref) + assert (out == stock).float().mean() > 0.99 + assert (out == ref).float().mean() > 0.99 + assert torch.equal(out.float().argmax(-1), stock.float().argmax(-1)) + + +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_integer_data_is_the_bf16_head_bit_for_bit(m) -> None: + """Small-integer weights and activations: every fp32 sum is exact in any order, so the coded head's logits are + the in-tree BF16 GEMM's bit for bit (exceptions included: zeros and integers outside the exponent window).""" + g = torch.Generator(device="cuda").manual_seed(_seed(f"integer {m}")) + n, k = 4096, HIDDEN + w = torch.randint(-3, 4, (n, k), generator=g, device="cuda").to(torch.bfloat16) + w[::7, ::5] = 96.0 # far above the window of the {1, 2, 3} exponents: exceptions + x = torch.randint(-2, 3, (m, k), generator=g, device="cuda").to(torch.bfloat16) + coded = encode_chunked(w) + assert coded[3].numel() > 0 + assert round_trip_ok(*coded, w) + with validating(op): + out = op(*_args(coded, x, n, _num_sms())) + assert torch.equal(out, torch.nn.functional.linear(x, w)) + + +def test_the_head_writes_only_the_shards_columns() -> None: + """Integer data over a 4000-row shard: every logit is the in-tree GEMM's bit for bit, and the columns past n of a + wider output are left untouched.""" + n, k, m = 4000, 1024, 2 + g = torch.Generator(device="cuda").manual_seed(_seed("partial last tile")) + # Small integers: every fp32 sum is exact in any order, so the logits are comparable bit for bit. + w = torch.randint(-3, 4, (n, k), generator=g, device="cuda").to(torch.bfloat16) + w[::7, ::5] = 96.0 # exceptions, including in the partial tile + x = torch.randint(-2, 3, (m, k), generator=g, device="cuda").to(torch.bfloat16) + coded = encode_chunked(w) + assert round_trip_ok(*coded, w) + args = list(_args(coded, x, n, _num_sms())) + args[2] = torch.full((m, n + 16), float("nan"), dtype=torch.bfloat16, device="cuda") + with validating(op): + op(*args) + out = args[2] + assert torch.isnan(out[:, n:]).all(), "the head wrote past the shard's last column" + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out[:, :n], stock) + + +def test_a_shard_of_partial_tiles_keeps_the_bf16_head() -> None: + """The kernel streams whole 16-row tiles: a shard whose rows are not a multiple of 16 is not taken (the fused + verify keeps the in-tree GEMM), and the op refuses it.""" + w = torch.zeros(4004, 1024, dtype=torch.bfloat16, device="cuda") + assert not target_head_kernels.head_supported(SimpleNamespace(weight=w, bias=None)) + assert target_head_kernels.head_supported(SimpleNamespace(weight=w[:4000], bias=None)) + + +@pytest.mark.parametrize("n,k", [(TARGET_SHARD, HIDDEN), (4004, 1024)]) +def test_round_trip_is_bitwise(n, k) -> None: + w, coded = _coded(n, k, 1, True) + assert coded[3].numel() > 0 + assert round_trip_ok(*coded, w) + assert torch.equal(decode(*coded, n, k).view(torch.int16), w.view(torch.int16)) + + +def test_a_corrupted_coding_fails_closed() -> None: + """A flipped sign/mantissa bit, a wrong exception value or a wrong base each fail the bitwise round trip.""" + w = _weight(4004, 1024, 2, exceptions=True) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + bad = packed.clone() + bad[100] ^= 1 + assert not round_trip_ok(bad, base, exc_off, exc_pos, exc_w, w) + bad_w = exc_w.clone() + bad_w[0, 0] += 1e-3 + assert not round_trip_ok(packed, base, exc_off, exc_pos, bad_w, w) + assert not round_trip_ok(packed, base + 1, exc_off, exc_pos, exc_w, w) + + +def test_prepare_head_keeps_only_a_bitwise_coding(monkeypatch) -> None: + """prepare_head keeps the coded shard when it round-trips and runs it as the entry; a coding that does not + round-trip is dropped (the in-tree GEMM stays the path).""" + w = _weight(4000, 1024, 3, exceptions=True) + head = target_head_kernels.prepare_head(SimpleNamespace(weight=w, bias=None)) + assert head.coded and head.weight is w + x = torch.randn(4, 1024, device="cuda").to(torch.bfloat16) + out = target_head_kernels.logits(head, x) + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out.float().argmax(-1), stock.float().argmax(-1)) + assert (out == stock).float().mean() > 0.99 + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[7] ^= 0x40 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(target_head_kernels, "encode_chunked", corrupt) + dropped = target_head_kernels.prepare_head(SimpleNamespace(weight=w, bias=None)) + assert not dropped.coded and dropped.w is None + + +def test_guard_refuses_a_dirty_counter() -> None: + w, coded = _coded(4000, 1024, 1, True) + args = list(_args(coded, torch.zeros(2, 1024, dtype=torch.bfloat16, device="cuda"), 4000, 8)) + args[4][3] = 1 + with pytest.raises(AssertionError, match="counter must be zero"): + with validating(op): + op(*args) + + +def _tiny_exceptions_head(n: int = 4096, k: int = 1024) -> torch.Tensor: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1: exceptions far + below the window whose exact products are the only non-zero logits.""" + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + return w + + +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_tiny_exception_products_are_exact(m) -> None: + """An exception adds exactly x * w: with a one-hot activation on column 0 the logits are 0, 2^-30 and 2^-29 bit + for bit as the BF16 head's, and row 1 wins (a correction against a non-zero placeholder cancels them to 0).""" + w = _tiny_exceptions_head() + coded = encode_chunked(w) + assert coded[1] == 113 and round_trip_ok(*coded, w) + x = torch.zeros((m, w.shape[1]), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + with validating(op): + out = op(*_args(coded, x, w.shape[0], _num_sms())) + stock = torch.nn.functional.linear(x, w) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)) + assert out.float().argmax(-1).tolist() == [1] * m + + +def test_a_shard_smaller_than_the_grid_is_stored() -> None: + """A shard with fewer coded units than warps of the requested grid: the launch narrows the grid so every warp owns + a unit, and every logit is stored.""" + w = _tiny_exceptions_head(n=256) + coded = encode_chunked(w) + x = torch.zeros((2, w.shape[1]), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + out = op(*_args(coded, x, w.shape[0], _num_sms())) + assert torch.equal(out.view(torch.int16), torch.nn.functional.linear(x, w).view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..a04f1da1ed81 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_moe_operands.py @@ -0,0 +1,170 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One tp4 rank's MiniMax-M3 MoE layer operands for the decode-engine entries' tests. + +The routed experts are random checkpoint-format NVFP4 tensors ([w3 (up); w1 (gate)] rows, e4m3 block scales) put +through TRT-LLM's own TRTLLM-Gen NVFP4 loader shuffle (``NVFP4TRTLLMGenFusedMoEMethod._shuffle_all_experts``), which +writes the buffers the engines read in place. The per-expert scalars follow the loader's conventions (act_beta and +act_clamp divided by fc31_alpha, fc31_scale_c = fc2_input_scale x fc31_alpha) and are sized so the FC1 sums land +around the SwiGLU-OAI clamp and the outputs near unit scale. The shared expert's MXFP8 weights come from +mxfp8_quantize with the swizzled scale layout, as the MXFP8 linear method stores them. +""" + +import functools +from types import SimpleNamespace + +import torch + +H, I_TP, NUM_EXPERTS, TOP_K = 6144, 768, 128, 4 +SHARED_I = 768 +SWIGLU_ALPHA, SWIGLU_LIMIT = 1.702, 7.0 +ROUTED_SCALING = 2.0 +MINIMAX2_ROUTING = 5 # RoutingMethodType.MiniMax2 +# NVFP4 global scales: 448 x 6 over the amax each side is sized for. +FC31_INPUT_SCALE = 448.0 * 6.0 / 8.0 +FC2_INPUT_SCALE = 448.0 * 6.0 / 24.0 + + +def _rand_e4m3_scales(shape, gen) -> torch.Tensor: + """E4M3 block scales in [0.5, 2): biased exponents 6 and 7, any mantissa.""" + return torch.randint(0x30, 0x40, shape, dtype=torch.uint8, device="cuda", generator=gen).view( + torch.float8_e4m3fn + ) + + +@functools.lru_cache(maxsize=None) +def experts(seed: int = 20261013): + """The routed experts and the router of one MoE layer, in the layout and conventions of the TRTLLM-Gen loader. + + Returns a SimpleNamespace: ``tensors`` (the ten expert operands in ``M3MoeState.layer``'s order), ``bias``, + ``router_weight``, ``fc31_input_scale`` and ``shared`` (gate_up, its scales, down, its scales). + """ + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + gen = torch.Generator(device="cuda").manual_seed(seed) + e = NUM_EXPERTS + w31 = torch.randint( + 0, 256, (e, 2 * I_TP, H // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + w2 = torch.randint(0, 256, (e, H, I_TP // 2), dtype=torch.uint8, device="cuda", generator=gen) + w31s = _rand_e4m3_scales((e, 2 * I_TP, H // 16), gen) + w2s = _rand_e4m3_scales((e, H, I_TP // 16), gen) + module = SimpleNamespace( + w3_w1_weight=w31, + w2_weight=w2, + w3_w1_weight_scale=w31s, + w2_weight_scale=w2s, + is_gated_activation=True, + bias=False, + ) + NVFP4TRTLLMGenFusedMoEMethod()._shuffle_all_experts(module) + torch.cuda.synchronize() + + # Raw FC1 sums: rms(x) x FC31_INPUT_SCALE x rms(E2M1 value) x rms(block scale) x sqrt(H) ~ 1.0e5 at rms(x) = 1; + # fc31_alpha brings them to ~1.5 in the dequantized domain, the FC2 alpha the outputs to ~1. + base = 1.5 / (FC31_INPUT_SCALE * 2.93 * 1.3 * H**0.5) + fc31_alpha = base * (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) + fc2_alpha = (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) / ( + FC2_INPUT_SCALE * 2.93 * 1.3 * I_TP**0.5 + ) + tensors = ( + w31, + w31s, + w2, + w2s, + fc31_alpha, + FC2_INPUT_SCALE * fc31_alpha, # fc31_scale_c + torch.full((e,), SWIGLU_ALPHA, device="cuda"), # act_alpha + 1.0 / fc31_alpha, # act_beta: the (up + 1) offset over fc31_alpha + SWIGLU_LIMIT / fc31_alpha, # act_clamp over fc31_alpha + fc2_alpha, + ) + bias = (0.02 * torch.randn(e, device="cuda", generator=gen)).float() + router_weight = (torch.randn(e, H, device="cuda", generator=gen) / H**0.5).float().contiguous() + fc31_input_scale = torch.tensor([FC31_INPUT_SCALE], dtype=torch.float32, device="cuda") + + gu = (0.02 * torch.randn(2 * SHARED_I, H, device="cuda", generator=gen)).to(torch.bfloat16) + dn = (0.05 * torch.randn(H, SHARED_I, device="cuda", generator=gen)).to(torch.bfloat16) + gu_q, gu_sf = torch.ops.trtllm.mxfp8_quantize(gu, True) + dn_q, dn_sf = torch.ops.trtllm.mxfp8_quantize(dn, True) + shared = (gu_q.view(2 * SHARED_I, H), gu_sf, dn_q.view(H, SHARED_I), dn_sf) + return SimpleNamespace( + tensors=tensors, + bias=bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared=shared, + ) + + +# Routing cases: random (up to 4 M distinct experts), and every token on the same four experts (the bias outweighs +# any sigmoid difference), which leaves later slot groups empty. +CONCENTRATED = (5, 40, 77, 126) + + +def layer_for(state, case: str, shared: bool = True): + """One layer on ``state`` over :func:`experts` for a routing case. ``shared=False``: the shared expert's weights + are zero, so the engine returns the routed part alone (bf16(0 + bf16(routed))).""" + ops = experts() + bias = ops.bias.clone() + if case == "concentrated": + bias[list(CONCENTRATED)] += 1.0 + shared_weights = ops.shared + if not shared: + gu, gu_sf, dn, dn_sf = ops.shared + shared_weights = (torch.zeros_like(gu), gu_sf, torch.zeros_like(dn), dn_sf) + return state.layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=shared_weights, + ) + + +def stock_routed( + hidden: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """The stock TRTLLM-Gen NVFP4 MoE (fp4_quantize, then fp4_block_scale_moe_runner with SwiGLU-OAI) on the given + routing, over :func:`experts`: what the built-in MiniMaxM3MoE computes for the routed part.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + + ops = experts() + w31, w31s, w2, w2s, fc31_alpha, scale_c, act_alpha, act_beta, act_clamp, fc2_alpha = ops.tensors + xq, xsf = torch.ops.trtllm.fp4_quantize(hidden, ops.fc31_input_scale, 16, False, False) + out = torch.ops.trtllm.fp4_block_scale_moe_runner( + None, None, xq, xsf.view(torch.float8_e4m3fn).flatten(), + w31, w31s, None, act_alpha, act_beta, act_clamp, w2, w2s, None, + scale_c, fc31_alpha, fc2_alpha, + NUM_EXPERTS, TOP_K, None, None, I_TP, 0, NUM_EXPERTS, ROUTED_SCALING, MINIMAX2_ROUTING, True, 0, + topk_weights=topk_weights.to(torch.bfloat16), topk_ids=topk_ids.to(torch.int32), + ) # fmt: skip + return out[0] if isinstance(out, (list, tuple)) else out + + +def hidden_states(m: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(m, H, device="cuda", generator=gen).to(torch.bfloat16) + + +def check_routed_part(op, state, m: int, seed: int) -> None: + """Exact, independent of the reference: the engine's routing (its route_dbg record: ids, then the bf16 weights' + fp32 bits) is the reference routing bit for bit, and its routed part is the stock TRTLLM-Gen MoE's output on that + routing bit for bit -- the claim the engines are built on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=False) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + stock = stock_routed(hidden, ids, weights) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)), ( + "routed part differs from TRTLLM-Gen's" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py new file mode 100644 index 000000000000..724a2979c234 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared operands and gates of the router entry's tests (m3_router_tc).""" + +import zlib + +import torch + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__all__ = ["_build", "_compare_logits", "_seed", "tc"] + + +def _build(m: int, e: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(e, k, generator=g, device="cuda") * 0.02 + return x, w + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +def _compare_logits(out: torch.Tensor, ref: torch.Tensor) -> None: + """fp32 logits that differ from the reference only in summation order. + + Over K = 6144 products a reordered fp32 sum moves by a few ulps of the sum of magnitudes, which near + cancellation is ~sqrt(K) times the row's largest logit's ulp at most. Gated in fp32 ulps of each row's largest + logit: 64 element-wise, 8 in relative RMS. The tensor-core logits against a true fp32 product (the TF32 rounding + the band must see) sit at 2200+ / 1700 ulp. + """ + assert out.dtype == ref.dtype == torch.float32 and out.shape == ref.shape + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=8.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py new file mode 100644 index 000000000000..4ad3e870f33a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe16_engine entry: the 9..13-row MoE engine on its caller-owned state. + +One GB300 rank of the tp4 deployment, the target's build (128 routed CTAs + 24 shared-expert CTAs). Operands and the +band: ``_m3_moe_operands`` and ``test_modeling_v2_m3_moe_engine`` (the same math). +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe16_engine import ( # noqa: E501 + M3Moe16State, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe16_engine import ( # noqa: E501 + m3_moe16_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import H, check_routed_part, hidden_states, layer_for # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_moe16_engine requires a CUDA device" + +# As test_modeling_v2_m3_moe_engine.compare (same arithmetic, same quantization boundaries). +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + +_STATES = {} + + +def _state(m: int) -> M3Moe16State: + if m not in _STATES: + _STATES[m] = M3Moe16State(torch.device("cuda"), num_tokens=m) + return _STATES[m] + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + assert_within_ulp(out, op.reference(hidden, layer), element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +@pytest.mark.parametrize("m", (9, 13)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..c77b345c4bc7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe_engine entry: the 1..8-token MoE engine on its caller-owned state. + +One GB300 rank of the tp4 deployment: 128 experts (all local, intermediate 768 per rank), top 4, hidden 6144, the +shared expert folded in, ``num_ctas`` = SMs - 24 as the target builds it. Operands: ``_m3_moe_operands``. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 # noqa: E501 + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import H, check_routed_part, hidden_states, layer_for # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +_STATES = {} + + +def _state(m: int) -> M3MoeState: + if m not in _STATES: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + _STATES[m] = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + return _STATES[m] + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """The engine against the plain-torch MoE: BAND_ELEMENT bf16 ulp of the row's largest magnitude per element, + BAND_RMS ulp relative RMS. + + Both sides quantize the same operands with the same recipes (the activation's and the intermediate's NVFP4 codes + and block scales, the shared expert's MXFP8); they differ in summation order (tensor-core fp32 against fp64) and in + the PTX approximations (rcp / ex2 / tanh). Such a difference changes the output only where it moves a value across + a rounding boundary: one bf16 ulp of a slot output or of the sum, or, rarely, one intermediate NVFP4 code, which + reaches the output as one weight column's share. Measured over every cell, both routing cases and 3 seeds each + (GB300): at most 1.48 ulp element, 0.18 ulp RMS. The band leaves ~3x on both; a misread recipe (the activation's + 1 / (sf x (1 / sc)) spelling of the reciprocal on ~0.06% of its codes) measured 12 ulp element, 7 ulp RMS. + """ + assert_within_ulp(out, ref, element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + # The workspace is tagged per call and never reset: a repeated call is the same call. + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + compare(out, op.reference(hidden, layer)) + + +@pytest.mark.parametrize("m", (1, 2, 8)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +_GB_STATES = {} + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("m", range(1, 9)) +def test_engine_is_the_gb_tp4_engine_bit_for_bit(m, case) -> None: + """The target's engine (the Rubin m3_moe at the TP4 widths) against the GB TP4 engine it replaces, same operands + and CTA count: the per-expert accumulation order is the same, so the outputs are equal bit for bit.""" + # The shared GB300 engine leaves the tree once no target imports it. + gb_op = pytest.importorskip("tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.moe.op") + if m not in _GB_STATES: + sms = torch.cuda.get_device_properties(0).multi_processor_count + _GB_STATES[m] = gb_op.M3MoeState( + torch.device("cuda"), num_tokens=m, num_ctas=sms - 24, shared=True + ) + hidden = hidden_states(m, seed=300 + m) + out = layer_for(_state(m), case)(hidden) + ref = layer_for(_GB_STATES[m], case)(hidden) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)) + + +def test_layers_on_one_state_do_not_leak_into_each_other() -> None: + """Two layers share one state's workspace, in stream order: A, B, A gives A's output twice, bit for bit, and the + caller's ``out`` receives the same rows as a fresh tensor.""" + state = _state(4) + a, b = layer_for(state, "random"), layer_for(state, "concentrated") + hidden = hidden_states(4, seed=7) + first = op(hidden, a) + op(hidden, b) + out = torch.empty(4, H, dtype=torch.bfloat16, device="cuda") + second = op(hidden, a, out=out) + assert second.data_ptr() == out.data_ptr() + assert torch.equal(first.view(torch.int16), out.view(torch.int16)) + + +@pytest.mark.parametrize( + "bad", + ("misaligned hidden", "fp32 out", "short out", "long out", "int32 arp"), +) +def test_malformed_buffers_are_rejected_by_the_call(bad) -> None: + """Why the entry has no is_valid: the compiled call checks each buffer against its signature and raises.""" + state = _state(2) + layer = layer_for(state, "random") + hidden, kw = hidden_states(2, seed=3), {} + if bad == "misaligned hidden": + hidden = torch.randn(2 * H + 8, device="cuda").to(torch.bfloat16)[2 : 2 + 2 * H].view(2, H) + elif bad == "fp32 out": + kw["out"] = torch.empty(2, H, dtype=torch.float32, device="cuda") + elif bad == "short out": + kw["out"] = torch.empty(1, H, dtype=torch.bfloat16, device="cuda") + elif bad == "long out": + kw["out"] = torch.empty(3, H, dtype=torch.bfloat16, device="cuda") + else: + kw["arp"] = torch.zeros(8, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + op(hidden, layer, **kw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py new file mode 100644 index 000000000000..d999e3e1e83e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_flat_gate_tp4.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The TP4 MoE engine's FLAT stack is applied only up to FLAT_MAX_TOKENS (4) tokens per call. + +At 8 tokens FLAT was +2.27 us/call slower than the per-m tables, so above the cap the engine keeps the per-m +configuration (bitwise the same outputs). CPU only: M3MoeState is constructed with its device +queries stubbed and stopped at the kernel build, so the test reads the engine config and the engagement lines +without a GPU. +""" + +import types + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_moe import ( # noqa: E501 + op, +) + + +class _Built(Exception): + """Raised by the stubbed kernel build: the config is final by then.""" + + +def _engine(monkeypatch, num_tokens): + """(config, info lines) of M3MoeState(num_tokens) up to its kernel build.""" + lines = [] + seen = {} + + def build(config, path): + seen["config"] = dict(config) + raise _Built() + + monkeypatch.setattr(op, "is_supported", lambda device=None: (True, "")) + monkeypatch.setattr(op, "_kernel_module", build) + monkeypatch.setattr(op.torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr( + op.torch.cuda, + "get_device_properties", + lambda device: types.SimpleNamespace(multi_processor_count=op.HIDDEN // 32), + ) + from tensorrt_llm.logger import logger + + monkeypatch.setattr(logger, "info", lambda msg, *a, **k: lines.append(str(msg))) + with pytest.raises(_Built): + op.M3MoeState("cuda", num_tokens, num_ctas=op.HIDDEN // 32) + return seen["config"], [line for line in lines if "MiniMax-M3 MoE engine FLAT (TP4)" in line] + + +def _has_flat(config): + return all(config.get(k) == v for k, v in op.FLAT_STACK.items()) + + +@pytest.mark.parametrize("num_tokens", range(1, 9)) +def test_flat_applies_only_up_to_the_cap(monkeypatch, num_tokens): + config, lines = _engine(monkeypatch, num_tokens) + assert op.FLAT_MAX_TOKENS == 4 + if num_tokens <= 4: + assert _has_flat(config), config + assert lines == [ + "MiniMax-M3 MoE engine FLAT (TP4): b2img=0 lg_repl=16 ar_pre=0 late_trigger=0 " + f"tokens={num_tokens}" + ] + else: + # above the cap: the per-m tables + assert config.get("b2img") == op.B2IMG_BY_M[num_tokens - 1] + assert config.get("ar_pre") == op.AR_PRE_BY_M[num_tokens - 1] + assert lines == [ + f"MiniMax-M3 MoE engine FLAT (TP4) off above 4 tokens (per-m tables): tokens={num_tokens}" + ] + + +def test_flat_applies_helper(): + """op.flat_applies: the FLAT stack up to FLAT_MAX_TOKENS tokens, the per-m tables above.""" + assert [op.flat_applies(m) for m in range(1, 9)] == [True] * 4 + [False] * 4 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..e43c57f109a4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_tc entry of the MiniMax-M3 tp4 target: the TF32 tensor-core router of the unfused MoE's +9..16-row pure generation steps (the arithmetic the built-in gate takes at these rows).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import true_fp32_matmul + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import _build, _compare_logits, _seed, tc # noqa: E402 +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + + +@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +def test_tc_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) + with validating(tc): + outs = tc(x, w, None, False) + ref = tc.reference(x, w, None, False) + assert len(outs) == len(ref) == 1 + _compare_logits(outs[0], ref[0]) + + +def test_tc_is_the_built_in_gates_router_call() -> None: + """The target's entry is the op the built-in gate calls at these rows, with the same arguments: the same logits + bit for bit.""" + x, w = _build(14, 128, 6144, 7) + ours = tc(x, w, None, True)[0] + gate = torch.ops.trtllm.minimax_m3_router_tc(x, w, None, True)[0] + assert torch.equal(ours.view(torch.int32), gate.view(torch.int32)) + + +def test_tc_reference_is_tf32_and_the_gate_sees_it() -> None: + """The tensor-core router rounds the weight to TF32: against a true fp32 product its logits sit far outside + the band, so the TF32 rounding in ``reference`` is load-bearing rather than decoration.""" + x, w = _build(16, 128, 6144, 5) + out = tc(x, w, None, False)[0] + with true_fp32_matmul(): + fp32 = x.float() @ w.t() + with pytest.raises(AssertionError): + _compare_logits(out, fp32) + + +def test_guard_refuses_a_misaligned_activation() -> None: + """A row view starting 8 bytes into the buffer keeps a legal row stride but not a legal bulk-copy source.""" + x, w = _build(9, 128, 6144, 3) + flat = torch.empty(9 * 6144 + 4, dtype=torch.bfloat16, device="cuda") + shifted = flat[4:].view(9, 6144) + shifted.copy_(x) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(tc): + tc(shifted, w, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..c207756c1b49 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/m3_embedding_firstnorm entry: the decode embedding gather + the first layer's RMSNorm. + +The table is the model's full bf16 [200064, 6144] embedding shape (random values at the checkpoint's scale), eps the +checkpoint's 1e-6. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + m3_embedding_firstnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" + +H = 6144 +EPS = 1e-6 + + +@functools.lru_cache(maxsize=None) +def _operands(): + g = torch.Generator(device="cuda").manual_seed(20261013) + table = (0.02 * torch.randn(VOCAB, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return table, weight + + +def _buffers(rows): + return (torch.empty(rows, H, dtype=torch.bfloat16, device="cuda"), + torch.empty(rows, H, dtype=torch.bfloat16, device="cuda")) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + rows, id_dtype = cell.spec["rows"], cell.spec["id_dtype"] + table, weight = _operands() + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device="cuda").to(id_dtype) + raw, out = _buffers(rows) + with validating(op): + normed, residual = op(ids, table, weight, EPS, raw, out) + assert normed.data_ptr() == out.data_ptr() and residual.data_ptr() == raw.data_ptr() + ref_normed, ref_raw = op.reference(ids, table, weight, EPS, raw, out) + # The gather is a copy: exact. + assert torch.equal(raw.view(torch.int16), ref_raw.view(torch.int16)) + # The norm: the kernel's fp32 sum of squares (warp butterflies, then across 4 warps) and fast rsqrt against an + # fp64 reference, then x * rstd * w rounded once to bf16. The fp32 side is ~1e-6 relative off, so the two roundings + # disagree only next to a bf16 rounding boundary: one ulp at most per element, a small fraction in aggregate. + # Measured over these cells (3 seeds each, GB300): at most 0.51 ulp element, 0.026 ulp RMS. + assert_within_ulp(out, ref_normed, element_ulp=1.0, rms_ulp=0.1) + + +def test_ids_outside_the_table_give_zero_rows() -> None: + """Documented, not refused: an id outside [0, V) gathers a zero row whose norm is zero.""" + table, weight = _operands() + ids = torch.tensor([-1, 5, VOCAB], dtype=torch.int64, device="cuda") + raw, out = _buffers(3) + op(ids, table, weight, EPS, raw, out) + assert not raw[0].any() and not raw[2].any() and not out[0].any() and not out[2].any() + assert torch.equal(raw[1], table[5]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..27157668c588 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs shared by the tp4 target's spec entry tests: TP4 dimensions, single-device stand-ins for the cross-rank +exchanges, and the comparison bands.""" + +from __future__ import annotations + +import zlib + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) + +#: MiniMax-M3 at TP4: hidden size, vocabulary, the target head's shard and RMSNorm epsilon. +HIDDEN = 6144 +VOCAB = 200064 +WORLD = 4 +TARGET_SHARD = VOCAB // WORLD +EPS = 1e-6 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): + """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" + uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") + keys = torch.zeros(slot_rows, dtype=torch.int64, device="cuda") + counters = torch.tensor([0, call], dtype=torch.int32, device="cuda") + return uc, keys, counters + + +def plant_peers( + uc: torch.Tensor, + counters: torch.Tensor, + own_best: torch.Tensor, + rank: int, + world: int, + shard: int, + g: torch.Generator, + slot_rows: int = SHARD_ROWS, + ties: bool = False, +) -> None: + """Write the other ranks' keys for each row into this rank's slots, as their multicast pushes would. + + Each peer's maximum is the own row maximum shifted by +-0.25 (well clear of the accumulation-order band), so + some rows are won by a peer and some by this rank; with ``ties`` a peer on an even row reports exactly the own + maximum, and the lower global index must win. + """ + if world == 1: + return + m = own_best.shape[0] + half = int(counters[1].item()) & 1 + for r in range(world): + if r == rank: + continue + sign = torch.where(torch.rand(m, generator=g, device="cuda") < 0.5, -1.0, 1.0) + value = own_best.float() + 0.25 * sign + if ties: + value = torch.where(torch.arange(m, device="cuda") % 2 == 0, own_best.float(), value) + index = r * shard + torch.randint(0, shard, (m,), generator=g, device="cuda") + keys = make_keys(value, index) + for t in range(m): + uc[(half * slot_rows + t) * world + r] = keys[t] + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..d31aebfe2610 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_capture_norm entry (Eagle3 capture with fc_norm folded in).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_capture_norm import ( # noqa: E501 + m3_capture_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" + +LAYERS = 3 + + +def _case(spec, seed): + g = gen(seed) + m, slot = spec["m"], spec["slot"] + h = randn((m, HIDDEN), g, 2.0) + r = randn((m, HIDDEN), g, 2.0) if spec["residual"] else None + w = (1.0 + randn((HIDDEN,), g, 0.1).float()).to(torch.bfloat16) + buf = torch.full((m, LAYERS * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + dst = buf[:, slot * HIDDEN : (slot + 1) * HIDDEN] + return (h, r, w, EPS, dst), buf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer + untouched.""" + args, buf = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + slot = cell.spec["slot"] + others = torch.cat([buf[:, : slot * HIDDEN], buf[:, (slot + 1) * HIDDEN :]], dim=1) + assert bool((others == 7.0).all()) + + +def test_guard_refuses_a_strided_weight() -> None: + (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + w2 = torch.stack([w, w], dim=1)[:, 0] + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(h, r, w2, eps, dst) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_normed.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_normed.py new file mode 100644 index 000000000000..e524f094101e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_normed.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU test: a first-draft-inputs reader that skips the drafter's fc_norm (the capture norm is armed) gets the captured +rows unchanged when every captured layer was written with the step's norm, and otherwise a copy with only the raw +layers normalized by their own fc_norm (M3Eagle3Worker._m3_renorm_raw_captures); the per-forward reset clears the +record and the unconsumed fc fork result.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) + +M3Eagle3Worker = modeling.M3Eagle3Worker + +H = 4 +LAYERS = (1, 29, 56) + + +def _worker(normed): + w = object.__new__(M3Eagle3Worker) + w._m3_capture_normed = set(normed) + return w + + +@pytest.fixture(autouse=True) +def _scaling_norm(monkeypatch): + """The drafter's RMSNorm as a stand-in that multiplies by its weight (CPU, no FlashInfer).""" + monkeypatch.setattr(modeling, "flashinfer_rmsnorm", lambda x, w, eps: x * w) + + +def _drafter(): + # fc_norm i scales by (i + 2): enough to tell which blocks were normalized here. + return SimpleNamespace(w={f"fc_norm{i}": i + 2 for i in range(len(LAYERS))}, rms_eps=1e-6) + + +def test_all_captured_layers_normed_returns_the_rows_as_they_lie() -> None: + hs = torch.randn(4, H * len(LAYERS)) + out = _worker(set(LAYERS))._m3_renorm_raw_captures( + hs, SimpleNamespace(layers_to_capture=LAYERS, hidden_size=H), _drafter() + ) + assert out is hs + + +@pytest.mark.parametrize("normed", [set(), {1, 29}, {29, 56}, {1, 56}]) +def test_raw_layers_get_their_own_fc_norm(normed) -> None: + hs = torch.randn(4, H * len(LAYERS)) + ref = hs.clone() + out = _worker(normed)._m3_renorm_raw_captures( + hs, SimpleNamespace(layers_to_capture=LAYERS, hidden_size=H), _drafter() + ) + assert torch.equal(hs, ref) # the capture buffer itself is not modified + for i, layer_id in enumerate(LAYERS): + block, src = out[:, i * H : (i + 1) * H], ref[:, i * H : (i + 1) * H] + assert torch.equal(block, src if layer_id in normed else src * (i + 2)), layer_id + + +def test_arm_resets_the_record_and_the_fork_result() -> None: + w = _worker(set(LAYERS)) + w._m3_fc_hidden = object() + w._m3_fc_early_stream = None # no fork pending + w._m3_join_pending = lambda: None # (tp8) nothing else pending + w.arm_capture_norm( + None, None, None, None + ) # no drafter: resets the per-forward state and returns + assert w._m3_capture_normed == set() + assert w._m3_fc_hidden is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..2caf1da9304d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_chunk_norm entry (the Eagle3 fc_norm over the captured layers).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_chunk_norm import ( # noqa: E501 + m3_chunk_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, c, h = spec["m"], spec["c"], spec["h"] + buf = randn((m, c * h + spec["extra_cols"]), g, 2.0) + src = buf[:, : c * h + spec["extra_cols"]] + w = (1.0 + randn((c, h), g, 0.1).float()).to(torch.bfloat16) + dst = torch.full((m, c * h), 7.0, dtype=torch.bfloat16, device="cuda") + return (src, w, EPS, dst, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + + +def test_guard_refuses_a_strided_weight() -> None: + src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() + wt.copy_(w) + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(src, wt, eps, dst, pdl) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..2f82a2a9d63b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_input entry (embedding + both RMSNorms + concat in one launch), against its reference and +against the in-tree drafter's modules.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_input import ( # noqa: E501 + m3_draft_input as op, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + VOCAB, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, h = spec["m"], spec["h"] + table = randn((VOCAB, h), g) + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + if spec["invalid_id"]: + ids[0] = -1 + hidden = randn((4 * m if spec["gather"] else m, h), g) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 2 if spec["gather"] else None + w_e = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + w_h = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + out = torch.full((m, 2 * h), 7.0, dtype=torch.bfloat16, device="cuda") + residual = torch.full((m, h), 7.0, dtype=torch.bfloat16, device="cuda") + return (ids, table, hidden, rows, w_e, w_h, EPS, EPS, out, residual, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Normalized halves within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic in a + different summation order, rsqrtf); the gathered residual and the zero embedding of an invalid id exact.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out, residual = op(*args) + ref_out, ref_res = op.reference(*args) + torch.cuda.synchronize() + assert_bf16_close(out, ref_out) + assert torch.equal(residual, ref_res) + if cell.spec["invalid_id"]: + assert not bool(out[0, : cell.spec["h"]].float().any()) + + +def test_guard_refuses_rows_past_hidden() -> None: + gather = next(c for c in op.CELLS if c.spec["gather"]) + args = list(_case(gather.spec, 1)) + args[3] = args[3] + 100 + with pytest.raises(AssertionError, match="rows must index hidden"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("m,gather", [(4, False), (32, False), (8, True)], ids=str) +def test_the_in_tree_draft_input(m, gather) -> None: + """Against the in-tree drafter's own modules: the Embedding lookup, then the layer's input_layernorm and + hidden_norm (in-tree RMSNorm) and the concat of Eagle3DecoderLayer.forward. Within one bf16 ulp on at most 0.1% of + the elements (the RMS statistic summed in another order); the residual (the carried hidden rows) exact.""" + g = gen(seed_of(f"in-tree draft input {m} {gather}")) + h = HIDDEN + embed = Embedding(VOCAB, h, dtype=torch.bfloat16).cuda() + embed.weight.data = randn((VOCAB, h), g) + norms = [] + for _ in range(2): + norm = RMSNorm(hidden_size=h, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + norms.append(norm) + input_layernorm, hidden_norm = norms + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + hidden = randn((4 * m if gather else m, h), g, 2.0) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 3 if gather else None + carried = hidden[rows] if gather else hidden + stock = torch.cat([input_layernorm(embed(ids)), hidden_norm(carried)], dim=-1) + out, residual = m3_spec_tail.draft_input( + ids, embed.weight, hidden, rows, input_layernorm.weight, EPS, hidden_norm.weight, EPS + ) + torch.cuda.synchronize() + assert_bf16_close(out, stock) + assert torch.equal(residual, carried) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py new file mode 100644 index 000000000000..53aaf105b5a3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_push_take.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter boundary push bookkeeping (ops/m3_draft_lowm arm_push / take_pushed, dispatch.DRAFT_AR_PUSH): the +boundary all-reduce runs prepushed only when the armed call pushed, and only on that call's own output; any other +input raises instead of reducing rows the exchange does not hold. Host state only (no GPU).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_draft_lowm, +) + + +@pytest.fixture(autouse=True) +def _clean(): + m3_draft_lowm._PUSH_ARM, m3_draft_lowm._PUSHED = None, 0 + yield + m3_draft_lowm._PUSH_ARM, m3_draft_lowm._PUSHED = None, 0 + + +def test_not_pushed_is_false_and_disarms() -> None: + m3_draft_lowm._PUSH_ARM = (1, 2, 3, 0, 8, 8) + assert m3_draft_lowm.take_pushed(torch.zeros(4)) is False + assert m3_draft_lowm._PUSH_ARM is None + + +def test_pushed_output_is_true_once() -> None: + out = torch.zeros(4, 8) + m3_draft_lowm._PUSHED = out.data_ptr() + assert m3_draft_lowm.take_pushed(out) is True + assert m3_draft_lowm._PUSHED == 0 + assert m3_draft_lowm.take_pushed(out) is False + + +def test_other_input_raises() -> None: + out = torch.zeros(4, 8) + other = out.clone() + m3_draft_lowm._PUSHED = out.data_ptr() + with pytest.raises(RuntimeError, match="not the boundary all-reduce's input"): + m3_draft_lowm.take_pushed(other) + assert m3_draft_lowm._PUSHED == 0 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py new file mode 100644 index 000000000000..6d1183397871 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_silu entry (the drafter MLP's plain SwiGLU as one C++ launch): bit for bit the Triton +silu_and_mul at its defaults, over every bf16 gate and every bf16 up pattern.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_silu import ( # noqa: E501 + m3_draft_silu as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" + +ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) + + +def _rows(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """[rows, 2d] gate | up from equally long flat gate and up values, d = 4608 (the TP4 drafter shard).""" + d = 4608 + n = gate.numel() + rows = -(-n // d) + pad = rows * d - n + g = torch.cat([gate, gate[:pad]]).view(rows, d) + u = torch.cat([up, up[:pad]]).view(rows, d) + return torch.cat([g, u], dim=1).contiguous() + + +def _assert_bitwise(x: torch.Tensor, what: str) -> None: + got = op(x) + assert got is not None, what + ref = op.reference(x) + torch.cuda.synchronize() + diff = got.view(torch.int16) != ref.view(torch.int16) + assert not bool(diff.any()), f"{what}: {int(diff.sum())} of {diff.numel()} differ" + + +@pytest.mark.parametrize("up", [0.0, 1.0, -1.0, 3.5, -0.0, 1e-3], ids=lambda v: f"up={v}") +def test_every_gate_pattern_is_bitwise_the_triton_kernel(up) -> None: + gate = ALL.cuda() + _assert_bitwise(_rows(gate, torch.full_like(gate, up)), f"gate sweep, up {up}") + + +@pytest.mark.parametrize("gate", [1.0, -1.0, 0.5, -3.0, 20.0, -0.0], ids=lambda v: f"gate={v}") +def test_every_up_pattern_is_bitwise_the_triton_kernel(gate) -> None: + up = ALL.cuda() + _assert_bitwise(_rows(torch.full_like(up, gate), up), f"up sweep, gate {gate}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed) -> None: + g = gen(100 + seed) + n = 1 << 22 + gate = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + up = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + _assert_bitwise(_rows(gate.view(torch.bfloat16), up.view(torch.bfloat16)), f"seed {seed}") + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The drafter's gate_up output range (and past it), in a CUDA graph as the drafter runs it.""" + g = gen(seed_of(cell.why)) + rows, d = cell.spec["rows"], cell.spec["d"] + x = (4.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :6] = torch.tensor([60.0, -60.0, 0.0, -0.0, 88.0, -88.0], dtype=torch.bfloat16) + certifying(op) + ref = op.reference(x) + eager = op(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(x) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(eager.view(torch.int16), ref.view(torch.int16)) + assert torch.equal(captured.view(torch.int16), ref.view(torch.int16)) + + +def test_declines_what_it_does_not_cover() -> None: + x = torch.zeros(2, 2 * 4608, dtype=torch.bfloat16, device="cuda") + assert op(x[:, 8:]) is None # strided rows + assert ( + op(torch.zeros(2, 2 * 12, dtype=torch.bfloat16, device="cuda")) is None + ) # d not a multiple of 8 + assert op(x.float()) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..8bb087bf3cd4 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,694 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_step entry (one fused greedy draft pass: the BF16 draft head, streamed from its lossless +12-bit coding or as BF16 fragments, the cross-rank argmax and the draft loop's bookkeeping), and its tokens against the +in-tree drafter's head and greedy draft sampling. + +The in-tree path: each rank's LM head shard GEMM (BF16 out), ``LogitsProcessor``'s float cast, then either the argmax +of the gathered logits (``torch.argmax``) or ``greedy_sample_draft_with_tp_gather``'s per-rank ``torch.max`` and the +pick among the ranks' (index, value) pairs. The edge cases use integer data, whose fp32 sums are exact in any order, +so the kernel and the in-tree path see the same BF16 logits and must agree bit for bit; every in-tree comparison runs +with the coded head (the default) and with the BF16 fragments, and on integer data the coded head's winning keys are +the BF16 fragments' bit for bit. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_draft_step as draft_step_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_LEN, + DRAFT_SHARD, + HIDDEN, + MAX_TOKENS, + bf16_fragments, + from_fragments, + head_weight, + shard_logits, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_step import ( # noqa: E501 + m3_draft_step as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + WORLD, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" + +T = DRAFT_LEN + 1 +_LOW32 = 0xFFFFFFFF + +# ------------------------------------------------------------------ the per-row rule against the in-tree token +# +# The kernel and the in-tree head (a cuBLAS BF16 GEMM) both round an FP32 sum of the same exact BF16 products to BF16, +# in different summation orders. A row's token can differ from the in-tree one only where some candidate's FP32 sum +# can land on either side of a BF16 rounding midpoint. Each row is held to: +# +# kernel token == in-tree token, or the row is provably undecidable: with the fp64 reference logit L_i (exact BF16 +# products, fp64 sum) and the FP32-order bound E_i = gamma_K * sum_k |w_ik * x_k|, the BF16 value of any FP32 sum of +# candidate i lies in [lo_i, hi_i] = [bf16(L_i - E_i), bf16(L_i + E_i)]; the kernel's token must win for some +# choice of values within those intervals (larger value, then the lower global id), and a BF16 rounding midpoint +# must lie within E of the kernel's or the in-tree token's L. +# +# gamma_K = K * 2^-22 is la-xh2's bound (4x the round-to-nearest unit roundoff 2^-24 per addition: any summation +# order, and a truncating tensor-core accumulator); K counts the row's non-zero products (adding an exact zero does +# not round). lo / hi take one more FP32 step outward before the BF16 rounding, which covers rounding L -+ E to FP32. +# Each undecidable row is reported as a warning with its midpoint distances, and each test reports its row count; any +# other mismatch fails. +GAMMA_UNIT = 2.0**-22 + + +def _bf16_outward(v: torch.Tensor, direction: float) -> torch.Tensor: + f = v.float() + f = torch.nextafter(f, torch.full_like(f, direction)) + b = f.to(torch.bfloat16).float() + return torch.where(b == 0, torch.zeros_like(b), b) + + +def bounds(h: torch.Tensor, w: torch.Tensor): + """(fp64 logits L, bound E, lo, hi) [rows, N] of the head input rows h against the head rows w.""" + cols = (h != 0).any(dim=0) # zero columns add exact zeros + hd, wd = h[:, cols].double(), w[:, cols].double() + logits = hd @ wd.t() + bound = (h != 0).sum(dim=1, keepdim=True).double() * GAMMA_UNIT * (hd.abs() @ wd.abs().t()) + del wd + inf = float("inf") + return logits, bound, _bf16_outward(logits - bound, -inf), _bf16_outward(logits + bound, inf) + + +def midpoint_distance(v: float) -> float: + """Distance of v to the nearest BF16 rounding midpoint.""" + b = torch.tensor(v, dtype=torch.float64).float().to(torch.bfloat16).float() + if float(b) == 0.0: + return abs(v) + ulp = 2.0 ** (int(torch.frexp(b).exponent) - 8) + return min(abs(v - (float(b) + ulp / 2)), abs(v - (float(b) - ulp / 2))) + + +def _winnable(tok: int, ids: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Whether candidate ``tok`` wins the row for some values within the candidates' [lo, hi].""" + k = int((ids == tok).nonzero()[0, 0]) + others = ids != tok + beat = (hi[k] > lo[others]) | ((hi[k] == lo[others]) & (tok < ids[others])) + return bool(beat.all()) + + +def check_row(name: str, kernel_tok: int, stock_tok: int, ids, logits, bound, lo, hi) -> bool: + """The per-row rule for one row over the candidates ``ids`` [C] (logits / bound / lo / hi [C]); True when the + row is undecidable (reported as a warning), False when the tokens match; a mismatch outside the bound fails.""" + if kernel_tok == stock_tok: + return False + assert _winnable(kernel_tok, ids, lo, hi), ( + f"{name}: token {kernel_tok} differs from the in-tree {stock_tok} outside the FP32-order bound" + ) + ik = int((ids == kernel_tok).nonzero()[0, 0]) + i_stock = int((ids == stock_tok).nonzero()[0, 0]) + lk, ls, ek, es = ( + float(logits[ik]), + float(logits[i_stock]), + float(bound[ik]), + float(bound[i_stock]), + ) + dk, ds = midpoint_distance(lk), midpoint_distance(ls) + assert dk <= ek or ds <= es, f"{name}: neither logit is within its bound of a BF16 midpoint" + warnings.warn( + f"[draft-step undecidable] {name}: kernel {kernel_tok} (L {lk:.9g}, midpoint distance {dk:.3e}, bound " + f"{ek:.3e}) vs in-tree {stock_tok} (L {ls:.9g}, midpoint distance {ds:.3e}, bound {es:.3e})" + ) + return True + + +def report(name: str, rows: int, undecidable: int) -> None: + warnings.warn(f"[draft-step rows] {name}: {rows} rows, {undecidable} undecidable") + + +def key_value(keys: torch.Tensor) -> torch.Tensor: + """The fp32 value a key orders by (the inverse of the kernel's order-preserving image).""" + u = (keys >> 32) & _LOW32 + bits = torch.where(u >= 0x80000000, u - 0x80000000, _LOW32 - u) + bits = torch.where(bits >= 2**31, bits - 2**32, bits) # as a signed 32-bit pattern + return bits.to(torch.int32).view(torch.float32) + + +def key_index(keys: torch.Tensor) -> torch.Tensor: + return _LOW32 - (keys & _LOW32) + + +def _bookkeeping(nreq: int, step: int, n_x_rows: int, g: torch.Generator): + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if step == 0 + else None + ) + seq_lens = torch.full((nreq,), T, dtype=torch.int32, device="cuda") if step == 0 else None + next_new = torch.full((nreq, T), -5, dtype=torch.int32, device="cuda") + pos_src = torch.randint(100, 3000, (n_x_rows,), generator=g, device="cuda", dtype=torch.int32) + kv_lens = torch.randint(200, 3000, (nreq,), generator=g, device="cuda", dtype=torch.int32) + ids_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + pos_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + return num_acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next + + +def _head(w: torch.Tensor, coded: bool): + """(wf, (exc_off, exc_pos, exc_w, exp_base)) of a BF16 shard: its 12-bit coding (checked to round-trip) or its + BF16 fragments.""" + if not coded: + return bf16_fragments(w), (None, None, None, 0) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + return packed, (exc_off, exc_pos, exc_w.view(-1), base) + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=(None, None, None, 0), +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + coded head's (exc_off, exc_pos, exc_w, exp_base).""" + acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next = _bookkeeping( + nreq, step, x.shape[0], g + ) + if num_acc is None: + num_acc = acc + if uc is None: + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + out_keys = torch.zeros(nreq, dtype=torch.int64, device="cuda") + tiles = wf.shape[0] if coded[0] is None else coded[0].numel() - 1 + return (x, x_rows, nreq, wf, rank * tiles * 16, keys, counters, uc, uc.data_ptr(), rank, world, + next_new[:, 1:], step, ids_next, pos_src, pos_next, kv_lens, num_acc, seq_lens, DRAFT_LEN, num_sms(), + True, out_keys, *coded) # fmt: skip + + +def _case(spec, seed): + g = gen(seed) + nreq, step, rank, world, vs, k = ( + spec["nreq"], + spec["step"], + spec["rank"], + spec["world"], + spec["vs"], + spec["k"], + ) + pass0 = step == 0 + x = randn((nreq * T if pass0 else nreq, k), g) + wf, coded = _head(randn((vs, k), g, 0.02), spec["coded"]) + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if pass0 + else None + ) + x_rows = (torch.arange(nreq, device="cuda") * T + num_acc.long() - 1) if pass0 else None + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + h = x[x_rows] if pass0 else x + own = shard_logits(h, wf, *coded).max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=MAX_TOKENS, ties=spec.get("ties", False) + ) + return _args(x, x_rows, nreq, wf, rank, world, step, g, uc, keys, counters, num_acc, coded) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The bookkeeping is exact on every row (next positions, KV and sequence lengths). Each row's token is the + in-tree one (this rank's BF16 shard logits from the head's GEMM, their first maximum, against the peers' keys) or + the row is provably undecidable (check_row); the winning key's index is the token, and its value lies in the + token's [lo, hi] (or is the peer's pushed key when a peer wins).""" + args = _case(cell.spec, seed_of(cell.why)) + x, x_rows, nreq, wf, offset, _, counters, uc = args[:8] + rank, world = args[9], args[10] + half = int(counters[1].item()) & 1 + pushed = ( + uc[half * MAX_TOKENS * world : (half * MAX_TOKENS + nreq) * world].view(nreq, world).clone() + ) + h = x[x_rows] if x_rows is not None else x[:nreq] + w = head_weight(wf, x.shape[1], *args[23:27]) + best, idx = local_argmax(torch.nn.functional.linear(h, w).float(), offset) + stock = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, MAX_TOKENS) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert args[6].tolist() == [0, 1] and not bool(args[5].any()), "counters / keys not re-armed" + step = args[12] + others = [c for c in range(args[11].shape[1]) if c != step] + assert torch.equal(args[11][:, others], ref["draft_out"][:, others]) + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + assert torch.equal(args[15], ref["pos_next"]) + assert torch.equal(args[16], ref["kv_lens"]) + if args[18] is not None: + assert torch.equal(args[18], ref["seq_lens"]) + # Candidates per row: this rank's head rows (fp64 bounds) and each peer's pushed key (an exact value). + logits, bound, lo, hi = bounds(h, w) + own_ids = offset + torch.arange(w.shape[0], device="cuda") + peers = [p for p in range(world) if p != rank] + undecidable = 0 + for r in range(nreq): + pv = key_value(pushed[r, peers]).double() + ids = torch.cat([own_ids, key_index(pushed[r, peers])]) + undecidable += check_row( + f"{cell.why[:40]} row {r}", + int(args[13][r]), + int(stock[r]), + ids, + torch.cat([logits[r], pv]), + torch.cat([bound[r], torch.zeros_like(pv)]), + torch.cat([lo[r], pv.float()]), + torch.cat([hi[r], pv.float()]), + ) + report(cell.why[:40], nreq, undecidable) + out_keys = args[22] + assert torch.equal(key_index(out_keys), args[13].long()), "the winning key's index is the token" + got = key_value(out_keys) + for r in range(nreq): + c = int(args[13][r]) - offset + if 0 <= c < w.shape[0]: + assert float(lo[r, c]) <= float(got[r]) <= float(hi[r, c]), ( + f"row {r}: winning value outside [lo, hi]" + ) + else: + assert int(out_keys[r]) in pushed[r].tolist(), ( + f"row {r}: a peer won with a key it did not push" + ) + + +def test_fragments_round_trip() -> None: + g = gen(seed_of("fragments")) + w = randn((64, 96), g) + wf = bf16_fragments(w) + assert wf.shape == (4, 6, 256) + assert torch.equal(from_fragments(wf), w) + head = SimpleNamespace(weight=w, bias=None) + # K 96: never coded, the BF16 fragments. + prepared = draft_step_kernels.prepare_head(head, 3, True) + assert not prepared.coded + assert torch.equal(prepared.wf, wf) and prepared.vocab_offset == 3 * 64 and prepared.weight is w + + +# ------------------------------------------------------------------ against the in-tree drafter, bit for bit + + +def _stock_tokens(shards: list, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The in-tree greedy draft tokens of the rows x over the vocab-sharded BF16 head: (argmax of the gathered float + logits, greedy_sample_draft_with_tp_gather's per-rank (index, value) pick).""" + logits = [torch.nn.functional.linear(x, w).float() for w in shards] + gathered = torch.argmax(torch.cat(logits, dim=-1), dim=-1).to(torch.int32) + combined = [] + for rank, shard_logit in enumerate(logits): + worker = SimpleNamespace(mapping=SimpleNamespace(tp_rank=rank)) + combined.append(Eagle3OneModelWorker._get_local_max_and_combined(worker, shard_logit)) + pick = Eagle3OneModelWorker._get_draft_tokens_from_gathered + pick = getattr(pick, "_torchdynamo_orig_callable", pick) + tp_gather = pick(None, torch.cat(combined, dim=-1)).to(torch.int32) + return gathered, tp_gather + + +def _fused_tokens(shards: list, x: torch.Tensor, coded: bool, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's four-rank pass with + its peers' keys in its slots, as their multicast pushes would put them.""" + g = gen(seed_of("fused tokens")) + nreq = x.shape[0] + heads = [_head(w, coded) for w in shards] + own_keys = [] + for rank, (wf, cd) in enumerate(heads): + args = list(_args(x, None, nreq, wf, rank, 1, step, g, coded=cd)) + args[9] = 0 # one-rank run: this rank's slot only, global ids of `rank` + op(*args) + own_keys.append(args[22].clone()) + tokens = [] + for rank, (wf, cd) in enumerate(heads): + uc, keys, counters = head_workspace(WORLD, slot_rows=MAX_TOKENS) + for peer in range(WORLD): + if peer != rank: + for t in range(nreq): + uc[t * WORLD + peer] = own_keys[peer][t] + args = _args(x, None, nreq, wf, rank, WORLD, step, g, uc, keys, counters, coded=cd) + with validating(op): + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + tokens.append(args[13].clone()) + return tokens + + +def _integer_head(g: torch.Generator, rows: int = WORLD * DRAFT_SHARD) -> torch.Tensor: + """A [rows, 6144] head of small integers over the first 4 columns (zero elsewhere): its logits against integer + rows are exact integers in any summation order, many of which share a BF16 value.""" + w = torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, :4] = torch.randint(-60, 61, (rows, 4), generator=g, device="cuda").to(torch.bfloat16) + return w + + +def _shards(w: torch.Tensor) -> list: + return [w[r * DRAFT_SHARD : (r + 1) * DRAFT_SHARD].contiguous() for r in range(WORLD)] + + +def _assert_stock(w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None) -> None: + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + if expect is not None: + assert gathered.tolist() == expect, (gathered.tolist(), expect) + for rank, tokens in enumerate(_fused_tokens(shards, x, coded)): + assert torch.equal(tokens, gathered), ( + f"rank {rank}: {tokens.tolist()} vs stock {gathered.tolist()}" + ) + + +def _ones_rows(nreq: int) -> torch.Tensor: + x = torch.zeros((nreq, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[:, :4] = 1.0 + return x + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_forced_bf16_ties_take_the_lowest_global_id(coded) -> None: + """A head row copied onto every rank that wins every row of ones (logit 244, every other at most 240): the lowest + global id must win, inside a shard and across ranks. Then integer rows in [-2, 2], whose logits (up to 480 in + magnitude, BF16 spacing 2 above 256) collide often, exactly or after BF16 rounding.""" + g = gen(seed_of("forced ties")) + w = _integer_head(g) + win = torch.tensor([61.0, 61.0, 61.0, 61.0], device="cuda").to(torch.bfloat16) + for gid in (3 * DRAFT_SHARD + 5, DRAFT_SHARD + 77, DRAFT_SHARD + 9000, 2 * DRAFT_SHARD + 1): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), coded, expect=[DRAFT_SHARD + 77] * 5) + # Random integer rows: ties after BF16 rounding happen naturally. + x = torch.randint(-2, 3, (8, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + x[:, 4:] = 0 + _assert_stock(w, x, coded) + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_signed_zero_maxima_take_the_lowest_global_id(coded) -> None: + """Rows whose maximum is zero: an all-zero head input (every logit +-0) and logits <= 0 with zeros by cancellation + on several ranks. -0 and +0 compare equal, so the lowest global id reaching zero wins.""" + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 0] = -1.0 + w[:, 1] = -1.0 + for gid in (2 * DRAFT_SHARD + 5, DRAFT_SHARD + 40000, 3 * DRAFT_SHARD + 7, 49999): + w[gid, 0] = 1.0 + w[gid, 1] = -1.0 + x = torch.zeros((3, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[1, :2] = 1.0 + x[2, :2] = 2.0 + x[2, 2] = -1.0 # w[:, 2] is zero: a -0 product on every row + _assert_stock(w, x, coded, expect=[0, 49999, 49999]) + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + (1000, 1001, "lower"), # 1001 rounds to 1000 in BF16: a tie, the lower id wins + (1001, 1002, "lower"), # 1001 and the tie-to-even midpoint 1002 both round to 1000 + (1000, 1004, "higher"), # adjacent BF16 values: the larger wins + (1003, 1000, "lower"), # 1003 rounds up to 1004 > 1000 + ], + ids=["one_fp32_step", "midpoint", "adjacent_bf16", "rounds_up"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where, coded) -> None: + """Top-2 logits one to four fp32 steps apart: the winner is the one the BF16 logits decide (the in-tree logits' + dtype), not the fp32 sum. The pair sits inside rank 3's shard or on ranks 1 and 2; every other logit is at most + 240 in magnitude.""" + g = gen(seed_of(f"near top2 {lower} {higher} {where}")) + w = _integer_head(g) + if where == "in_shard": + lo_id, hi_id = 3 * DRAFT_SHARD + 100, 3 * DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD + 49000, 2 * DRAFT_SHARD + 3 + for gid, value in ((lo_id, lower), (hi_id, higher)): + q, r = divmod(value, 4) # four integers <= 251, exact in BF16 + w[gid, :4] = torch.tensor([q + (i < r) for i in range(4)], device="cuda").to(torch.bfloat16) + expect = lo_id if winner == "lower" else hi_id + _assert_stock(w, _ones_rows(2), coded, expect=[expect, expect]) + + +# ------------------------------------------------------------------ non-integer rows at a BF16 rounding midpoint + +#: A BF16 rounding midpoint (between 6.0 and 6.03125; a tie there rounds to the even 6.0) and the FP32 spacing there. +MIDPOINT = 6.0 + 2.0**-6 +FP32_ULP = 2.0**-21 + + +def _near_midpoint_head(g: torch.Generator, targets: dict): + """A full-vocabulary head [4 * 50016, 6144] and one head-input row x with 65 non-zero entries: 1, 2^-12, 2^-20 + and 2^-28 on columns 0..3 and 61 random non-integer BF16 values. Every head row carries random non-integer BF16 + values on those 61 columns (logits about 0.2 in magnitude); each row of ``targets`` (global id -> fp64 value) also + gets columns 0..3 set so its fp64 logit is the target to within 2^-30 (a coarse value, then three finer + corrections, each BF16-rounded).""" + x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device="cuda") + scales = (1.0, 2.0**-12, 2.0**-20, 2.0**-28) + for col, scale in enumerate(scales): + x[0, col] = scale + x[0, 4:65] = randn((61,), g, 0.5) + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 4:65] = randn((WORLD * DRAFT_SHARD, 61), g, 0.05) + xd = x[0].double() + for gid, target in targets.items(): + rest = float(target - (w[gid].double() * xd).sum()) + for col, scale in enumerate(scales): + v = torch.tensor(rest / scale, dtype=torch.float64).to(torch.bfloat16) + w[gid, col] = v + rest -= float(v) * scale + assert abs(rest) <= 2.0**-30, (gid, rest) + return w, x + + +def _check_fused_against_stock( + name: str, w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None +) -> int: + """Each rank's fused tokens against the in-tree tokens under the per-row rule (full-vocabulary candidates); + ``expect`` (a token per row) additionally pins rows the bound decides. Returns the undecidable rows.""" + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + fused = _fused_tokens(shards, x, coded) + for tokens in fused[1:]: + assert torch.equal(tokens, fused[0]), "the ranks disagree on the global token" + logits, bound, lo, hi = bounds(x, w) + ids = torch.arange(w.shape[0], device="cuda") + undecidable = 0 + for r in range(x.shape[0]): + undecidable += check_row( + f"{name} row {r}", + int(fused[0][r]), + int(gathered[r]), + ids, + logits[r], + bound[r], + lo[r], + hi[r], + ) + if expect is not None: + assert int(gathered[r]) == expect[r] and int(fused[0][r]) == expect[r], ( + name, + int(gathered[r]), + int(fused[0][r]), + expect[r], + ) + others = ids != expect[r] + assert _winnable(expect[r], ids, lo[r], hi[r]), ( + f"{name}: the bound does not let {expect[r]} win" + ) + rivals = [int(i) for i in ids[others][hi[r][others] >= lo[r][expect[r]]]] + assert not any(_winnable(i, ids, lo[r], hi[r]) for i in rivals), ( + f"{name}: the row is not decided" + ) + report(name, x.shape[0], undecidable) + return undecidable + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + ( + -2, + 2, + None, + ), # straddling the midpoint by 2 FP32 ulps: either side is in-tree up to the order + (2, -2, None), + (0, -2, None), # exactly on the midpoint (ties to the even 6.0) against just below it + ( + -2048, + 2048, + "higher", + ), # ~1e-3 either side, well outside the bound: decided by BF16 rounding + ( + -4096, + -2048, + "lower", + ), # both below the midpoint, ~1e-3 apart: a BF16 tie the lower id wins + ], + ids=["straddle", "straddle_reversed", "on_midpoint", "decided_up", "decided_tie"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where, coded) -> None: + """Non-integer rows whose top-2 fp64 logits sit the given number of FP32 ulps from a BF16 rounding midpoint, inside + rank 3's shard or on ranks 1 and 2, every other logit far below. The per-row rule holds on every row; outside the + bound the token is the one the BF16 rounding decides, bit for bit.""" + g = gen(seed_of(f"midpoint {lower} {higher} {where}")) + if where == "in_shard": + lo_id, hi_id = 3 * DRAFT_SHARD + 100, 3 * DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD + 49000, 2 * DRAFT_SHARD + 3 + w, x = _near_midpoint_head( + g, {lo_id: MIDPOINT + lower * FP32_ULP, hi_id: MIDPOINT + higher * FP32_ULP} + ) + expect = None if winner is None else [lo_id if winner == "lower" else hi_id] + _check_fused_against_stock(f"midpoint {where} {lower:+d}/{higher:+d} ulp", w, x, coded, expect) + + +# ------------------------------------------------------------------ the coded head against the BF16 fragments + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_coded_head_is_the_bf16_head_bit_for_bit(nreq) -> None: + """Integer data (fp32 sums exact in any order) with exceptions on both sides of the exponent window: the coded + head's winning keys (value and global id) and its bookkeeping are the BF16 fragments' bit for bit, on one rank and + across the ranks.""" + g = gen(seed_of(f"coded vs bf16 {nreq}")) + sign = torch.randint(0, 2, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * 2 - 1 + w = (torch.randint(1, 41, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * sign).to( + torch.bfloat16 + ) + # Exceptions on both sides of the window of 1..40: exact zeros, and 2^15 (fp32 sums stay below 2^24). + w[::5, ::512] = 0.0 + w[::13, 7::600] = 32768.0 + x = torch.randint(-2, 3, (nreq, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + wf_c, cd = _head(w, True) + assert cd[1].numel() > 0, "the shard must carry exceptions" + wf_b, cb = _head(w, False) + outs = [] + for wf, c in ((wf_c, cd), (wf_b, cb)): + args = _args(x, None, nreq, wf, 2, 1, 1, gen(seed_of("bookkeeping")), coded=c) + args = list(args) + args[9] = 0 + with validating(op): + op(*args) + torch.cuda.synchronize() + outs.append(args) + coded_args, bf16_args = outs + for i in (11, 13, 15, 16, 22): # draft_out, ids_next, pos_next, kv_lens, out_keys + assert torch.equal(coded_args[i], bf16_args[i]), i + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 2 * DRAFT_SHARD) + assert torch.equal(coded_args[22], make_keys(best, idx)), ( + "the winning keys are not the in-tree ones" + ) + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_tiny_exception_products_are_exact(nreq) -> None: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1, and a one-hot + activation on column 0: the exceptions' exact products are the only non-zero logits, so the coded head must pick + row 1 with the BF16 head's key bit for bit (a correction against a non-zero placeholder cancels them to 0).""" + n, k = 256, 1024 + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + x = torch.zeros((nreq, k), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + wf_c, cd = _head(w, True) + assert cd[3] == 113 + keys = [] + for wf, c in ((wf_c, cd), _head(w, False)): + args = _args(x, None, nreq, wf, 0, 1, 1, gen(seed_of("tiny exceptions")), coded=c) + with validating(op): + # 16 tiles on an all-SM grid: every tile is streamed by all 16 warps, which only the fixed-order sum + # adds deterministically (is_valid rejects the atomics there). + op(*args, fixed_order=True) + torch.cuda.synchronize() + assert args[13].tolist() == [1] * nreq + keys.append(args[22]) + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 0) + assert torch.equal(keys[0], keys[1]) and torch.equal(keys[0], make_keys(best, idx)) + + +def test_prepare_head_codes_the_shard_and_fails_closed(monkeypatch) -> None: + """prepare_head keeps the coded shard when it decodes to the weight bit for bit; a coding that does not is + dropped for the BF16 fragments.""" + g = gen(seed_of("prepare coded")) + w = randn((4096, 1024), g, 0.02) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 2, True) + assert prepared.coded and prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(head_weight(prepared.wf, 1024, prepared.exc_off, prepared.exc_pos, prepared.exc_w, + prepared.exp_base), w) # fmt: skip + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[11] ^= 0x10 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(draft_step_kernels, "encode_chunked", corrupt) + dropped = draft_step_kernels.prepare_head(head, 2, True) + assert not dropped.coded and dropped.exc_off is None + assert torch.equal(dropped.wf, bf16_fragments(w)) + + +def test_prepare_head_keeps_bf16_when_coding_is_off(monkeypatch) -> None: + """``coded`` False (``DRAFT_HEAD_CODED`` off) keeps the BF16 fragments even for a shard whose coding would decode + bit for bit.""" + g = gen(seed_of("prepare bf16")) + w = randn((4096, 1024), g, 0.02) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2, False) + assert not prepared.coded and prepared.exc_off is None + assert torch.equal(prepared.wf, bf16_fragments(w)) + + +def test_guard_refuses_head_rows_past_x() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[1] = args[1] + 1000 + with pytest.raises(AssertionError, match="x_rows must index x"): + with validating(op): + op(*args) + + +def test_guard_refuses_too_few_ctas() -> None: + args = list(_case(op.CELLS[1].spec, 2)) + args[20] = 8 + with pytest.raises(AssertionError, match="tiles per CTA"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py new file mode 100644 index 000000000000..f8ffd8e47635 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step_fixed_order.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The draft step's tile pieces and their summation order (ops/m3_draft_step: the fixed warp order). + +A CTA's (tile, chunk) range is cut into 16 equal warp segments; the shared-memory atomics add a tile's pieces in +arrival order, which is deterministic only while no tile takes more than 2 pieces. ``fixed_order`` sums them in warp +order instead. The geometry tests enumerate the real grid (the shard on 212 SMs) so a grid or shard change cannot +silently bring back an order-dependent sum; the GPU tests hold the fixed order to the atomics on exact data and to the +warp-order oracle on a planted tile whose sum depends on the order. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( + m3_draft_step as kds, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_SHARD, + MAX_TOKENS, + bf16_fragments, + max_tile_pieces, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, head_workspace # noqa: E402 + +OP = "minimax_m3_sm107tp4_draft_step" +SMS = 212 # Rubin's SMs: the grid the draft step launches on +K = 6144 +CHUNKS = K // 16 +# This target's shard on Rubin: 14-15 tiles per CTA, at most 2 pieces per tile (the atomics are deterministic and the +# fixed order stays off). +ORDER_DEPENDENT = False + +gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + + +@pytest.mark.parametrize("grid", [SMS, kds.GRID]) +def test_real_grid_pieces(grid) -> None: + """One CTA per SM and the launch the target uses (ops/m3_draft_step.GRID).""" + pieces = max_tile_pieces(DRAFT_SHARD // 16, grid, CHUNKS) + assert (pieces > 2) == ORDER_DEPENDENT, (grid, pieces) + + +def test_ds_grid_launches_the_configured_grid(monkeypatch) -> None: + monkeypatch.setattr(kds, "_num_sms", lambda: SMS) + assert kds.GRID == 209 + assert kds._grid() == min(kds.GRID, SMS) + + +def test_max_tile_pieces_counts_warp_segments() -> None: + # One tile per CTA: all 16 warps stream it. 7 tiles per CTA: a 168-chunk segment can sit inside a tile (TP8). + assert max_tile_pieces(16, 16, CHUNKS) == 16 + assert max_tile_pieces(7, 1, CHUNKS) == 4 + assert max_tile_pieces(8, 1, CHUNKS) == 2 + assert max_tile_pieces(15, 1, CHUNKS) == 2 + + +def test_the_fixed_order_engages_where_the_geometry_needs_it() -> None: + flag = True + x = torch.empty(1, K) + + def head(n_tiles): + return types.SimpleNamespace(wf=torch.empty(n_tiles, 1), exc_off=None, bnd=None) + + assert kds._fixed_order(1, head(DRAFT_SHARD // 16), x, SMS) == (flag and ORDER_DEPENDENT) + assert kds._fixed_order(1, head(DRAFT_SHARD // 16), x, kds.GRID) == (flag and ORDER_DEPENDENT) + assert kds._fixed_order(1, head(16), x, 16) == flag # 16 pieces per tile + + +def _launcher(wf, x, fixed, nctas): + """(launch, ids, keys): one pass's op launch on buffers made here (outside any capture), world 1.""" + uc, keys, counters = head_workspace(1, slot_rows=MAX_TOKENS) + m = x.shape[0] + i32 = dict(dtype=torch.int32, device="cuda") + draft_out = torch.zeros(m, 4, **i32) + ids = torch.zeros(m, **i32) + pos_src, pos_next, kv_lens = ( + torch.full((m,), 100, **i32), + torch.zeros(m, **i32), + torch.full((m,), 500, **i32), + ) + out = torch.zeros(m, dtype=torch.int64, device="cuda") + + def launch(): + getattr(torch.ops.trtllm, OP)(x, None, m, wf, 0, keys, counters, uc, uc.data_ptr(), 0, 1, draft_out, 1, ids, + pos_src, pos_next, kv_lens, None, None, 3, nctas, False, out, + fixed_order=fixed) # fmt: skip + + return launch, ids, out + + +@gpu +@pytest.mark.parametrize("nctas,vs", [(SMS, DRAFT_SHARD), (16, 256)]) +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_fixed_order_equals_the_atomics_on_exact_data(nctas, vs, m) -> None: + """Integer data: every fp32 sum is exact in any order, so both summations give the same bits.""" + g = gen(17 + m) + w = torch.randint(-1, 2, (vs, K), generator=g, device="cuda").to(torch.bfloat16) + x = torch.randint(-1, 2, (m, K), generator=g, device="cuda").to(torch.bfloat16) + wf = bf16_fragments(w) + out = {} + for fixed in (False, True): + launch, ids, keys = _launcher(wf, x, fixed, nctas) + launch() + torch.cuda.synchronize() + out[fixed] = (keys, ids) + assert torch.equal(out[False][0], out[True][0]) + ref = (x.double() @ w.double().t()).float().bfloat16().float().argmax(dim=1).int() + assert torch.equal(out[True][1], ref) + + +@gpu +@pytest.mark.parametrize("graph", [False, True]) +def test_fixed_order_is_the_warp_order_sum_on_a_planted_tile(graph) -> None: + """One tile per CTA (16 pieces): row 0's pieces are +2^25 (warp 0), +1 (warps 1..14), -2^25 (warp 15). In warp + order the +1s are absorbed by 2^25 and the row sums to 0, so the competitor (row 255, logit 0.5) wins; arrival + order could give up to 14. Every launch must give the warp-order result.""" + w = torch.zeros(256, K, device="cuda") + seg = CHUNKS // 16 * 16 # one warp segment in K + w[0, 0] = 2.0**25 + for s in range(1, 15): + w[0, s * seg] = 1.0 + w[0, 15 * seg] = -(2.0**25) + w[255, 0] = 0.5 + wf = bf16_fragments(w.to(torch.bfloat16)) + x = torch.ones(1, K, device="cuda", dtype=torch.bfloat16) + launch, ids, _ = _launcher(wf, x, True, 16) + run = launch + if graph: + launch() + torch.cuda.synchronize() + gr = torch.cuda.CUDAGraph() + cs = torch.cuda.Stream() + with torch.cuda.stream(cs), torch.cuda.graph(gr, stream=cs): + launch() + run = gr.replay + for _ in range(16): + ids.zero_() + run() + torch.cuda.synchronize() + assert int(ids[0]) == 255 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py new file mode 100644 index 000000000000..ef33aecf3a6a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -0,0 +1,289 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp4 target's speculative step against the in-tree (stock) Eagle3 numerics, at the real TP4 shapes. + +* Verify: the target logits are the in-tree BF16 LM head's own shard GEMM, and the fused shard argmax + acceptance + gives the in-tree greedy acceptance of the gathered logits, bit for bit, on every rank (peers' keys are the real + per-shard maxima, computed by the same kernel). +* Drafter: the worker runs the in-tree drafter's modules and weights; its own drafter arithmetic is the first pass's + fc_norm (chunk_norm, or capture_norm folded into the capture), held here to the in-tree ``apply_eagle3_fc`` RMSNorm + modules within one bf16 ulp (its effect on the fc output and on the full-vocabulary BF16 head's draft tokens is + measured too), and the fused draft loop's draft input and BF16 head, held to the in-tree modules and greedy draft + sampling by their own tests (test_modeling_v2_m3_draft_input.py, test_modeling_v2_m3_draft_step.py). + +Synthetic weights: hidden 6144, vocabulary 200064 (4 x 50016 shards), 3 captured layers, draft length 3. +""" + +import importlib.util +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as shard_accept_op, +) +from tensorrt_llm._torch.models.modeling_speculative import Eagle3ForCausalLM +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.rms_norm import RMSNorm +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + EPS, + HIDDEN, + TARGET_SHARD, + VOCAB, + WORLD, + assert_bf16_close, + craft_drafts, + gen, + head_workspace, + num_sms, + randn, + seed_of, +) + +assert torch.cuda.is_available(), "the exact-numerics checks require a CUDA device" + +DRAFT_LEN = 3 +CAPTURED_LAYERS = 3 +#: Largest share of fc_norm outputs one bf16 ulp off the in-tree RMSNorm (fp32 statistic summed in another order). +NORM_ULP_FRAC = 1e-3 + + +def _stock_accept(logits: torch.Tensor, draft: torch.Tensor): + """The in-tree greedy strict acceptance (``SpecWorkerBase._sample_and_accept_draft_tokens_base``) of the gathered + float logits of a decode-only batch, no forcing.""" + md = SimpleNamespace(is_all_greedy_sample=True, is_ragged_verify=False, is_cuda_graph=False) + worker = SimpleNamespace(force_num_accepted_tokens=0.0) + worker._sample_tokens_for_batch = lambda *a: Eagle3OneModelWorker._sample_tokens_for_batch( + worker, *a + ) + worker._apply_force_accepted_tokens = lambda *a, **k: ( + Eagle3OneModelWorker._apply_force_accepted_tokens(worker, *a, **k) + ) + nreq = draft.shape[0] + return Eagle3OneModelWorker._sample_and_accept_draft_tokens_base( + worker, logits, draft, 0, nreq, md + ) + + +def _lm_head_shard(weight: torch.Tensor, rank: int) -> LMHead: + """Rank ``rank``'s local part of the target's vocab-sharded BF16 LM head at TP4: an LM head over that rank's + 50016 rows (the TP4 module's local GEMM is this module's; its all-gather is the caller's).""" + head = LMHead(num_embeddings=TARGET_SHARD, embedding_dim=HIDDEN, dtype=torch.bfloat16) + head.weight.data = weight[rank * TARGET_SHARD : (rank + 1) * TARGET_SHARD].contiguous() + return head + + +def _run_shard_accept(logits, rank, world, draft, keys_of_peers=None, call=0): + """The fused verify of one rank: its shard logits, the peers' keys planted in its slots.""" + nreq = draft.shape[0] + rows = logits.shape[0] + uc, keys, counters = head_workspace(world, call=call, slot_rows=SHARD_ROWS) + half = call & 1 + for peer, peer_keys in (keys_of_peers or {}).items(): + for t in range(rows): + uc[(half * SHARD_ROWS + t) * world + peer] = peer_keys[t] + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + shard_accept_op( + logits, rank * TARGET_SHARD if world > 1 else 0, keys, counters, uc, uc.data_ptr(), rank, world, + draft, batch_idx, acc, num, gather, 0, next_new, num_sms(), None, + ) # fmt: skip + torch.cuda.synchronize() + return acc, num, gather, next_new + + +def _local_keys(logits: torch.Tensor, rank: int) -> torch.Tensor: + """What rank ``rank`` pushes: the key of its own shard maximum per row, found by the same kernel at world 1.""" + rows = logits.shape[0] + nreq = rows // (DRAFT_LEN + 1) + dummy = torch.full((nreq, DRAFT_LEN), -1, dtype=torch.int32, device="cuda") + acc, _, _, _ = _run_shard_accept(logits, 0, 1, dummy) + local = acc.view(-1).to(torch.int64) + value = logits.float().gather(1, local[:, None]).squeeze(1) + return make_keys(value, local + rank * TARGET_SHARD) + + +@pytest.mark.parametrize( + "nreq,ties", [(1, False), (2, False), (4, False), (8, True), (12, False)], ids=str +) +def test_fused_verify_is_the_stock_greedy_acceptance(nreq, ties) -> None: + """Target logits (in-tree BF16 head) and verify outputs bit-identical to the in-tree path, on all four ranks.""" + g = gen(seed_of(f"verify {nreq} {ties}")) + rows = nreq * (DRAFT_LEN + 1) + x = randn((rows, HIDDEN), g, 1.0) + weight = randn((VOCAB, HIDDEN), g, 0.02) + if ties: + # Two equal head rows on ranks 1 and 3 that win every row: the lower global id must be the token. + win = (x.float().sum(0) / x.float().sum(0).norm() * 4.0).to(torch.bfloat16) + weight[TARGET_SHARD + 77] = win + weight[3 * TARGET_SHARD + 5] = win + heads = [_lm_head_shard(weight, r) for r in range(WORLD)] + # Stock: each rank's LMHead forward, all-gathered, then LogitsProcessor's float cast. + shard_logits = [head(x) for head in heads] + for head, stock in zip(heads, shard_logits): + assert torch.equal(head.apply_linear(x, None), stock), ( + "the verify's shard GEMM is not the head's" + ) + gathered = torch.cat(shard_logits, dim=1).float() + target_tokens = gathered.argmax(dim=-1).to(torch.int32) + if ties: + assert bool((target_tokens == TARGET_SHARD + 77).all()) + draft = craft_drafts(target_tokens, nreq, DRAFT_LEN) + stock_acc, stock_num = _stock_accept(gathered, draft) + peer_keys = {r: _local_keys(shard_logits[r], r) for r in range(WORLD)} + for rank in range(WORLD): + peers = {r: k for r, k in peer_keys.items() if r != rank} + acc, num, gather, next_new = _run_shard_accept( + shard_logits[rank], rank, WORLD, draft, peers + ) + assert torch.equal(acc, stock_acc.to(torch.int32)), f"rank {rank}: accepted tokens" + assert torch.equal(num, stock_num.to(torch.int32)), f"rank {rank}: accepted counts" + expect_gather = torch.arange(nreq, device="cuda") * (DRAFT_LEN + 1) + num.long() - 1 + assert torch.equal(gather, expect_gather), f"rank {rank}: first-pass gather rows" + assert torch.equal(next_new[:, 0], target_tokens[expect_gather]), f"rank {rank}: next_new" + + +class _Drafter: + """The in-tree Eagle3 first-pass pieces at the target's shapes: per-layer fc_norm RMSNorms and the fc.""" + + def __init__(self, g: torch.Generator): + self.norms = [] + for _ in range(CAPTURED_LAYERS): + norm = RMSNorm(hidden_size=HIDDEN, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + 0.1 * torch.randn(HIDDEN, generator=g, device="cuda")).to( + torch.bfloat16 + ) + self.norms.append(norm) + self.fc = Linear(HIDDEN * CAPTURED_LAYERS, HIDDEN, bias=False, dtype=torch.bfloat16).cuda() + self.fc.weight.data = randn((HIDDEN, HIDDEN * CAPTURED_LAYERS), g, 0.02) + self.model = SimpleNamespace( + dtype=torch.bfloat16, + hidden_size=HIDDEN, + fc_norm=torch.nn.ModuleList(self.norms), + _norm_before_fc=False, + fc=self.fc, + ) + self.stacked = torch.stack([n.weight.detach() for n in self.norms]).contiguous() + + def stock_normed(self, hs: torch.Tensor) -> torch.Tensor: + """``apply_eagle3_fc``'s fc input: each captured layer through its own in-tree RMSNorm.""" + chunks = hs.chunk(CAPTURED_LAYERS, dim=-1) + return torch.cat([norm(c) for norm, c in zip(self.norms, chunks)], dim=-1) + + def stock_fc(self, hs: torch.Tensor) -> torch.Tensor: + return Eagle3ForCausalLM.apply_eagle3_fc(self, hs) + + +def _captures(rows: int, g: torch.Generator): + """Per captured layer (hidden, residual) and the in-tree capture buffer of bf16(hidden + residual).""" + layers = [(randn((rows, HIDDEN), g, 1.5), randn((rows, HIDDEN), g, 3.0)) for _ in range(3)] + buf = torch.cat([h + r for h, r in layers], dim=-1) + return layers, buf + + +def _row_identical(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return (a == b).all(dim=-1) + + +@pytest.mark.parametrize("rows", (4, 8, 32, 48), ids=lambda r: f"rows{r}") +def test_first_drafter_inputs_are_the_stock_fc_norm(rows) -> None: + """chunk_norm and capture_norm against ``apply_eagle3_fc``'s in-tree RMSNorms: within one bf16 ulp on at most + NORM_ULP_FRAC of the elements; the fc (the in-tree module in both paths) then differs only on rows whose input + did.""" + g = gen(seed_of(f"fc_norm {rows}")) + drafter = _Drafter(g) + layers, buf = _captures(rows, g) + stock_in = drafter.stock_normed(buf) + stock_out = drafter.stock_fc(buf) + assert torch.equal(drafter.fc(stock_in), stock_out) + + chunk_in = m3_spec_tail.chunk_norm(buf, drafter.stacked, EPS) + capture_in = torch.empty_like(buf) + for i, (h, r) in enumerate(layers): + m3_spec_tail.capture_norm( + h, r, drafter.stacked[i], EPS, capture_in[:, i * HIDDEN : (i + 1) * HIDDEN] + ) + torch.cuda.synchronize() + for name, ours in (("chunk_norm", chunk_in), ("capture_norm", capture_in)): + assert_bf16_close(ours, stock_in, max_frac=NORM_ULP_FRAC) + out = drafter.fc(ours) + same_rows = _row_identical(ours, stock_in) + assert torch.equal(out[same_rows], stock_out[same_rows]), ( + f"{name}: fc moved an unchanged row" + ) + moved = (ours != stock_in).float().mean().item() + print( + f"[exact-numerics] rows={rows} {name}: {moved:.2e} of fc inputs one ulp off, " + f"{int((~same_rows).sum())}/{rows} rows differ, fc max |diff| " + f"{(out.float() - stock_out.float()).abs().max().item():.3e}" + ) + assert torch.equal(chunk_in, capture_in), "capture_norm and chunk_norm disagree" + + +@pytest.mark.parametrize("rows", (4, 32), ids=lambda r: f"rows{r}") +def test_draft_tokens_through_the_full_vocab_head(rows) -> None: + """The fc_norm's ulp-level differences carried through the full-vocabulary BF16 head (a proxy of the draft head + input): the greedy draft token equals the stock one on every row whose stock top-2 margin exceeds the logit + perturbation.""" + g = gen(seed_of(f"draft head {rows}")) + drafter = _Drafter(g) + _, buf = _captures(rows, g) + weight = randn((VOCAB, HIDDEN), g, 0.02) + head = LMHead(num_embeddings=VOCAB, embedding_dim=HIDDEN, dtype=torch.bfloat16) + head.weight.data = weight + stock_logits = head(drafter.stock_fc(buf)).float() + ours_logits = head(drafter.fc(m3_spec_tail.chunk_norm(buf, drafter.stacked, EPS))).float() + perturb = (ours_logits - stock_logits).abs().max(dim=-1).values + top2 = stock_logits.topk(2, dim=-1).values + decidable = (top2[:, 0] - top2[:, 1]) > 2 * perturb + stock_tok, ours_tok = stock_logits.argmax(-1), ours_logits.argmax(-1) + assert torch.equal(ours_tok[decidable], stock_tok[decidable]) + print( + f"[exact-numerics] rows={rows} draft head: {int((ours_tok == stock_tok).sum())}/{rows} tokens equal, " + f"{int(decidable.sum())} decidable, max logit |diff| {perturb.max().item():.3e}" + ) + + +def test_the_drafter_is_the_in_tree_drafter() -> None: + """Outside the fused draft loop's steps the worker drafts with the in-tree loop, forward and sampling, and nothing + re-quantizes the drafter's linears or the LM heads: no hook of a quantized / pruned head or drafter projection is + left in the target, and the fused head streams the BF16 head weight itself.""" + assert ( + modeling.M3Eagle3Worker._forward_draft_loop is not Eagle3OneModelWorker._forward_draft_loop + ) + for name in ( + "_forward_linear_draft_loop", + "_run_draft_forward", + "sample_draft_tokens", + "sample_and_accept_draft_tokens", + ): + assert getattr(modeling.M3Eagle3Worker, name) is getattr(Eagle3OneModelWorker, name), name + for module, name in ( + (modeling, "M3DraftGemvLinearMethod"), + (modeling, "adopt_draft_linears"), + (modeling, "greedy_head_outputs"), + (modeling, "M3FP8LMHeadMethod"), + (modeling, "adopt_fp8_lm_heads"), + ): + assert not hasattr(module, name), f"{module.__name__}.{name}" + pkg = modeling.__name__.rsplit(".", 1)[0] + for gone in ("ops.m3_draft_head", "ops.spec.m3_draft_head", "ops.gemm.m3_draft_gemv"): + assert importlib.util.find_spec(f"{pkg}.{gone}") is None, gone diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py new file mode 100644 index 000000000000..d3e117c326e7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the forward-tail spec fold (ops/m3_fold_tail and M3SpecStepFold.in_forward). + +The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers +(step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing +when the step was not armed. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( + modeling as spec_fold, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( + m3_fold_tail, + m3_spec_tail, +) + +assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" + +VOCAB = 200064 +NUM_SLOTS = 9 +MAX_ROWS = 8 +STALE = -4 + + +def _store(widths, g): + wn, wx, wd = widths + i32 = dict(dtype=torch.int32, device="cuda") + return types.SimpleNamespace( + new_tokens=torch.randint( + 0, VOCAB, (wn, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_new_tokens=torch.randint( + 0, VOCAB, (wx, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_draft_tokens=torch.randint( + 0, VOCAB, (NUM_SLOTS, wd), generator=g, device="cuda", dtype=torch.int32 + ), + new_tokens_lens=torch.randint(1, 5, (NUM_SLOTS,), generator=g, **i32), + ) + + +def _clone(ns): + return types.SimpleNamespace(**{k: v.clone() for k, v in vars(ns).items()}) + + +def _outputs(rows, t, g): + d = t - 1 + return { + "new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_draft_tokens": torch.randint( + 0, VOCAB, (rows, d), generator=g, device="cuda", dtype=torch.int32 + ), + "new_tokens_lens": torch.randint( + 1, t + 1, (rows,), generator=g, device="cuda", dtype=torch.int32 + ), + } + + +def _engine(t): + i32 = dict(dtype=torch.int32, device="cuda") + return ( + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS * (t - 1),), STALE, **i32), + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS,), STALE, **i32), + ) + + +def _reference(store, engine, slots, out, t): + """slot_scatter into the store, then step_gather of the next step's inputs from it.""" + assert m3_spec_tail.slot_scatter( + slots, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["new_tokens_lens"], + store.new_tokens_lens, + out["next_draft_tokens"], + store.next_draft_tokens, + ) + p = slots.numel() + ids, draft, pos_off, kv_off = engine + slots32 = slots.to(torch.int32) + assert m3_spec_tail.step_gather( + store.next_new_tokens, + store.next_draft_tokens, + store.new_tokens_lens, + slots32, + slots32.repeat_interleave(t), + p, + t, + t - 1, + ids, + 0, + draft, + 0, + pos_off, + kv_off, + 0, + ) + + +def _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=False): + ids, draft, pos_off, kv_off = engine + assert m3_fold_tail.fold_tail( + ctrl, + slots_buf, + rows, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["next_draft_tokens"], + store.next_draft_tokens, + out["new_tokens_lens"], + store.new_tokens_lens, + ids, + draft, + pos_off, + kv_off, + t, + cpp=cpp, + ) + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("widths", [(4, 4, 3), (8, 8, 7)], ids=["store=step", "store wider"]) +def test_tail_matches_scatter_then_gather(rows, widths, cpp) -> None: + """Store bit for bit; engine rows [0, P*T) / [0, P) bit for bit (the forward reads only those rows).""" + t = 4 + g = torch.Generator(device="cuda").manual_seed(1000 * rows + widths[0]) + store = _store(widths, g) + out = _outputs(rows, t, g) + slots = torch.randperm(NUM_SLOTS, generator=g, device="cuda")[:rows] + ref_store, tail_store = _clone(store), _clone(store) + ref_engine, tail_engine = _engine(t), _engine(t) + _reference(ref_store, ref_engine, slots, out, t) + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + slots_buf = torch.zeros(MAX_ROWS, dtype=torch.long, device="cuda") + slots_buf[:rows] = slots + _tail(tail_store, tail_engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(ref_store, name), getattr(tail_store, name)), name + spans = (rows * t, rows * (t - 1), rows * t, rows) + for k, (ref, got, n) in enumerate(zip(ref_engine, tail_engine, spans)): + assert torch.equal(ref[:n], got[:n]), k + assert torch.all(got[n:] == STALE), k + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +def test_unarmed_tail_writes_nothing(cpp) -> None: + """ctrl[0] = 0 (the step was not armed): the store and the engine buffers are untouched.""" + t, rows = 4, 2 + g = torch.Generator(device="cuda").manual_seed(7) + store = _store((4, 4, 3), g) + out = _outputs(rows, t, g) + before, engine = _clone(store), _engine(t) + ctrl = torch.zeros(1, dtype=torch.int32, device="cuda") + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(before, name), getattr(store, name)), name + for buf in engine: + assert torch.all(buf == STALE) + + +def _fold(): + fold = spec_fold.M3SpecStepFold(in_forward=True) + t = 4 + engine = _engine(t) + fold.bind_engine_buffers(*engine, max_batch_size=MAX_ROWS) + g = torch.Generator(device="cuda").manual_seed(11) + fold.bind_sampler_store(_store((4, 4, 3), g)) + return fold, g, t + + +def test_fold_handshake_stages_only_confirmed_rows() -> None: + """begin_step arms; the sampler confirms only the armed rows from the tail's outputs; the next step is staged only + for exactly those rows in that order.""" + fold, g, t = _fold() + assert fold.begin_step([3, 5], 2) is False # nothing staged before the first forward + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [3, 5]) is True + assert fold.begin_step([3, 5], 2) is True + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert ( + fold.sampler_rows_stored(out, 0, [5, 3]) is False + ) # other order: the sampler stores itself + assert fold.begin_step([5, 3], 2) is False + + +def test_fold_handshake_needs_the_tail_outputs() -> None: + """Outputs the tail did not run on (another dict, a mixed step, an unarmed step) are never reported stored.""" + fold, g, t = _fold() + fold.begin_step([1], 1) + fold.store_in_forward(_outputs(1, t, g)) + assert fold.sampler_rows_stored(_outputs(1, t, g), 0, [1]) is False + fold.begin_step(None, 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [1]) is False + fold.begin_step([1], 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 1, [1]) is False # rows were skipped (chunked context) + + +def test_a_padded_step_after_an_armed_one_writes_nothing() -> None: + """K3 #150 check #5: CUDA-graph padding rows carry no seq slot, so the engine passes rows=None for a padded step. + Even right after an armed step, that step uploads ctrl 0 and its forward tail leaves the store and the engine + buffers untouched (a dummy row never writes a live slot).""" + fold, g, t = _fold() + fold.begin_step([2, 4], 2) + fold.store_in_forward(_outputs(2, t, g)) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 2 + + fold.begin_step(None, 4) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 0 + store_before = _clone(fold.store) + engine_before = [buf.clone() for buf in fold.step_inputs] + fold.store_in_forward(_outputs(4, t, g)) + torch.cuda.synchronize() + for name in vars(store_before): + assert torch.equal(getattr(store_before, name), getattr(fold.store, name)), name + for before, buf in zip(engine_before, fold.step_inputs): + assert torch.equal(before, buf) + + +def test_fold_off_by_default() -> None: + fold = spec_fold.M3SpecStepFold() + assert not fold.in_forward + assert fold.begin_step([0], 1) is False + + +@pytest.mark.parametrize( + "in_forward, store_copy_wait, expected", + [(True, True, True), (True, False, False), (False, True, False)], +) +def test_the_store_copy_event_exists_only_with_the_tail_and_the_wait( + in_forward, store_copy_wait, expected +): + """STEP_SPEC_STORE_COPY_SIDE: an external event the sampler records after its side-stream store copies and the + forward's tail waits on; only a fold that stores in the forward needs it.""" + fold = spec_fold.M3SpecStepFold(in_forward=in_forward, store_copy_wait=store_copy_wait) + device = torch.device("cuda") + fold.bind_engine_buffers( + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + 4, + ) + assert (fold.store_copies_done is not None) == expected + if expected: + fold.store_copies_done.synchronize() + + +def test_the_store_wait_precedes_only_a_plain_launch_under_capture(): + """The store wait is an external-event wait node under graph capture, which drops the next kernel's PDL edge: + it may precede the fold tail (a plain launch); any other launch fails the capture (eager steps are not + checked).""" + fold = spec_fold.M3SpecStepFold(in_forward=True, store_copy_wait=True) + device = torch.device("cuda") + fold.bind_engine_buffers( + *(torch.zeros(8, dtype=torch.int32, device=device) for _ in range(4)), 4 + ) + assert fold.STORE_WAIT_BEFORE == frozenset({"fold_tail"}) + fold._wait_store_copies(before="draft_step") # eager: allowed + x = torch.zeros(4, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + fold._wait_store_copies(before="fold_tail") + x.add_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(x, torch.ones(4, device=device)) + bad = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="PDL edge"): + with torch.cuda.graph(bad, stream=stream): + fold._wait_store_copies(before="draft_step") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..717deed4c1aa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec entry of the MiniMax-M3 tp4 target: the fractional forced-acceptance draw.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(5) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + counter = torch.tensor([41], dtype=torch.int64, device="cuda") + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + counter, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert int(counter.item()) == 42, "the draw bumps the step counter once" + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py new file mode 100644 index 000000000000..8a4d03c2fd03 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec_req entry of the MiniMax-M3 tp4 target: the per-request fractional forced-acceptance +draw (FORCE_AL_PER_REQUEST: one draw per request at its own key, its KV length at the verify, no shared counter), so +a request's forced count does not depend on the call history, on its slot in the batch, or on the other requests.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_force_vec_req import ( # noqa: E501 + m3_force_vec_req as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SEED, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(7) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + # Each request's key is its KV length at the verify: long, distinct, arbitrary. + keys = torch.randint( + 1, 1 << 20, (spec["batch"],), generator=g, device="cuda", dtype=torch.int64 + ).to(torch.int32) + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + keys, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert torch.equal(out, ref) + + +def _pool(): + g = torch.Generator(device="cpu") + g.manual_seed(_FORCE_ACCEPT_RNG_SEED) + return torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, dtype=torch.float32, generator=g).cuda() + + +def _draw(pool, keys, base_total=2, max_total=4, frac=0.6): + out = torch.empty(keys.numel(), dtype=torch.int32, device="cuda") + return op( + pool, + keys, + out, + keys.numel(), + base_total, + max_total, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (keys.numel() - 1).bit_length()), + ) + + +def test_a_request_draws_the_same_whatever_ran_before_or_beside_it() -> None: + pool = _pool() + keys = torch.tensor([165_521, 40_003, 323_007, 7], dtype=torch.int32, device="cuda") + first = _draw(pool, keys) + for _ in range(5): # no shared counter: repeated calls draw the same + assert torch.equal(_draw(pool, keys), first) + perm = torch.tensor([2, 0, 3, 1], device="cuda") + assert torch.equal(_draw(pool, keys[perm]), first[perm]) # slot-independent + alone = torch.cat([_draw(pool, keys[i : i + 1]) for i in range(4)]) + assert torch.equal(alone, first) # independent of the other requests + + +def test_the_counts_follow_the_forced_fraction() -> None: + pool = _pool() + keys = torch.arange(100_000, 100_000 + 64 * 400, dtype=torch.int32, device="cuda") + counts = torch.cat([_draw(pool, keys[i : i + 64]) for i in range(0, keys.numel(), 64)]) + assert set(counts.tolist()) <= {2, 3} + extra = (counts == 3).float().mean().item() + assert abs(extra - 0.6) < 0.02, extra diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py new file mode 100644 index 000000000000..01aa7eb2af53 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (ops/m3_lowm_splitk, +ops/m3_draft_lowm): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's +linears, eager and in a CUDA graph, with and without PDL.""" + +import pytest +import torch +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, +) + +assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" + +# Every certified (N, K, rows) of the TP4 drafter table that a split-K tactic serves: (N, K, rows, tactic). +CELLS = [ + (nk[0], nk[1], m, tactic_for(nk[0], nk[1], m)[1]) + for nk, rows in sorted(ROWS.items()) + for m in rows + if tactic_for(nk[0], nk[1], m)[0] == "s" +] + + +def test_the_table_has_the_certified_split_k_cells() -> None: + assert sorted((n, k, m) for n, k, m, _ in CELLS) == [ + (2304, 12288, 4), + (6144, 2048, 4), + (6144, 4608, 1), + (6144, 4608, 4), + (9216, 6144, 4), + ] + + +@pytest.mark.parametrize("use_pdl", [False, True], ids=["plain", "pdl"]) +@pytest.mark.parametrize("cell", CELLS, ids=[f"{n}x{k}@{m}" for n, k, m, _ in CELLS]) +def test_bitwise_flashinfer(cell, use_pdl) -> None: + n, k, m, tactic = cell + g = torch.Generator(device="cuda").manual_seed(n + 7 * k + m) + for trial in range(3): + x = (torch.randn(m, k, generator=g, device="cuda") * (1.0 + trial)).to(torch.bfloat16) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + ref = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + got = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + fi_run(x, w.t(), None, ref, use_pdl, FiTactic(*tactic)) + run_splitk_dense(x, w.t(), None, got, use_pdl, SplitKTactic(*tactic)) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), f"trial {trial}" + # In a CUDA graph, as the drafter runs it. + cap = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_splitk_dense(x, w.t(), None, cap, use_pdl, SplitKTactic(*tactic)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(cap.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..4ed33c8a7251 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py @@ -0,0 +1,297 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sampled_accept entry (fused temperature + top-p verify) and its workspace size.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_sampled_accept import ( # noqa: E501 + m3_sampled_accept as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_sampled_accept_ws_words import ( # noqa: E501 + m3_sampled_accept_ws_words as ws_words, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" + +DRAFT_LEN = 3 +#: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). +HEAD = 12 +#: Draws per distribution cell, and the total-variation band against the reference distribution. With 12 +#: categories and >= 4000 draws the expected TV of an exact sampler is ~0.02 (sd ~0.006); 0.06 is ~6 sd above it +#: while a wrong nucleus cut (one category too many or too few, >= 1 / 12 of the mass moved) lands at >= 0.08. +DRAWS = 4000 +TV_BAND = 0.06 + + +def _logits(spec, g, rows): + vs = spec["vs"] + if spec["check"] == "flat": + return (torch.randn((rows, vs), generator=g, device="cuda") * 0.3).to(torch.bfloat16) + if spec["check"] == "argmax": + return (torch.randn((rows, vs), generator=g, device="cuda") * 3).to(torch.bfloat16) + row = torch.full((vs,), -20.0, device="cuda") + hot = torch.randperm(vs, generator=g, device="cuda")[:HEAD] + row[hot] = torch.linspace(2.0, 0.2, HEAD, device="cuda") + return row.to(torch.bfloat16).expand(rows, vs).contiguous() + + +def _case(spec, seed): + g = gen(seed) + nreq = spec["nreq"] + rows = nreq * (DRAFT_LEN + 1) + logits = _logits(spec, g, rows) + words = ws_words(1) + uc = torch.zeros(words, dtype=torch.int64, device="cuda") + counters = torch.zeros(2, dtype=torch.int32, device="cuda") + tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + extra = torch.zeros(1, dtype=torch.int32, device="cuda") + temperatures = torch.full((rows,), spec["temperature"], dtype=torch.float32, device="cuda") + top_ps = torch.full((rows,), spec["top_p"], dtype=torch.float32, device="cuda") + seeds = torch.randint(0, 2**62, (rows,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.zeros(rows, dtype=torch.int64, device="cuda") + # Drafts drawn from the likeliest tokens, so some prefixes match. + top = logits[0].float().topk(3).indices.to(torch.int32) + spec["vocab_offset"] + draft = top[torch.randint(0, 3, (nreq, DRAFT_LEN), generator=g, device="cuda")].contiguous() + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + force = spec.get("force", "none") + force_total = 3 if force == "total" else 0 + force_vec = (torch.randint(1, DRAFT_LEN + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if force == "vec" else None) # fmt: skip + return [logits, spec["vocab_offset"], temperatures, top_ps, seeds, offsets, uc, uc.data_ptr(), counters, tokens, + 0, 1, draft, batch_idx, acc, num, gather, force_total, next_new, force_vec, extra, True] # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Statistical for the draw, exact for everything derived from it. + + Every drawn token has nonzero reference probability (lies in the nucleus); the acceptance outputs equal the + reference's strict acceptance of the drawn tokens, every call; the call counter advances once per call; and, + for the distribution cells, the draws over many calls (fresh Philox offsets) are within TV_BAND of the + reference distribution (see DRAWS / TV_BAND). + """ + spec = cell.spec + args = _case(spec, seed_of(cell.why)) + logits, offset, rows = args[0], spec["vocab_offset"], args[0].shape[0] + certifying(op) + ref = op.reference(*args) + probs = ref["probs"] + calls = 1 if spec["check"] == "argmax" else -(-DRAWS // rows) + counts = torch.zeros(logits.shape[1], dtype=torch.float64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + drawn = args[9][:rows].clone() + local = drawn.long() - offset + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + assert bool((probs.gather(1, local[:, None]) > 0).all()), ( + "a drawn token outside the nucleus" + ) + acc = ref["accept"](drawn) + assert torch.equal(args[14], acc["accepted"]) + assert torch.equal(args[15], acc["num_accepted"]) + assert torch.equal(args[16], acc["gather_ids"]) + assert torch.equal(args[18], acc["next_new"]) + assert args[8].tolist() == [0, c + 1] + counts.index_add_(0, local, torch.ones(rows, dtype=torch.float64, device="cuda")) + if spec["check"] == "distribution": + emp = counts / counts.sum() + tv = 0.5 * (emp - probs[0].double()).abs().sum().item() + assert tv <= TV_BAND, f"total variation {tv:.3f} against the reference nucleus distribution" + + +@pytest.mark.parametrize("top_p", [0.001, 0.0002]) +def test_long_rejection_chains_stay_in_the_nucleus(top_p) -> None: + """A tiny nucleus over a flat tail: rows that need more exchange rounds than the workspace has slots. + + Each round's candidate is drawn over the tokens heavier than the last one, so a row takes ~ln(1 / top_p) + 1 + rounds on average, with a tail past the 16 exchange slots (about 1% of the rows at top_p 0.001). Rounds reuse the + slots cyclically, told apart by a lap-parity bit in the wire tags, and no round accepts unconditionally: every + draw, however many rounds it took, lies in the reference nucleus. World 1 (one rank's 8 CTAs per row exchange + through the same slot ring); the cross-rank path needs the multi-device job. + """ + spec = dict(nreq=1, vs=50016, vocab_offset=50016, check="flat", top_p=top_p, temperature=1.0) + args = _case(spec, seed_of(f"flat tail top_p {top_p}")) + logits, rows = args[0], args[0].shape[0] + probs = op.reference(*args)["probs"] + calls = -(-DRAWS // rows) + outside = torch.zeros((), dtype=torch.int64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + local = args[9][:rows].long() - spec["vocab_offset"] + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + outside += (probs.gather(1, local[:, None]) == 0).sum() + assert args[8].tolist() == [0, c + 1] + draws = calls * rows + rounds = 2.0 + args[20].item() / draws + kept = int((probs > 0).sum(1).min().item()) + assert outside.item() == 0, ( + f"{outside.item()} of {draws} draws outside the nucleus (kept >= {kept} of {logits.shape[1]}, " + f"{rounds:.2f} rounds per row on average)" + ) + assert rounds > 4.0, ( + f"only {rounds:.2f} rounds per row: the rows no longer stress the slot ring" + ) + + +@pytest.mark.parametrize("cell", ws_words.CELLS, ids=[c.why[:44] for c in ws_words.CELLS]) +def test_workspace_words(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(ws_words) + world = cell.spec["world"] + assert ws_words(world) == ws_words.reference(world) + + +def test_guard_refuses_top_k_rows() -> None: + args = _case(op.CELLS[0].spec, 1) + top_ks = torch.full((args[0].shape[0],), 20, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="top-k rows"): + with validating(op): + op(*args, top_ks=top_ks) + + +def test_guard_refuses_fp64_temperatures() -> None: + args = _case(op.CELLS[0].spec, 2) + args[2] = args[2].double() + with pytest.raises(AssertionError, match="fp32 sampling parameters"): + with validating(op): + op(*args) + + +def test_the_verify_calls_this_targets_op_pair() -> None: + """The verify calls this target's op pair (the sm_107 tp4 kernel), whose workspace sizes as the shared one.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_sampled_accept as kernels, + ) + + assert kernels._sampled_accept_op is op + assert kernels._ws_words_op is ws_words + assert int(torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept_ws_words(4, 0)) == int( + torch.ops.trtllm.minimax_m3_sampled_accept_ws_words(4) + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_prewait_parameter_loads_change_nothing(cell) -> None: + """dispatch.SAMPLED_ACCEPT_PREWAIT_PARAMS moves only where the sampling parameters load (before or after the PDL + wait): the same draws and acceptance outputs, bit for bit, call after call.""" + outs = [] + for prewait in (True, False): + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + got = [] + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept(*args, prewait) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + outs.append(got) + for call, (a, b) in enumerate(zip(*outs)): + for x, y in zip(a, b): + assert torch.equal(x, y), f"call {call}: prewait changes {x.tolist()} -> {y.tolist()}" + + +def _walk_outputs(args, one_round_k, ilp, calls=3, offset_step=1000): + """The draws and acceptance outputs of `calls` calls with the owner walk one token (ilp False) or 8 tokens at a + time (ilp True), at SA-1R variant `one_round_k`, each call at its own Philox offsets.""" + if one_round_k: + words = torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept_ws_words(1, one_round_k) + args[6] = torch.zeros(words, dtype=torch.int64, device="cuda") + args[7] = args[6].data_ptr() + rows = args[0].shape[0] + got = [] + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * offset_step) + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept( + *args, True, False, one_round_k, None, ilp + ) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + return got + + +@pytest.mark.parametrize("one_round_k", [0, 82]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_ilp_walk_changes_nothing(cell, one_round_k) -> None: + """The owner walk's 8-token form (``ilp_walk``; off in this target) only computes the owner's token weights 8 at a + time; the walk adds them in the same order with the same early exit, so the draws and acceptance outputs are the + one-token walk's, bit for bit, call after call, with SA-1R off and on.""" + one = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, False) + ilp = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, True) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.95]) +def test_the_ilp_walk_changes_nothing_over_many_rounds(top_p) -> None: + """A flat tail (many tokens of equal weight) under a tiny and a wide nucleus: rows take many exchange rounds and + each round's walk exits at a different token; the 8-token walk draws the same tokens as the one-token walk.""" + spec = dict(nreq=1, vs=50016, vocab_offset=50016, check="flat", top_p=top_p, temperature=1.0) + one = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), 0, False, calls=24) + ilp = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), 0, True, calls=24) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +def test_the_ilp_walk_is_the_same_under_graph_replay() -> None: + """The 8-token walk captured in a CUDA graph and replayed draws what the one-token walk draws eagerly.""" + cell = op.CELLS[0] + eager = _walk_outputs(_case(cell.spec, seed_of(cell.why)), 0, False) + _walk_outputs( + _case(cell.spec, seed_of("ilp walk warm-up")), 0, True, calls=1 + ) # first launch outside the capture + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept(*args, True, False, 0, None, True) + torch.cuda.current_stream().wait_stream(stream) + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + graph.replay() + torch.cuda.synchronize() + for x, y in zip(eager[c], [args[i] for i in (9, 14, 15, 16, 18)]): + assert torch.equal(x, y), ( + f"replay {c}: the graphed 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_early_trigger_changes_nothing(cell) -> None: + """dispatch.SAMPLED_ACCEPT_EARLY_TRIGGER moves only where the grid releases its dependents: the same draws and + acceptance outputs, bit for bit, call after call, PDL launch on.""" + outs = [] + for early in (True, False): + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + got = [] + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept(*args, True, early) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + outs.append(got) + for call, (a, b) in enumerate(zip(*outs)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the early trigger changes {x.tolist()} -> {y.tolist()}" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept_1r.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept_1r.py new file mode 100644 index 000000000000..42bacb1467be --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept_1r.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SA-1R (dispatch.SAMPLED_ACCEPT_1R_K) on trtllm::minimax_m3_sm107tp4_sampled_accept, world 1, the TP4 shard. + +The one-round variants (one_round_k 9: exact top-9 list; 82: pooled top-8 + bound) must return the two-round +kernel's (one_round_k 0) sampled tokens, accepted tokens, accepted counts, gather ids and extra-attempt count BIT FOR +BIT on the same Philox stream, over random AgentX-like rows (several temperatures / top-p, broad and heavy-tailed +heads) and over rows built so the decision sits within ulps of top_p * Z (the margin fallback) or a slice holds more +heavy tokens than its list (the incomplete-list fallback), eager and under CUDA-graph replay. The round-0 decided +counter must engage on the AgentX-like rows and stay below the row count on the boundary rows (fallbacks fire). +Ported from the prototype test (exact on 39,968 rows). +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + +VS, D = 50016, 3 +KX = (9, 82) + +if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the sm_107 tp4 sampled accept runs on SM 10.7", allow_module_level=True) + + +class Runner: + """One variant with its own workspace / counters (call tag sequence) and outputs.""" + + def __init__(self, k, rows): + self.k, self.rows = k, rows + words = int(torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept_ws_words(1, k)) + self.ws = torch.zeros(words, dtype=torch.int64, device="cuda") + self.counters = torch.zeros(2, dtype=torch.int32, device="cuda") + self.tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + self.extra = torch.zeros(1, dtype=torch.int32, device="cuda") + self.decided = torch.zeros(1, dtype=torch.int32, device="cuda") + nreq = rows // (D + 1) + self.acc = torch.zeros((nreq, D + 1), dtype=torch.int32, device="cuda") + self.nacc = torch.zeros(nreq, dtype=torch.int32, device="cuda") + self.gids = torch.zeros(nreq, dtype=torch.int64, device="cuda") + + def __call__(self, lg, temps, tps, seeds, offs, draft, bidx): + torch.ops.trtllm.minimax_m3_sm107tp4_sampled_accept( + lg, 0, temps, tps, seeds, offs, self.ws, self.ws.data_ptr(), self.counters, self.tokens, 0, 1, draft, + bidx, self.acc, self.nacc, self.gids, 0, None, None, self.extra, False, True, False, self.k, + self.decided) # fmt: skip + + def out(self): + return (self.tokens[: self.rows].clone(), self.acc.clone(), self.nacc.clone(), self.gids.clone(), + self.extra.clone()) # fmt: skip + + +def _params(rows, temp, top_p, seed0): + temps = torch.full((rows,), temp, dtype=torch.float32, device="cuda") + tps = ( + (top_p if torch.is_tensor(top_p) else torch.full((rows,), top_p)) + .float() + .cuda() + .contiguous() + ) + seeds = torch.arange(rows, dtype=torch.int64, device="cuda") * 7919 + seed0 + offs = torch.zeros(rows, dtype=torch.int64, device="cuda") + return temps, tps, seeds, offs + + +def _sets(n, rows, g, kind): + out = [] + for _ in range(n): + lg = torch.randn(rows, VS, device="cuda", generator=g) * (1.0 if kind == "flat" else 2.0) + if kind == "lm": # a few heavy tokens over a broad tail + hot = torch.randint(0, VS, (rows, 8), device="cuda", generator=g) + lg.scatter_(1, hot, 12.0 + 3.0 * torch.rand(rows, 8, device="cuda", generator=g)) + elif kind == "tail": # 64 near-equal top tokens: lists overflow + hot = torch.randint(0, VS, (rows, 64), device="cuda", generator=g) + lg.scatter_(1, hot, 9.0 + 0.5 * torch.rand(rows, 64, device="cuda", generator=g)) + out.append(lg.to(torch.bfloat16).contiguous()) + return out + + +def _draft_from(lg, g): + nreq = lg.shape[0] // (D + 1) + am = lg.float().argmax(1).view(nreq, D + 1)[:, :D].int() + rnd = torch.randint(0, VS, (nreq, D), device="cuda", generator=g, dtype=torch.int32) + return torch.where(torch.rand(nreq, D, device="cuda", generator=g) < 0.7, am, rnd).contiguous() + + +def _compare(ref, got, what): + for n, x, y in zip( + ("tokens", "acc_tokens", "num_acc", "gather_ids", "extra_attempts"), ref, got + ): + assert torch.equal(x, y), f"{what}: {n} differs" + + +CASES = [(1.0, 0.95, "lm"), (0.7, 0.8, "lm"), (1.3, 0.5, "lm"), (1.0, 1.0, "lm"), (1.0, 0.95, "flat"), + (1.0, 0.95, "tail"), (0.6, 0.99, "tail")] # fmt: skip + + +def test_one_round_is_two_round_bitwise(): + g = torch.Generator(device="cuda").manual_seed(2) + total = 0 + decided = {k: 0 for k in KX} + for rows, calls in ((32, 20), (4, 60)): + for ci, (t, p, kind) in enumerate(CASES): + runners = {k: Runner(k, rows) for k in (0, *KX)} + temps, tps, seeds, offs = _params(rows, t, p, 1000 + ci) + sets = _sets(4, rows, g, kind) + bidx = torch.arange(rows // (D + 1), dtype=torch.int32, device="cuda") + for i in range(calls): + lg = sets[i % 4] + draft = _draft_from(lg, g) + offs.fill_(i) + for r in runners.values(): + r(lg, temps, tps, seeds, offs, draft, bidx) + for k in KX: + _compare( + runners[0].out(), runners[k].out(), f"K={k} rows={rows} case={ci} call={i}" + ) + total += rows + if kind == "lm" and p < 1.0: + for k in KX: + decided[k] += int(runners[k].decided.item()) + assert int(runners[0].decided.item()) == 0 + assert total >= 5000 + for k in KX: + assert decided[k] > 0, ( + f"K={k}: no row decided in round 0 on the AgentX-like rows (SA-1R not engaged)" + ) + + +def _boundary_rows(rows, n_heavy_same_slice): + """Heavy tokens heavier than a 50-token candidate group; top_p swept across ulps of S / Z (margin fallback), or + n heavy tokens in one slice (more than its list: incomplete fallback).""" + lg = torch.full((rows, VS), -40.0, device="cuda") + sl = ((VS + 7) // 8 + 7) // 8 * 8 + heavy = list(range(5, 5 + n_heavy_same_slice)) if n_heavy_same_slice else [5, 3 * sl + 11] + grp = [2 * sl + 3 + 37 * j for j in range(25)] + [5 * sl + 7 + 41 * j for j in range(25)] + lg[:, heavy] = 5.0 + lg[:, grp] = 4.0 + lg = lg.to(torch.bfloat16).contiguous() + w = torch.exp(lg[0].double() - lg[0].double().max()) + crit = float(w[heavy].sum()) / float(w.sum()) + ks = torch.arange(rows, dtype=torch.float64) - rows // 2 + return lg, (crit * (1.0 + ks * 2.0**-24)).float() + + +@pytest.mark.parametrize("graph", [False, True]) +def test_one_round_boundary_and_incomplete(graph): + rows = 32 + g = torch.Generator(device="cuda").manual_seed(3) + seen = {k: [0, 0] for k in KX} # decided rows, rows + for n_same in (0, 12): + lg, tps_rows = _boundary_rows(rows, n_same) + for shift in range(-4, 4): + tps = (tps_rows.double() * (1.0 + shift * 32 * 2.0**-24)).float() + runners = {k: Runner(k, rows) for k in (0, *KX)} + temps, tp_t, seeds, offs = _params(rows, 1.0, tps, 5000 + shift) + bidx = torch.arange(rows // (D + 1), dtype=torch.int32, device="cuda") + draft = _draft_from(lg, g) + graphs = {} + if graph: + for r in runners.values(): + r(lg, temps, tp_t, seeds, offs, draft, bidx) + torch.cuda.synchronize() + for k, r in runners.items(): + graphs[k] = torch.cuda.CUDAGraph() + with torch.cuda.graph(graphs[k]): + r(lg, temps, tp_t, seeds, offs, draft, bidx) + for r in runners.values(): + r.decided.zero_() + for i in range(12): + offs.fill_(i) + for k, r in runners.items(): + if graph: + graphs[k].replay() + else: + r(lg, temps, tp_t, seeds, offs, draft, bidx) + for k in KX: + _compare(runners[0].out(), runners[k].out(), + f"boundary n_same={n_same} shift={shift} call={i} K={k}") # fmt: skip + for k in KX: + seen[k][0] += int(runners[k].decided.item()) + seen[k][1] += 12 * rows + for k in KX: + assert 0 < seen[k][0] < seen[k][1], ( + f"K={k}: decided {seen[k]}: the boundary sweep missed both sides" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..0c90ceac2e70 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py @@ -0,0 +1,125 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_shard_accept entry (greedy verify from this rank's logits shard).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + craft_drafts, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" + + +def _case(spec, seed, edit_logits=None): + g = gen(seed) + nreq, d, vs, world, rank = ( + spec["nreq"], + spec["draft_len"], + spec["vs"], + spec["world"], + spec["rank"], + ) + m = nreq * (d + 1) + if spec["strided"]: + logits = randn((m, vs + 72), g, 3.0)[:, :vs] + else: + logits = randn((m, vs), g, 3.0) + if edit_logits is not None: + edit_logits(logits) + uc, keys, counters = head_workspace(world, slot_rows=SHARD_ROWS) + own = logits.float().max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=SHARD_ROWS, ties=spec.get("ties", False) + ) + offset = rank * vs + best, idx = local_argmax(logits.float(), offset) + tokens = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, SHARD_ROWS) + draft = craft_drafts(tokens, nreq, d) + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + force_total = 3 if spec["force"] == "total" else 0 + force_vec = (torch.randint(1, d + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if spec["force"] == "vec" else None) # fmt: skip + return (logits, offset, keys, counters, uc, uc.data_ptr(), rank, world, draft, batch_idx, acc, num, gather, + force_total, next_new, num_sms(), force_vec) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bit-exact: the op reduces BF16 logits without arithmetic, ties (within the shard and across ranks) go to + the lowest global index on both sides.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) + assert args[3].tolist() == [0, 1] and not bool(args[2].any()) + + +def test_guard_refuses_nan_logits() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0][1, 7] = float("nan") + with pytest.raises(AssertionError, match="NaN"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_batch_index_past_the_requests() -> None: + args = list(_case(op.CELLS[0].spec, 2)) + args[9] = args[9] + 1 + with pytest.raises(AssertionError, match="batch_idx"): + with validating(op): + op(*args) + + +def _signed_zero_maxima(logits: torch.Tensor) -> None: + """Every row's maximum is zero, reached as -0.0 at a lower index and +0.0 at a higher one.""" + logits.fill_(-1.0) + logits[:, 3] = -0.0 + logits[:, 9] = 0.0 + + +@pytest.mark.parametrize("ties", [False, True], ids=["peer_rows", "peer_ties"]) +def test_signed_zero_maxima_take_the_lower_index(ties) -> None: + """-0.0 and +0.0 compare equal, so the lower index wins as with torch.argmax (also across ranks).""" + spec = dict(op.CELLS[1].spec, ties=ties) + args = _case(spec, seed_of(f"signed-zero-{ties}"), edit_logits=_signed_zero_maxima) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..4b27e42af50b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_slot_scatter entry (the spec sampler's per-slot stores in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_slot_scatter import ( # noqa: E501 + m3_slot_scatter as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" + +STALE = -4 + + +def _case(spec, seed): + """Worker outputs and the store's transposed [slots, width] views of its [width, slots, 1] buffers.""" + g = gen(seed) + n, nslots = spec["n"], spec["slots"] + + def rnd(shape): + return torch.randint(0, VOCAB, shape, generator=g, device="cuda", dtype=torch.int32) + + store_new = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_next = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_lens = torch.full((nslots,), STALE, dtype=torch.int32, device="cuda") + store_draft = torch.full((nslots, spec["store_d"]), STALE, dtype=torch.int32, device="cuda") + slots = torch.randperm(nslots, generator=g, device="cuda")[:n] + lens = torch.randint(1, 5, (n,), generator=g, device="cuda", dtype=torch.int32) + return (slots, rnd((n, spec["src_w"])), store_new.squeeze(-1).T, rnd((n, spec["src_w"])), + store_next.squeeze(-1).T, lens, store_lens, rnd((n, spec["src_d"])), store_draft) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, zero padding included, other slots untouched.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in (("da", args[2]), ("db", args[4]), ("dl", args[6]), ("dd", args[8])): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_repeated_slots() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0] = args[0].clone() + args[0][1] = args[0][0] + with pytest.raises(AssertionError, match="distinct"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py new file mode 100644 index 000000000000..4a2a242816d3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for m3_spec_offsets (RuntimeHooks.spec_offsets_op): the overlap scheduler's two speculative step corrections +in one plain launch, bit for bit the in-place ATen adds, eager and in a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_spec_offsets import ( # noqa: E501 + m3_spec_offsets as op, +) + +assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" + +I32 = torch.iinfo(torch.int32) + + +def _inputs(tokens, seqs, seed, wrap=False): + g = torch.Generator(device="cuda").manual_seed(seed) + lo, hi = (I32.min, I32.max) if wrap else (0, 1 << 20) + pos = torch.randint(lo, hi, (tokens,), generator=g, device="cuda", dtype=torch.int32) + kv = torch.randint(lo, hi, (seqs,), generator=g, device="cuda", dtype=torch.int32) + pos_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (tokens,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + kv_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (seqs,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + if wrap: + pos[:2], pos_off[:2] = I32.max, 1 # overflow wraps, as the ATen add does + kv[:1], kv_off[:1] = I32.min, -1 + return pos, pos_off, kv, kv_off + + +CELLS = [c.spec for c in op.CELLS] + [ + dict(tokens=1, seqs=1), + dict(tokens=300, seqs=75), + dict(tokens=4, seqs=0), +] + + +@pytest.mark.parametrize("spec", CELLS, ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items())) +def test_bitwise_aten_adds(spec) -> None: + for trial in range(3): + pos, pos_off, kv, kv_off = _inputs( + spec["tokens"], spec["seqs"], 97 * trial + spec["tokens"], spec.get("wrap", False) + ) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + assert op(pos, pos_off, kv, kv_off) is True + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"trial {trial}" + + +def test_graph_replays_read_the_live_offsets() -> None: + pos, pos_off, kv, kv_off = _inputs(32, 8, 5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + for step in range(3): + fresh = _inputs(32, 8, 11 + step) + for dst, src in zip((pos, pos_off, kv, kv_off), fresh): + dst.copy_(src) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"step {step}" + + +def test_runner_views_are_accepted() -> None: + """The runner passes position_ids[0, ctx:] and kv_lens_cuda[ctx_reqs:seqs]: 1-D slices of contiguous buffers.""" + position_ids = torch.arange(40, dtype=torch.int32, device="cuda").view(1, 40) + kv_lens = torch.arange(12, dtype=torch.int32, device="cuda") + pos_off = torch.full((40,), 2, dtype=torch.int32, device="cuda") + kv_off = torch.full((12,), 3, dtype=torch.int32, device="cuda") + assert op(position_ids[0, 8:], pos_off[:32], kv_lens[2:10], kv_off[:8]) is True + torch.cuda.synchronize() + assert torch.equal( + position_ids[0], torch.cat([torch.arange(8), torch.arange(8, 40) + 2]).int().cuda() + ) + assert torch.equal( + kv_lens, + torch.cat([torch.arange(2), torch.arange(2, 10) + 3, torch.arange(10, 12)]).int().cuda(), + ) + + +@pytest.mark.parametrize( + "bad", + ["int64", "strided", "length", "2d"], +) +def test_declines_other_layouts_without_writing(bad) -> None: + pos, pos_off, kv, kv_off = _inputs(8, 2, 3) + if bad == "int64": + kv = kv.long() + elif bad == "strided": + pos = torch.zeros(16, dtype=torch.int32, device="cuda")[::2] + elif bad == "length": + pos_off = pos_off[:7] + else: + pos = pos.view(2, 4) + before = (pos.clone(), kv.clone()) + assert op(pos, pos_off, kv, kv_off) is False + torch.cuda.synchronize() + assert torch.equal(pos, before[0]) and torch.equal(kv, before[1]) + + +def test_launches_plainly() -> None: + """It writes step inputs that pre-wait readers take before their grid wait (MSA contract, dispatch + PRODUCER_PREWAIT_STEP_INPUTS): its kernel node carries no programmatic (PDL) serialization.""" + from cuda.bindings import runtime as rt + + pos, pos_off, kv, kv_off = _inputs(32, 8, 7) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + err, _, n = rt.cudaGraphGetNodes(raw, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + flags = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind == rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + flags.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + assert flags == [0], flags diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..c57435ce1f30 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_gather entry (the overlap scheduler's step-input gathers in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_step_gather import ( # noqa: E501 + m3_step_gather as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP4 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" + +#: Rows of the engine's flat buffers (the max batch), and their stale fill. +MAX_ROWS = 8 +STALE = -4 + + +def _case(spec, seed): + g = gen(seed) + p, t, d, nslots, b0 = spec["p"], spec["t"], spec["d"], spec["slots"], spec["b0"] + i32 = dict(dtype=torch.int32, device="cuda") + new_tok = torch.randint(0, VOCAB, (t, nslots, 1), generator=g, device="cuda", dtype=torch.int32) + if spec.get("two_d"): + new_tok = new_tok.squeeze(-1) + next_draft = torch.randint(0, VOCAB, (nslots, d), generator=g, device="cuda", dtype=torch.int32) + lens = torch.randint(1, t + 1, (nslots,), generator=g, device="cuda", dtype=torch.int32) + slots = torch.randperm(nslots, generator=g, device="cuda")[:p].to(torch.int32) + ppi = slots.repeat_interleave(t) + ids = torch.full((MAX_ROWS * t,), STALE, **i32) + draft = torch.full((MAX_ROWS * d,), STALE, **i32) + pos_off = torch.full((MAX_ROWS * t,), STALE, **i32) + kv_off = torch.full((MAX_ROWS,), STALE, **i32) + return ( + new_tok, + next_draft, + lens, + slots, + ppi, + p, + t, + d, + ids, + spec["ids0"], + draft, + spec["d0"], + pos_off, + kv_off, + b0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in ( + ("ids", args[8]), + ("draft", args[10]), + ("pos_off", args[12]), + ("kv_off", args[13]), + ): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_a_slot_past_the_store() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[3] = args[3].clone() + args[3][0] = 8 + with pytest.raises(AssertionError, match="slots must index the store"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py new file mode 100644 index 000000000000..a5d0cec0ed6c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_prologue entry (the decode graph's step-input prologue in one launch).""" + +import random +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_step_prologue import ( # noqa: E501 + m3_step_prologue as op, +) +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" + +MAX_DESC = StepInputPack.MAX_DESC +WINDOW = StepInputPack.GRAPH_WINDOW +DATA_START = StepInputPack.HEADER_BYTES + StepInputPack.DESC_BYTES * MAX_DESC +ARENA = 1 << 20 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _layout(spec, seed): + """Copies as (dst, src, bytes) arena offsets (a host copy's src: a window offset) and the host copies' bytes.""" + rng = random.Random(seed) + h2d, d2d, data = [], [], [] + used = DATA_START + pos = 0 + for _ in range(spec["h2d"]): + n = rng.choice([1, 3, 4, 7, 16, 24, 32, 48, 64, 5, 12]) + pos += rng.randint(0, 40) # gaps and every alignment + h2d.append((pos, used, n)) + data.append(bytes(rng.getrandbits(8) for _ in range(n))) + pos += n + used += -(-n // 16) * 16 + assert pos < ARENA // 2 and used <= WINDOW + # Device copies write the last quarter; unchained ones read the third. + dpos, spos = 3 * ARENA // 4, ARENA // 2 + for i in range(spec["d2d"]): + if spec.get("chained") and h2d: + dst_h, _, n_h = h2d[i % len(h2d)] + off = rng.randint(0, n_h - 1) + src, n = dst_h + off, n_h - off + else: + n = rng.choice([4, 8, 16, 32, 6, 64]) + spos += rng.randint(0, 24) + src = spos + spos += n + dpos += rng.randint(0, 24) + d2d.append((dpos, src, n)) + dpos += n + assert dpos < ARENA and spos < 3 * ARENA // 4 + return h2d, d2d, data + + +def _window(layout, base: int) -> torch.Tensor: + """The window of ``layout`` for an arena at device address ``base``.""" + h2d, d2d, data = layout + w = torch.zeros(WINDOW, dtype=torch.uint8) + words = [len(h2d), len(d2d)] + for dst, src, n in h2d: + words += [base + dst, src, n] + for dst, src, n in d2d: + words += [base + dst, base + src, n] + w[: 8 * len(words)] = torch.tensor(words, dtype=torch.int64).view(torch.uint8) + for (_, off, n), b in zip(h2d, data): + w[off : off + n] = torch.frombuffer(bytearray(b), dtype=torch.uint8) + return w.cuda() + + +def _arenas(seed): + a = torch.randint(0, 256, (ARENA,), generator=gen(seed), device="cuda", dtype=torch.int32) + a = a.to(torch.uint8) + return a, a.clone() + + +@pytest.mark.parametrize("with_epoch", [True, False], ids=["epoch", "no_epoch"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, with_epoch, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Byte for byte the default prologue's destinations, header clear and epoch.""" + seed = seed_of(cell.why) + layout = _layout(cell.spec, seed) + arena, ref_arena = _arenas(seed) + window, ref_window = _window(layout, arena.data_ptr()), _window(layout, ref_arena.data_ptr()) + epoch = torch.full((1,), 41, dtype=torch.int64, device="cuda") if with_epoch else None + ref_epoch = epoch.clone() if with_epoch else None + certifying(op) + op(window, epoch, MAX_DESC) + op.reference(ref_window, ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena) + # Each window holds its own arena's addresses: past the consumed header, both are left as staged. + for w, a in ((window, arena), (ref_window, ref_arena)): + assert not w[:16].any(), "the header is consumed" + assert torch.equal(w[16:], _window(layout, a.data_ptr())[16:]) + if with_epoch: + assert int(epoch.item()) == int(ref_epoch.item()) == 42 + + +def test_graph_replays_scatter_each_new_window() -> None: + """Captured once on a static window: every replay writes that replay's window and bumps the epoch once.""" + arena, ref_arena = _arenas(7) + window = torch.zeros(WINDOW, dtype=torch.uint8, device="cuda") + epoch = torch.zeros(1, dtype=torch.int64, device="cuda") + ref_epoch = epoch.clone() + op(window, epoch, MAX_DESC) # the empty window: writes nothing + op.reference(window.clone(), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(window, epoch, MAX_DESC) + for step, spec in enumerate( + [dict(h2d=24, d2d=8, chained=True), dict(h2d=0, d2d=0), dict(h2d=40, d2d=4)] + ): + layout = _layout(spec, 100 + step) + window.copy_(_window(layout, arena.data_ptr())) + graph.replay() + op.reference(_window(layout, ref_arena.data_ptr()), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena), step + assert not window[:16].any(), step + assert int(epoch.item()) == int(ref_epoch.item()) == step + 2, step + + +def _pack_case(prologue_op): + pack = StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=True, + in_graph=True, + early_copy=True, + step_epoch=True, + prologue_op=prologue_op, + ) + a, b, c = (torch.zeros(16, dtype=torch.int32, device="cuda") for _ in range(3)) + out = torch.zeros(48, dtype=torch.int32, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _pack_step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + + def host(v): + return torch.tensor(v, dtype=torch.int32).pin_memory() + + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, host(values[0])) + step_inputs.copy_(b, host(values[1])) + if overlap: + step_inputs.copy_(a[:4], host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if not flush: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_pack_with_the_one_launch_prologue_matches_the_default() -> None: + """The same pack steps through the default prologue and the one-launch one: the same outputs and epochs, and the + one-launch prologue is engaged (the default's warm-up compile is not what bumped its epoch).""" + packs = [_pack_case(None), _pack_case(op)] + assert packs[0][0].prologue_op is None and packs[1][0].prologue_op is op + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + outs = [] + for pack, graph, dsts, out in packs: + # A replay with a new pack, one with an eagerly issued pack, and one with nothing to scatter. + _pack_step(pack, graph, dsts, (va, vb), overlap=step == 2, flush=step != 3) + torch.cuda.synchronize() + outs.append((out.tolist(), int(pack.epoch.item()))) + assert outs[0] == outs[1], step + assert outs[1][1] == step + 2, step # one bump at warm-up, then one per replay diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..d8eb4f29e589 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The feed-forward exit on Eagle3 capture layers: the generic layer in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``), and the decode fast path's capture layers (``CAPTURE_LAYERS_FAST``). + +Every layer's MoE / MLP output goes unreduced to the fused boundary (all-reduce + residual + the next layer's input +norm, the final norm after the last layer), and a capture layer captures the reduced residual sum that boundary +writes, with no residual of its own. CPU only: the attention, the feed-forward and the boundaries are stand-ins. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) + +LAYERS = 4 + + +class _Spec: + def __init__(self, layers): + self.layers = set(layers) + + def is_layer_capture(self, i): + return i in self.layers + + +@pytest.fixture(params=["prefill", "decode"]) +def layer(request, monkeypatch): + """(core, the target's generic layer bound to a target over ``core``).""" + calls = [] + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + c = SimpleNamespace( + w=w, num_layers=LAYERS, moe_layers=frozenset({2, 3}), rms_eps=1e-6, calls=calls + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, + "_native_allreduce_norm", + lambda core, x, r, norm: (("normed", x, norm), ("updated", x, r)), + ) + + def boundary(self, x, x2, residual, norm_w): + calls.append((x, x2, norm_w)) + return ("next", norm_w), ("sum", x, x2) + + if request.param == "prefill": + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i))) + monkeypatch.setattr( + cls, "_swiglu_mlp", lambda self, h, prefix, low_m, dense=False: ("mlp", prefix) + ) + else: + cls = modeling.DecodeTarget + monkeypatch.setattr( + cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i), False, False) + ) + monkeypatch.setattr( + cls, + "_swiglu_mlp", + lambda self, h, prefix, low_m, push=None, dense=False: (("mlp", prefix), False), + ) + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + monkeypatch.setattr(cls, "_generic_boundary", boundary) + target = cls(c) + return c, target._layer if request.param == "prefill" else target._generic_layer + + +@pytest.mark.parametrize("i", range(LAYERS)) +def test_a_capture_layer_captures_the_boundary_residual_sum(layer, i) -> None: + core, run = layer + captured = [] + out = run( + i, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([i]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + norm = core.w[f"l{i + 1}_in_norm"] if i + 1 < LAYERS else "final" + assert core.calls[-1][2] == norm, ( + "the boundary applies the next layer's norm (the final norm after the last)" + ) + assert captured == [(i, out[1], None)] + if i in core.moe_layers: + assert core.calls[-1][1] == ("routed", i), "the routed part folds into the boundary" + + +def test_a_layer_without_capture_captures_nothing(layer) -> None: + _, run = layer + captured = [] + run( + 1, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([]), + lambda *a: captured.append(a), + ) + assert not captured + + +class _FastCore: + num_layers = 3 + w = {"embed": None} + + +@pytest.mark.parametrize("written", [True, False]) +def test_a_decode_capture_layer_takes_the_fast_layer_with_the_capture_norm( + monkeypatch, written +) -> None: + """A decode step's capture layer (past the first) runs the fast decode layer with the worker's capture target, and + captures the boundary's updated residual; ``written=True`` when the boundary all-reduce wrote the capture norm + (``CAPTURE_NORM_IN_BOUNDARY``).""" + assert modeling.CAPTURE_LAYERS_FAST and modeling.CAPTURE_FUSED_BOUNDARY + assert modeling.CAPTURE_NORM_IN_BOUNDARY + target = modeling.DecodeTarget(_FastCore()) + seen, captured = [], [] + monkeypatch.setattr(target, "_check_step_contract", lambda md: None) + monkeypatch.setattr(target, "step_args", lambda md: {"num_tokens": 4}) + monkeypatch.setattr(target, "_embed", lambda ids, n: ("h", "r")) + monkeypatch.setattr(modeling, "plan_for", lambda n: SimpleNamespace(name="plan")) + monkeypatch.setattr(modeling, "_record", lambda *a: None) + + def fast(i, plan, pos, h, r, step, capture=None): + seen.append((i, capture)) + if capture is not None: + target._capture_written = written + return ("h", i), ("r", i) + + def generic(*a): + raise AssertionError("a capture layer ran the generic layer on a decode step") + + monkeypatch.setattr(target, "_layer", fast) + monkeypatch.setattr(target, "_generic_layer", generic) + out = target.forward( + None, + torch.zeros(4, dtype=torch.int32), + None, + None, + _Spec([1]), + lambda spec, i, h, r, written=False: captured.append((i, h, r, written)), + lambda spec, i, n: ("capture_out", i, n), + ) + assert out == ("h", 2) + assert seen == [(0, None), (1, ("capture_out", 1, 4)), (2, None)] + assert captured == [(1, ("r", 1), None, written)] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..8a70e7e9947a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,286 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp4 (Rubin) decode target's fast-path table (``modeling.py``: the row ranges and the selection functions): its +row ranges against the catalog cells that certify them, the selection for every decode row count, and the phase split. + +The paths whose cells pass on SM 10.7 take the sm_103 tp4 rows; the ones built on kernels with no SM 10.7 build (the +the fp32 routers, the low-M GEMM, the fused / mega / c1 attention) and the ones not yet driven on +SM 10.7 (the long-row select, the capped dense decode split, the early drafter fc) cover nothing. The dense layers' and +the drafter's full-context decode take the one-request shapes measured at TP4. + +CPU only: nothing here launches a kernel. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + triton_sparse_decode as tsd, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm import ( # noqa: E501 + m3_mx_gemv, +) + +ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 2) # every decode row count and the first one past it + +# The TP4 per-rank MXFP8 projections (N, K): 16 query heads and 1 KV head, expert width 768 per rank. +TP4_PROJECTIONS = ( + m3_mx_gemv.QKV_SPARSE, + m3_mx_gemv.QKV, + m3_mx_gemv.O_PROJ, + m3_mx_gemv.SHARED_GATE_UP, + m3_mx_gemv.SHARED_DOWN, + m3_mx_gemv.DENSE_GATE_UP, + m3_mx_gemv.DENSE_DOWN, +) + + +@pytest.fixture(autouse=True) +def _policy(): + """Every test starts outside require mode with nothing logged, and leaves it so.""" + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_every_covered_shape_is_a_certified_cell(fast_path) -> None: + uncertified = fast_path.keys - fast_path.certified() + assert not uncertified, ( + f"{fast_path.path} covers {sorted(uncertified, key=str)} but " + f"{type(fast_path.entry).__name__}.CELLS certifies none of them: add a cell or narrow the range" + ) + + +def test_the_fast_path_names_are_unique() -> None: + names = [r.path for r in dispatch.FAST_PATHS] + assert len(names) == len(set(names)) + + +def test_the_decode_target_keeps_the_tp4_geometry() -> None: + """Up to 8 requests verifying 1 + 3 drafts and their padding rows; a wider step is the prefill target's.""" + assert dispatch.DECODE_MAX_TOKENS == 48 + assert [dispatch.is_decode_step(n) for n in ROWS] == [True] * 48 + [False] + + +# ------------------------------------------------------------------ the paths still off on SM 10.7 + +OFF_ON_SM107 = ( + dispatch.DENSE_DECODE_SPLIT, + dispatch.LONG_SELECT, + dispatch.MC_FINALIZE, +) + + +@pytest.mark.parametrize("fast_path", OFF_ON_SM107, ids=lambda r: r.path) +def test_paths_without_an_sm107_build_cover_nothing(fast_path) -> None: + assert not fast_path.keys, f"{fast_path.path} covers {sorted(fast_path.keys, key=str)}" + + +def test_decode_boundaries_follow_the_multicast_all_reduce_kill() -> None: + """MC_AR_TP4 moves every decode boundary (and the o_proj / MoE pushes into it) onto the SM 10.7 TP4 multicast + all-reduce; off, the built-in all-reduce runs and nothing pushes.""" + on = dispatch.MC_AR_TP4 + for plan in dispatch.DECODE_PLANS: + assert (plan.attn_mc_ar, plan.moe_mc_ar, plan.qkv_prequant) == (on, on, on), plan + assert tuple(p.max_tokens for p in dispatch.DECODE_PLANS) == (dispatch.DECODE_MAX_TOKENS,) + assert bool(dispatch.MC_ALLREDUCE.keys) == on + assert dispatch.DRAFT_MC_AR.keys == (frozenset(range(1, 9)) if on else frozenset()) + assert dispatch.ATTN_PUSH_MAX_TOKENS == (8 if on else 0) + assert dispatch.DENSE_PUSH_MAX_TOKENS == (8 if on and dispatch.DENSE_PUSH_TP4 else 0) + assert dispatch.MOE_PUSH_MAX_TOKENS == (dispatch.MOE_ENGINE_MAX_TOKENS if on else 0) + assert dispatch.GENERIC_MC_AR_MAX_TOKENS == 0 + assert dispatch.GENERIC_MC_AR_FOLDED_MAX_TOKENS == 0 + + +def test_the_dense_decode_split_is_the_built_in_one() -> None: + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {} + + +# ------------------------------------------------------------------ on: the dense layers' full-context decode + + +def test_the_dense_layers_take_the_fullctx_decode_at_the_measured_tp4_shapes() -> None: + """One or two requests, 1 + 3 verify rows or one draft-pass row, each at its own pinned recipe.""" + assert dispatch.FULLCTX_DECODE_ENABLED + assert dispatch.FULLCTX_DECODE.keys == {(4, 1), (1, 1), (4, 2), (1, 2)} + assert set(dispatch.FULLCTX_DECODE_RECIPES) == set(dispatch.FULLCTX_DECODE.keys) + for recipe in dispatch.FULLCTX_DECODE_RECIPES.values(): + assert 2 <= recipe["splits"] <= 128 # the standalone reducer's range + + +def test_the_fullctx_plans_keep_the_tp4_geometry() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + + assert (kernel.HEADS, kernel.KV_HEADS, kernel.HEAD_DIM) == (16, 1, 128) + assert kernel.HEADS_PER_KV == 16 + + +def test_the_fullctx_entry_cells_run_at_the_forward_plan_capacity() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention import ( # noqa: E501 + m3_fullctx_decode as entry, + ) + + assert entry.MAX_KV_LEN == dispatch.FULLCTX_DECODE_MAX_KV_LEN + + +def test_the_fullctx_capacity_covers_the_agentx_cache_bound() -> None: + """The forward compares the cache's length bound (max_seq_len plus the speculative rows it reserves) with the plan + capacity; at the AgentX max_seq_len of 1 << 20 with up to 8 draft rows the dense layers must still be taken.""" + assert dispatch.FULLCTX_DECODE_MAX_KV_LEN >= (1 << 20) + 8 + + +def test_the_drafter_takes_the_fullctx_decode_at_the_same_shapes() -> None: + assert dispatch.DRAFTER_FULLCTX_ENABLED + assert dispatch.DRAFTER_FULLCTX.keys == dispatch.FULLCTX_DECODE.keys + + +def test_the_dense_fp8_qkv_and_kv_insert_kill_constants_default_on() -> None: + assert dispatch.DENSE_QKV_FP8_VIA_BF16 is True + assert dispatch.DENSE_KV_INSERT_FUSED is True + assert dispatch.FULLCTX_DECODE_OWN_REDUCER is True + assert dispatch.FULLCTX_DECODE_MXFP8 is True + + +def test_the_fullctx_taken_lines_log_without_raising_under_require() -> None: + dispatch.set_require(True) + dispatch.taken(dispatch.FULLCTX_DECODE.path, 4) + dispatch.taken(dispatch.DENSE_KV_INSERT_PATH, 4) + dispatch.taken(dispatch.FULLCTX_DECODE_MXFP8_PATH, 4) + + +# ------------------------------------------------------------------ the paths on + + +@pytest.mark.parametrize("rows", range(1, dispatch.DECODE_MAX_TOKENS + 1)) +def test_the_moe_takes_the_engines_and_the_generation_router(rows) -> None: + """1..8 rows the MoE engine, 9..13 the MoE16 engine (routers in-kernel); the unfused MoE of a pure generation step + takes the TF32 tensor-core router at 9..16 rows (14..16 always, 9..13 on an MoE16 miss), the fp32 router else.""" + expected = ( + dispatch.MOE_ENGINE.path + if rows <= 8 + else dispatch.MOE16_ENGINE.path + if rows <= 13 + else None + ) + assert dispatch.moe_path(rows) == expected + assert dispatch.ROUTER_TC_GEN.covers(rows) == (9 <= rows <= 16) + + +def _router_core(cluster16: bool): + import torch + + return SimpleNamespace(router_tc_cluster16=cluster16, w={"l3_router": torch.zeros(128, 64)}) + + +@pytest.mark.parametrize("rows", [1, 8, 9, 13, 14, 16, 17, 48]) +@pytest.mark.parametrize("cluster16", [True, False]) +def test_the_generation_router_takes_router_tc_only_on_its_rows( + monkeypatch, rows, cluster16 +) -> None: + """The built-in gate's choice: router_tc on 9..16 bf16 rows when the device co-schedules 16-CTA clusters; the fp32 + router otherwise. Context steps (``_router_logits``) never take it.""" + import torch + + calls = [] + monkeypatch.setattr(dispatch, "m3_router_tc", lambda *a: calls.append("tc") or [a[0]]) + monkeypatch.setattr(dispatch, "m3_fp32_linear", lambda *a: calls.append("fp32") or a[0]) + h = torch.zeros(rows, 64, dtype=torch.bfloat16) + dispatch._decode_router_logits(_router_core(cluster16), 3, h) + assert calls == ["tc" if cluster16 and 9 <= rows <= 16 else "fp32"] + calls.clear() + dispatch._router_logits(_router_core(cluster16), 3, h) + assert calls == ["fp32"] + + +def test_the_core_selects_the_target_on_the_context_count_alone() -> None: + """Every pure generation step is decode's, whatever its width; a step with context rows (a mixed one too) is + prefill's.""" + core = SimpleNamespace(_prefill="prefill", _decode="decode") + select = dispatch.MiniMaxM3Tp4ModelingV2._select_target + assert select(core, SimpleNamespace(num_contexts=0)) == "decode" + assert select(core, SimpleNamespace(num_contexts=1)) == "prefill" + assert select(core, SimpleNamespace(num_contexts=7)) == "prefill" + + +@pytest.mark.parametrize("rows", range(1, dispatch.DECODE_MAX_TOKENS + 1)) +def test_the_projections_take_the_gemvs_else_the_mxfp8_gemm(rows) -> None: + for nk in TP4_PROJECTIONS: + prequant = dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) + assert prequant == ( + dispatch.MX_GEMV_PATH + if rows in dispatch.MX_GEMV_ROWS[nk] + else dispatch.PREQUANT_GEMM_PATH + ) + bf16 = dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) + assert bf16 == ( + dispatch.MX_GEMV_BF16_PATH + if rows <= dispatch.MX_GEMV_BF16_MAX_TOKENS[nk] + else dispatch.GENERIC_GEMM_PATH + ) + + +def test_the_attention_producer_and_scorer_are_on() -> None: + assert dispatch.PRODUCER.covers("nvfp4_tp4") + assert dispatch.INDEX_SCORE.keys == {"early", "prefetch"} + assert dispatch.SPARSE_DECODE_PRODUCER.keys == {"per_row", "multi_query"} + assert dispatch.EMBED_FIRSTNORM.keys == frozenset(range(1, 9)) + + +def test_speculation_runs_the_fused_tail_with_the_bf16_draft_head() -> None: + requests = dispatch.DECODE_MAX_TOKENS // (dispatch.SHARD_ACCEPT_DRAFT_LEN + 1) + assert dispatch.SHARD_ACCEPT_MAX_REQUESTS == requests + assert dispatch.SAMPLED_VERIFY_MAX_REQUESTS == requests + assert dispatch.FIRST_DRAFTER_INPUTS_MAX_REQUESTS == requests + assert dispatch.TARGET_HEAD_MAX_ROWS == 0 # the in-tree BF16 head is faster at the TP4 shard + assert dispatch.FC_AFTER_HEAD_MAX_ROWS == 8 + assert dispatch.DRAFT_STEP_MAX_REQUESTS == 8 + assert dispatch.FC_EARLY_MAX_ROWS == 0 + assert not dispatch.DRAFT_HEAD_CODED + + +def test_a_miss_raises_only_under_require() -> None: + dispatch.missed("some_path", 1, "a precondition") + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="some_path"): + dispatch.missed("some_path", 1, "a precondition") + + +def test_taken_logs_without_raising_under_require() -> None: + dispatch.set_require(True) + dispatch.taken(dispatch.INDEX_SCORE.path, 1) + dispatch.taken(dispatch.MOE_ENGINE.path, 8) + + +def test_taken_logs_once_per_path_and_rows(monkeypatch) -> None: + lines = [] + monkeypatch.setattr(dispatch.logger, "info", lines.append) + for _ in range(3): + dispatch.taken(dispatch.MX_GEMV_PATH, 4, "[2560, 6144] mxfp8 in") + dispatch.taken(dispatch.MX_GEMV_PATH, 8, "[2560, 6144] mxfp8 in") + assert len(lines) == 2 and all("m3_mx_gemv taken at" in line for line in lines) + + +def test_the_sparse_decode_writes_the_mxfp8_handoff() -> None: + """Every decode step (at most DECODE_MAX_TOKENS rows) can take the MXFP8 o_proj hand-off of the sparse decode.""" + assert dispatch.MERGE_MXFP8_HANDOFF is True + assert tsd.MERGE_MXFP8_MIN_TOKENS <= 1 and dispatch.DECODE_MAX_TOKENS <= tsd.MergeMxfp8.MAX_ROWS + + +def test_the_decode_mxfp8_decision_is_pinned_to_cute_dsl() -> None: + from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, + ) + + assert dispatch.MX_TAC_TP4 + assert dispatch.MXFP8_DECODE_QUANT_TACTIC == MXFP8QuantizeRunner.CUTE_DSL + assert dispatch.MXFP8_DECODE_GEMM_TACTIC == FlashInferMXFP8GemmRunner.CUTE_DSL diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py new file mode 100644 index 000000000000..3046e9bc61bd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The decode kernels the core builds before the executor's warmup (``prebuild_kernels``), and that a forward after it +never compiles one: the MoE engines of every ``MOE_ENGINE`` / ``MOE16_ENGINE`` token count, the decode GEMVs with +their push variants, the index scorer with and without its early-issue schedule at every decode query length (and +the T2SC scorer, the long-row select and the full-context plans where they apply), and the 9..16-row generation +router's device decision. + +The engine builds are stubbed (a real one needs a loaded checkpoint's expert buffers; the engine kernels have their +own tests); the GEMVs and the scorer compile for real on the GPU. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers the trtllm ops +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_index_score as index_score_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_index_score_t2sc as index_score_t2sc_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mx_gemv import ( # noqa: E501 + O_PROJ, + SHARED_DOWN, + SHARED_GATE_UP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.mx_gemv import ( # noqa: E501 + op as mx_gemv, +) + +dispatch = modeling + +assert torch.cuda.is_available(), "the prebuild checks require a CUDA device" + +INDEX_HEADS = 4 +INDEX_DIM = 128 +DRAFT_LEN = 3 +_CORE_METHODS = ( + "prebuild_kernels", + "_prebuild_engines", + "_try_build", + "engine", + "moe16_engine", + "_lookup", + "_scorer_variants", + "_prebuild_index_scorer", + "_prebuild_index_scorer_t2sc", + "decode_query_lens", +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require, frozen = dispatch.require_mode(), mx_gemv.frozen() + yield + dispatch.set_require(require) + mx_gemv.freeze(frozen) + dispatch.reset_logged() + + +class _FakeSparseAttn: + def __init__(self, indexer_kv_dtype: str): + self.indexer_kv_dtype = indexer_kv_dtype + self.m3_config = SimpleNamespace(num_index_heads=INDEX_HEADS, sparse_index_dim=INDEX_DIM) + + +def _core(num_layers: int = 3, indexer_kv_dtype: str = "fp8", built=None, fail=False): + """The parts of ``MiniMaxM3Tp4ModelingV2`` the prebuild reads: one rank (the build agreement takes no + collective), MoE + sparse layers 1.. (layer 0 dense), no projections (the GEMV prebuild has its own test), the + speculative config; the engine builds stubbed.""" + core = SimpleNamespace( + mapping=SimpleNamespace(tp_size=1, enable_attention_dp=False), + num_experts=128, + hidden=6144, + moe_layers=frozenset(range(1, num_layers)), + sparse_layers=frozenset(range(1, num_layers)), + attn=[_FakeSparseAttn(indexer_kv_dtype) for _ in range(num_layers)], + w={"embed": torch.empty(1, device="cuda")}, + mxfp8=SimpleNamespace(quant_method=SimpleNamespace(use_cutlass=True)), + model_config=SimpleNamespace( + spec_config=SimpleNamespace(max_draft_len=DRAFT_LEN, max_total_draft_tokens=DRAFT_LEN) + ), + moe_engines={}, + moe16_engines={}, + projections=lambda: [], + ) + for name in _CORE_METHODS: + setattr(core, name, getattr(modeling.MiniMaxM3Tp4ModelingV2, name).__get__(core)) + built = [] if built is None else built + + def build(i, num_tokens, shared=True): + if fail: + raise ValueError("the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI") + built.append(("moe", i, num_tokens, shared)) + return ("engine", i, num_tokens) + + def build16(i, num_tokens): + if fail: + raise ValueError("the routed experts are not TRTLLM-Gen NVFP4 SwiGLU-OAI") + built.append(("moe16", i, num_tokens)) + return ("engine16", i, num_tokens) + + core._build_engine, core._build_moe16_engine = build, build16 + return core + + +def _forbid_compiles(monkeypatch, core=None) -> None: + """Every compile / build entry point a forward could reach raises.""" + import cutlass.cute as cute + + def boom(*args, **kwargs): + raise AssertionError("a compile entry point ran after prebuild") + + monkeypatch.setattr(cute, "compile", boom) + monkeypatch.setattr(mx_gemv, "ensure_compiled", boom) + monkeypatch.setattr(mx_gemv, "precompile", boom) + monkeypatch.setattr(index_score_kernel.IndexScoreRunner, "_compile", boom) + monkeypatch.setattr(index_score_t2sc_kernel.IndexScoreRunner, "_compile", boom) + if core is not None: + core._build_engine = core._build_moe16_engine = boom + + +def test_decode_query_lens_follow_the_speculative_config() -> None: + core = _core() + assert core.decode_query_lens() == [1, 1 + DRAFT_LEN] + core.model_config.spec_config = None + assert core.decode_query_lens() == [1] + core.model_config.spec_config = SimpleNamespace( + max_draft_len=4, max_total_draft_tokens=4, draft_len_schedule={1: 4, 8: 2} + ) + assert core.decode_query_lens() == [1, 3, 5] + + +def test_prebuild_builds_every_engine_count_and_scorer_variant() -> None: + built = [] + core = _core(built=built) + times = core.prebuild_kernels() + + assert mx_gemv.frozen(), "prebuild must close the GEMV compiles" + + # Layer 0 is dense; every MoE layer gets every count, in layer order. + def per_layer(i): + engines = [("moe", i, n, True) for n in range(1, dispatch.MOE_ENGINE_MAX_TOKENS + 1)] + return engines + [ + ("moe16", i, n) for n in range(dispatch.MOE16_MIN_TOKENS, dispatch.MOE16_MAX_TOKENS + 1) + ] + + assert built == per_layer(1) + per_layer(2) + assert set(core.moe_engines[2]) == {(n, True, True) for n in dispatch.MOE_ENGINE.keys} + assert set(core.moe16_engines[2]) == set(dispatch.MOE16_ENGINE.keys) + + want = { + (torch.float8_e4m3fn, INDEX_HEADS, qlen, INDEX_DIM, modeling.INDEX_SCORE_SPLIT_K, early) + for qlen in (1, 1 + DRAFT_LEN) + for early in (True, False) + } + assert set(times["scorer_keys"]) == want + assert want <= set(index_score_kernel.IndexScoreRunner.kernel_cache) + for key in ("gemv_s", "moe_s", "scorer_s"): + assert times[key] >= 0.0 + + # The generation router's device decision, as the built-in gate makes it. + assert core.router_tc_cluster16 == bool(torch.ops.trtllm.minimax_m3_router_tc_cluster16(6144)) + + # A second build builds no engine again. + built.clear() + core.prebuild_kernels() + assert built == [] + + +def test_engine_lookups_after_prebuild_never_build(monkeypatch) -> None: + core = _core() + core._prebuild_engines(5) + _forbid_compiles(monkeypatch, core) + for require in (False, True): + dispatch.set_require(require) + for n in sorted(dispatch.MOE_ENGINE.keys): + assert core.engine(5, n, shared=True) == ("engine", 5, n) + for n in sorted(dispatch.MOE16_ENGINE.keys): + assert core.moe16_engine(5, n) == ("engine16", 5, n) + + +def test_missing_engine_is_a_miss_not_a_build(monkeypatch) -> None: + core = _core() + core.moe_engines[7], core.moe16_engines[7] = {}, {} + _forbid_compiles(monkeypatch, core) + dispatch.set_require(False) + assert core.engine(7, 3, shared=True) is None + assert core.moe16_engine(7, 10) is None + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + core.engine(7, 3, shared=True) + with pytest.raises(dispatch.FastPathUnavailable, match="not built before warmup"): + core.moe16_engine(7, 10) + + +def test_failed_engine_build_is_logged_then_missed() -> None: + core = _core(fail=True) + core._prebuild_engines(9) + assert all(v is None for v in core.moe_engines[9].values()) and len(core.moe_engines[9]) == 8 + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="build failed"): + core.engine(9, 1, shared=True) + + +def _projection(nk, push): + n, k = nk + w = (torch.randn(n, k, device="cuda") / k**0.5).to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return SimpleNamespace(weight=wq.view(n, k), weight_scale=wsf, push=push) + + +@pytest.mark.parametrize( + "nk,push", + ((O_PROJ, "attn"), (SHARED_DOWN, "moe"), (SHARED_GATE_UP, None)), + ids=("o_proj", "shared_down", "shared_gate_up"), +) +def test_gemvs_after_prebuild_never_compile(monkeypatch, nk, push) -> None: + use_pdl = True + module = _projection(nk, push) + variants = modeling._gemv_variants(*nk, push, 4) + pushes = [v for v in variants if v[2] is not None] + if push is None: + assert not pushes + else: + covered = dispatch.ATTN_PUSH.keys if push == "attn" else dispatch.MOE_PUSH.keys + assert {rows for rows, _, _ in pushes} == set(covered) + assert {p[1] for _, _, p in pushes} == {b for b in (8, 16) if b <= max(covered)} + mx_gemv.freeze(False) + # As prebuild_kernels compiles one projection. + mx_gemv.ensure_compiled(module.weight, module.weight_scale, use_pdl) + for rows, x_bf16, ar in variants: + mx_gemv.precompile(module.weight, module.weight_scale, rows, x_bf16, use_pdl=use_pdl, ar=ar) + mx_gemv.freeze(True) + _forbid_compiles(monkeypatch) + + n, k = nk + for rows, x_bf16, ar in variants: + if ar is not None: + # A push needs a multicast exchange to launch: its kernel is compiled (the forward finds it). + resolved = mx_gemv._resolve(rows, n, k, x_bf16, True, None, use_pdl, (0, 0, 0, *ar)) + if resolved is not None: + assert mx_gemv.is_compiled(resolved[0]), (rows, x_bf16, ar) + continue + if x_bf16: + x = torch.randn(rows, k, device="cuda").to(torch.bfloat16) + out = mx_gemv.m3_mx_gemv_bf16(x, module.weight, module.weight_scale, use_pdl=use_pdl) + else: + xq, xsf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(rows, k, device="cuda").to(torch.bfloat16), True + ) + out = mx_gemv.m3_mx_gemv( + xq.view(rows, k), xsf, module.weight, module.weight_scale, use_pdl=use_pdl + ) + served = mx_gemv.serves(n, k, rows, x_bf16) + assert (out is not None) == served, (rows, x_bf16) + assert mx_gemv.last_uncompiled is None + torch.cuda.synchronize() + + +def test_frozen_gemv_declines_an_uncompiled_kernel(monkeypatch) -> None: + module = _projection(SHARED_DOWN, None) + mx_gemv.freeze(True) + _forbid_compiles(monkeypatch) + # A launch option nothing prebuilt (no early trigger, without a push): declined, never compiled. + xq, xsf = torch.ops.trtllm.mxfp8_quantize( + torch.randn(3, 768, device="cuda").to(torch.bfloat16), True + ) + key = mx_gemv._resolve(3, 6144, 768, False, False, None, True, None)[0] + if mx_gemv.is_compiled(key): + pytest.skip("this process already compiled the probe kernel") + out = mx_gemv.m3_mx_gemv( + xq.view(3, 768), xsf, module.weight, module.weight_scale, trigger_early=False, use_pdl=True + ) + assert out is None and mx_gemv.last_uncompiled == key + + +@pytest.mark.parametrize("indexer_kv_dtype", ("fp8", "bf16")) +def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> None: + core = _core(indexer_kv_dtype=indexer_kv_dtype) + keys = core._prebuild_index_scorer(core.decode_query_lens(), torch.device("cuda")) + assert len(keys) == 4 + _forbid_compiles(monkeypatch, core) + dtype = torch.float8_e4m3fn if indexer_kv_dtype == "fp8" else torch.bfloat16 + for _, heads, qlen, dim, split_k, early in keys: + rows = 1 if early else 2 + blocks = 3 + idx_q = torch.randn(rows * qlen, heads, dim, device="cuda").to(dtype) + idx_k = torch.randn(rows * blocks, 1, 128, dim, device="cuda").to(dtype) + block_table = torch.arange(rows * blocks, dtype=torch.int32, device="cuda").view( + rows, blocks + ) + seq_lens = torch.full((rows,), blocks * 128 - 5, dtype=torch.int32, device="cuda") + score = torch.empty(heads, blocks, rows * qlen, dtype=torch.float32, device="cuda") + m3_index_score( + idx_q, + idx_k.squeeze(1), + block_table, + seq_lens, + score.transpose(1, 2), + qlen, + split_k, + early, + ) + torch.cuda.synchronize() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..0b5d60b19523 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode robustness without a GPU kernel: the multicast all-reduce epilogue options keep the cluster-barrier bit, +an engine build that fails on one TP rank drops that engine on every rank (two ranks simulated with threads), and +the Eagle3 fc fork's output is recorded on the stream that reads it.""" + +import threading +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = modeling.require_mode() + yield + modeling.set_require(require) + modeling.reset_logged() + + +@pytest.mark.parametrize("finalize", [False, True]) +@pytest.mark.parametrize("bucket", [8, 16, 32, 48]) +def test_ar_opts_keep_cluster_barrier_bit(monkeypatch, bucket, finalize): + # Tables without bit 2 and with buckets missing: every lookup, fallback included, keeps the barrier. + monkeypatch.setattr(mc, "AR_OPTS_BY_M", (0, 1, 4, 5, 8, 9, 12, 13)) + monkeypatch.setattr(mc, "AR_OPTS_BY_BUCKET", {}) + monkeypatch.setattr(mc, "AR_FIN_OPTS_BY_BUCKET", {}) + for m in range(1, min(bucket, 8) + 1): + opts = mc._ar_opts(m, bucket, finalize=finalize) + assert opts & mc._OPT_CLUSTER_BARRIER and 0 <= opts <= 15, (m, bucket, finalize, opts) + + +def test_ar_opts_shipped_tables_unchanged(): + for m in range(1, 9): + assert mc._ar_opts(m, 8) == mc.AR_OPTS_BY_M[m - 1] + for b, o in mc.AR_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b) == o + for b, o in mc.AR_FIN_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b, finalize=True) == o + + +class _FakeComm: + """An MPI-like TP communicator for ``world`` threads: allreduce sums the ranks' ints.""" + + def __init__(self, world): + self.world = world + self.barrier = threading.Barrier(world) + self.lock = threading.Lock() + self.vals = [] + self.calls = 0 + + def allreduce(self, v): + with self.lock: + self.vals.append(v) + self.calls += 1 + self.barrier.wait() + total = sum(self.vals) + self.barrier.wait() + with self.lock: + self.vals.clear() + self.barrier.wait() + return total + + def Get_size(self): # noqa: N802 (MPI name) + return self.world + + +def _mapping(tp_size=2, adp=False): + return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) + + +def _run_ranks(fns): + out, errs = [None] * len(fns), [] + + def run(i): + try: + out[i] = fns[i]() + except BaseException as e: # noqa: BLE001 + errs.append(e) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errs, errs + return out + + +@pytest.fixture +def two_rank_comm(monkeypatch): + from tensorrt_llm._torch.distributed import ops + + comm = _FakeComm(2) + monkeypatch.setattr(mc, "_tp_comm", lambda mapping: comm) + monkeypatch.setattr(ops, "mpi_disabled", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return comm + + +def test_tp_agree_single_rank_and_adp_take_no_collective(monkeypatch): + monkeypatch.setattr(mc, "_tp_comm", mock.Mock(side_effect=AssertionError("no collective"))) + assert mc.tp_agree(_mapping(1), True) is True + assert mc.tp_agree(_mapping(1), False) is False + assert mc.tp_agree(_mapping(4, adp=True), False) is False + + +def test_tp_agree_refuses_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError): + mc.tp_agree(_mapping(2), True) + + +@pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) +def test_tp_agree_built_two_ranks(two_rank_comm, oks): + built = [object() if ok else None for ok in oks] + out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + assert out == (built if all(oks) else [None, None]) + + +def _core(fail: bool) -> SimpleNamespace: + """The engine-building half of the core (``MiniMaxM3Tp4ModelingV2``) on one of two simulated TP ranks, its + engine builds stubbed (and failing on one rank).""" + core = SimpleNamespace(mapping=_mapping(2), moe_engines={}, moe16_engines={}) + cls = modeling.MiniMaxM3Tp4ModelingV2 + + def build(i, num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + core._build_engine = build + core._build_moe16_engine = lambda i, num_tokens: build(i, num_tokens) + for name in ("_prebuild_engines", "_try_build", "engine", "moe16_engine", "_lookup"): + setattr(core, name, getattr(cls, name).__get__(core)) + return core + + +@pytest.mark.parametrize("fail", [(False, False), (False, True)]) +def test_engine_build_agreed_across_ranks(two_rank_comm, fail): + cores = [_core(f) for f in fail] + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + keys = sorted(cores[0].moe_engines[3]) + sorted(cores[0].moe16_engines[3]) + assert keys and two_rank_comm.calls == 2 * len(keys) + for engines in ("moe_engines", "moe16_engines"): + a, b = getattr(cores[0], engines)[3], getattr(cores[1], engines)[3] + assert a.keys() == b.keys() + for key in a: + assert (a[key] is None) == (b[key] is None) == any(fail), (engines, key) + # Both ranks then take the same path in the forward: the engine, or (after a miss) the generic MoE. + modeling.set_require(False) + for n in sorted(modeling.MOE_ENGINE.keys): + got = [c.engine(3, n, shared=True) for c in cores] + assert (got[0] is None) == (got[1] is None) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + M3Eagle3Worker = modeling.M3Eagle3Worker + main, fork = mock.Mock(name="main"), mock.Mock(name="fork") + hidden = mock.Mock(name="hidden", is_cuda=True) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *a: main) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + worker = SimpleNamespace(_m3_fc_hidden=hidden) + M3Eagle3Worker._m3_fc_join(worker, fork) + main.wait_stream.assert_called_once_with(fork) + if capturing: + hidden.record_stream.assert_not_called() + else: + hidden.record_stream.assert_called_once_with(main) + worker._m3_fc_hidden = None + M3Eagle3Worker._m3_fc_join(worker, fork) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py new file mode 100644 index 000000000000..1153eaee47bd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_silu_and_mul entry of the MiniMax-M3 tp8 target: SwiGLU-OAI and plain SwiGLU.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.activation.m3_silu_and_mul import ( # noqa: E501 + m3_silu_and_mul as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" + +OAI = (7.0, 1.702, 1.0) + + +def _build(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + # Gate and up spread past the clamp (|x| up to ~12) so both clamps engage. + x = (torch.randn(spec["m"], 2 * spec["inter"], generator=g, device="cuda") * 4.0).to( + torch.bfloat16 + ) + return (x, *OAI) if spec["oai"] else (x, None, None, None) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few + fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); + measured on GB300: worst element 0.5 unit, relative RMS 0.01.""" + args = _build(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(*args) + assert out.shape == (cell.spec["m"], cell.spec["inter"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=1.0, rms_ulp=0.1) + + +def test_the_gate_sees_a_missing_clamp() -> None: + """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" + spec = op.CELLS[0].spec + x, *oai = _build(dict(spec, m=64), 3) + out = op(x, *oai) + with pytest.raises(AssertionError): + assert_within_ulp(out, op.reference(x, None, oai[1], oai[2]), element_ulp=1.0, rms_ulp=0.1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_attn_runtime_tp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_attn_runtime_tp8.py new file mode 100644 index 000000000000..e9aaad5c65fa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_attn_runtime_tp8.py @@ -0,0 +1,274 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real MiniMax-M3 attention runtime state at the TP8 per-rank geometry, for the backend-step entries' GPU tests. + +One hybrid NVFP4 KV cache manager (``MiniMaxM3KVCacheManagerV2``) laid out as the target's: a dense layer in the FP8 +half, a sparse MSA layer in the NVFP4 half with its FP8 index-K cache, and (``draft=True``) the one-model Eagle3 +drafter's shared layer, read through the manager's FP8 draft view as the runtime's draft swap reads it. Metadata is +the runtime's own type, filled by its own ``prepare()``; the backends are built the way the target builds them. + +The geometry is one rank's: 8 query heads over one KV head (M3's 4 KV heads, each duplicated on a rank pair) and +one index head of dim 128. The mapping is a single rank holding exactly those heads, so nothing is sharded again. + +History rows are written by the in-tree scatters (``fused_write_layer_caches`` for the FP8 layers, +``fused_write_layer_caches_nvfp4`` for the NVFP4 layer and its index-K), at the slots the manager assigned, which are +not in request order. +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import ( + MiniMaxM3KVCacheManagerV2, + MiniMaxM3MsaSparseAttention, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches, + fused_write_layer_caches_nvfp4, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_backend import ( + MiniMaxM3MsaSparseAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.utils import create_attention +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.pyexecutor.runtime_hooks import MsaRuntimePolicy +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + KvCacheConfig, + MiniMaxM3SparseAttentionConfig, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.quantization.mode import QuantAlgo + +HEADS = 8 +KV_HEADS = 1 +INDEX_HEADS = 1 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +DENSE, SPARSE, DRAFT = 0, 1, 2 +MAX_BATCH = 8 +MAX_SEQ_LEN = 8192 +MAX_NUM_TOKENS = 8192 +#: The Eagle3 drafter checkpoint's RoPE: NeoX over the whole head, theta 5e6. +DRAFT_ROPE_THETA = 5.0e6 + +#: E4M3-exact planted index levels, 1/8 apart in relative terms (0.25 .. 15): 48 distinct block scores. +LEVELS = tuple(2.0**e * (1 + m / 8) for e in range(-2, 4) for m in range(8)) +#: The index value of a block's non-peak rows; its score (8) is below every level's (>= 16). +LOW = 0.125 + + +def mapping() -> Mapping: + return Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + +def sparse_attention_config() -> MiniMaxM3SparseAttentionConfig: + """The target's runtime: MSA, the FP8 indexer, the fused five-way projection (M3's checkpoint selection rule: + no forced initial blocks, the query's own block forced).""" + return MiniMaxM3SparseAttentionConfig( + implementation="msa", + indexer_kv_dtype="fp8", + fuse_qkv_index_projection=True, + sparse_num_index_heads=INDEX_HEADS, + num_attention_heads=HEADS, + num_key_value_heads=KV_HEADS, + ) + + +def make_manager(*, draft: bool) -> MiniMaxM3KVCacheManagerV2: + torch.cuda.init() + return MiniMaxM3KVCacheManagerV2( + KvCacheConfig( + max_tokens=4 * MAX_SEQ_LEN, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + host_cache_size=0, + event_buffer_max_size=0, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=[KV_HEADS] * (3 if draft else 2), + head_dim=HEAD_DIM, + tokens_per_block=PAGE, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + mapping=mapping(), + dtype=DataType.NVFP4, + vocab_size=16, + spec_config=( + Eagle3DecodingConfig(max_draft_len=3, speculative_model="unused") if draft else None + ), + sparse_layer_ids=[SPARSE], + sparse_attention_config=sparse_attention_config(), + ) + + +def make_msa_backend(layer: int) -> MiniMaxM3MsaSparseAttention: + """One layer's MSA backend as ``MiniMaxM3Tp8ModelingV2`` builds it (NVFP4 KV quant config on every layer).""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.NVFP4) + backend = create_attention( + "TRTLLM", + layer, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=None, + quant_config=quant, + q_scaling=1.0, + sparse_params=sparse_attention_config().to_sparse_params(layer_idx=layer), + ) + backend.update_quant_config(quant) + assert isinstance(backend, MiniMaxM3MsaSparseAttention), type(backend).__name__ + return backend + + +def make_draft_backend() -> TrtllmAttention: + """The drafter layer's attention as ``M3Eagle3Draft`` builds it: fused NeoX RoPE, an FP8 KV cache (the target + overrides the drafter's KV dtype to FP8), no other quantization.""" + quant = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + backend = create_attention( + "TRTLLM", + DRAFT, + HEADS, + HEAD_DIM, + KV_HEADS, + pos_embd_params=PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=HEAD_DIM, theta=DRAFT_ROPE_THETA, max_positions=MAX_SEQ_LEN), + ), + quant_config=quant, + q_scaling=1.0, + attention_chunk_size=None, + sparse_params=None, + ) + backend.update_quant_config(quant) + assert type(backend) is TrtllmAttention, type(backend).__name__ + return backend + + +def make_msa_metadata(manager) -> MiniMaxM3MsaSparseAttentionMetadata: + metadata = MiniMaxM3MsaSparseAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + max_num_sequences=MAX_BATCH, + kv_cache_manager=manager, + mapping=mapping(), + sparse_metadata_params=sparse_attention_config().to_sparse_metadata_params(), + ) + # The target's MSA preparation choice that changes what prepare() stages (the identity tables of short steps). + metadata.msa_policy = MsaRuntimePolicy(short_sequence_identity=True) + return metadata + + +def make_draft_metadata(manager) -> TrtllmAttentionMetadata: + """Drafter metadata over the shared manager's FP8 draft view, the manager the runtime's draft swap installs.""" + view = manager.get_draft_kv_cache_view() + assert view is not None and view.dtype == DataType.FP8 + return TrtllmAttentionMetadata( + max_num_requests=MAX_BATCH, + max_num_tokens=MAX_NUM_TOKENS, + kv_cache_manager=view, + mapping=mapping(), + ) + + +def allocate(manager, request_ids: Sequence[int], qlens: Sequence[int], cached: Sequence[int]): + manager.add_dummy_requests(list(request_ids), token_nums=[c + n for c, n in zip(cached, qlens)]) + + +def prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: + metadata.seq_lens = torch.tensor(list(qlens), dtype=torch.int) + metadata.num_contexts = num_contexts + metadata.request_ids = list(request_ids) + metadata.prompt_lens = [ + c + n if row < num_contexts else c for row, (c, n) in enumerate(zip(cached, qlens)) + ] + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=list(cached)) + metadata.prepare() + + +def flat_slots(manager, layer: int, request_id: int, positions: torch.Tensor) -> torch.Tensor: + """The cache slots (page slot id * PAGE + offset) of a request's positions in ``layer``.""" + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0], + dtype=torch.int64, + device="cuda", + ) + positions = positions.to(device="cuda", dtype=torch.int64) + return (pages[positions // PAGE] * PAGE + positions % PAGE).to(torch.int32) + + +def write_fp8_rows(manager, layer, request_id, positions, k, v) -> None: + """K / V rows [n, 128] (E4M3) into an FP8 layer's cache at ``positions``.""" + view = manager.get_buffers(layer, "HND") + assert view.dtype == torch.float8_e4m3fn + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches(view[:, 0], view[:, 1], None, slots, k, v, None) + + +def write_nvfp4_rows(manager, metadata, layer, request_id, positions, k, v, idx_k, quant) -> None: + """BF16 K / V rows quantized into the NVFP4 layer's cache, and index-K rows into its FP8 index cache.""" + data = manager.get_buffers(layer, "HND") + scales = manager.get_block_scale_buffers(layer, "HND") + slots = flat_slots(manager, layer, request_id, positions) + assert fused_write_layer_caches_nvfp4( + data[:, 0], + data[:, 1], + scales[:, 0], + scales[:, 1], + metadata.msa_idx_k_cache(layer), + slots, + k, + v, + idx_k, + quant, + ) + + +def index_directions(gen: torch.Generator) -> tuple[torch.Tensor, torch.Tensor]: + """Two orthogonal ±1 index directions (the halves of the 128 dims), E4M3-exact.""" + signs = torch.randint(0, 2, (HEAD_DIM,), generator=gen, device="cuda").float() * 2 - 1 + u0 = torch.zeros(HEAD_DIM, device="cuda") + u1 = torch.zeros(HEAD_DIM, device="cuda") + u0[: HEAD_DIM // 2] = signs[: HEAD_DIM // 2] + u1[HEAD_DIM // 2 :] = signs[HEAD_DIM // 2 :] + return u0, u1 + + +def planted_index_keys(length: int, u0, u1, gen: torch.Generator) -> torch.Tensor: + """Index-K rows [length, 128] (bf16, E4M3-exact) whose per-block maximum score is unambiguous. + + Every block has one peak row at a random position, ``a_b u0 + c_b u1`` with ``a`` and ``c`` two independent + permutations of ``LEVELS``; its other rows are ``LOW (u0 + u1)``. A query along ``u0`` then ranks the blocks by + ``a``, one along ``u1`` by ``c``, every score 1/8 or more apart from the next.""" + blocks = (length + PAGE - 1) // PAGE + assert blocks <= len(LEVELS) + levels = torch.tensor(LEVELS, device="cuda") + a = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + c = levels[torch.randperm(len(LEVELS), generator=gen, device="cuda")[:blocks]] + rows = (LOW * (u0 + u1)).expand(length, HEAD_DIM).clone() + for b in range(blocks): + span = min(PAGE, length - b * PAGE) + peak = b * PAGE + int(torch.randint(0, span, (1,), generator=gen, device="cuda")) + rows[peak] = a[b] * u0 + c[b] * u1 + return rows.to(torch.bfloat16) + + +def nvfp4_quant(*values: torch.Tensor) -> torch.Tensor: + """The (Q, K, V) quantization scales ``fused_write_layer_caches_nvfp4`` takes: each role's amax at the top of the + E2M1 x E4M3 range.""" + quant = torch.ones(3, dtype=torch.float32, device="cuda") + for role, x in enumerate(values, start=1): + quant[role] = (448.0 * 6.0) / x.abs().max().float() + return quant diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_nvfp4_cache.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_nvfp4_cache.py new file mode 100644 index 000000000000..bfa03c09da03 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/_m3_nvfp4_cache.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A MiniMax-M3 TP8 sparse-layer decode step over a packed NVFP4 cache, for the attention entries' tests. + +The cache is written by the in-tree scatter (``fused_write_layer_caches_nvfp4``, +which quantizes with ``torch.ops.trtllm.fp4_quantize``), into one +``[pages, 2 (K, V), 1, 128, .]`` pool per role exactly as ``msa_paged_kv`` and +``get_block_scale_buffers`` hand the K and V views over. Pages are shuffled so +a kernel that ignored the block table would read the wrong ones. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Sequence + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_scatter import ( + fused_write_layer_caches_nvfp4, +) + +#: Query heads per rank at TP8 (on one KV head). +NUM_HEADS = 8 +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 + + +def token_kv_len(seq_len: int, decode_query_len: int, offset: int) -> int: + """The KV extent query token ``offset`` of a request attends, as every decode kernel derives it.""" + return max(seq_len - decode_query_len + offset + 1, 0) + + +def build_step( + seq_lens: Sequence[int], decode_query_len: int, *, identity: bool, seed: int +) -> SimpleNamespace: + """One decode step: E4M3 q [T, 8, 128], the packed NVFP4 K/V views and their scales, a [1, T, 16] + head-major top-k table (the identity table, or a random ascending selection with a -1 tail), the block + table and per-request lengths.""" + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + total_q = batch * decode_query_len + max_blocks = max(1, max((s + PAGE - 1) // PAGE for s in seq_lens)) + num_pages = batch * max_blocks + num_slots = num_pages * PAGE + block_table = ( + torch.randperm(num_pages, device="cuda", generator=g) + .to(torch.int32) + .reshape(batch, max_blocks) + ) + data_pool = torch.zeros(num_pages, 2, 1, PAGE, HEAD_DIM // 2, dtype=torch.uint8, device="cuda") + scale_pool = torch.zeros( + num_pages, 2, 1, PAGE, HEAD_DIM // 16, dtype=torch.uint8, device="cuda" + ) + quant = torch.ones(3, dtype=torch.float32, device="cuda") + values = [] + for role in (0, 1): + x = torch.randn(num_slots, HEAD_DIM, device="cuda", generator=g, dtype=torch.float32).to( + torch.bfloat16 + ) + quant[role + 1] = (448.0 * 6.0) / x.abs().max().float().item() + values.append(x) + assert fused_write_layer_caches_nvfp4( + data_pool[:, 0], + data_pool[:, 1], + scale_pool[:, 0], + scale_pool[:, 1], + None, + torch.arange(num_slots, device="cuda", dtype=torch.int32), + values[0], + values[1], + None, + quant, + ) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, device="cuda", generator=g).to( + torch.float8_e4m3fn + ) + topk = torch.full((1, total_q, TOPK), -1, dtype=torch.int32, device="cuda") + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + nblocks = (token_kv_len(seq_lens[req], decode_query_len, offset) + PAGE - 1) // PAGE + if identity: + assert nblocks <= TOPK, "an identity step has every request within 16 blocks" + chosen = torch.arange(nblocks, device="cuda") + else: + chosen = torch.randperm(nblocks, device="cuda", generator=g)[:TOPK].sort().values + topk[0, tok, : chosen.numel()] = chosen.to(torch.int32) + return SimpleNamespace( + q=q, + k_paged=data_pool[:, 0], + v_paged=data_pool[:, 1], + k_block_scale=scale_pool[:, 0], + v_block_scale=scale_pool[:, 1], + k_global_scale=(1.0 / quant[1:2]).contiguous(), + v_global_scale=(1.0 / quant[2:3]).contiguous(), + topk=topk, + block_table=block_table, + seq_lens=torch.tensor(list(seq_lens), dtype=torch.int32, device="cuda"), + decode_query_len=decode_query_len, + sm_scale=HEAD_DIM**-0.5, + source=values, + quant=quant, + num_slots=num_slots, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..bef0db5ce41b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7): the TP8 dense layers' FP8 q/k/v (minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8, head-count +generic, at the TP8 per-rank 8 query heads over 1 KV head) carries the bits of the stock path, the BF16 +fused_qk_norm_rope output cast to E4M3, for every norm / RoPE variant: Q and K are rounded to BF16 before E4M3 (K is +stored in the FP8 cache), V is cast.""" + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers the trtllm ops) + +assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) + +HQ, HK, HV, D, ROT = 8, 1, 1, 128, 64 + + +@pytest.mark.parametrize("tokens", [1, 4, 8, 4096]) +@pytest.mark.parametrize("use_gemma", [False, True], ids=["rmsnorm", "gemma"]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_dense_fp8_qkv_is_the_bf16_kernel_cast(tokens, use_gemma, is_neox) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8") + g = torch.Generator(device="cuda").manual_seed(tokens * 4 + int(use_gemma) * 2 + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + ours = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8( + qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, use_gemma + ) + stock = qkv.clone() + torch.ops.trtllm.fused_qk_norm_rope( + stock, HQ, HK, HV, D, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, + 1.0, 0.0, 0.0, 1.0, True, use_gemma, False, 0, 0, + ) # fmt: skip + assert torch.equal(ours.view(torch.uint8), stock.to(torch.float8_e4m3fn).view(torch.uint8)) + + +@pytest.mark.parametrize("tokens,live", [(1, 1), (4, 4), (8, 5), (64, 64)]) +@pytest.mark.parametrize("tokens_per_block", [32, 128]) +@pytest.mark.parametrize("is_neox", [True, False], ids=["neox", "interleaved"]) +def test_the_fused_kv_insert_stores_the_bytes_of_write_kv_slots_and_the_compact_q( + tokens, live, tokens_per_block, is_neox +) -> None: + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + g = torch.Generator(device="cuda").manual_seed(tokens * 1000 + tokens_per_block + int(is_neox)) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + pages, layers, layer = 16, 3, 1 + # The layer's HND view of a multi-layer pool, strided as get_buffers returns it. + raw = torch.randint( + 0, + 256, + (pages, layers, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.randperm(pages * tokens_per_block, device="cuda", generator=g)[:tokens].to( + torch.int32 + ) + slots[live:] = -1 # padded rows own no slot + + def view(t): + return t.view(torch.float8_e4m3fn)[:, layer].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, is_neox, pos, False) + q_compact = torch.empty(tokens, HQ, D, device="cuda", dtype=torch.float8_e4m3fn) + ours = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, live, q_compact + ) + plain = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + assert torch.equal(ours.view(torch.uint8), plain.view(torch.uint8)) + assert torch.equal( + q_compact.view(tokens, HQ * D).view(torch.uint8), plain[:, : HQ * D].view(torch.uint8) + ) + + ref_raw = raw.clone() + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D) + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D) + ref = view(ref_raw) + write_kv_slots(ref[:, 0], slots, k, live, layout="HND") + write_kv_slots(ref[:, 1], slots, v, live, layout="HND") + assert torch.equal(ours_raw, ref_raw) + + +def test_the_fused_kv_insert_skips_a_negative_slot() -> None: + """A live row whose slot is the -1 sentinel stores nothing (no write before the pool row).""" + if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv"): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv") + from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + write_kv_slots, + ) + + tokens, tokens_per_block, pages = 4, 32, 8 + g = torch.Generator(device="cuda").manual_seed(7) + qkv = (torch.randn(tokens, (HQ + HK + HV) * D, device="cuda", generator=g) * 2).to( + torch.bfloat16 + ) + pos = torch.randint(0, 1 << 20, (tokens,), device="cuda", generator=g, dtype=torch.int32) + qw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + kw = (1 + 0.1 * torch.randn(D, device="cuda", generator=g)).to(torch.bfloat16) + raw = torch.randint( + 0, + 256, + (pages, 3, 2, HK * tokens_per_block * D), + device="cuda", + generator=g, + dtype=torch.uint8, + ) + slots = torch.tensor([5, -1, 40, 77], device="cuda", dtype=torch.int32) + + def view(t): + return t.view(torch.float8_e4m3fn)[:, 1].view(pages, 2, HK, tokens_per_block, D) + + ours_raw = raw.clone() + args = (qkv, HQ, HK, HV, ROT, 1e-6, qw, kw, 5e6, True, pos, False) + torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8_kv( + *args, view(ours_raw), slots, tokens + ) + plain = torch.ops.trtllm.minimax_m3_sm107_tp2_dense_qk_norm_rope_fp8(*args) + keep = torch.tensor([0, 2, 3], device="cuda") + ref_raw = raw.clone() + ref = view(ref_raw) + k = plain[:, HQ * D : (HQ + HK) * D].reshape(tokens, HK, D)[keep] + v = plain[:, (HQ + HK) * D :].reshape(tokens, HV, D)[keep] + write_kv_slots(ref[:, 0], slots[keep], k, 3, layout="HND") + write_kv_slots(ref[:, 1], slots[keep], v, 3, layout="HND") + assert torch.equal(ours_raw, ref_raw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py new file mode 100644 index 000000000000..c89183995648 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_attention entry: the Eagle3 drafter layer's attention against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp8``): the target's hybrid NVFP4 +``MiniMaxM3KVCacheManagerV2`` with its one-model Eagle3 draft layer, the requests allocated in it, drafter +``TrtllmAttentionMetadata`` over the manager's FP8 draft view (the manager the runtime's draft swap installs) filled +by its ``prepare()``, and the drafter's ``TrtllmAttention`` built as ``M3Eagle3Draft`` builds it (fused NeoX RoPE at +the drafter checkpoint's theta, FP8 KV). Earlier rows go into the draft layer's cache through the in-tree scatter at +the slots the manager assigned; the call rotates and appends the step's own. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_draft_attention import ( # noqa: E501 + m3_draft_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp8 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" + +# Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. +# The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: +# it takes the rotated Q in E4M3 (rounding the reference's Q to E4M3 takes the RMS gap from ~9 to ~7 ulps) and rounds +# the softmax probabilities before the PV product. Measured on GB300 over these cells (max element / RMS): 14.8 / 8.7 +# (4 x 4 rows), 13.3 / 7.4 (8 x 1), 19.1 / 7.7 (512 context rows). The band leaves about 1.5x; a reference shifted +# by one row is far outside it. +_ELEMENT_ULP = 30.0 +_RMS_ULP = 12.0 + +# Cached tokens per request of each cell; the new rows per request come from the cell. +_CACHED = { + (2, 4): (200, 1000), + (2, 1): (127, 2047), + (1, None): (0,), +} +#: Rows past the live ones in the fused QKV, as a padded step hands them over. +_PAD = 3 + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + +# Non-paged context FMHA on SM107 does not apply RoPE (this fixture's path; the served drafter runs the paged context +# FMHA, which rotates): the context cell is a strict expected failure until the runtime rotates there too. +_UNROTATED_CONTEXT = pytest.mark.xfail( + strict=True, + reason="non-paged context FMHA on SM107 does not apply RoPE (shared runtime)", +) +_CELLS = [ + pytest.param(c, marks=_UNROTATED_CONTEXT) if c.spec.get("context") else c for c in op.CELLS +] + + +@pytest.mark.parametrize("cell", _CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + context = spec.get("context") + cached = _CACHED[spec["reqs"], spec.get("qlen")] + qlens = (context,) if context else (spec["qlen"],) * spec["reqs"] + assert len(cached) == spec["reqs"] + contexts = 1 if context else 0 + rids = list(range(spec["reqs"])) + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=True) + try: + backend = rt.make_draft_backend() + metadata = rt.make_draft_metadata(manager) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") + hist = hist.to(torch.float8_e4m3fn) + rt.write_fp8_rows(manager, rt.DRAFT, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + assert int(metadata.num_tokens) == total + width = (rt.HEADS + 2 * rt.KV_HEADS) * rt.HEAD_DIM + qkv = torch.randn(total + _PAD, width, generator=gen, device="cuda").to(torch.bfloat16) + + certifying(op) + out = op(backend, qkv, metadata) + ref = op.reference(backend, qkv, metadata) + assert out.shape == (total, rt.HEADS * rt.HEAD_DIM) and out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + element, rms = _ulp_stats(out, ref) + print(f"\n[{cell.why}] max element {element:.2f} ulp, rms {rms:.3f} ulp") + group = rt.HEADS // rt.KV_HEADS + for h in range(rt.KV_HEADS): + cols = slice(h * group * rt.HEAD_DIM, (h + 1) * group * rt.HEAD_DIM) + element, rms = _ulp_stats(out[:, cols], ref[:, cols]) + print(f"[{cell.why}] KV head {h}: max element {element:.2f} ulp, rms {rms:.3f} ulp") + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + # The call's append: each live row's rotated K and its V at the row's position, E4M3 (unit KV scale), to + # within one E4M3 step of the float64 rotation (the kernel rotates in FP32). + pool = manager.get_buffers(rt.DRAFT, "HND") + q_size = rt.HEADS * rt.HEAD_DIM + first = 0 + for rid, c, n in zip(rids, cached, qlens): + positions = torch.arange(c, c + n, device="cuda") + slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() + rows = qkv[first : first + n] + for h in range(rt.KV_HEADS): + k_cols = q_size + h * rt.HEAD_DIM + v_cols = q_size + (rt.KV_HEADS + h) * rt.HEAD_DIM + want_k = _rope(rows[:, k_cols : k_cols + rt.HEAD_DIM].float(), positions) + want_v = rows[:, v_cols : v_cols + rt.HEAD_DIM].to(torch.float8_e4m3fn) + got_k = pool[slots // rt.PAGE, 0, h, slots % rt.PAGE].to(torch.float64) + got_v = pool[slots // rt.PAGE, 1, h, slots % rt.PAGE] + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), ( + f"K head {h} is appended rotated, in E4M3" + ) + assert torch.equal(got_v.view(torch.uint8), want_v.view(torch.uint8)), f"V head {h}" + first += n + + # The band discriminates: a reference whose rows are shifted by one token is outside it. + shifted = op.reference(backend, qkv.roll(1, 0), metadata) + with pytest.raises(AssertionError): + assert_within_ulp(out, shifted, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py new file mode 100644 index 000000000000..174be4bef5a1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for ops/m3_drafter_attn: the Eagle3 drafter layer's generation pass on the full-context +PrimTS decode, driven through ``M3DrafterFullCtxFmha.run_generation`` (in-tree preprocess with the step's FP8 K/V +append, the K / V sub-page views of the drafter cache, the pinned plan) at every recipe the target pins: the 4-row +first draft pass and the 1-row later passes of 1-2 requests. + +Lengths sweep the KV tile edges each query row in turn ends on (the rows of a Q CTA with no key in its trailing +splits), with mixed per-request lengths; every call is held to the fp32 reference over the cache after the append +(E4M3 query, unit KV scales, no RoPE) within the FP8-KV band, finite, for both split-KV reducers. +""" + +import math +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) + +assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +FP8_TOL = 6e-2 +# Persistent like the FMHA manager's workspace and the spec-decoding metadata of a real step: the in-tree preprocess +# kernels must not see these move between calls. +_PERSISTENT = {} + + +def _persistent(name: str, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _edge_lens(plan, qlen: int): + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +def _library(qlen: int, nreq: int, own_reducer: bool): + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = own_reducer + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = torch.cuda.get_device_properties(0).multi_processor_count + lib._on_taken = lambda rows: None + # The in-tree preprocess pass (its RoPE-free cache append); the one-launch preprocess has its own bitwise test. + lib._fused_rope_mode = None + lib._pdl = False + lib._on_rope_taken = None + return lib + + +def _attn(qlen: int): + return SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=0, + rotary_inv_freq=None, + rotary_cos_sin=None, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=0, theta=10000.0, scale_type=0, scale=1.0, max_positions=1024 + ), + ) + + +def _run(lib, qlen: int, kv_lens, seed: int, spec: bool = True): + """One drafter generation pass over a fresh FP8 cache; returns (output, fp32 reference). ``spec`` passes the + spec-decoding lengths/positions (a multi-token generation step); without them a multi-row pass is the drafter's + plain generation step of input_seq_length rows, as it runs in-model.""" + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g).clamp_(-4, 4) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + for r, n in enumerate( + kv_lens + ): # shift the tail third of V, so a missed trailing split fails the band + start = 2 * n // 3 + for pg in range(start // PAGE, math.ceil(n / PAGE)): + pool[int(perm[r, pg]), 1, :, (start % PAGE if pg == start // PAGE else 0) :] += 1.0 + pool = pool.to(torch.float8_e4m3fn) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + pointers = torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64) + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=dev, generator=g) + seq = torch.tensor(kv_lens, dtype=torch.int32, device=dev) + one = torch.ones(1, dtype=torch.float32, device=dev) + multi = qlen > 1 and spec + out = torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=_attn(qlen), + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=pointers, + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=seq, + spec_decoding_generation_lengths=( + _persistent( + f"gen_lengths{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=( + _persistent( + f"position_offsets{qlen}x{nreq}", + lambda: ( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous() + ), + ) + if multi + else None + ), + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=out, + ) + lib.run_generation(params) + torch.cuda.synchronize() + # fp32 reference over the cache after the append: E4M3 query, unit scales, causal over the step's rows. + q = qkv[:, : H * D].float().view(nreq, qlen, H, D).to(torch.float8_e4m3fn).float() + flat = pool.float() + ref = torch.empty(nreq, qlen, H, D, device=dev) + for r, n in enumerate(kv_lens): + blocks = perm[r, : math.ceil(n / PAGE)] + k = flat[blocks, 0].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + v = flat[blocks, 1].permute(1, 0, 2, 3).reshape(KVH, -1, D)[:, :n] + k = k.repeat_interleave(H // KVH, dim=0) + v = v.repeat_interleave(H // KVH, dim=0) + s = torch.einsum("qhd,hkd->hqk", q[r], k) * D**-0.5 + visible = ( + torch.arange(n, device=dev)[None, :] + < (n - qlen + 1 + torch.arange(qlen, device=dev))[:, None] + ) + s = s.masked_fill(~visible[None], float("-inf")) + ref[r] = torch.einsum("hqk,hkd->qhd", torch.softmax(s, dim=-1), v) + return out.float().view(nreq, qlen, H, D), ref + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_the_drafter_pass_at_tile_edges(qlen, nreq) -> None: + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if fullctx.has_own_reducer() else (False,) + for own in reducers: + fullctx.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_LEN, own, get_env_enable_pdl() + ) + lens = _edge_lens(fullctx.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + lens = [n for n in lens if n > qlen] + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for i, kv_lens in enumerate(batches): + # The first draft pass runs in-model as a plain multi-row generation step; every other batch also takes the + # spec-decoding form. + spec = qlen > 1 and i % 2 == 1 + for own in reducers: + got, ref = _run(_library(qlen, nreq, own), qlen, kv_lens, sum(kv_lens), spec=spec) + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +def _drafter_model(scale: float): + """A drafter whose one attention module holds a TrtllmAttention backend as ``attn`` (the backend is no module).""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + backend = TrtllmAttention.__new__(TrtllmAttention) + backend.kv_cache_scaling_factor = torch.full((1,), scale, device="cuda") + backend.is_mla_enable, backend.kv_lora_rank, backend.v_head_dim, backend.head_dim = ( + False, + None, + None, + D, + ) + backend._fmha_manager = SimpleNamespace(fmha_libs=[]) + model = torch.nn.Module() + model.weight = torch.nn.Parameter(torch.zeros(1, device="cuda")) + model.self_attn = torch.nn.Module() + model.self_attn.attn = backend + return model, backend + + +def test_install_declines_a_scaled_kv_cache() -> None: + """At unit KV scales the arithmetic is the stock decode's; a scaled cache keeps the stock path.""" + reasons = [] + model, layer = _drafter_model(0.5) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 0 + assert layer._fmha_manager.fmha_libs == [] and reasons == ["a scaled FP8 KV cache"] + model, layer = _drafter_model(1.0) + assert drafter.install(model, {}, MAX_LEN, False, False, None, reasons.append, None) == 1 + assert isinstance(layer._fmha_manager.fmha_libs[0], drafter.M3DrafterFullCtxFmha) + + +def test_a_cache_bound_past_the_plan_capacity_keeps_the_stock_path() -> None: + """The plans are compiled for ``max_kv_len`` of K/V: a step whose cache length bound is longer falls through.""" + + from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase + + class _Backend: + num_heads, num_kv_heads, head_dim, quant_mode = H, KVH, D, FP8_KV_QUANT_MODE + kv_scale_orig_quant = kv_scale_quant_orig = None + + backend = _Backend() + backend.kv_scale_orig_quant, backend.kv_scale_quant_orig = object(), object() + reasons = [] + lib = _library(1, 1, False) + lib._on_uncovered = reasons.append + lib._attn_ref = weakref.ref(backend) + fwd = SimpleNamespace( + attention_input_type=None, + kv_scale_orig_quant=backend.kv_scale_orig_quant, + kv_scale_quant_orig=backend.kv_scale_quant_orig, + ) + meta = SimpleNamespace(num_contexts=0, num_generations=1, max_seq_len=MAX_LEN + 8) + q = torch.empty(1, H * D, device="cuda") + assert not lib._is_supported(q, None, None, meta, fwd, phase=FmhaPhase.GENERATION) + assert reasons == ["the cache length bound is beyond the plans' K/V capacity"] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py new file mode 100644 index 000000000000..b1759f370b7a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -0,0 +1,371 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for minimax_m3_sm107tp8_drafter_rope_fp8_kv, the TP8 drafter layer's one-launch generation +preprocess: bitwise against the in-tree ``thop.trtllm_gen_generation_preprocess`` followed by the torch E4M3 cast +of its BF16 query, on the query bytes and on every byte of the paged cache, with GPT-NeoX RoPE from a cos/sin table +(full and partial rotary width), multi-token (spec-decoding positions) and plain passes, one and two requests, and +query values past the E4M3 range. The library's fused pass is also held bitwise to its preprocess pass end to end. +""" + +import math +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_drafter_attn as drafter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_fullctx_decode as fullctx, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp8_drafter_rope_fp8_kv"): + pytest.skip("minimax_m3_sm107tp8_drafter_rope_fp8_kv is not built", allow_module_level=True) + +H, KVH, D, PAGE = fullctx.HEADS, fullctx.KV_HEADS, fullctx.HEAD_DIM, 128 +MAX_LEN = 1 << 20 +FP8_KV_QUANT_MODE = 128 +ROPE_GPT_NEOX = 2 +_PERSISTENT = {} + + +def _persistent(name, make): + if name not in _PERSISTENT: + _PERSISTENT[name] = make() + return _PERSISTENT[name] + + +def _case(qlen, kv_lens, rotary_dim, multi, seed): + dev = torch.device("cuda") + nreq = len(kv_lens) + g = torch.Generator(device=dev).manual_seed(seed) + pages = math.ceil(max(kv_lens) / PAGE) + nblocks = nreq * pages + 2 + pool = (torch.randn(nblocks, 2, KVH, PAGE, D, device=dev, generator=g) * 2).to( + torch.float8_e4m3fn + ) + perm = torch.randperm(nblocks, device=dev, generator=g)[: nreq * pages].view(nreq, pages) + offsets = torch.zeros(1, nreq, 2, pages, dtype=torch.int32, device=dev) + offsets[0, :, 0] = 2 * perm + offsets[0, :, 1] = 2 * perm + 1 + ntok = nreq * qlen + qkv = torch.randn(ntok, (H + 2 * KVH) * D, device=dev, generator=g) * 4 + # A few query / key values near and past the E4M3 range (both casts saturate at 448). + hot = torch.rand(qkv.shape, device=dev, generator=g) < 0.01 + qkv = torch.where( + hot, qkv.sign() * (430 + 80 * torch.rand(qkv.shape, device=dev, generator=g)), qkv + ) + qkv = qkv.to(torch.bfloat16).contiguous() + max_pos = max(kv_lens) + 64 + cos_sin = torch.randn(max_pos, rotary_dim // 2, 2, device=dev, generator=g).float().contiguous() + attn = SimpleNamespace( + num_heads=H, + num_kv_heads=KVH, + head_dim=D, + q_scaling=1.0, + attention_chunk_size=None, + quant_mode=FP8_KV_QUANT_MODE, + local_layer_idx=0, + position_embedding_type=ROPE_GPT_NEOX, + rotary_inv_freq=None, + rotary_cos_sin=cos_sin, + predicted_tokens_per_seq=qlen, + rope_params=SimpleNamespace( + dim=rotary_dim, theta=10000.0, scale_type=0, scale=1.0, max_positions=max_pos + ), + ) + one = torch.ones(1, dtype=torch.float32, device=dev) + spec_offsets = None + if multi: + spec_offsets = _persistent( + f"pos{qlen}x{nreq}", + lambda: torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1).contiguous(), + ) + # Spec positions are not the plain row index: shift them per request. + spec_offsets.copy_( + torch.arange(qlen, dtype=torch.int32, device=dev).repeat(nreq, 1) + + torch.arange(nreq, dtype=torch.int32, device=dev)[:, None] + ) + params = SimpleNamespace( + qkv_input=qkv, + query_input=None, + attn=attn, + meta=SimpleNamespace( + kv_cache_block_offsets=offsets, + host_kv_cache_pool_pointers=torch.tensor([[pool.data_ptr(), 0]], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ), + fwd=SimpleNamespace( + kv_scale_orig_quant=one, + kv_scale_quant_orig=one, + out_scale=None, + mrope_position_deltas=None, + ), + workspace=_persistent( + "workspace", lambda: torch.empty(512 << 20, dtype=torch.uint8, device=dev) + ), + sequence_lengths=torch.tensor(kv_lens, dtype=torch.int32, device=dev), + spec_decoding_generation_lengths=( + _persistent( + f"len{qlen}x{nreq}", + lambda: torch.full((nreq,), qlen, dtype=torch.int32, device=dev), + ) + if multi + else None + ), + spec_decoding_position_offsets=spec_offsets, + seq_offset=0, + tokens_per_block=PAGE, + max_attention_window_size=MAX_LEN, + cyclic_attention_window_size=MAX_LEN, + num_tokens=ntok, + batch_size=nreq, + input_seq_length=qlen, + max_past_kv_length=max(kv_lens), + total_num_blocks=nblocks, + kv_factor=2, + is_cross=False, + output=torch.zeros(ntok, H * D, dtype=torch.bfloat16, device=dev), + ) + return params, pool + + +def _preprocess(params, mpc): + """run_generation's in-tree preprocess call and the E4M3 cast of its query.""" + attn, meta, fwd, rope = params.attn, params.meta, params.fwd, params.attn.rope_params + q_processed = thop.trtllm_gen_generation_preprocess( + params.qkv_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope.dim, + rope.theta, + int(rope.scale_type), + rope.scale, + rope.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + False, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + mpc, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + )[0] + rows = params.num_tokens * attn.num_heads * attn.head_dim + return ( + q_processed.flatten()[:rows] + .to(torch.float8_e4m3fn) + .view(params.num_tokens, attn.num_heads, attn.head_dim) + ) + + +def _op(params, pool, mode): + attn = params.attn + nreq, qlen = params.batch_size, params.input_seq_length + kv_pages = pool.view(-1, KVH, PAGE, D) + tables = params.meta.kv_cache_block_offsets[0, :nreq] + q = torch.empty( + params.num_tokens, + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=pool.device, + ) + multi = params.spec_decoding_generation_lengths is not None and qlen > 1 + torch.ops.trtllm.minimax_m3_sm107tp8_drafter_rope_fp8_kv( + params.qkv_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + attn.rope_params.dim, + PAGE, + mode, + True, + ) + return q + + +# One-request multi-token passes keep the in-tree preprocess (it reads an unbuilt token map), so the spec cases hold +# two requests. +CASES = [ + (qlen, kv_lens, rotary, multi) + for qlen, multi in ((1, False), (4, True), (4, False)) + for kv_lens in (((4, 130) if multi else (4,)) if qlen > 1 else (130,), (8191, 131073)) + for rotary in (128, 64) +] + + +def _bytes(t): + return t.contiguous().view(torch.uint8) + + +def _matches(case, seed): + """Per rope mode: (query bytes equal, cache bytes equal) against the in-tree preprocess + cast.""" + qlen, kv_lens, rotary, multi = case + kv_lens = [max(n, qlen) for n in kv_lens] + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(qlen, kv_lens, rotary, multi, seed) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + pool_ref = pool.clone() + out = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + out[mode] = ( + bool(torch.equal(_bytes(q), _bytes(q_ref))), + bool(torch.equal(_bytes(pool), _bytes(pool_ref))), + ) + return out + + +@pytest.mark.parametrize( + "case", + CASES, + ids=lambda c: f"q{c[0]}-{'x'.join(map(str, c[1]))}-r{c[2]}-{'spec' if c[3] else 'plain'}", +) +def test_the_pinned_rope_mode_is_bitwise(case) -> None: + result = _matches(case, seed=7) + print(f"[rope modes] {case}: {result}", flush=True) + assert result[modeling.DRAFTER_ROPE_MODE] == (True, True), result + + +def test_the_fused_pass_matches_the_preprocess_pass() -> None: + mpc = torch.cuda.get_device_properties(0).multi_processor_count + for qlen, nreq in sorted(modeling.FULLCTX_DECODE_RECIPES): + for multi in (False, True) if qlen > 1 else (False,): + kv_lens = [70001 + 517 * r for r in range(nreq)] + outs = [] + for fused in (False, True): + params, pool = _case(qlen, kv_lens, 128, multi, seed=11) + lib = drafter.M3DrafterFullCtxFmha.__new__(drafter.M3DrafterFullCtxFmha) + lib._recipes = {(qlen, nreq): modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)]} + lib._own_reducer = True + lib._max_kv_len = MAX_LEN + lib._multi_processor_count = mpc + lib._fused_rope_mode = modeling.DRAFTER_ROPE_MODE if fused else None + lib._pdl = True + taken = [] + lib._on_rope_taken = lambda rows, what: taken.append(what) + lib._on_taken = lambda rows: None + fullctx.prebuild(torch.device("cuda"), lib._recipes, MAX_LEN, True, True) + lib.run_generation(params) + torch.cuda.synchronize() + assert bool(taken) == (fused and not (multi and nreq == 1)) + outs.append((params.output.clone(), pool.clone())) + assert torch.equal(_bytes(outs[0][1]), _bytes(outs[1][1])), (qlen, nreq, multi, "cache") + assert torch.equal(_bytes(outs[0][0]), _bytes(outs[1][0])), ( + qlen, + nreq, + multi, + "output", + ) + + +def _round_bf16(x): + u = x.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000 + return u.to(torch.int64).to(torch.int32).view(torch.float32) + + +def _discriminating(n, seed): + """(cos, sin, x, pair) quadruples whose rotated first-half output's E4M3 byte differs between the rope modes + (fp32 CPU emulation; the GPU comparison is the verdict).""" + g = torch.Generator().manual_seed(seed) + keep = [] + while sum(k.shape[0] for k in keep) < n: + m = 1 << 20 + c, s = torch.randn(m, generator=g), torch.randn(m, generator=g) + x = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + xp = (torch.randn(m, generator=g) * 4).to(torch.bfloat16).float() + sn = -s + r0 = (c * x) + (sn * xp) + r1 = ((c.double() * x.double()) + (sn * xp).double()).float() + r2 = ((sn.double() * xp.double()) + (c * x).double()).float() + b = [ + _round_bf16(r.contiguous()).to(torch.float8_e4m3fn).view(torch.uint8) + for r in (r0, r1, r2) + ] + diff = (b[0] != b[1]) | (b[0] != b[2]) | (b[1] != b[2]) + diff &= (r0.abs() < 400) & (r1.abs() < 400) + keep.append(torch.stack([c, s, x, xp], 1)[diff]) + return torch.cat(keep)[:n] + + +def test_crafted_inputs_pin_the_rope_mode() -> None: + """Rotations whose E4M3 byte depends on how the two products are summed: the in-tree preprocess and exactly the + pinned mode agree on every one.""" + mpc = torch.cuda.get_device_properties(0).multi_processor_count + params, pool = _case(1, [130], 128, False, seed=3) + half = 64 + quads = _discriminating(H * half, seed=5).cuda() + pos = 129 # the one row's rope position: cache length 130, one query row + table = params.attn.rotary_cos_sin.view(-1, half, 2) + # One (cos, sin) per dim: every head shares the position's table row, so the heads take the same dims' quads with + # their own (x, pair). + table[pos, :, 0] = quads[:half, 0] + table[pos, :, 1] = quads[:half, 1] + qkv = params.qkv_input.view(-1, H + 2 * KVH, D) + for h in range(H): + qkv[0, h, :half] = quads[:half, 2].to(torch.bfloat16) + qkv[0, h, half:] = quads[:half, 3].to(torch.bfloat16) + start = pool.clone() + q_ref = _preprocess(params, mpc) + torch.cuda.synchronize() + agree = {} + for mode in (0, 1, 2): + pool.copy_(start) + q = _op(params, pool, mode) + torch.cuda.synchronize() + agree[mode] = int((_bytes(q)[0, :, :half] == _bytes(q_ref)[0, :, :half]).sum()) + print( + f"[rope modes crafted] agreeing first-half query bytes of {H * half}: {agree}", flush=True + ) + assert agree[modeling.DRAFTER_ROPE_MODE] == H * half, agree diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py new file mode 100644 index 000000000000..2179f1720e0d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp8_indexer_qk_norm_rope entry of the MiniMax-M3 tp8 target: index-Q out and index-K in the +cache against the in-place fused norm + RoPE followed by the E4M3 cast, byte for byte up to the 0.1% budget the +in-tree test of the kernel allows.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 + m3_fp8_indexer_qk_norm_rope as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" + +DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 + + +def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: + assert a.shape == b.shape and a.dtype == b.dtype == torch.float8_e4m3fn + differ = int((a.view(torch.uint8) != b.view(torch.uint8)).sum()) + assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + m, heads = cell.spec["m"], cell.spec["heads"] + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + qk = torch.randn(m, (heads + 1) * DIM, generator=g, device="cuda").to(torch.bfloat16) + qw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (m,), generator=g, device="cuda", dtype=torch.int32) + pages = (m + PAGE - 1) // PAGE + 1 + cache = torch.zeros(pages, 1, PAGE, DIM, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(pages * PAGE, generator=g, device="cuda")[:m].to(torch.int32) + before = qk.clone() + args = (qk, cache, slots, heads, DIM, ROTARY, EPS, qw, kw, THETA, pos) + with validating(op): + q = op(*args) + assert torch.equal(qk, before), "the BF16 input is left untouched" + q_ref, k_ref = op.reference(*args) + _bytes_close(q.reshape(m, heads, DIM), q_ref) + flat = cache.view(-1, DIM) + _bytes_close(flat[slots.long()], k_ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_saturate.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_saturate.py new file mode 100644 index 000000000000..900e6c7843f2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_saturate.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""paged_cache.fp8_e4m3_saturate (FP8_SATURATE): the E4M3 cache writes and the transient FP8 query casts saturate to ++-448 like the CUDA writers (__NV_SATFINITE); every finite in-range value keeps the plain cast's bits, values already +in E4M3 pass through, and the kill switch is exactly the plain cast. Where torch's cast already saturates (this image) +the conversion is the cast alone, one kernel; otherwise it clamps first, the same bytes.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels import paged_cache + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fp8_e4m3_saturate_matches_satfinite(dtype, monkeypatch) -> None: + g = torch.Generator(device="cuda").manual_seed(0) + inside = (torch.randn(4096, generator=g, device="cuda") * 100).clamp(-448, 448).to(dtype) + plain = inside.to(torch.float8_e4m3fn) + assert torch.equal( + paged_cache.fp8_e4m3_saturate(inside).view(torch.uint8), plain.view(torch.uint8) + ) + edge = torch.tensor( + [449.0, 463.0, 464.0, 1e4, float("inf"), -500.0, -float("inf")], dtype=dtype, device="cuda" + ) + want = torch.tensor([448.0, 448.0, 448.0, 448.0, 448.0, -448.0, -448.0], device="cuda") + assert torch.equal(paged_cache.fp8_e4m3_saturate(edge).float(), want) + nan = paged_cache.fp8_e4m3_saturate(torch.tensor([float("nan")], dtype=dtype, device="cuda")) + assert torch.isnan(nan.float()).all() + already = edge.clamp(-448, 448).to(torch.float8_e4m3fn) + assert paged_cache.fp8_e4m3_saturate(already) is already + monkeypatch.setattr(paged_cache, "FP8_SATURATE", False) + killed = paged_cache.fp8_e4m3_saturate(edge) + assert torch.equal(killed.view(torch.uint8), edge.to(torch.float8_e4m3fn).view(torch.uint8)) + + +def test_cache_write_saturates() -> None: + """The cache-dtype conversion every write_kv_slots branch applies (compared as bytes: torch has no float8 + index_copy_ kernel), and an E4M3 source passing through untouched.""" + values = torch.tensor([[1.0, 500.0, -1e5, float("inf")]], dtype=torch.bfloat16, device="cuda") + out = paged_cache._to_cache_dtype(values, torch.float8_e4m3fn) + assert out.dtype == torch.float8_e4m3fn + want = torch.tensor([[1.0, 448.0, -448.0, 448.0]], device="cuda").to(torch.float8_e4m3fn) + assert torch.equal(out.view(torch.uint8), want.view(torch.uint8)) + assert paged_cache._to_cache_dtype(out, torch.float8_e4m3fn) is out + bf = paged_cache._to_cache_dtype(values, torch.bfloat16) + assert bf.dtype == torch.bfloat16 and torch.equal(bf, values) + + +_EDGES = [ + 448.0, + 449.0, + 463.0, + 464.0, + 1e4, + float("inf"), + -448.0, + -449.0, + -464.0, + -1e4, + -float("inf"), + float("nan"), + 0.0, +] + + +def _cuda_kernels(fn) -> int: + from torch.profiler import ProfilerActivity, profile + + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CUDA]) as prof: + fn() + torch.cuda.synchronize() + return sum(1 for e in prof.events() if e.device_type.name == "CUDA") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("cast_saturates", [True, False], ids=["cast_only", "clamp_fallback"]) +def test_fp8_e4m3_saturate_bytes_equal_clamp_then_cast(dtype, cast_saturates, monkeypatch) -> None: + """Both paths give the bytes of clamp(+-448) followed by the cast, edge values included.""" + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", cast_saturates) + x = torch.tensor(_EDGES, dtype=dtype, device="cuda") + want = x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + assert torch.equal(paged_cache.fp8_e4m3_saturate(x).view(torch.uint8), want.view(torch.uint8)) + + +def test_fp8_e4m3_saturate_takes_the_cast_only_path_here(monkeypatch) -> None: + """This image's torch cast saturates, on the CPU (the import-time probe) and on CUDA, so the conversion is one + kernel; with the probe failing it falls back to clamp + cast, two kernels.""" + assert paged_cache._CAST_SATURATES + x = torch.tensor(_EDGES, dtype=torch.bfloat16, device="cuda") + assert torch.equal( + x.to(torch.float8_e4m3fn).view(torch.uint8), + x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).view(torch.uint8), + ) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 1 + monkeypatch.setattr(paged_cache, "_CAST_SATURATES", False) + assert _cuda_kernels(lambda: paged_cache.fp8_e4m3_saturate(x)) == 2 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py new file mode 100644 index 000000000000..48d1bba922ac --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for the attention/m3_fullctx_decode entry: the dense layers' full-context FP8 decode on PrimTS +split-KV decode at the recipe the target pins per (query rows per request, requests). + +Every cell runs against the fp32 reference (q rounded to E4M3 first, as the kernel runs it), eager and replayed from a +CUDA graph (the same bits), and twice (the same bits). The K/V pages are read through the strided views of a flat +sub-page pool, as the KV cache manager hands them out. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_fullctx_decode import ( # noqa: E501 + HEAD_DIM, + HEADS, + KV_HEADS, + MAX_KV_LEN, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_fullctx_decode import ( # noqa: E501 + m3_fullctx_decode as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) + +PAGE = 128 +SM_SCALE = HEAD_DIM**-0.5 +# The FP8-KV band of the in-tree dense decode test (test_minimax_m3_dense_decode.test_matches_reference). +FP8_TOL = 6e-2 + + +class _Pool: + """The one method the entry asks of a KV cache manager.""" + + def __init__(self, pool: torch.Tensor, subpages: int): + self.pool = pool + self.subpages = subpages + + def get_kv_subpage_pool(self, layer_idx: int, kv_layout: str = "HND"): + assert kv_layout == "HND" + return self.pool, self.subpages + + +def _inputs(kv_lens, qlen: int, seed: int, subpages: int): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(kv_lens) + pages = max((n + PAGE - 1) // PAGE for n in kv_lens) + slots = batch * pages + 3 + pool = torch.randn( + (slots * subpages, KV_HEADS, PAGE, HEAD_DIM), device="cuda", generator=g + ).clamp(-4, 4) + perm = torch.randperm(slots, generator=g, device="cuda")[: batch * pages] + block_table = perm.reshape(batch, pages).to(torch.int32) + # The last third of every request's values is shifted by +1, so a split that missed the tail fails the band. + for b, n in enumerate(kv_lens): + start = 2 * n // 3 + for page in range(start // PAGE, (n + PAGE - 1) // PAGE): + v = int(block_table[b, page]) * subpages + 1 + pool[v, :, (start % PAGE if page == start // PAGE else 0) :] += 1.0 + # q as the layer hands it over on decode steps: the FP8 query columns of the fused norm/RoPE kernel's + # [q | k | v] output (row stride wider than q), so no cast makes it compact. + qkv = torch.randn( + (batch * qlen, (HEADS + 2 * KV_HEADS) * HEAD_DIM), device="cuda", generator=g + ).to(torch.float8_e4m3fn) + q = qkv[:, : HEADS * HEAD_DIM].view(batch * qlen, HEADS, HEAD_DIM) + seq_lens = torch.tensor(kv_lens, device="cuda", dtype=torch.int32) + return q, _Pool(pool.to(torch.float8_e4m3fn), subpages), block_table, seq_lens + + +@pytest.mark.parametrize("subpages", [2, 3], ids=["kv", "kv+index"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, subpages, certifying) -> None: # noqa: F811 + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs( + spec["kv_lens"], qlen, zlib.crc32(cell.why.encode()) % 2**31, subpages + ) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + certifying(op) + ref = op.reference(*args) + eager = op(*args).clone() + torch.testing.assert_close(eager.float(), ref, rtol=FP8_TOL, atol=FP8_TOL) + again = op(*args).clone() + assert torch.equal(again, eager), "not deterministic" + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*args) + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, eager), "the CUDA-graph replay differs from the eager call" + + +def test_the_cells_cover_the_target_rows() -> None: + certified = {(c.spec["qlen"], len(c.spec["kv_lens"])) for c in op.CELLS} + assert set(modeling.FULLCTX_DECODE.keys) <= certified + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_bf16_merge_is_the_primts_reducer_to_one_ulp(cell) -> None: + """The BF16 split-KV merge op against PrimTS's own FP16 reducer and the cast: the same partials, merged once + in fp32 and rounded once, so within one BF16 spacing of the FP16-then-BF16 result.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + + if not kernel.has_own_reducer(): + pytest.skip("this build does not carry minimax_m3_sm107_tp2_fullctx_reduce") + spec = cell.spec + qlen, nreq = spec["qlen"], len(spec["kv_lens"]) + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + q, manager, block_table, seq_lens = _inputs(spec["kv_lens"], qlen, 11, 2) + outs = [] + for own in (True, False): + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + op(q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe, own_reducer=own) + outs.append(out.float()) + ours, primts = outs + _, exp = torch.frexp(primts.abs()) + spacing = torch.where( + primts == 0, + torch.full_like(primts, 2.0**-133), + torch.ldexp(torch.ones_like(primts), exp - 8), + ) + assert bool(((ours - primts).abs() <= spacing).all()) + + +def _tile_edge_kv_lens(plan, qlen: int): + """KV lengths around every kind of boundary the split ranges have, so each query row of the Q CTA in turn ends + its keys on a KV tile edge (k * tile + d, d in -1 .. qlen - 1): the first tiles, the reducer's 8-group stride, the + lengths where a split's range grows (splits * insts tiles) and a split edge past it.""" + tile, insts, splits = plan.tile_kv, plan.insts_kv, int(plan.recipe["splits"]) + grow = splits * insts + tiles = sorted({1, 2, 8, 9, grow - 1, grow, grow + 1, grow + 2 * insts, 2 * grow}) + return [k * tile + d for k in tiles for d in range(-1, qlen)] + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_rows_with_no_key_in_a_trailing_split(qlen, nreq) -> None: + """A row whose keys end before its Q CTA's causal union does has trailing splits holding none of its keys + (log2-LSE -inf); both reducers must leave them out. [1300, 1281] at 4 query rows is the batch a GSM8K decode + step first hit (NaN rows from exp2(-inf - -inf)).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, False, get_env_enable_pdl() + ) + lens = _tile_edge_kv_lens(kernel.plan_for(torch.device("cuda"), qlen, nreq, False), qlen) + batches = [[n] for n in lens] if nreq == 1 else [[a, b] for a, b in zip(lens, lens[::-1])] + batches.append([1300, 1281][:nreq]) + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, sum(kv_lens), 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +@pytest.mark.parametrize("qlen,nreq", sorted(modeling.FULLCTX_DECODE_RECIPES)) +def test_kv_lengths_past_the_policy_length(qlen, nreq) -> None: + """The plans are compiled for MAX_KV_LEN of K/V with PrimTS's policy sized at 1 << 20 (POLICY_KV_LEN): caches just + past the policy length (the AgentX bound is 1 << 20 plus the speculative rows) and well past it stay exact.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_fullctx_decode as kernel, + ) + from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + + recipe = modeling.FULLCTX_DECODE_RECIPES[(qlen, nreq)] + reducers = (True, False) if kernel.has_own_reducer() else (False,) + lengths = [(1 << 20) + 1, (1 << 20) + 8] + ([1_500_001] if nreq == 1 else []) + batches = [[n] * nreq if nreq == 1 else [n, (1 << 20) - 3] for n in lengths] + failures = [] + for kv_lens in batches: + q, manager, block_table, seq_lens = _inputs(kv_lens, qlen, 13, 2) + out = torch.zeros(q.shape, dtype=torch.bfloat16, device="cuda") + args = (q, manager, 0, block_table, seq_lens, SM_SCALE, out, qlen, recipe) + ref = op.reference(*args) + for own in reducers: + kernel.prebuild( + torch.device("cuda"), {(qlen, nreq): recipe}, MAX_KV_LEN, own, get_env_enable_pdl() + ) + out.zero_() + op(*args, own_reducer=own) + got = out.float() + if not bool(torch.isfinite(got).all()) or not torch.allclose( + got, ref, rtol=FP8_TOL, atol=FP8_TOL + ): + failures.append((kv_lens, "bf16-merge" if own else "primts-reducer")) + assert not failures, failures + + +@pytest.mark.parametrize( + "splits,rows,seq_len_q", [(108, 8, 1), (64, 8, 1), (96, 32, 4), (96, 32, 4)] +) +def test_the_tp8_merge_is_the_tp2_merge_bitwise(splits, rows, seq_len_q) -> None: + """The TP8 copy of the BF16 split-KV merge (one row per CTA, early dependent release) returns the image TP2 op's + bits on random 16-bit partials with empty (-inf log2-LSE) splits, both partial dtypes, two requests.""" + ops = torch.ops.trtllm + if not ( + hasattr(ops, "minimax_m3_sm107tp8_fullctx_reduce") + and hasattr(ops, "minimax_m3_sm107_tp2_fullctx_reduce") + ): + pytest.skip("needs both the TP8 merge copy and the image's TP2 merge op") + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(splits * 131 + rows) + nreq, kvh = 2, 1 + for dtype in (torch.float16, torch.bfloat16): + po = ( + (torch.randn(nreq, kvh, splits, rows, 128, device=dev, generator=g) * 3) + .to(dtype) + .contiguous() + ) + ps = torch.randn(nreq, kvh, splits, rows, device=dev, generator=g) * 4 + ps[torch.rand(ps.shape, device=dev, generator=g) < 0.05] = float("-inf") + seq = torch.tensor([165521, 120011], dtype=torch.int32, device=dev) + outs = [] + for reduce in ( + ops.minimax_m3_sm107_tp2_fullctx_reduce, + ops.minimax_m3_sm107tp8_fullctx_reduce, + ): + out = torch.zeros( + nreq * seq_len_q, kvh * rows // seq_len_q, 128, dtype=torch.bfloat16, device=dev + ) + reduce( + po, ps.contiguous(), seq, out, seq_len_q, seq_len_q, 128, 1, seq_len_q > 1, False + ) + outs.append(out) + torch.cuda.synchronize() + assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)), ( + splits, + rows, + seq_len_q, + dtype, + ) + + +def _swizzled_scale_offsets(rows: int, width: int) -> torch.Tensor: + """Offsets of rows [0, rows)'s UE8M0 scales in the 128x4-swizzled layout of mxfp8_quantize(x, True).""" + k_sf = width // 32 + k_tiles = (k_sf + 3) // 4 + r = torch.arange(rows, device="cuda").view(-1, 1) + k = torch.arange(k_sf, device="cuda").view(1, -1) + offsets = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + return offsets.reshape(-1) + + +@pytest.mark.parametrize("splits,rows,seq_len_q", [(108, 8, 1), (64, 8, 1), (96, 32, 4)]) +def test_the_tp8_merge_mxfp8_copy_is_stock_mxfp8_quantize(splits, rows, seq_len_q) -> None: + """The TP8 merge's MXFP8 copy (the dense o_proj hand-off) equals mxfp8_quantize(out, True) bit for bit: E4M3 rows + and the real rows' swizzled UE8M0 scales, with an all-zero head row (amax 0) and saturating outliers; the BF16 + output is the merge without the copy.""" + ops = torch.ops.trtllm + if not hasattr(ops, "minimax_m3_sm107tp8_fullctx_reduce"): + pytest.skip("needs the TP8 merge copy") + dev = torch.device("cuda") + g = torch.Generator(device=dev).manual_seed(splits * 7 + rows) + nreq, kvh = 2, 1 + tokens = nreq * seq_len_q + width = kvh * rows // seq_len_q * 128 + for dtype in (torch.float16, torch.bfloat16): + po = torch.randn(nreq, kvh, splits, rows, 128, device=dev, generator=g) * 3 + po[0, 0, :, 1] = 0.0 + po[1, 0, :, 0, 5] = 6.0e4 + po = po.to(dtype).contiguous() + ps = (torch.randn(nreq, kvh, splits, rows, device=dev, generator=g) * 4).contiguous() + seq = torch.tensor([165521, 120011], dtype=torch.int32, device=dev) + args = (po, ps, seq) + tail = (seq_len_q, seq_len_q, 128, 1, seq_len_q > 1, False) + plain = torch.zeros(tokens, width, dtype=torch.bfloat16, device=dev) + ops.minimax_m3_sm107tp8_fullctx_reduce(*args, plain, *tail) + out = torch.zeros(tokens, width, dtype=torch.bfloat16, device=dev) + q8 = torch.empty(128, width, dtype=torch.float8_e4m3fn, device=dev) + sf = torch.zeros(128 * (width // 32), dtype=torch.uint8, device=dev) + ops.minimax_m3_sm107tp8_fullctx_reduce(*args, out, *tail, q8, sf) + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), plain.view(torch.int16)) + stock_q, stock_sf = ops.mxfp8_quantize(out, True) + assert torch.equal( + q8[:tokens].view(torch.uint8), stock_q.view(torch.uint8).reshape(tokens, -1) + ), dtype + at = _swizzled_scale_offsets(tokens, width) + assert torch.equal(sf[at], stock_sf.view(torch.uint8).reshape(-1)[at]), dtype diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_identity_selection.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_identity_selection.py new file mode 100644 index 000000000000..8819b1fcfb73 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_identity_selection.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test: the short-sequence identity selection (runtime policy ``short_sequence_identity``) is the stock table. + +On a step whose every query token has at most 16 valid blocks, this target skips the index scorer and the select and +stages the table (i < n_valid ? i : -1). Stock runs the in-tree CuTe DSL scorer and ``select_blocks_from_maxscore``, +which returns ascending ids and drops a valid block only when its max score is -inf. The scorer's scores are FP8 +products summed in FP32 (|score| <= 448^2 * 128), so -inf comes only from the causal mask, which n_valid excludes. +This test drives the real scorer with extreme FP8 operands at the TP8 geometry (1 index head per rank) and checks +both claims. +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) +from tensorrt_llm._torch.custom_ops import ( + cute_dsl_custom_ops, # noqa: F401 — registers the scorer op +) + +PAGE, DIM, TOPK, HEADS, QLEN = 128, 128, 16, 1, 4 +E4M3_MAX = 448.0 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + + +def _fill(kind, shape, gen): + if kind == "random": + x = torch.randn(shape, generator=gen, device="cuda") * 64 + elif kind == "max": + x = torch.full(shape, E4M3_MAX, device="cuda") + elif kind == "negmax": + x = torch.full(shape, -E4M3_MAX, device="cuda") + elif kind == "alternating": + x = torch.full(shape, E4M3_MAX, device="cuda") + x[..., 1::2] = -E4M3_MAX + elif kind == "zero": + x = torch.zeros(shape, device="cuda") + else: # the smallest E4M3 subnormal + x = torch.full(shape, 2.0**-9, device="cuda") + return x.clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + + +_KINDS = ("random", "max", "negmax", "alternating", "zero", "tiny") + + +@pytest.mark.parametrize( + "seq_lens", [(4,), (129,), (2048,), (700, 2048), (5, 1500, 2047, 300), (1, 128, 256, 1024)] +) +@pytest.mark.parametrize("q_kind", _KINDS) +@pytest.mark.parametrize("k_kind", _KINDS) +def test_stock_selection_of_a_short_step_is_the_identity_table(seq_lens, q_kind, k_kind) -> None: + gen = torch.Generator(device="cuda").manual_seed(0) + reqs, total_q = len(seq_lens), len(seq_lens) * QLEN + nblk = [(s + PAGE - 1) // PAGE for s in seq_lens] + pages = sum(nblk) + 4 + perm = torch.randperm(pages, generator=gen, device="cuda").int() + block_table = torch.zeros((reqs, max(nblk)), dtype=torch.int32, device="cuda") + start = 0 + for r, n in enumerate(nblk): + block_table[r, :n] = perm[start : start + n] + start += n + n_valid = torch.tensor( + [(max(s - QLEN + o + 1, 0) + PAGE - 1) // PAGE for s in seq_lens for o in range(QLEN)], + dtype=torch.int32, + device="cuda", + ) + idx_q = _fill(q_kind, (total_q, HEADS, DIM), gen) + idx_k = _fill(k_kind, (pages, PAGE, DIM), gen) + max_score = torch.full((HEADS, max(nblk), total_q), float("-inf"), device="cuda") + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, + idx_k, + block_table, + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + max_score.transpose(1, 2), + QLEN, + ) + for tok, nv in enumerate(n_valid.tolist()): + assert bool(torch.isfinite(max_score[:, :nv, tok]).all()), ( + f"token {tok}: a causally valid block scored non-finite" + ) + table = select_blocks_from_maxscore( + max_score, + topk=TOPK, + n_valid_blocks=n_valid, + init_blocks=1, + local_blocks=4, + head_major_output=True, + ).to(torch.int32) + # The selector returns [tokens, kv_heads, topk]; the decode reads it as [kv_heads, tokens, topk] (layers.py). + table = table.permute(1, 0, 2) + i = torch.arange(TOPK, dtype=torch.int32, device="cuda") + identity = torch.where( + i[None, :] < n_valid[:, None], i[None, :], torch.full_like(i[None, :], -1) + ) + assert torch.equal(table, identity[None].expand(HEADS, total_q, TOPK)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py new file mode 100644 index 000000000000..e768b7442a15 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_index_score entry: the MSA index scorer with its pre-wait schedule. + +Every cell is held to the fp32 reference (a few fp32 ulps: the kernel sums the 128 exact products in its k-split order) +and to the in-tree scorer (``cute_dsl_minimax_m3_index_decode_score`` at its own 256-CTA split) bit for bit, which the +pre-wait prefetch, the early tile issue and the cell's split must not change. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_index_score import ( # noqa: E501 + PAGE, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_index_score import ( # noqa: E501 + m3_index_score as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" + +HEADS = 1 # index heads per rank at TP8 (4 over 4 KV groups, each group on a rank pair) +HEAD_DIM = 128 + + +def _inputs(seq_lens, qlen: int, seed: int, dtype=torch.float8_e4m3fn): + g = torch.Generator(device="cuda").manual_seed(seed) + batch = len(seq_lens) + max_blocks = (max(seq_lens) + PAGE - 1) // PAGE + stride = ((max_blocks + 15) // 16) * 16 + pages = batch * max_blocks + block_table = ( + torch.randperm(pages, device="cuda", generator=g).to(torch.int32).view(batch, max_blocks) + ) + idx_q = torch.randn(batch * qlen, HEADS, HEAD_DIM, device="cuda", generator=g).to(dtype) + index_k = torch.randn(pages, PAGE, HEAD_DIM, device="cuda", generator=g).to(dtype) + seq = torch.tensor(seq_lens, device="cuda", dtype=torch.int32) + return idx_q, index_k, block_table, seq, stride + + +def _score(stride: int, total_q: int): + backing = torch.full( + (HEADS, stride, total_q), -float("inf"), device="cuda", dtype=torch.float32 + ) + return backing.transpose(1, 2) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + qlen = spec["qlen"] + idx_q, index_k, block_table, seq, stride = _inputs( + spec["seq_lens"], qlen, zlib.crc32(cell.why.encode()) + ) + score = _score(stride, idx_q.shape[0]) + assert spec["heads"] == HEADS + args = (idx_q, index_k, block_table, seq, score, qlen, spec["split_k"], spec["early"]) + certifying(op) + ref = op.reference(*args) + op(*args) + stock = _score(stride, idx_q.shape[0]) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + torch.cuda.synchronize() + torch.testing.assert_close(score, ref, rtol=1e-5, atol=1e-4) + assert torch.equal(score, stock), "the scores differ from the in-tree scorer" + + +@pytest.mark.parametrize("early", [True, False], ids=["early", "prefetch"]) +def test_bf16_index_is_the_in_tree_scorer(early) -> None: + idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) + score, stock = _score(stride, 4), _score(stride, 4) + op(idx_q, index_k, block_table, seq, score, 4, 1536, early) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, 4 + ) + assert torch.equal(score, stock) + + +@pytest.mark.parametrize("qlen", [1, 2, 3, 4]) +def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> None: + """Lengths around the edge between pre-staged (all-old) and post-wait first blocks, 1 and 2 requests.""" + for k in (1, 2, 1294): + for length in (k * PAGE - qlen - 1, k * PAGE - qlen, k * PAGE - qlen + 1, k * PAGE + 1): + for seq_lens, split in ( + ((length,), 1536), + ((length,), 1272), + ((length, max(qlen, length // 3)), 768), + ): + idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) + score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) + op(idx_q, index_k, block_table, seq, score, 4, split, True) + torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( + idx_q, index_k, block_table, seq, stock, qlen + ) + assert torch.equal(score, stock), (qlen, seq_lens) + + +def test_guard_refuses_a_request_past_its_block_table() -> None: + idx_q, index_k, block_table, seq, stride = _inputs((1000,), 1, 5) + seq = seq + 4 * PAGE + with pytest.raises(AssertionError, match="past its block-table row"): + with validating(op): + op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, 1536, True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py new file mode 100644 index 000000000000..0539d9639125 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_msa_attention entry: one layer's MSA backend step against its float64 reference. + +Every cell runs on the runtime's own objects (``_m3_attn_runtime_tp8``): a hybrid NVFP4 ``MiniMaxM3KVCacheManagerV2`` +with the requests allocated in it, ``MiniMaxM3MsaSparseAttentionMetadata`` filled by its ``prepare()``, and the layer's +``MiniMaxM3MsaSparseAttention`` built as the target builds it. Earlier rows (and, on a sparse layer, the step's own +rows, which the target's producer writes before the call) go into the cache through the in-tree scatters, at the +slots the manager assigned. The sparse layers' index-K rows are planted so every block's score is 1/8 or more apart +from every other's: the top-k selection the reference recomputes is then unambiguous. + +The cells are the target's calls (``_attention_core`` in modeling.py): a dense layer with live bf16 K / V in a +prefill chunk and live FP8 K / V in a decode step (the dense layers sit in the hybrid cache's FP8 half), a sparse +layer whose producer already wrote the caches (``k = v = None``) in a prefill chunk beyond 16 blocks and in a mixed +step, both with the NVFP4 scales the target passes. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_msa_attention import ( # noqa: E501 + m3_msa_attention as op, +) + +__extra_import_path__ = ["../../../..", "."] +import _m3_attn_runtime_tp8 as rt # noqa: E402 — needs the path declared above +from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" + +# Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, +# per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 +# FMHA does), so what is left is each FMHA's own arithmetic, measured on GB300 over these cells (max element / RMS): +# * fmha_sm100 over the FP8 pool (dense context rows): 10.7 / 6.0. It runs E4M3 x E4M3 and its softmax +# probabilities are rounded before the PV product. +# * trtllm-gen over the FP8 pool (dense generation rows): 7.7 / 6.1. It rounds the probabilities to E4M3 (scaled by +# 448 against the running maximum): emulating that brings a single-tile request to 0.6 RMS. +# * the NVFP4 CSR prefill (sparse context rows): 33.6 / 11.8 at 4096 rows, 24.3 / 9.1 at 512. It stages the +# reconstructed K / V in E4M3 (the in-tree A/B test, test_minimax_m3_sparse_attn_decode, notes it and holds it +# to 3e-2). +# * the Triton NVFP4 sparse decode (sparse generation rows): 1.0 / 0.19 -- the band of its own entry +# (m3_sparse_decode): exact operands, FP16 probabilities, BF16 rounding. +# Each band leaves about 1.5x the measured maximum; a shifted or wrongly selected reference is far outside them. +_FP8_CONTEXT = dict(element_ulp=16.0, rms_ulp=9.0) +_FP8_DECODE = dict(element_ulp=12.0, rms_ulp=9.0) +_NVFP4_CONTEXT = dict(element_ulp=48.0, rms_ulp=16.0) +_NVFP4_DECODE = dict(element_ulp=2.0, rms_ulp=0.5) +_BANDS = { + ("dense", "context"): (_FP8_CONTEXT, None), + ("dense", "generation"): (None, _FP8_DECODE), + ("sparse", "context"): (_NVFP4_CONTEXT, None), + ("sparse", "mixed"): (_NVFP4_CONTEXT, _NVFP4_DECODE), +} + + +def _assert_in_bands(output, ref, bands, context_rows: int) -> None: + """Gate the context rows and the generation rows each on their own kernel's band.""" + context_band, decode_band = bands + if context_rows: + assert_within_ulp(output[:context_rows], ref[:context_rows], **context_band) + if context_rows < output.shape[0]: + assert_within_ulp(output[context_rows:], ref[context_rows:], **decode_band) + + +# (request rows, cached tokens per request, new tokens per request, context rows) per cell. +_SHAPES = { + ("dense", "context"): ((384,), (1024,), 1), + ("dense", "generation"): ((128, 2500), (1, 1), 0), + ("sparse", "context"): ((0,), (4096,), 1), + ("sparse", "mixed"): ((0, 2999, 4500), (512, 1, 1), 1), +} + + +def _ulp_stats(out: torch.Tensor, ref: torch.Tensor) -> tuple[float, float]: + """The two quantities ``assert_within_ulp`` gates, in BF16 ulps.""" + o, r = out.float(), ref.float() + scale = r.abs().amax(dim=1, keepdim=True).clamp_min(1e-9) + element = ((o - r).abs() / scale).max().item() / ULP[torch.bfloat16] + rms = ((o - r).pow(2).mean().sqrt() / r.pow(2).mean().sqrt()).item() / ULP[torch.bfloat16] + return element, rms + + +def _shift(x: torch.Tensor) -> torch.Tensor: + """``x``'s rows moved down by one (byte-wise, so E4M3 rows move too).""" + return x.view(torch.uint8).roll(1, 0).view(x.dtype) if x.element_size() == 1 else x.roll(1, 0) + + +def _shape(cell): + spec = cell.spec + cached, qlens, contexts = _SHAPES[spec["layer"], spec["phase"]] + if spec["phase"] == "context": + assert sum(qlens) == spec["tokens"] + elif spec["phase"] == "generation": + assert len(qlens) == spec["reqs"] + else: + assert sum(qlens[:contexts]) == spec["tokens"] and len(qlens) - contexts == spec["reqs"] + return cached, qlens, contexts + + +def _run_dense(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + for rid, c in zip(rids, cached): + if c: + hist = torch.randn(2, c, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + rt.write_fp8_rows(manager, rt.DENSE, rid, torch.arange(c), hist[0], hist[1]) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + # A prefill chunk arrives in BF16 (the hybrid cache's dense layers take FP8 Q / K / V only on decode steps). + dtype = torch.bfloat16 if contexts else torch.float8_e4m3fn + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + k = torch.randn(total, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + v = torch.randn(total, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to(dtype) + # The NVFP4 target passes its layer scales on every layer; an FP8 layer's write and FMHA take none. + inv = torch.tensor([1.0, 0.5, 2.0], device="cuda") + fwd = 1.0 / inv + return (q, k, v, None, inv, fwd), dict(rids=rids, cached=cached, qlens=qlens, contexts=contexts) + + +def _run_sparse(cell, gen, manager, backend, metadata): + cached, qlens, contexts = _shape(cell) + rids = list(range(len(qlens))) + rt.allocate(manager, rids, qlens, cached) + u0, u1 = rt.index_directions(gen) + lengths = [c + n for c, n in zip(cached, qlens)] + keys = [ + torch.randn(n, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths + ] + values = [ + torch.randn(n, rt.KV_HEADS * rt.HEAD_DIM, generator=gen, device="cuda") for n in lengths + ] + quant = rt.nvfp4_quant(torch.cat(keys), torch.cat(values)) + for rid, n, kk, vv in zip(rids, lengths, keys, values): + # Every row, the step's own included: the producer writes a sparse layer's caches before the call. + idx_k = rt.planted_index_keys(n, u0, u1, gen) + rt.write_nvfp4_rows( + manager, + metadata, + rt.SPARSE, + rid, + torch.arange(n), + kk.to(torch.bfloat16), + vv.to(torch.bfloat16), + idx_k, + quant, + ) + rt.prepare(metadata, rids, qlens, cached, contexts) + total = sum(qlens) + q = torch.randn(total, rt.HEADS * rt.HEAD_DIM, generator=gen, device="cuda").to( + torch.float8_e4m3fn + ) + # Query tokens alternate between the two planted directions, so neighbouring tokens select different blocks; the + # rank's index heads take alternating directions, so with more than one KV head they select different blocks too. + turn = torch.arange(total, device="cuda") % 2 + directions = torch.stack([u0, u1]) + idx_q = torch.cat([directions[(turn + h) % 2] for h in range(rt.INDEX_HEADS)], dim=1) + idx_q = idx_q.to(torch.float8_e4m3fn) + return (q, None, None, idx_q, quant, 1.0 / quant), dict( + rids=rids, cached=cached, qlens=qlens, contexts=contexts + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + sparse = spec["layer"] == "sparse" + gen = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + manager = rt.make_manager(draft=False) + try: + backend = rt.make_msa_backend(rt.SPARSE if sparse else rt.DENSE) + metadata = rt.make_msa_metadata(manager) + run = _run_sparse if sparse else _run_dense + (q, k, v, idx_q, inv, fwd), step = run(cell, gen, manager, backend, metadata) + output = torch.full( + (q.shape[0], rt.HEADS * rt.HEAD_DIM), float("nan"), dtype=torch.bfloat16, device="cuda" + ) + certifying(op) + result = op(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert result is output + ref = op.reference(backend, q, k, v, idx_q, metadata, output, sparse, inv, fwd) + assert torch.isfinite(output.float()).all(), "every live row is written" + bands = _BANDS[spec["layer"], spec["phase"]] + context_rows = sum(step["qlens"][: step["contexts"]]) + for name, rows in ( + ("context", slice(0, context_rows)), + ("generation", slice(context_rows, None)), + ): + if output[rows].numel(): + element, rms = _ulp_stats(output[rows], ref[rows]) + print( + f"\n[{cell.why}] {name} rows: max element {element:.2f} ulp, rms {rms:.3f} ulp" + ) + _assert_in_bands(output, ref, bands, context_rows) + + if k is not None: + # The call's own cache write: the live K / V rows land, as a plain E4M3 cast, at their positions. + pool = manager.get_buffers(rt.DENSE, "HND") + first = 0 + for rid, c, n in zip(step["rids"], step["cached"], step["qlens"]): + slots = rt.flat_slots(manager, rt.DENSE, rid, torch.arange(c, c + n)).long() + for role, src in ((0, k), (1, v)): + for h in range(rt.KV_HEADS): + got = pool[slots // rt.PAGE, role, h, slots % rt.PAGE] + want = src[first : first + n, h * rt.HEAD_DIM : (h + 1) * rt.HEAD_DIM] + want = want.to(torch.float8_e4m3fn) + assert torch.equal(got.view(torch.uint8), want.view(torch.uint8)) + first += n + + # The band discriminates: a reference whose query rows are shifted by one token is outside it. + shifted = op.reference(backend, _shift(q), k, v, idx_q, metadata, output, sparse, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, shifted, bands, context_rows) + if sparse: + # And the selection is load-bearing: attending every causal block instead is outside it too. + dense = op.reference(backend, q, k, v, idx_q, metadata, output, False, inv, fwd) + with pytest.raises(AssertionError): + _assert_in_bands(output, dense, bands, context_rows) + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py new file mode 100644 index 000000000000..b883a15acd38 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_n_valid_contract.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The per-token valid-block count is written only before the step's first layer. + +The long-row select (``ops/m3_select_long_dsl.py``, ``SELECT_N_VALID_PREWAIT``) reads its row's entry of +``MsaMetadata.msa_n_valid_blocks`` BEFORE its PDL grid-dependency wait. That is exact only while every write to the +buffer happens before the step's first layer: + +- the allocation (``_alloc_msa_proxy_scratch``); +- the host-to-device copy in ``prepare()`` (``_build_step_plans``); +- ``on_update_kv_lens``, which the runner calls before the forward. It writes through the plain + ``<<<>>>`` launch of ``msa_kv_update_kernel`` or through torch copies, so any PDL grid of the step starts only + after it completes. + +This test scans ``tensorrt_llm/_torch`` and fails on any other write site (including writes through a local alias of +the buffer), so a new writer has to revisit the pre-wait read. A buffer handed to an opaque kernel launch is not +tracked; such a launch inside a layer must be reviewed by hand. + +CPU only: it parses the sources. +""" + +import ast +from pathlib import Path + +import pytest + +import tensorrt_llm + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(tensorrt_llm.__file__).resolve().parent / "_torch" +_NAME = "msa_n_valid_blocks" +# In-place tensor methods that write their receiver. +_MUTATORS = { + "copy_", + "fill_", + "zero_", + "index_put_", + "index_copy_", + "scatter_", + "masked_fill_", + "add_", + "sub_", +} +# Calls whose first argument is the destination they write: ``step_inputs.copy_(dst, src)`` and the H2D helpers. +_COPY_HELPERS = {"copy_", "copy_h2d_if_changed", "copy_h2d_pool_shared"} +_BACKEND = "attention/backends/sparse/minimax_m3/msa_backend.py" +# (file relative to tensorrt_llm/_torch, enclosing function) -> writes allowed there +_ALLOWED = { + (_BACKEND, "_alloc_msa_proxy_scratch"), # the allocation + (_BACKEND, "_build_step_plans"), # the H2D copy in prepare() + (_BACKEND, "on_update_kv_lens"), # the device-side correction before the forward +} + + +def _base(node): + """The innermost receiver of an attribute / subscript / call chain.""" + while isinstance(node, (ast.Subscript, ast.Attribute, ast.Call)): + if isinstance(node, ast.Attribute) and node.attr == _NAME: + return node + node = node.func if isinstance(node, ast.Call) else node.value + return node + + +def _mentions(node, aliases) -> bool: + """Whether the expression's receiver chain is ``*.msa_n_valid_blocks`` or a local alias of it.""" + b = _base(node) + if isinstance(b, ast.Attribute): + return b.attr == _NAME + return isinstance(b, ast.Name) and (b.id == _NAME or b.id in aliases) + + +def _refers(node, aliases) -> bool: + """Whether any sub-expression (e.g. a branch of a conditional) names the buffer or an alias.""" + return any( + (isinstance(n, ast.Attribute) and n.attr == _NAME) + or (isinstance(n, ast.Name) and (n.id == _NAME or n.id in aliases)) + for n in ast.walk(node) + ) + + +def _write_sites(tree: ast.AST): + sites = [] + + def visit(node, func, aliases): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + func, aliases = node.name, set() + if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + value = node.value + for t in targets: + if isinstance(t, ast.Name) and value is not None and _refers(value, aliases): + aliases.add(t.id) # a local view of the buffer + elif ( + not isinstance(t, ast.Name) + and _mentions(t, aliases) + and not isinstance(node, ast.AnnAssign) + ): + sites.append((func, node.lineno)) + elif ( + isinstance(t, ast.Attribute) + and t.attr == _NAME + and isinstance(node, ast.AnnAssign) + ): + pass # a dataclass field declaration + if isinstance(node, ast.Call): + f = node.func + if isinstance(f, ast.Attribute) and f.attr in _MUTATORS and _mentions(f.value, aliases): + sites.append((func, node.lineno)) + name = ( + f.id + if isinstance(f, ast.Name) + else (f.attr if isinstance(f, ast.Attribute) else "") + ) + if name in _COPY_HELPERS and node.args and _mentions(node.args[0], aliases): + sites.append((func, node.lineno)) + for child in ast.iter_child_nodes(node): + visit(child, func, aliases) + + visit(tree, None, set()) + return sites + + +def test_only_the_pre_forward_paths_write_the_valid_block_count() -> None: + found = set() + for path in sorted(_ROOT.rglob("*.py")): + src = path.read_text() + if _NAME not in src: + continue + rel = path.relative_to(_ROOT).as_posix() + for func, line in _write_sites(ast.parse(src)): + assert (rel, func) in _ALLOWED, ( + f"{rel}:{line} ({func}) writes msa_n_valid_blocks; the long-row select reads it before its PDL grid " + "wait (ops/m3_select_long_dsl.py, SELECT_N_VALID_PREWAIT), so a writer inside the " + "step must be ordered before the select's grid or the read moved after the wait" + ) + found.add((rel, func)) + assert found == _ALLOWED, f"expected write sites {_ALLOWED}, found {found}" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py new file mode 100644 index 000000000000..124942104b82 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_producer entry: the pre-wait-staged NVFP4 QKV + indexer producer. + +Every cell is held to the native producer (``minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert``) bit for bit on q, +index-q and every K / V / scale / index-K cache byte, and to the fp32 torch reference: the FP8 outputs within one E4M3 +step, the decoded NVFP4 K / V within one E2M1 step of their block scale. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_producer import ( # noqa: E501 + FP4_BLOCK, + HEADS_Q, + PAGE, + decode_nvfp4_token, + scale_index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_producer import ( # noqa: E501 + m3_producer as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_producer requires a CUDA device" + +PAGES = 1024 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _caches(): + return ( + torch.zeros(PAGES, 2, 1, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, 1, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _inputs(tokens: int, scale: float, bad_slots: bool, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + weights = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + packed = ( + torch.randn(tokens, (HEADS_Q + 4) * 128, device="cuda", generator=g) * scale + ).bfloat16() + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return packed, slots, quant, weights, rotary, positions + + +def _args(inputs, caches, early: bool = True): + packed, slots, quant, weights, rotary, positions = inputs + return (packed, *caches, slots, quant, EPS, *weights, rotary, positions, HEADS_Q, True, early) + + +def _native(inputs, caches): + packed, slots, quant, weights, rotary, positions = inputs + return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( + packed, *caches, slots, quant, HEADS_Q, 1, 1, 128, 64, EPS, *weights, rotary, positions + ) + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + early = spec.get("early", True) + inputs = _inputs( + spec["tokens"], spec["scale"], spec.get("bad_slots", False), zlib.crc32(cell.why.encode()) + ) + stock_caches, caches = _caches(), _caches() + stock = _native(inputs, stock_caches) + args = _args(inputs, caches, early) + certifying(op) + ref = op.reference(*args) + out = op(*args) + torch.cuda.synchronize() + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, s, o in zip(names, (*stock, *stock_caches), (*out, *caches)): + assert torch.equal(s.view(torch.uint8), o.view(torch.uint8)), ( + f"{name} differs from the native producer" + ) + + _within_fp8(out[0], ref["q"], "q") + _within_fp8(out[1][:, 0], ref["index_q"], "index_q") + kv_data, kv_scale, index_k = caches + slots = inputs[1] + checked = 0 + for t in range(spec["tokens"]): + slot = int(slots[t]) + if slot < 0 or slot >= PAGES * PAGE: + continue + page, within = divmod(slot, PAGE) + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") + for plane, name in ((0, "k"), (1, "v")): + scales = kv_scale[page, plane, 0].reshape(-1)[scale_index(within, plane).cuda()] + dq = decode_nvfp4_token(kv_data[page, plane, 0, within], scales, KV_QUANT[1 + plane]) + r = ref[name][t].bfloat16().float() + # Round to nearest E2M1 under the block's E4M3-rounded scale s (in value units s / quant_scale): at most + # half the widest E2M1 step (one unit), plus the clipping at 6 units when s rounded down; a scale that + # underflows E4M3 zeroes the block. A slack of one E4M3 step covers the kernel's approximate reciprocals. + qs = KV_QUANT[1 + plane] + amax = r.view(-1, FP4_BLOCK).abs().amax(dim=1) + unit = (qs * amax / 6.0).to(torch.float8_e4m3fn).float() / qs + bound = torch.where( + unit > 0, 1.15 * unit + (amax - 6.0 * unit / 1.15).clamp_min(0.0), amax + ) + bound = bound.repeat_interleave(FP4_BLOCK) + 1e-12 + assert bool(((dq - r).abs() <= bound).all()), (name, t) + checked += 1 + assert checked >= spec["tokens"] - 2 + + +def test_early_trigger_does_not_change_the_bits() -> None: + inputs = _inputs(48, 1.0, False, 7) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=True)) + out_b = op(*_args(inputs, b, early=False)) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +@pytest.mark.parametrize( + "tokens, scale, bad_slots, early", + [ + (1, 1.0, False, True), + (4, 3.0, False, True), + (16, 1e-3, False, True), + (32, 40.0, True, True), + (48, 1.0, False, True), + (1000, 1.0, False, False), + ], +) +def test_prewait_step_inputs_is_bitwise_the_shared_op(tokens, scale, bad_slots, early) -> None: + """modeling's PRODUCER_PREWAIT_STEP_INPUTS: this target's copy (minimax_m3_sm107_tp8_producer_pw, the position ids, + slots and RoPE rows read before its grid wait) writes exactly the shared op's q, index-q and cache bytes.""" + assert hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_producer_pw"), ( + "the target's producer op is not loaded (its TP8 extension)" + ) + inputs = _inputs(tokens, scale, bad_slots, 11) + a, b = _caches(), _caches() + out_a = op(*_args(inputs, a, early=early)) + out_b = op(*_args(inputs, b, early=early), prewait_step_inputs=True) + for x, y in zip((*out_a, *a), (*out_b, *b)): + assert torch.equal(x.view(torch.uint8), y.view(torch.uint8)) + + +def test_guard_refuses_a_repeated_slot() -> None: + inputs = list(_inputs(4, 1.0, False, 3)) + inputs[1] = inputs[1].clone() + inputs[1][1] = inputs[1][0] + with pytest.raises(AssertionError, match="must not repeat a slot"): + with validating(op): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_operand_off_the_device() -> None: + inputs = list(_inputs(4, 1.0, False, 9)) + inputs[5] = inputs[5].cpu() # position ids on the host + with pytest.raises(RuntimeError, match="packed's CUDA device"): + op(*_args(tuple(inputs), _caches())) + + +def test_op_refuses_an_unaligned_index_k_page_stride() -> None: + """The index-K rows are written as 32-bit words: a page stride that is not a multiple of 4 is refused.""" + kv_data, kv_scale, _ = _caches() + backing = torch.zeros(PAGES * (128 * 128 + 1), dtype=torch.float8_e4m3fn, device="cuda") + index_k = backing.as_strided((PAGES, 1, PAGE, 128), (128 * 128 + 1, 128 * 128, 128, 1)) + with pytest.raises(RuntimeError, match="index_k"): + op(*_args(_inputs(4, 1.0, False, 13), (kv_data, kv_scale, index_k))) + + +def test_op_refuses_unaligned_weights_or_rope_table() -> None: + """The norm weights are read as 8-byte words and the RoPE rows as float4: offset views are refused.""" + packed, slots, quant, weights, rotary, positions = _inputs(4, 1.0, False, 17) + offset_weight = torch.zeros(129, dtype=torch.bfloat16, device="cuda")[1:] + with pytest.raises(RuntimeError, match="norm weights"): + op( + *_args( + (packed, slots, quant, [offset_weight, *weights[1:]], rotary, positions), _caches() + ) + ) + flat = torch.zeros(rotary.numel() + 1, device="cuda") + offset_rope = flat[1:].view(rotary.shape) + with pytest.raises(RuntimeError, match="rotary_cos_sin"): + op(*_args((packed, slots, quant, weights, offset_rope, positions), _caches())) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py new file mode 100644 index 000000000000..6e58bacee51c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_qk_norm_rope_fp8 entry of the MiniMax-M3 tp8 target: the out-of-place FP8 form of the fused +per-head Gemma RMSNorm + partial NeoX RoPE against the in-place kernel followed by the E4M3 cast. V is a cast, so it +matches bit for bit; Q and K round once from fp32 where the composition rounds through bf16 first, so each element is +within one E4M3 step of it.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 + m3_qk_norm_rope_fp8 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" + +HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 + + +def _args(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q, kv = spec["q"], spec["kv"] + qkv = torch.randn(spec["m"], (q + 2 * kv) * HEAD_DIM, generator=g, device="cuda").to( + torch.bfloat16 + ) + qw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + kw = (torch.randn(HEAD_DIM, generator=g, device="cuda") * 0.1).to(torch.bfloat16) + pos = torch.randint(0, 200000, (spec["m"],), generator=g, device="cuda", dtype=torch.int32) + return ( + qkv, + q, + kv, + kv, + HEAD_DIM, + ROTARY_DIM, + EPS, + qw, + kw, + THETA, + True, + pos, + 1.0, + 0.0, + 0.0, + 1.0, + True, + True, + False, + 0, + 0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) + before = args[0].clone() + with validating(op): + out = op(*args) + assert torch.equal(args[0], before), "the FP8 form must leave qkv untouched" + ref = op.reference(*args) + assert out.dtype == torch.float8_e4m3fn and out.shape == ref.shape + v_start = (cell.spec["q"] + cell.spec["kv"]) * HEAD_DIM + assert torch.equal(out[:, v_start:].view(torch.uint8), ref[:, v_start:].view(torch.uint8)) + assert _e4m3_steps(out[:, :v_start], ref[:, :v_start]) <= 1 + + +def _e4m3_steps(a, b) -> int: + """The largest distance between two E4M3 tensors in representable values (E4M3 codes order by magnitude).""" + + def signed(t): + c = t.view(torch.uint8).int() + # Sign-magnitude to a line: -0 and +0 both map to 0. + return torch.where(c >= 128, 128 - c, c) + + return int((signed(a) - signed(b)).abs().max()) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py new file mode 100644 index 000000000000..514567365f8c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the attention/m3_qkv_insert entry: the generic horizontal producer on an FP8 main cache. + +Q, index-Q and every written K / V / index-K cache row are held within one E4M3 step of the fp32 reference. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_qkv_insert import ( # noqa: E501 + m3_qkv_insert as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" + +PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 8, 1, 1, 128, 64 +MAX_POSITIONS = 262144 +EPS = 1e-6 + + +def _within_fp8(out: torch.Tensor, ref: torch.Tensor, name: str) -> None: + # One E4M3 step (3 mantissa bits) of the value, or its smallest subnormal step. + err = (out.float() - ref).abs() + assert bool((err <= ref.abs() * 2.0**-3 + 2.0**-9).all()), (name, float(err.max())) + + +def _args(tokens: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + tokens, (HEADS + 2 * KV + INDEX + 1) * D, device="cuda", generator=g + ).bfloat16() + kv = torch.zeros(PAGES, 2, KV, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + index_k = torch.zeros(PAGES, 1, PAGE, D, dtype=torch.float8_e4m3fn, device="cuda") + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + weights = [(torch.randn(D, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + rotary = torch.randn(MAX_POSITIONS, 2, ROTARY // 2, device="cuda", generator=g) + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + return ( + packed, + kv, + None, + index_k, + slots, + None, + HEADS, + KV, + INDEX, + D, + ROTARY, + EPS, + *weights, + rotary, + positions, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _args(cell.spec["m"], zlib.crc32(cell.why.encode())) + certifying(op) + ref = op.reference(*args) + q, index_q = op(*args) + torch.cuda.synchronize() + _within_fp8(q.view(ref["q"].shape), ref["q"], "q") + _within_fp8(index_q.view(ref["index_q"].shape), ref["index_q"], "index_q") + kv, index_k, slots = args[1], args[3], args[4] + for t in range(cell.spec["m"]): + page, within = divmod(int(slots[t]), PAGE) + _within_fp8(kv[page, 0, :, within], ref["k"][t], f"k token {t}") + _within_fp8(kv[page, 1, :, within], ref["v"][t], f"v token {t}") + _within_fp8(index_k[page, 0, within], ref["index_k"][t], f"index_k token {t}") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py new file mode 100644 index 000000000000..bf3c3c271148 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_select_blocks entry of the MiniMax-M3 tp8 target: MSA top-16 block selection with the forced +init / local blocks and per-row valid-block masks.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_blocks import ( # noqa: E501 + m3_select_blocks as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" + +TOPK, INIT, LOCAL = 16, 0, 1 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + n_blocks, total_q = spec["n_blocks"], spec["total_q"] + # Distinct scores (a permutation) so the selection has no ties to break. + scores = ( + torch.randperm(n_blocks * total_q, generator=g, device="cuda") + .float() + .view(1, n_blocks, total_q) + ) + if spec["mixed"]: + n_valid = torch.randint( + 17, n_blocks + 1, (total_q,), generator=g, device="cuda", dtype=torch.int32 + ) + else: + n_valid = torch.full((total_q,), n_blocks, dtype=torch.int32, device="cuda") + with validating(op): + out = op(scores, TOPK, n_valid, INIT, LOCAL, True) + ref = op.reference(scores, TOPK, n_valid, INIT, LOCAL, True) + assert out.shape == ref.shape + assert torch.equal(out, ref) + assert out.permute(1, 0, 2).is_contiguous(), "head-major backing" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py new file mode 100644 index 000000000000..945025a650ca --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp8 (Rubin) m3_select_long entry: top-16 MSA block selection for long decode rows, 1 KV head.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long import ( # noqa: E501 + m3_select_long as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( + select_blocks_from_maxscore, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key +# widths, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 2048, 2049, 4096, 4097, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # re-read path). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + # Every row: 10 winners above, ~1500 blocks tied at the 16th score, the rest below; the in-tree op stages up to + # 2048 boundary candidates, so it must break this tie to the smaller block like this op. + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + # Odd tokens: almost everything -inf, so picks run past the finite scores. + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +def _mass_tie_rows(scores, nvb, init_blocks, local_blocks): + """[total_q, num_kv_heads] mask of rows where more than 2048 valid blocks share the 16th-largest score.""" + num_kv_heads, num_blocks, total_q = scores.shape + blocks = torch.arange(num_blocks, device=scores.device) + raw = nvb[:total_q].long() + valid = raw.clamp(0, num_blocks) + forced = scores.permute(2, 0, 1).clone() # [total_q, heads, blocks] + forced = torch.where(blocks < init_blocks, 1.0e30, forced) + forced = torch.where(blocks >= (raw - local_blocks).clamp_min(0)[:, None, None], 1.0e29, forced) + in_row = (blocks[None, :] < valid[:, None])[:, None, :] + ranked = torch.where(in_row, forced, float("-inf")).sort(dim=-1, descending=True).values + kth = ranked[..., min(TOPK, num_blocks) - 1 : min(TOPK, num_blocks)] + return ((forced == kth) & in_row).sum(dim=-1) > 2048 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + + # Selection is an exact claim: the same blocks or a different attention. + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # And it is the in-tree selector's answer, bit for bit (the op the target falls back to) -- except on rows + # where more than 2048 blocks tie at the 16th-largest score: there the in-tree op does not break the tie to + # the smallest blocks, while this op and the reference do. + scores, nvb, _, init_blocks, local_blocks, _ = args + in_tree = select_blocks_from_maxscore( + scores, + topk=TOPK, + n_valid_blocks=nvb, + init_blocks=init_blocks, + local_blocks=local_blocks, + head_major_output=True, + ) + mass_tie = _mass_tie_rows(scores, nvb, init_blocks, local_blocks) + assert torch.equal(out[~mass_tie], in_tree[~mass_tie]), ( + "selected blocks differ from the in-tree minimax_m3_select_blocks" + ) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + """Releasing the dependent grid right after the select's own wait changes only when the next kernel may launch.""" + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + """Releasing the dependent grid once the row's scores are loaded changes only when the next kernel may launch.""" + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) + + +def test_guard_refuses_out_aliasing_scores() -> None: + scores, nvb, _, init_blocks, local_blocks, _ = _case(op.CELLS[0].spec, seed=3) + aliased = ( + scores.view(torch.int32) + .reshape(-1)[: scores.shape[2] * TOPK] + .view(scores.shape[2], 1, TOPK) + ) + with pytest.raises(AssertionError, match="alias"): + with validating(op): + op(scores, nvb, aliased, init_blocks, local_blocks, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py new file mode 100644 index 000000000000..4022b2cfb133 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py @@ -0,0 +1,191 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp8 (Rubin) m3_select_long_dsl entry: the long-row top-16 select from one 384-thread CTA per row, +against the reference and bit for bit against the 1024-thread m3_select_long on every cell.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( + m3_select_long_dsl as select_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long import ( # noqa: E501 + TOPK, + m3_select_long, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_long_dsl import ( # noqa: E501 + m3_select_long_dsl as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" + +# Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key +# widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). +_MIXED_VALID = (0, 5, 16, 17, 200, 1536, 1537, 4224, 4225, 8193, 9000, 1, 300, 6000, 129, 3000) + + +def _scores(spec, g): + nb, tq = spec["num_blocks"], spec["total_q"] + s = torch.randn(spec["num_kv_heads"], nb, tq, generator=g, device="cuda", dtype=torch.float32) + kind = spec["scores"] + if kind == "ties": + # Even tokens: a flat row with ten winners, so thousands of keys tie at the 16th-largest score (the + # whole-row re-rank). Odd tokens: scores on a coarse grid, so ties are broken by block index. + s[:, :, 0::2] = 0.25 + winners = torch.randint(0, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda") + s[:, :, 0::2].scatter_(1, winners[:, :, 0::2], 3.0) + s[:, :, 1::2] = (s[:, :, 1::2] * 2).round() / 2 + elif kind == "ties1500": + s = torch.full_like(s, -1.0) + s[:, :1500, :] = 0.25 + winners = torch.randint( + 1500, nb, (spec["num_kv_heads"], 10, tq), generator=g, device="cuda" + ) + s.scatter_(1, winners, 3.0) + elif kind == "nonfinite": + s[:, 7::50, :] = float("inf") + s[:, 11::70, :] = float("nan") + s[:, 3::40, :] = -0.0 + s[:, 4::40, :] = 0.0 + s[:, :, 1::2] = float("-inf") + s[:, 0:5, 1::2] = 1.0 + return s + + +def _valid(spec, g): + tq, nb = spec["total_q"], spec["num_blocks"] + if spec["valid"] == "mixed": + vals = [_MIXED_VALID[i % len(_MIXED_VALID)] for i in range(tq)] + else: + vals = (nb - torch.randint(0, 12, (tq,), generator=g, device="cuda")).tolist() + vals[0] = nb + return torch.tensor(vals, dtype=torch.int32, device="cuda") + + +def _case(spec, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + scores = _scores(spec, g) + nvb = _valid(spec, g) + out = torch.empty( + spec["num_kv_heads"], spec["total_q"], TOPK, dtype=torch.int32, device="cuda" + ).permute(1, 0, 2) + return scores, nvb, out, spec["init_blocks"], spec["local_blocks"], False + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + out = op(*args) + ref = op.reference(*args) + assert out.dtype is torch.int32 + assert torch.equal(out, ref), "selected blocks differ from the reference" + # Every row, mass ties included, is the 1024-thread kernel's answer. + scores, nvb, _, init_blocks, local_blocks, _ = args + long_out = torch.empty_like(out) + m3_select_long(scores, nvb, long_out, init_blocks, local_blocks, False) + assert torch.equal(out, long_out), "selected blocks differ from m3_select_long" + + +@pytest.mark.parametrize("layout", ["kv_blocks_q", "kv_q_blocks"]) +@pytest.mark.parametrize("valid", [17, 1294, 2524, 3907, 9216]) +def test_layouts_and_widths_match_the_1024_thread_select(layout, valid) -> None: + """Both score layouts the scorer can hand over, at a valid count in each key width.""" + g = torch.Generator(device="cuda").manual_seed(valid) + nb = max(valid, 129) + s = torch.randn(1, nb, 8, generator=g, device="cuda") + if layout == "kv_q_blocks": + s = s.permute(0, 2, 1).contiguous().permute(0, 2, 1) + nvb = torch.tensor([valid - (t % 2) for t in range(8)], dtype=torch.int32, device="cuda") + out = torch.empty(1, 8, TOPK, dtype=torch.int32, device="cuda").permute(1, 0, 2) + ref = torch.empty_like(out) + op(s, nvb, out, 0, 1, False) + m3_select_long(s, nvb, ref, 0, 1, False) + assert torch.equal(out, ref) + + +def test_pdl_launch_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=1) + pdl_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, False) + op(scores, nvb, pdl_out, init_blocks, local_blocks, True) + assert torch.equal(out, pdl_out) + + +def test_early_trigger_selects_the_same() -> None: + spec = op.CELLS[2].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=2) + early_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False) + op(scores, nvb, early_out, init_blocks, local_blocks, True, True) + assert torch.equal(out, early_out) + + +@pytest.mark.parametrize("cell", [0, 1, 2]) +def test_trigger_phase_selects_the_same(cell: int) -> None: + spec = op.CELLS[cell].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=3 + cell) + loaded_out = torch.empty_like(out) + op(scores, nvb, out, init_blocks, local_blocks, True, False, 0) + op(scores, nvb, loaded_out, init_blocks, local_blocks, True, False, 1) + assert torch.equal(out, loaded_out) + + +@pytest.mark.parametrize("early_trigger", [False, True]) +@pytest.mark.parametrize("cell", range(len(op.CELLS))) +def test_valid_count_read_before_the_wait_selects_the_same(cell: int, early_trigger: bool) -> None: + """ops/m3_select_long_dsl.SELECT_N_VALID_PREWAIT: the kernel that reads its row's valid-block count before the + grid wait selects exactly the blocks of the one that reads it after, on every certified cell (padding, clamped and + tied rows).""" + phase = 0 if early_trigger else 1 + spec = op.CELLS[cell].spec + scores, nvb, before, init_blocks, local_blocks, _ = _case(spec, seed=11 + cell) + after = torch.empty_like(before) + select_kernels._compile((True, early_trigger, phase, True))( + scores, nvb, before, init_blocks, local_blocks + ) + select_kernels._compile((True, early_trigger, phase, False))( + scores, nvb, after, init_blocks, local_blocks + ) + assert torch.equal(before, after) + + +def test_graph_replay_selects_the_same() -> None: + spec = op.CELLS[1].spec + scores, nvb, out, init_blocks, local_blocks, _ = _case(spec, seed=9) + ref = op.reference(scores, nvb, out, init_blocks, local_blocks) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(scores, nvb, out, init_blocks, local_blocks, True, False, 1) + for _ in range(3): + out.fill_(-7) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, ref) + + +def test_op_refuses_a_row_past_its_width() -> None: + scores = torch.zeros(1, 9217, 1, device="cuda") + nvb = torch.ones(1, dtype=torch.int32, device="cuda") + out = torch.empty(1, 1, TOPK, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="at most 9216"): + op(scores, nvb, out, 0, 1, False) + + +def test_guard_refuses_negative_forcing() -> None: + args = list(_case(op.CELLS[0].spec, seed=2)) + args[4] = -1 + with pytest.raises(AssertionError, match="non-negative"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py new file mode 100644 index 000000000000..36dbdd14d592 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py @@ -0,0 +1,626 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode entry (the target's Triton block-sparse decode and its routes) at the TP8 heads. + +Besides the band of every route around a float64 reference, the Triton merge's MXFP8 o_proj hand-off (published to +``m3_prequant.SLOT``) is held to ``torch.ops.trtllm.mxfp8_quantize(output, True)`` bit for bit: on identity steps and on +selected steps with the long-context range killed (the routes the o_proj consumes when LC8 does not run). +""" + +import zlib +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( + modeling, # noqa: E501 +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_attn_lc8, + m3_prequant, + m3_report, + triton_sparse_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_sparse_decode import ( # noqa: E501 + m3_sparse_decode as op, +) +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( + minimax_m3_sparse_attn_decode as stock_decode, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import ( # noqa: E402 — needs the path declared above + HEAD_DIM, + NUM_HEADS, + build_step, +) +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +# Band, in BF16 ulps (2^-8) of each [8, 128] row's largest magnitude, and of the relative RMS. Every route +# computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference +# is the BF16 rounding on both sides (one ulp of the row at most), the FP16 probabilities before PV, the split-K +# partials kept in FP32 and merged by their log2-sum-exp, and approximate exp2/log2: 2 ulps per element, the two +# roundings (each at most half an ulp) dominating the RMS, so half an ulp. (GB300 TP4: at most 1.80 ulps per element +# and 0.20 ulps RMS over the cells.) +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _call(case, output, identity): + return op( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + identity, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = build_step( + spec["seq_lens"], + spec["decode_query_len"], + identity=spec["identity"], + seed=zlib.crc32(cell.why.encode()) % 2**31, + ) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + certifying(op) + _call(case, output, spec["identity"]) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + spec["identity"], + ) + assert torch.isfinite(output.float()).all(), "padding rows must be zeros, not NaN" + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_flag_does_not_change_the_answer() -> None: + """On an identity table the identity route (the per-row producer) and the selected-table route (the + long-context kernel) compute the same attention.""" + case = build_step((2048, 130), 4, identity=True, seed=17) + outs = [] + for identity in (True, False): + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, identity) + outs.append(output) + assert_within_ulp(outs[0], outs[1], element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_guard_refuses_a_padded_selection_slot() -> None: + """The Triton producers index the block table with every one of a token's first min(16, blocks) ids, + -1 included.""" + case = build_step((1500,), 1, identity=True, seed=23) + case.topk[0, 0, 3] = -1 + output = torch.empty(1, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="blocks of the request"): + with validating(op): + _call(case, output, True) + + +def test_guard_refuses_a_false_identity_promise() -> None: + """identity_blocks lets the multi-query producer read one table for the whole request.""" + case = build_step((1500, 1600, 1700, 1800, 1500, 1600, 1700, 1800), 4, identity=True, seed=29) + case.topk[0, 5, :2] = case.topk[0, 5, :2].flip(0) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + with pytest.raises(AssertionError, match="identity_blocks promises"): + with validating(op): + _call(case, output, True) + + +def _swizzled_scale_offsets(rows: int, width: int) -> torch.Tensor: + """Offsets of rows [0, rows)'s UE8M0 scales in the 128x4-swizzled layout of mxfp8_quantize(x, True).""" + k_sf = width // 32 + k_tiles = (k_sf + 3) // 4 + r = torch.arange(rows, device="cuda").view(-1, 1) + k = torch.arange(k_sf, device="cuda").view(1, -1) + offsets = ( + (r // 128) * k_tiles * 512 + (k // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + k % 4 + ) + return offsets.reshape(-1) + + +def _mxfp8_handoff_matches_stock(output) -> None: + """The real rows' E4M3 bytes and UE8M0 scales equal mxfp8_quantize's. The merge writes only the step's rows into + its persistent 128-row scale tile; the padding rows' scales are never read by the o_proj and are not compared.""" + rows = output.shape[0] + flat = output.view(rows, -1) + taken = m3_prequant.SLOT.take(flat) + assert taken is not None, "the route published no MXFP8 hand-off for its output" + q8, sf = taken + stock_q, stock_sf = torch.ops.trtllm.mxfp8_quantize(flat, True) + assert torch.equal( + q8.view(torch.uint8).reshape(rows, -1), stock_q.view(torch.uint8).reshape(rows, -1) + ) + at = _swizzled_scale_offsets(rows, flat.shape[1]) + assert torch.equal( + sf.view(torch.uint8).reshape(-1)[at], stock_sf.view(torch.uint8).reshape(-1)[at] + ) + + +@pytest.mark.parametrize("seq_lens", [(1900,), (2048, 0), (700, 1500)]) +def test_identity_merge_mxfp8_handoff_is_stock_mxfp8_quantize(seq_lens) -> None: + """The Triton merge writes the o_proj hand-off (E4M3 rows, 128x4-swizzled UE8M0 scales) as mxfp8_quantize.""" + case = build_step(seq_lens, 4, identity=True, seed=31 + len(seq_lens)) + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(case, output, True) + _mxfp8_handoff_matches_stock(output) + + +@pytest.mark.parametrize("seq_lens", [(40000,), (9000, 40000)]) +def test_killed_long_context_route(seq_lens) -> None: + """With SPARSE_DECODE_LC8 killed a selected table runs the Triton producer + merge: in the band of the reference, + and its MXFP8 hand-off is mxfp8_quantize of the rows it wrote.""" + case = build_step(seq_lens, 4, identity=False, seed=41 + len(seq_lens)) + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + output = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + _call(case, output, False) + _mxfp8_handoff_matches_stock(output) + finally: + triton_sparse_decode.LC8_KEYS = landed + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_identity_route_is_the_certified_producer_range() -> None: + """Every identity step at the TP8 heads (one KV head) takes the per-row producer below 17 rows and the multi-query + producer from 17 rows; modeling.SPARSE_DECODE_PRODUCER covers both, and the cells certify both.""" + assert modeling.SPARSE_DECODE_PRODUCER.covers("per_row") + assert modeling.SPARSE_DECODE_PRODUCER.covers("multi_query") + seen = set() + for cell in op.CELLS: + keys = tuple(modeling._sparse_decode_producer(cell.spec)) + rows = cell.spec["decode_query_len"] * len(cell.spec["seq_lens"]) + want = (("multi_query" if rows >= 17 else "per_row"),) if cell.spec["identity"] else () + assert keys == want, (cell.why, keys) + seen.update(keys) + assert seen == {"per_row", "multi_query"} + + +_IDENTITY_LENS = (1900, 2048, 700, 1500, 1200, 2000, 300, 1024) +_SELECTED_LENS = (9000, 40000, 20000, 5000, 7000, 12000, 30000, 3000) + + +def _stock(case, output): + stock_decode( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + sm_scale=case.sm_scale, + output=output, + decode_query_len=case.decode_query_len, + k_block_scale=case.k_block_scale, + v_block_scale=case.v_block_scale, + k_global_scale=case.k_global_scale, + v_global_scale=case.v_global_scale, + ) + + +@pytest.mark.parametrize("identity", [True, False], ids=["identity", "selected_lc_killed"]) +@pytest.mark.parametrize("requests", [1, 2, 4, 8]) +def test_triton_route_is_bitwise_the_stock_decode(requests, identity) -> None: + """The target's per-row Triton producer + split-K merge writes the stock in-tree MSA decode's rows bit for bit: on + SM107 it takes stock's split-K factor at every row count (4 / 8 / 16 / 32 verify rows here). An identity step of + at least SPARSE_DECODE_MQ_MIN_TOKENS rows takes the multi-query producer (one KV head), which stock does not run: + exact-class instead, in the band of the float64 reference and the same bits on every run.""" + lens = (_IDENTITY_LENS if identity else _SELECTED_LENS)[:requests] + case = build_step(lens, 4, identity=identity, seed=50 + requests + 100 * identity) + rows = case.q.shape[0] + multi_query = identity and rows >= modeling.SPARSE_DECODE_MQ_MIN_TOKENS + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + again = torch.empty_like(target) + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + _call(case, target, identity) + _call(case, again, identity) + finally: + triton_sparse_decode.LC8_KEYS = landed + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), again.view(torch.int16)), "run to run" + if not multi_query: + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)) + return + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + case.topk, + case.block_table, + case.seq_lens, + case.sm_scale, + target, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + True, + ) + assert_within_ulp(target, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def test_plain_route_is_bitwise_the_stock_decode() -> None: + """Above MergeMxfp8.MAX_ROWS (here 160 verify rows) the Triton producer + merge writes no hand-off; with + SPARSE_DECODE_LC8 killed (its "plain" key covers these rows) it writes the stock decode's rows bit for bit.""" + case = build_step((3000,) * 40, 4, identity=False, seed=77) + target = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + _call(case, target, False) + finally: + triton_sparse_decode.LC8_KEYS = landed + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)) + + +@pytest.mark.parametrize("identity", [True, False], ids=["triton_identity", "lc8_selected"]) +def test_handoff_kill_publishes_nothing_and_keeps_the_rows(identity) -> None: + """MERGE_MXFP8_HANDOFF False: no route publishes the MXFP8 hand-off (the o_proj then quantizes its BF16 input), and + the attention rows are the same bits as with the hand-off.""" + case = build_step((1900,) if identity else (40000,), 4, identity=identity, seed=91 + identity) + on = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + off = torch.empty_like(on) + _call(case, on, identity) + assert m3_prequant.SLOT.take(on.view(on.shape[0], -1)) is not None + triton_sparse_decode.MERGE_MXFP8_HANDOFF = False + try: + _call(case, off, identity) + assert m3_prequant.SLOT.take(off.view(off.shape[0], -1)) is None + finally: + triton_sparse_decode.MERGE_MXFP8_HANDOFF = True + assert torch.equal(on.view(torch.int16), off.view(torch.int16)) + + +def _short_lens(requests: int, g: torch.Generator) -> tuple: + """Short contexts of 1..16 blocks of 128 tokens (the GSM / identity-route regime), a few exactly on a block edge.""" + lens = torch.randint(4, 16 * 128 + 1, (requests,), generator=g).tolist() + if requests > 1: + lens[0] = 128 * int(torch.randint(1, 17, (1,), generator=g)) + return tuple(lens) + + +@pytest.mark.parametrize("seed", range(12)) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_short_table_stress_is_bitwise_the_stock_decode(requests, seed) -> None: + """Short contexts (1..16 blocks) at 1..4 requests x 4 verify rows, the identity route and a random selected route + with -1 tails past each token's blocks: the target's Triton route writes the stock decode's rows bit for bit.""" + g = torch.Generator().manual_seed(1000 * requests + seed) + for identity in (True, False): + case = build_step(_short_lens(requests, g), 4, identity=identity, seed=7 * seed + requests) + target = torch.empty( + case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + stock = torch.empty_like(target) + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + _call(case, target, identity) + finally: + triton_sparse_decode.LC8_KEYS = landed + _stock(case, stock) + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + identity, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("fault", ["neg_in_window", "id_past_row", "page_out_of_pool"]) +def test_invalid_slot_fails_closed(fault) -> None: + """A corrupt decode input (a -1 inside a token's block window, an id past the block-table row, or a block-table + page outside the pool) is read as an absent block instead of a wild KV read: the rows equal the reference over the + remaining blocks.""" + case = build_step((1500, 2048, 700), 4, identity=False, seed=303) + tok, slot = 5, 2 + if fault == "neg_in_window": + case.topk[0, tok, slot] = -1 + elif fault == "id_past_row": + case.topk[0, tok, slot] = case.block_table.shape[1] + 7 + else: + req = tok // 4 + blk = int(case.topk[0, tok, slot]) + bad_page = case.k_paged.shape[0] + 1000 + case.block_table = case.block_table.clone() + case.block_table[req, blk] = bad_page + # Every token of the request that selected that block loses it. + hits = (case.topk[:, req * 4 : req * 4 + 4] == blk).nonzero().tolist() + output = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + landed = triton_sparse_decode.LC8_KEYS + triton_sparse_decode.LC8_KEYS = frozenset() + try: + _call(case, output, False) + finally: + triton_sparse_decode.LC8_KEYS = landed + torch.cuda.synchronize() + # The reference walks each token's first min(16, blocks) entries and drops -1: the corrupt slots as -1. + ref_topk = case.topk.clone() + if fault == "page_out_of_pool": + for h, t, k in hits: + ref_topk[h, req * 4 + t, k] = -1 + else: + ref_topk[0, tok, slot] = -1 + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + ref_topk, + case.block_table, + case.seq_lens, + case.sm_scale, + output, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert torch.isfinite(output.float()).all() + assert_within_ulp(output, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +def _identity_table( + seq_lens: torch.Tensor, decode_query_len: int, heads: int, lag: int = 0 +) -> torch.Tensor: + """The identity selection of every verify token, staged from lengths `lag` tokens behind seq_lens.""" + seq = (seq_lens.cpu().to(torch.int64) - lag).clamp(min=0).repeat_interleave(decode_query_len) + intra = torch.arange(seq.shape[0]) % decode_query_len + n_valid = ((seq - decode_query_len + intra + 1).clamp(min=0) + 127) // 128 + ar = torch.arange(16) + table = torch.where(ar[None, :] < n_valid[:, None], ar[None, :], -1).to(torch.int32) + return table.unsqueeze(0).expand(heads, -1, -1).contiguous().cuda() + + +# Each crosses a 128-token edge within its four verify tokens or between the stale and the true lengths. +_EDGE_LENS = (1284, 1287, 1281, 129, 132, 2048, 1028, 5, 643, 1920, 256, 1) + + +@pytest.mark.parametrize("fault", ["stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_identity_route_reads_its_own_lengths_not_the_table(requests, fault) -> None: + """The identity route derives block i for slot i from the kv lengths the kernel reads: a staged table that lags + the lengths by a verify step (a 128-token edge crossed), or is garbage, gives the stock decode's rows over the + true identity selection bit for bit.""" + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=400 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + assert torch.equal(case.topk, true_table), ( + "build_step's identity table is the identity selection" + ) + if fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + stock = torch.empty_like(target) + _stock(case, stock) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + # The Triton identity route (triton_sparse_decode.LC8_IDENTITY off); the LC8 identity route has its own test below. + triton_sparse_decode.LC8_IDENTITY = False + try: + _call(case_staged, target, True) + finally: + triton_sparse_decode.LC8_IDENTITY = True + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), stock.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) + + +@pytest.mark.parametrize("fault", ["none", "stale_one_step", "all_minus_one", "random_ids"]) +@pytest.mark.parametrize("requests", [1, 2, 3, 4]) +def test_lc8_identity_route_reads_its_own_lengths_not_the_table( + requests, fault, monkeypatch +) -> None: + """triton_sparse_decode.LC8_IDENTITY: an identity step runs the LC8, which takes block i for CTA i from the kv + lengths it reads. Whatever the staged table holds (true, lagging a verify step, all -1, garbage), the rows and the + MXFP8 hand-off are the LC8's on the true identity table (the selected-table route) bit for bit, and in the band of + the float64 reference.""" + lens = tuple(_EDGE_LENS[(requests * 5 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=700 + requests) + heads = case.topk.shape[0] + true_table = _identity_table(case.seq_lens, 4, heads) + if fault == "none": + staged = true_table.clone() + elif fault == "stale_one_step": + staged = _identity_table(case.seq_lens, 4, heads, lag=4) + elif fault == "all_minus_one": + staged = torch.full_like(true_table, -1) + else: + g = torch.Generator().manual_seed(requests) + staged = ( + torch.randint(-1, case.block_table.shape[1], true_table.shape, generator=g) + .to(torch.int32) + .cuda() + ) + rows = case.q.shape[0] + target = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + selected = torch.empty_like(target) + case_staged = SimpleNamespace(**{**vars(case), "topk": staged}) + case_true = SimpleNamespace(**{**vars(case), "topk": true_table}) + taken = [] + monkeypatch.setattr(m3_report, "taken", lambda path, rows, what="": taken.append(path)) + _call(case_staged, target, True) + monkeypatch.undo() + # The LC8 engagement, plus its one-division merge's (m3_attn_lc8.ONE_DIV_MERGE) when on. + assert sorted(taken) == sorted( + [triton_sparse_decode.LC8_PATH] + + ([triton_sparse_decode.ONE_DIV_MERGE_PATH] if m3_attn_lc8.ONE_DIV_MERGE else []) + ), taken + target_handoff = m3_prequant.SLOT.take(target.view(rows, -1)) + _call(case_true, selected, False) + selected_handoff = m3_prequant.SLOT.take(selected.view(rows, -1)) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert torch.equal(target.view(torch.int16), selected.view(torch.int16)), ( + fault, + case.seq_lens.tolist(), + ) + assert (target_handoff is None) == (selected_handoff is None) + if target_handoff is not None: + for a, b in zip(target_handoff, selected_handoff): + assert torch.equal(a.view(torch.uint8), b.view(torch.uint8)) + ref = op.reference( + case.q, + case.k_paged, + case.v_paged, + true_table, + case.block_table, + case.seq_lens, + case.sm_scale, + target, + case.decode_query_len, + None, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + None, + False, + ) + assert_within_ulp(target, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + + +@pytest.mark.parametrize("requests", [1, 2]) +def test_lc8_identity_step_reads_q_after_the_wait(requests, monkeypatch) -> None: + """An identity step has no select, so the LC8 must read q after its grid wait (783877: a route that skipped it + read q before the producer wrote it): the op gets identity and identity_step on every identity step, and the rows + are the same call after call.""" + lens = tuple(_EDGE_LENS[(requests * 3 + k) % len(_EDGE_LENS)] for k in range(requests)) + case = build_step(lens, 4, identity=True, seed=900 + requests) + true_table = _identity_table(case.seq_lens, 4, case.topk.shape[0]) + case_true = SimpleNamespace(**{**vars(case), "topk": true_table}) + rows = case.q.shape[0] + seen = [] + real = m3_attn_lc8.m3_sparse_decode_lc8 + + def spy(*args, **kwargs): + seen.append((kwargs.get("identity", False), kwargs.get("identity_step", False))) + return real(*args, **kwargs) + + monkeypatch.setattr(m3_attn_lc8, "m3_sparse_decode_lc8", spy) + out = torch.empty(rows, NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + ref = torch.empty_like(out) + _call(case_true, out, True) + _call(case_true, ref, True) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert seen[0] == (True, True), seen + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)) + + +@pytest.mark.parametrize("requests", [2, 3]) +def test_lc_fails_closed_on_ids_past_the_row(requests) -> None: + """A selected id past the token's row (a stale table, e.g. a padded row keeping a 500K request's n_valid) is + skipped, as the Triton producer skips it: the rows equal those with that slot at -1, bit for bit, and stay finite + even though the block-table tail it would reach points at a page whose block scales are NaN (0xFF).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_prequant, + ) + + lens = (165521, 700, 3000)[:requests] + case = build_step(lens, 4, identity=False, seed=1300 + requests) + dql = case.decode_query_len + good = case.topk.clone() + bad = case.topk.clone() + for tok in range(case.q.shape[0]): + req, qoff = divmod(tok, dql) + kv_len = max(int(case.seq_lens[req]) - dql + qoff + 1, 0) + nblk = (kv_len + 127) // 128 + past = nblk + 1 + if 2 <= nblk < 16 and past < case.block_table.shape[1]: + bad[:, tok, nblk - 1] = past + good[:, tok, nblk - 1] = -1 + page = int(case.block_table[req, past]) + case.k_block_scale[page].fill_(0xFF) + case.v_block_scale[page].fill_(0xFF) + assert not torch.equal(bad, good) + outs = [] + for table in (bad, good): + out = torch.empty(case.q.shape[0], NUM_HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda") + _call(SimpleNamespace(**{**vars(case), "topk": table}), out, False) + outs.append(out) + m3_prequant.SLOT.clear() + torch.cuda.synchronize() + assert torch.isfinite(outs[0].float()).all() + assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py new file mode 100644 index 000000000000..ac61687d6968 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_sparse_decode_lc8 entry (one-launch long-context sparse decode at the TP8 8 query heads, +``torch.ops.trtllm.minimax_m3_tp8_sparse_decode_lc8``). + +Every cell is held to the float64 reference within the derived band, its MXFP8 hand-off to the recipe bit for bit, and +the two ``prewait`` variants (staging before or after the grid wait) to each other bit for bit. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_sparse_decode_lc8 import ( # noqa: E501 + NUM_HEADS, + WIDTH, + dequant_nvfp4_pages, + mxfp8_handoff_reference, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_sparse_decode_lc8 import ( # noqa: E501 + m3_sparse_decode_lc8 as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_nvfp4_cache import NUM_HEADS as CASE_HEADS # noqa: E402 — needs the path declared above +from _m3_nvfp4_cache import build_step # noqa: E402 +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +assert torch.cuda.is_available(), "m3_sparse_decode_lc8 requires a CUDA device" + +#: Hand-off buffers are sized as MERGE_MXFP8 sizes them. +_HANDOFF_ROWS = 128 + +# Band, in BF16 ulps (2^-8) of each output row's largest magnitude, and of the relative RMS. +# +# The reference is float64 attention over the exactly dequantized blocks, rounded once to BF16. The kernel's +# operands are exact too (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so what separates the two is: the final +# BF16 rounding on both sides (up to one ulp of the element, at most one of the row's scale), the softmax +# probabilities rounded to FP16 before PV (2^-11 relative each, averaging down over the block), ex2.approx / +# lg2.approx (~2^-22) and FP32 accumulation order. Hence 2 ulps per element; the RMS is dominated by the two +# roundings (each at most half an ulp, about 0.3 in aggregate), so half an ulp. None of these depends on the head +# count or on the MMA orientation: the tp4 target's 16-head kernel measured at most 1.91 ulps per element and 0.25 +# ulps RMS on GB300 over the same cells. +_ELEMENT_ULP = 2.0 +_RMS_ULP = 0.5 + + +def _args(case, out, q8, sf, linear_softmax, prewait, num_sf_k_tiles=WIDTH // 128): + return ( + case.q.view(torch.uint8), + case.k_paged, + case.v_paged, + case.k_block_scale, + case.v_block_scale, + case.k_global_scale, + case.v_global_scale, + case.topk[0], + case.block_table, + case.seq_lens, + out, + q8, + sf, + num_sf_k_tiles, + case.decode_query_len, + case.sm_scale, + linear_softmax, + False, + prewait, + ) + + +def _buffers(total_q, mxfp8): + out = torch.empty(total_q, WIDTH, dtype=torch.bfloat16, device="cuda") + if not mxfp8: + return out, None, None + q8 = torch.zeros(_HANDOFF_ROWS, WIDTH, dtype=torch.uint8, device="cuda") + sf = torch.zeros(_HANDOFF_ROWS * (WIDTH // 32), dtype=torch.uint8, device="cuda") + return out, q8, sf + + +def _case(spec, seed): + return build_step(spec["seq_lens"], spec["decode_query_len"], identity=False, seed=seed) + + +def test_the_cases_are_the_tp8_geometry() -> None: + assert CASE_HEADS == NUM_HEADS == 8 and WIDTH == 1024 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + spec = cell.spec + case = _case(spec, zlib.crc32(cell.why.encode()) % 2**31) + total_q = case.q.shape[0] + out, q8, sf = _buffers(total_q, spec["mxfp8"]) + args = _args(case, out, q8, sf, spec["linear_softmax"], spec["prewait"]) + certifying(op) + op(*args) + ref = op.reference(*args) + assert_within_ulp(out, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + # Padding rows attend nothing and must be exact zeros, not NaN. + assert torch.isfinite(out.float()).all() + if spec["mxfp8"]: + # The hand-off is a deterministic function of the op's own BF16 rows: exact. + ref_q8, ref_sf = mxfp8_handoff_reference(out, _HANDOFF_ROWS) + assert torch.equal(q8[:total_q], ref_q8), ( + "MXFP8 rows differ from the recipe applied to the BF16 rows" + ) + assert torch.equal(sf, ref_sf), "UE8M0 scales differ from the recipe / swizzle" + + # The other staging variant: the same arithmetic, so the same bits (rows and hand-off). + other_out, other_q8, other_sf = _buffers(total_q, spec["mxfp8"]) + op(*_args(case, other_out, other_q8, other_sf, spec["linear_softmax"], not spec["prewait"])) + assert torch.equal(other_out.view(torch.int16), out.view(torch.int16)), ( + "prewait and plain staging wrote different rows" + ) + if spec["mxfp8"]: + assert torch.equal(other_q8, q8) and torch.equal(other_sf, sf) + + +def test_reference_dequant_matches_the_quantizer() -> None: + """The reference reads the paged bytes through its own layout claims (K scales token-major, V scales in + 4x4 token quads). Hold them to the quantizer's own flat output: every slot was written in order, so slot + ``s`` is page ``s // 128``, token ``s % 128``, and the flat scales are token-major by construction.""" + case = build_step((3000,), 1, identity=False, seed=5) + pages = torch.arange(case.k_paged.shape[0], device="cuda") + views = ((case.k_paged, case.k_block_scale, False), (case.v_paged, case.v_block_scale, True)) + for role, (data, scale, quad) in enumerate(views): + flat_data, flat_scale = torch.ops.trtllm.fp4_quantize( + case.source[role].reshape(case.num_slots, 1, -1).contiguous(), + case.quant[role + 1 : role + 2], + 16, + False, + False, + ) + expected = dequant_nvfp4_pages( + flat_data.view(torch.uint8).reshape(-1, 1, 128, 64), + flat_scale.view(torch.uint8).reshape(-1, 1, 128, 8), + pages, + quad_scales=False, + ) + got = dequant_nvfp4_pages(data, scale, pages, quad_scales=quad) + assert torch.equal(got, expected), ( + f"role {role}: paged dequantization disagrees with the quantizer" + ) + + +def test_the_variants_agree() -> None: + """The two softmax recurrences, the hand-off switch and the staging must not change the BF16 rows beyond + rounding; the hand-off switch and the staging not at all.""" + case = build_step((5000, 12345), 4, identity=False, seed=11) + outs = {} + for linear in (False, True): + for mxfp8 in (False, True): + for prewait in (False, True): + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op(*_args(case, out, q8, sf, linear, prewait)) + outs[(linear, mxfp8, prewait)] = out + base = outs[(False, False, False)] + for key, other in outs.items(): + assert_within_ulp(other, base, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + same_recurrence = outs[(key[0], False, False)] + torch.testing.assert_close(other, same_recurrence, rtol=0, atol=0) + + +def test_guard_refuses_a_block_past_the_request() -> None: + """A selected id at or past the request's block count reads a page the token cannot attend; with every + position masked the CTA's softmax is NaN.""" + case = build_step((3000,), 1, identity=False, seed=4) + case.topk[0, 0, 0] = 40 + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="block count"): + with validating(op): + op(*_args(case, out, q8, sf, False, True)) + + +def test_guard_refuses_strided_page_rows() -> None: + """The kernel walks a page's token rows at a fixed 64-byte pitch; a padded row stride is misread.""" + case = build_step((3000,), 1, identity=False, seed=6) + wide = torch.zeros(*case.k_paged.shape[:3], 80, dtype=torch.uint8, device="cuda")[..., :64] + wide.copy_(case.k_paged) + case.k_paged = wide + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="64-byte rows"): + with validating(op): + op(*_args(case, out, q8, sf, False, True)) + + +def test_guard_refuses_a_narrow_handoff() -> None: + """The hand-off rows are 1024 wide (8 scale k-tiles); a smaller count overlaps them.""" + case = build_step((3000,), 1, identity=False, seed=7) + out, q8, sf = _buffers(1, True) + with pytest.raises(AssertionError, match="num_sf_k_tiles"): + with validating(op): + op(*_args(case, out, q8, sf, False, True, num_sf_k_tiles=4)) + + +@pytest.mark.parametrize("prewait", [True, False], ids=["prewait", "plain"]) +@pytest.mark.parametrize("mxfp8", [True, False], ids=["handoff", "no_handoff"]) +def test_run_to_run_deterministic(prewait, mxfp8) -> None: + """The exact-class reference for the TP8 sparse decode (band vs fp64, as the TP4 LC): every launch of the same + step writes the same bits, in eager repeats and in a CUDA-graph replay.""" + spec = dict(seq_lens=(165521, 5000), decode_query_len=4) + case = _case(spec, seed=91) + for linear in (False, True): + runs = [] + for _ in range(3): + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op(*_args(case, out, q8, sf, linear, prewait)) + runs.append((out, q8, sf)) + g_out, g_q8, g_sf = _buffers(case.q.shape[0], mxfp8) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(*_args(case, g_out, g_q8, g_sf, linear, prewait)) + graph.replay() + torch.cuda.synchronize() + runs.append((g_out, g_q8, g_sf)) + base = runs[0] + for other in runs[1:]: + assert torch.equal(other[0].view(torch.int16), base[0].view(torch.int16)), ( + linear, + prewait, + mxfp8, + ) + if mxfp8: + assert torch.equal(other[1], base[1]) and torch.equal(other[2], base[2]), ( + linear, + prewait, + ) + + +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + [((165521,), 4), ((5000, 12345), 4), ((700,), 1), ((200000, 3000, 64, 70000), 4)], + ids=["1x165k", "2req", "short-1row", "4req-incl-64"], +) +def test_merge_once_is_bitwise_the_image_op(seq_lens, decode_query_len) -> None: + """ops/m3_attn_lc8.MERGE_ONCE: the target's op (merge weights formed once per CTA) writes exactly the + image op's rows and hand-off, at both softmax recurrences and both staging variants, including requests with fewer + than 16 selected blocks (empty chunks, -inf LSE). With SPARSE_DECODE_LC8_PREFETCH_OPROJ it also pulls two ranges + the size of the TP8 o_proj weight and scales (one not a multiple of 16 bytes) into L2: the same bits again. With + SPARSE_DECODE_LC8_PREFETCH_NEXT_KV it records each row's selection (the first min(16, blocks) ids, -1 after) and + prefetches a next layer's recorded blocks, stale ids past a request's blocks included: the same bits, and the + record.""" + assert hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_sparse_decode_lc8"), ( + "the target's LC8 op is not loaded (its TP8 extension)" + ) + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + ) + o_proj_weight = torch.empty(6144, 1024, dtype=torch.uint8, device="cuda") + o_proj_scale = torch.empty(6144 * 32 + 9, dtype=torch.uint8, device="cuda") + rows = case.q.shape[0] + n1_next = torch.randint( + -1, 2 * (max(seq_lens) // 128 + 1), (128, 16), dtype=torch.int32, device="cuda" + ) + for linear in (False, True): + for prewait in (False, True): + for mxfp8 in (False, True): + n1_self = torch.full((128, 16), -7, dtype=torch.int32, device="cuda") + next_kv = ( + n1_self, + n1_next, + case.k_paged.view(torch.uint8), + case.v_paged.view(torch.uint8), + case.k_block_scale.view(torch.uint8), + case.v_block_scale.view(torch.uint8), + ) + variants = ( + (False, None, None), + (True, None, None), + (True, (o_proj_weight, o_proj_scale), None), + (True, (o_proj_weight, o_proj_scale), next_kv), + ) + got = [] + for merge_once, prefetch, n1 in variants: + out, q8, sf = _buffers(case.q.shape[0], mxfp8) + op( + *_args(case, out, q8, sf, linear, prewait), + merge_once=merge_once, + prefetch=prefetch, + next_kv=n1, + ) + got.append((out, q8, sf)) + torch.cuda.synchronize() + a, aq, asf = got[0] + for i, (b, bq, bsf) in enumerate(got[1:], 1): + assert torch.equal(a.view(torch.int16), b.view(torch.int16)), ( + linear, + prewait, + mxfp8, + i, + ) + if mxfp8: + assert torch.equal(aq, bq) and torch.equal(asf, bsf), (linear, prewait, i) + # The record: each row's first min(16, blocks) selected ids, -1 past them; rows past the step untouched. + want = case.topk[0].clone() + for tok in range(rows): + req, qoff = divmod(tok, case.decode_query_len) + kv_len = max(int(case.seq_lens[req]) - case.decode_query_len + qoff + 1, 0) + want[tok, min(16, (kv_len + 127) // 128) :] = -1 + assert torch.equal(n1_self[:rows], want), "the recorded selection" + assert bool((n1_self[rows:] == -7).all()), "rows past the step" + + +@pytest.mark.parametrize("linear", [False, True]) +@pytest.mark.parametrize("prewait", [True, False], ids=["prewait", "plain"]) +@pytest.mark.parametrize( + "seq_lens, decode_query_len", + [((165521,), 4), ((5000, 12345), 4), ((700,), 1), ((200000, 3000, 64, 70000), 4)], + ids=["1x165k", "2req", "short-1row", "4req-incl-64"], +) +def test_one_div_merge_is_deterministic_and_in_band( + seq_lens, decode_query_len, prewait, linear +) -> None: + """ops/m3_attn_lc8.ONE_DIV_MERGE (numerics rule 1b, gate a): the one-division merge writes the same rows on + every run (5 eager runs, PDL on and off) and under CUDA-graph replay (3 replays == eager), MXFP8 hand-off included; + the rows are within the float64 reference's band, and within summation-order distance of the per-chunk merge.""" + case = build_step( + seq_lens, decode_query_len, identity=False, seed=zlib.crc32(repr(seq_lens).encode()) + 3 + ) + total_q = case.q.shape[0] + + def run(out, q8, sf, use_pdl, one_div): + args = list(_args(case, out, q8, sf, linear, prewait)) + args[18] = use_pdl + op(*args, merge_once=True, one_div=one_div) + + runs = [] + for i in range(6): + out, q8, sf = _buffers(total_q, True) + run(out, q8, sf, i == 5, True) + runs.append((out, q8, sf)) + torch.cuda.synchronize() + a, aq, asf = runs[0] + for b, bq, bsf in runs[1:]: + assert torch.equal(a.view(torch.int16), b.view(torch.int16)) + assert torch.equal(aq, bq) and torch.equal(asf, bsf) + out = torch.empty_like(a) + run(out, None, None, True, True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run(out, None, None, True, True) + for _ in range(3): + out.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out.view(torch.int16), a.view(torch.int16)) + per, _, _ = _buffers(total_q, False) + run(per, None, None, False, False) + ref = op.reference(*_args(case, a, None, None, linear, prewait)) + assert torch.isfinite(a.float()).all() + assert_within_ulp(a, ref, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) + assert_within_ulp(a, per, element_ulp=_ELEMENT_ULP, rms_ulp=_RMS_ULP) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py new file mode 100644 index 000000000000..e21326054b2e --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test (SM 10.7) for ops/m3_tokmap_guard: a one-request multi-token call of +thop.trtllm_gen_generation_preprocess (rows padded to the bucket, the request's generation length 1..bucket) with a +stale token map in the workspace gives, through the guard, the bytes the op gives with a freshly built map: the +appended FP8 K/V pages and the processed query's live rows, bit for bit, eager and from one captured graph replayed for +every length. Without the guard the stale map moves rows (the negative control), and a several-request call passes +through with its own map untouched. + +The fresh-map reference forces the op's own build (a scaled RoPE type with no inv-freq cache, at rotary dim 0, so no +row is rotated either way). +""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_tokmap_guard as guard, +) +from tensorrt_llm.bindings.internal import thop + +assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) + +H, KVH, D, PAGE = 8, 1, 128, 128 +BUCKET = 4 +_PEER_EXTRA = 777 # the second request's extra length behind the first +FP8_KV_QUANT_MODE = 128 +_DEV = torch.device("cuda") # unindexed on purpose: install() and the lookup normalize it +_TAKEN = [] + + +@pytest.fixture(scope="module", autouse=True) +def _installed(): + guard.install(_DEV, 48, _TAKEN.append) + yield + + +class _Step: + """One request (or ``nreq``) of BUCKET padded rows over a fresh FP8 paged cache, persistent across calls like a + real step's buffers.""" + + def __init__(self, kv_len: int, nreq: int = 1, seed: int = 0): + g = torch.Generator(device=_DEV).manual_seed(seed + kv_len) + self.nreq = nreq + # Per-request inputs (page tables, generation lengths, position offsets) for at least two requests, the page + # tables at the peer's longer length (_seq_with_peer): the stale map sends rows to request 1, which must have + # real entries behind it, also on the unguarded calls. + tables = max(nreq, 2) + pages = math.ceil((kv_len + _PEER_EXTRA + BUCKET) / PAGE) + self.nblocks = tables * pages + 2 + self.pool0 = ( + torch.randn(self.nblocks, 2, KVH, PAGE, D, device=_DEV, generator=g) + .clamp_(-4, 4) + .to(torch.float8_e4m3fn) + ) + self.pool = self.pool0.clone() + perm = torch.randperm(self.nblocks, device=_DEV, generator=g)[: tables * pages].view( + tables, pages + ) + self.offsets = torch.zeros(1, tables, 2, pages, dtype=torch.int32, device=_DEV) + self.offsets[0, :, 0] = 2 * perm + self.offsets[0, :, 1] = 2 * perm + 1 + self.pointers = torch.tensor([[self.pool.data_ptr(), 0]], dtype=torch.int64) + self.mapping = torch.zeros(1, 2, dtype=torch.int32) + self.qkv = torch.randn( + nreq * BUCKET, (H + 2 * KVH) * D, dtype=torch.bfloat16, device=_DEV, generator=g + ) + # Request lengths after this step's append; generation lengths (accepted rows) set per case. + self.seq = torch.full((nreq,), kv_len, dtype=torch.int32, device=_DEV) + self.gen = torch.full((tables,), BUCKET, dtype=torch.int32, device=_DEV) + self.pos = ( + torch.arange(BUCKET, dtype=torch.int32, device=_DEV).repeat(tables, 1).contiguous() + ) + self.one = torch.ones(1, dtype=torch.float32, device=_DEV) + self.ws = torch.empty(64 << 20, dtype=torch.uint8, device=_DEV) + self.kv_len = kv_len + + def call(self, op, fresh_map: bool = False): + """One generation preprocess; returns a copy of the processed query.""" + out = op( + self.qkv, + self.ws, + self.seq, + self.gen, + self.pos, + self.offsets, + self.pointers, + self.mapping, + self.one, + self.one, + None, + None, + None, + None, + 0, + 0, + H, + KVH, + D, + PAGE, + FP8_KV_QUANT_MODE, + 1 << 20, + 1 << 20, + self.nreq * BUCKET, + self.nreq, + BUCKET, + self.kv_len, + 0, + 10000.0, + 1 if fresh_map else 0, + 1.0, + 1 << 20, + 0, + 1.0, + 1.0, + False, + BUCKET, + 0, + torch.cuda.get_device_properties(0).multi_processor_count, + self.nblocks, + 2, + True, + False, + ) + return out[0].clone() + + def poison_map(self) -> None: + """A stale multi-request map in the op's one-request map slice: rows 1.. on request 1.""" + layout = thop.get_trtllm_gen_generation_workspace_layout( + self.qkv.dtype, 1, self.nreq * BUCKET, H, D, 0, KVH, 0, False, False + ) + off = int(layout["tokens_info_offset"]) + stale = torch.tensor( + [[0, 0]] + [[1, t] for t in range(BUCKET - 1)], dtype=torch.int32, device=_DEV + ) + self.ws[off : off + 8 * BUCKET].view(torch.int32).view(BUCKET, 2).copy_(stale) + + def reset(self) -> None: + self.pool.copy_(self.pool0) + + +def _seq_with_peer(step: _Step) -> None: + # A second request's length behind the first, so a row the stale map moves to request 1 reads a real length. + step.seq = torch.tensor( + [step.kv_len, step.kv_len + _PEER_EXTRA], dtype=torch.int32, device=_DEV + ) + + +@pytest.mark.parametrize("kv_len", [130, 8191, 34117]) +@pytest.mark.parametrize("n", [1, 2, 3, 4]) +def test_guard_equals_a_fresh_map(kv_len, n) -> None: + step = _Step(kv_len) + _seq_with_peer(step) + step.gen.fill_(n) + step.poison_map() + step.reset() + q_ref = step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + + step.poison_map() + step.reset() + q = step.call(thop.trtllm_gen_generation_preprocess) + # The processed query's n live rows (rows past the request's length are padding the op leaves unwritten). + assert torch.equal(q[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)) + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + assert BUCKET in _TAKEN + + +def test_stale_map_moves_rows_without_the_guard() -> None: + step = _Step(34117) + _seq_with_peer(step) + step.gen.fill_(BUCKET) + step.poison_map() + step.reset() + step.call(guard._original, fresh_map=True) + pool_ref = step.pool.clone() + step.poison_map() + step.reset() + step.call(guard._original) + assert not torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)) + + +def test_one_captured_graph_serves_every_length() -> None: + step = _Step(34000) + _seq_with_peer(step) + refs = {} + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + q = step.call(guard._original, fresh_map=True) + refs[n] = (q, step.pool.clone()) + step.gen.fill_(BUCKET) + step.call(thop.trtllm_gen_generation_preprocess) # eager first, as the warmup runs it + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out = thop.trtllm_gen_generation_preprocess + q_static = step.call(out) + for n in range(1, BUCKET + 1): + step.gen.fill_(n) + step.poison_map() + step.reset() + g.replay() + torch.cuda.synchronize() + q_ref, pool_ref = refs[n] + assert torch.equal(q_static[:n].view(torch.uint8), q_ref[:n].view(torch.uint8)), n + assert torch.equal(step.pool.view(torch.uint8), pool_ref.view(torch.uint8)), n + + +def test_several_requests_pass_through() -> None: + step = _Step(8191, nreq=2) + layout = thop.get_trtllm_gen_generation_workspace_layout( + step.qkv.dtype, 2, 2 * BUCKET, H, D, 0, KVH, 0, False, False + ) + step.reset() + step.call(thop.trtllm_gen_generation_preprocess) + pool_guarded = step.pool.clone() + step.reset() + step.call(guard._original) + assert torch.equal(step.pool.view(torch.uint8), pool_guarded.view(torch.uint8)) + assert int(layout["tokens_info_size"]) >= 8 * 2 * BUCKET diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py new file mode 100644 index 000000000000..c29fd0085966 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the msa_kv_update entry (the one-launch MSA KV-length correction).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.msa_kv_update import ( # noqa: E501 + TOKENS_PER_BLOCK, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.msa_kv_update import ( # noqa: E501 + msa_kv_update as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.msa_kv_update import ( # noqa: E501 + msa_kv_update as hook, +) + +__extra_import_path__ = ["../../../.."] +from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above + +assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" + +#: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. +_SPARE = 3 + + +def _case(spec, seed): + """The metadata's buffers for one step, with stale values in every output so a skipped write shows.""" + g = torch.Generator(device="cpu").manual_seed(seed) + batch, qlen = spec["batch"], spec["query_len"] + total_q = batch * qlen + staged = torch.randint(qlen, 5000, (batch,), generator=g, dtype=torch.int32) + accept = spec["accept"] + if accept == "full": + kv = staged.clone() + elif accept == "page_edge": + staged = torch.tensor([130, 257, 2 * TOKENS_PER_BLOCK + qlen][:batch], dtype=torch.int32) + kv = staged - torch.tensor([0, 1, qlen - 1][:batch], dtype=torch.int32) + else: + kv = staged - torch.randint(0, qlen, (batch,), generator=g, dtype=torch.int32) + if accept == "padded": + kv[batch // 2 :] = 0 + staged[batch // 2 :] = 0 + ncols = (5000 + qlen) // TOKENS_PER_BLOCK + 2 + table = ( + torch.randperm(batch * ncols + 7, generator=g)[: batch * ncols] + .to(torch.int32) + .reshape(batch, ncols) + ) + dev = dict(device="cuda", dtype=torch.int32) + stale = -12345 + return ( + torch.cat([kv, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.cat([staged, torch.full((_SPARE,), 99, dtype=torch.int32)]).cuda(), + torch.full((batch + _SPARE,), qlen, **dev), + torch.arange(batch, **dev).repeat_interleave(qlen), + torch.arange(qlen, **dev).repeat(batch), + table.cuda(), + TOKENS_PER_BLOCK, + batch, + total_q, + torch.full((batch + _SPARE,), stale, **dev), + torch.full((batch + 1 + _SPARE,), stale, **dev) if spec["cu_kv"] else None, + torch.full((total_q + _SPARE,), stale, **dev), + torch.full((total_q + _SPARE,), stale, **dev) if spec["n_valid"] else None, + ) + + +def _check(args, ref): + seq, cu_kv, loc, n_valid = args[9], args[10], args[11], args[12] + batch, total_q = args[7], args[8] + ref_seq, ref_cu, ref_loc, ref_nv = ref + # Integer bookkeeping: exact, and nothing past the live rows is touched. + assert torch.equal(seq[:batch], ref_seq) + assert torch.equal(loc[:total_q], ref_loc) + assert bool((seq[batch:] == -12345).all()) and bool((loc[total_q:] == -12345).all()) + assert (cu_kv is None) == (ref_cu is None) and (n_valid is None) == (ref_nv is None) + if cu_kv is not None: + assert torch.equal(cu_kv[: batch + 1], ref_cu) + assert bool((cu_kv[batch + 1 :] == -12345).all()) + if n_valid is not None: + assert torch.equal(n_valid[:total_q], ref_nv) + assert bool((n_valid[total_q:] == -12345).all()) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31) + certifying(op) + op(*args) + _check(args, op.reference(*args)) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_hook_takes_every_cell(cell) -> None: + """The target's hook (what the metadata calls) takes each certified step and writes what the op does.""" + args = _case(cell.spec, seed=zlib.crc32(cell.why.encode()) % 2**31 + 1) + assert hook(*args) is True + _check(args, op.reference(*args)) + + +def test_the_hook_declines_what_the_op_would_misread() -> None: + """A strided operand is read wrongly by the op, so the hook declines it (the metadata then runs its own + torch update) and writes nothing.""" + args = list(_case(op.CELLS[2].spec, seed=5)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + before = args[11].clone() + assert hook(*args) is False + assert torch.equal(args[11], before) + + +def test_guard_refuses_a_strided_operand() -> None: + args = list(_case(op.CELLS[2].spec, seed=6)) + args[0] = torch.stack([args[0], args[0]], dim=1)[:, 0] + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_request_index_past_the_batch() -> None: + args = list(_case(op.CELLS[2].spec, seed=7)) + args[3] = args[3].clone() + args[3][-1] = args[7] + with pytest.raises(AssertionError, match="q_batch_row"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py new file mode 100644 index 000000000000..45b043d24e95 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py @@ -0,0 +1,454 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for the ``m3_allreduce`` entry, 2 ranks (SM107): its cells against the entry's reference, +and the fused draft loop's layer boundaries built on it (``M3Eagle3Draft.ar_norm``: the plain all-reduce, then +FlashInfer's fused add + RMSNorm) against stock Eagle3's, bit for bit. + +Stock ``Eagle3DecoderLayer`` all-reduces each projection plainly (BF16 out) and runs the ``RMSNorm`` module with the +residual: FlashInfer's ``fused_add_rmsnorm``, which adds in FP32 and normalizes the FP32 sum. The draft loop's +boundaries must be that path exactly, under every all-reduce strategy a single-node job can run (AUTO, ONESHOT, +TWOSHOT, NCCL; MNNVL, which the tp8 drafter's two-tray group may take, is multi-node only and not driven here), at the +drafter's shapes (H = 6144, eps = 1e-6, 1 to 32 rows), at the job's world size (2, 4 or 8 ranks on one node; the tp8 +drafter runs at 8 over two trays, which one node cannot form, and nothing here depends on the rank count). The inputs +are +random rows and rows built so a BF16 rounding of the residual sum would move the norm the most, including the +review's case: unit weights, sums 1 + 2^-8 on 6143 coordinates and 2 - 2^-8 on the last, both BF16 rounding +midpoints that round to even in opposite directions. + +Asserts the normalized rows and the residual sum bitwise equal to the stock path's, and reports the max ulp of each +against an fp64 reference (ulps of BF16 at the larger magnitude; a correctly rounded output is within 0.5). + +Its own launcher, one rank per device in CUDA_VISIBLE_DEVICES, as ``_m3_mc_allreduce_norm_op_matrix.py`` is; the +collected entry point is ``test_modeling_v2_m3_allreduce.py``. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_allreduce_op_matrix.py +""" + +import math +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +TOKENS = (1, 4, 8, 32) +RESULT_TAG = "[ulp]" +"""Prefix of rank 0's result lines; the pytest entry reports them.""" + +#: World sizes the matrix runs at (one node): see ``_rank_job.WORLD_SIZES``. +WORLD_SIZES = (2, 4, 8) +RANK = 0 +WORLD = 1 + +Inputs = Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor] +"""(every rank's [m, H] BF16 partial rows, the [m, H] BF16 residual, the [H] BF16 norm weight).""" + + +# ---------------------------------------------------------------- inputs (every rank builds every rank's rows) + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _spread(parts: List[torch.Tensor]) -> List[torch.Tensor]: + """Every rank's BF16 partial: ``parts`` on the first ranks, zeros on the rest.""" + assert len(parts) <= WORLD, (len(parts), WORLD) + zeros = [torch.zeros_like(parts[0]) for _ in range(WORLD - len(parts))] + return [p.bfloat16() for p in parts] + [z.bfloat16() for z in zeros] + + +def _random(seed: int, m: int, x_scale: float, r_scale: float, unit_w: bool) -> Inputs: + g = _gen(seed) + xs = [ + (torch.randn(m, HIDDEN, generator=g, device="cuda") * x_scale).bfloat16() + for _ in range(WORLD) + ] + residual = (torch.randn(m, HIDDEN, generator=g, device="cuda") * r_scale).bfloat16() + if unit_w: + w = torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16) + else: + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return xs, residual, w + + +def _on_rank0(m: int, x: torch.Tensor) -> List[torch.Tensor]: + """``x`` as rank 0's partial, zeros elsewhere (the cross-rank sum is exact).""" + return [x.bfloat16()] + [ + torch.zeros(m, HIDDEN, device="cuda", dtype=torch.bfloat16) for _ in range(WORLD - 1) + ] + + +def _review_case(m: int) -> Inputs: + """Sums 1 + 2^-8 on the first H - 1 coordinates and 2 - 2^-8 on the last, unit weights. + + Both sums are BF16 midpoints: 1 + 2^-8 rounds (to even) down to 1 and 2 - 2^-8 up to 2, so the fused path's + norm sees a residual moved in opposite directions on the bulk and on the last coordinate. + """ + residual = torch.ones(m, HIDDEN, device="cuda") + residual[:, -1] = 2.0 - 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _review_case_flipped(m: int) -> Inputs: + """The mirror: the bulk on odd-mantissa midpoints (1 + 3 * 2^-8 rounds up), the last one rounding down.""" + residual = torch.full((m, HIDDEN), 1.0 + 2.0**-7, device="cuda") + residual[:, -1] = 2.0 - 2.0 * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + return ( + _on_rank0(m, x), + residual.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _midpoints(seed: int, m: int, outliers: int) -> Inputs: + """Every sum a BF16 midpoint: residual r random in [1, 2) and x = half its ulp (+2^-8), split over the ranks. + + The bulk has an odd mantissa (rounds up); ``outliers`` coordinates per row carry an even one (rounds down) and + a larger magnitude, so they weigh in the RMS while rounding the other way. + """ + g = _gen(seed) + mant = ( + torch.randint(0, 64, (m, HIDDEN), generator=g, device="cuda") * 2 + 1 + ) # odd 7-bit mantissa + r = 1.0 + mant.float() * 2.0**-7 + x = torch.full((m, HIDDEN), 2.0**-8, device="cuda") + if outliers: + cols = torch.randperm(HIDDEN, generator=g, device="cuda")[:outliers] + even = torch.randint(0, 64, (m, outliers), generator=g, device="cuda") * 2 + r[:, cols] = 8.0 * (1.0 + even.float() * 2.0**-7) + x[:, cols] = 8.0 * 2.0**-8 + # Split x over the ranks in exact BF16 pieces: x/2 on rank 0, x/4 on ranks 1 and 2 (x/2 on rank 1 at two ranks), + # 0 on the rest. + parts = [x / 2, x / 2] if WORLD == 2 else [x / 2, x / 4, x / 4] + return ( + _spread(parts), + r.bfloat16(), + torch.ones(HIDDEN, device="cuda", dtype=torch.bfloat16), + ) + + +def _cancelling(seed: int, m: int) -> Inputs: + """Large partials that cancel across the ranks (+A, -A, +B, -B + d; +A, -A + d at two or three ranks) onto a unit + residual.""" + g = _gen(seed) + a = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + b = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 64).bfloat16().float() + d = (torch.randn(m, HIDDEN, generator=g, device="cuda") * 2.0**-6).bfloat16().float() + parts = [a, -a, b, (-b + d)] if WORLD >= 4 else [a, (-a + d)] + residual = (1.0 + torch.rand(m, HIDDEN, generator=g, device="cuda")).bfloat16() + w = (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).bfloat16() + return _spread(parts), residual, w + + +def _cases() -> List[Tuple[str, Callable[[int], Inputs]]]: + cases: List[Tuple[str, Callable[[int], Inputs]]] = [ + ("review: 1+2^-8 x6143, 2-2^-8 x1", _review_case), + ("review mirrored", _review_case_flipped), + ("all midpoints, no outliers", lambda m: _midpoints(11, m, 0)), + ("all midpoints, 1 outlier", lambda m: _midpoints(12, m, 1)), + ("all midpoints, 64 outliers", lambda m: _midpoints(13, m, 64)), + ("cancelling partials", lambda m: _cancelling(14, m)), + ] + for i, (xs, rs, unit) in enumerate( + ((1.0, 4.0, False), (1.0, 1.0, True), (0.05, 8.0, False), (8.0, 0.5, False)) + ): + cases.append( + ( + f"random x{xs} r{rs}{' unit-w' if unit else ''}", + lambda m, i=i, xs=xs, rs=rs, unit=unit: _random(100 + i, m, xs, rs, unit), + ) + ) + return cases + + +# ---------------------------------------------------------------- the three paths + + +def _fp64(xs: List[torch.Tensor], residual: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + s = residual.double() + for x in xs: + s = s + x.double() + rms = torch.rsqrt(s.square().mean(-1, keepdim=True) + EPS) + return s * rms * w.double() + + +def _ulp_at(v: torch.Tensor) -> torch.Tensor: + """BF16's ulp at |v| (fp64): 2^(e - 8) for |v| in [2^(e-1), 2^e), floored at the smallest normal's.""" + _, e = torch.frexp(v.abs().double()) + return torch.ldexp(torch.ones_like(v, dtype=torch.float64), torch.clamp(e - 8, min=-133)) + + +def _ulps(out: torch.Tensor, ref: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the larger magnitude of the two.""" + o, r = out.double(), ref.double() + return ((o - r).abs() / _ulp_at(torch.maximum(o.abs(), r.abs()))).max().item() + + +def _ulps_at(out: torch.Tensor, ref: torch.Tensor, scale: torch.Tensor) -> float: + """Max |out - ref| in BF16 ulps at the magnitude ``scale`` (the operands'), not at the result's.""" + return ((out.double() - ref.double()).abs() / _ulp_at(scale.abs())).max().item() + + +def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """The draft loop's boundary (``M3Eagle3Draft.ar_norm``): ``m3_allreduce`` plain, then the catalog's + ``flashinfer_fused_add_rmsnorm``.""" + from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 + flashinfer_fused_add_rmsnorm, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + + return flashinfer_fused_add_rmsnorm(m3_allreduce(x, ar, None), residual.clone(), w, EPS) + + +def _check_cells(mapping: Any) -> None: + """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum + to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude + u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final + roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). + The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + m3_allreduce, + ) + from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams + from tensorrt_llm.functional import AllReduceStrategy + + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating + + ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) + for i, cell in enumerate(m3_allreduce.CELLS): + spec = cell.spec + xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) + params = ( + None + if spec.get("plain") + else AllReduceParams( + fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM, + residual=residual.clone(), + norm_weight=w, + eps=EPS, + trigger_completion_at_end=False, + ) + ) + # The reference first: the fused path may write the updated residual into ``params.residual``. + ref = m3_allreduce.reference(xs, ar, params) + with validating(m3_allreduce): + out = m3_allreduce(xs[RANK], ar, params) + torch.cuda.synchronize() + total = sum(x.double() for x in xs) + if params is None: + err = _ulps_at(out, ref, total) + assert err <= 1.0, (cell.why, err) + else: + err = _ulps_at(out[1], ref[1], torch.maximum(total.abs(), residual.double().abs())) + assert err <= 2.5, (cell.why, "residual", err) + h = out[1].double() + own = h * torch.rsqrt(h.square().mean(-1, keepdim=True) + EPS) * w.double() + assert _ulps(out[0], own) <= 2.0, (cell.why, "normed", _ulps(out[0], own)) + dist.barrier() + if RANK == 0: + print(f"{RESULT_TAG} cell {cell.why} passed", flush=True) + + +def _stock(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): + """Stock Eagle3: ``o_proj``'s plain all-reduce (BF16 out), then ``RMSNorm(h, residual)``'s FlashInfer + ``fused_add_rmsnorm`` (FP32 add, norm of the FP32 sum; in place on both).""" + h = ar(x) + res = residual.clone() + torch.ops.trtllm.flashinfer_fused_add_rmsnorm(h, res, w, EPS) + return h, res + + +# ---------------------------------------------------------------- the matrix + + +def _single_node_cluster_info() -> None: + """``TorchDist`` (the communicator under ``TLLM_DISABLE_MPI``) asks Ray for each rank's node and GPU to build its + node-local group. These ranks are spawned without Ray on one node, so every rank reports the same node and its + own device.""" + from tensorrt_llm._torch.distributed.communicator import TorchDist + + def cluster_info(self): + if self.cluster_info is None: + self.cluster_info = [("127.0.0.1", r) for r in range(dist.get_world_size())] + return self.cluster_info + + TorchDist._get_cluster_info = cluster_info + + +def _run_matrix() -> None: + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers trtllm::flashinfer_fused_add_rmsnorm + from tensorrt_llm._torch.distributed import AllReduce + from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE + from tensorrt_llm.functional import AllReduceStrategy + from tensorrt_llm.mapping import Mapping + + assert IS_FLASHINFER_AVAILABLE, "stock Eagle3's add + RMSNorm is FlashInfer's fused_add_rmsnorm" + _single_node_cluster_info() + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=WORLD) + _check_cells(mapping) + strategies = ( + AllReduceStrategy.AUTO, + AllReduceStrategy.ONESHOT, + AllReduceStrategy.TWOSHOT, + AllReduceStrategy.NCCL, + ) + worst = 0.0 + mismatches: List[str] = [] + for strategy in strategies: + ar = AllReduce(mapping=mapping, strategy=strategy, dtype=torch.bfloat16) + assert ar.mnnvl_allreduce is None and ar.symm_mem_allreduce is None, strategy + for name, build in _cases(): + ulp, res_ulp = 0.0, 0.0 + for m in TOKENS: + xs, residual, w = build(m) + target_out, target_res = _target(ar, xs[RANK], residual, w) + stock_out, stock_res = _stock(ar, xs[RANK], residual, w) + torch.cuda.synchronize() + for what, a, b in ( + ("normed", target_out, stock_out), + ("residual", target_res, stock_res), + ): + bad = int((a.view(torch.int16) != b.view(torch.int16)).sum().item()) + if bad: + mismatches.append( + f"{strategy.name} {name} m={m}: {bad} {what} elements differ from stock" + ) + ref = _fp64(xs, residual, w) + ref_res = residual.double() + sum(x.double() for x in xs) + u, ur = _ulps(target_out, ref), _ulps(target_res, ref_res) + assert math.isfinite(u) and math.isfinite(ur), (strategy, name, m) + ulp, res_ulp = max(ulp, u), max(res_ulp, ur) + dist.barrier() + worst = max(worst, ulp) + if RANK == 0: + print( + f"{RESULT_TAG} {strategy.name:7s} | {name:34s} | target = stock vs fp64: normed {ulp:.3f}" + f" | residual {res_ulp:.3f}", + flush=True, + ) + if RANK == 0: + scope = f"all cases, m in {TOKENS}" + print( + f"{RESULT_TAG} MAX | {scope:34s} | target = stock vs fp64: normed {worst:.3f}", + flush=True, + ) + assert not mismatches, "\n".join(mismatches) + + +def _run_one_rank() -> int: + global RANK, WORLD + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD in WORLD_SIZES, f"the matrix runs at {WORLD_SIZES} ranks on one node, got {WORLD}" + try: + _run_matrix() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + os._exit(13) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker per claimed device, each in its own session; on the first nonzero exit or the deadline every + session is killed (as ``_m3_mc_allreduce_norm_op_matrix.py`` does).""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns (one rank each)" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size in WORLD_SIZES, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at {WORLD_SIZES}" + ) + 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", + ) + env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} + children = [ + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..6d7def71212a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,788 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification matrix for ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm``, at the job's ranks. + +Its own launcher, as the catalog's collective matrices are: run plainly it spawns +one rank process per device in CUDA_VISIBLE_DEVICES and kills them all on the +first failure or on the deadline. A rank whose peers never push polls forever, +so the deadline is what keeps a broken exchange from wedging the calling run. + + CUDA_VISIBLE_DEVICES=0,1,2,3 python _m3_mc_allreduce_norm_op_matrix.py + +World size: one rank per visible device, 2, 4 or 8. The tp8 target runs these +entries at 8 ranks over two trays, which one 4-device node cannot form; the +world size is a runtime argument of the kernel and no cell's spec names one, so +the cells are certified at the world size of the job (``_rank_job.WORLD_SIZES``). + +Op: the entries call the image's op by default and this target's own copy +(``trtllm::minimax_m3_sm107_tp8_mc_allreduce_norm``) with ``own_op``, which is +what the target's boundaries call when it is loaded. Every cell is driven +through each op the build has. + +Each rank builds the three multicast workspaces (8 / 16 / 48 tokens) the target +builds, over the torch.distributed group (the non-MPI path of +``_make_mnnvl_mcast_buffer``). Every rank can regenerate every rank's inputs from +the shared seed convention, which is what keeps each reference arithmetic rather +than a second collective. Producer pushes (``prepushed`` 2 / 3) are emulated by +storing the canonicalized rows through the multicast mapping at the call parity +the flags select, which is exactly what the producing kernels store. + +Not a pytest module, despite the ``check_*`` bodies: they are one fixed sequence +inside a single multi-rank job whose workspace state (call parities) carries from +one to the next. The collected entry point is +``test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py``. +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 900 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +BUCKETS = (8, 16, 48) +EMPTY_WORD = -2147483648 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +allreduce_entry: Any = None +finalize_entry: Any = None +ar_module: Any = None +validating: Any = None +#: World sizes the matrix runs at (one node): see ``_rank_job.WORLD_SIZES``. +WORLD_SIZES = (2, 4, 8) +#: The ``own_op`` values the cells are driven with: False (the image's op), then True when this target's op is +#: loaded. The last one is the op the target's boundaries call, which the other checks use. +OWN_OPS: Tuple[bool, ...] = (False,) +RANK = 0 +WORLD = 1 +WORKSPACES: Dict[int, Any] = {} + + +# ---------------------------------------------------------------- inputs + + +def _gen(seed: int, rank: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed * 1009 + rank + 1) + + +def _rows(seed: int, rank: int, m: int, scale: float = 1.0) -> torch.Tensor: + """One rank's [m, H] BF16 rows; ``rank`` -1 for tensors every rank holds alike.""" + g = _gen(seed, rank + 1) + return (torch.randn(m, HIDDEN, generator=g, device="cuda") * scale).to(torch.bfloat16) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return _rows(seed + 7, -1, m, 4.0) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11, 0) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _finalize_inputs(seed: int, rank: int, m: int, topk: int, dropped: int): + """One rank's unfinalized MoE outputs: fc2 [m * topk, H], perm [m, topk], fp32 scales, shared [m, H].""" + g = _gen(seed + 23, rank) + rows = m * topk + fc2 = (torch.randn(rows, HIDDEN, generator=g, device="cuda") * 0.5).to(torch.bfloat16) + perm = torch.randperm(rows, generator=g, device="cuda").to(torch.int32) + if dropped: + drop = torch.randperm(rows, generator=g, device="cuda")[:dropped] + perm[drop] = -1 + perm = perm.view(m, topk).contiguous() + scales = torch.rand(m, topk, generator=g, device="cuda", dtype=torch.float32) + shared = _rows(seed + 29, rank, m, 0.5) + return fc2, perm, scales, shared + + +# ---------------------------------------------------------------- workspace / producer emulation + + +def _make_workspace(maxt: int): + """One bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK) + words = 4 * maxt * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar_module.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=maxt, + ) + ws.flags_for(CLUSTER) + return ws + + +def _push(ws, rows: torch.Tensor, exchange_set: int) -> None: + """Store this rank's rows into every rank's exchange, as a producer kernel does before a prepushed call. + + Slot of (token t, cluster CTA c): set * 2 * half + parity(t, c) * half + (t * world + rank) * H / 2 + c * H / 16 + words, parity read from the CTA's flag; -0.0 halves are stored as +0.0 (0x80000000 is the empty word). + """ + m = rows.shape[0] + torch.cuda.synchronize() + parity = (ws.flags_for(CLUSTER).view(ws.max_tokens, CLUSTER)[:m] & 1).long() + half = ws.max_tokens * WORLD * (HIDDEN // 2) + chunk = HIDDEN // 2 // CLUSTER + h16 = rows.contiguous().view(torch.int16) + h16 = torch.where(h16 == -32768, torch.zeros_like(h16), h16) + words = h16.view(torch.int32).view(m, CLUSTER, chunk) + t = torch.arange(m, device="cuda").view(m, 1, 1) + c = torch.arange(CLUSTER, device="cuda").view(1, CLUSTER, 1) + w = torch.arange(chunk, device="cuda").view(1, 1, chunk) + index = ( + exchange_set * 2 * half + + parity.view(m, CLUSTER, 1) * half + + (t * WORLD + RANK) * (HIDDEN // 2) + + c * chunk + + w + ) + ws.mc[index.reshape(-1)] = words.reshape(-1) + torch.cuda.synchronize() + dist.barrier() + + +# ---------------------------------------------------------------- gates + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +#: Share of normed elements one BF16 step off the float64 reference: measured at most 1.1e-4 per call over +#: every call of this matrix on GB300 (see ``_assert_normed``). The gate sits ~18x above that and far below what a +#: wrong statistic produces (most elements move). +NORMED_STEP_FRACTION = 0.002 +NORMED_STEP_SEEN: List[float] = [] + + +def _assert_normed(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + """At most one BF16 rounding step from the float64 reference, and rarely that. + + The kernel's RMS statistic is a shuffle / warp / cluster tree of FP32 partial sums and its ``rsqrtf`` is + approximate (~2^-22 relative): the scaled value moves by ~1e-6 relative, which changes the BF16 rounding only + for values within that distance of a rounding midpoint -- one step, never two. The fraction bound catches a + statistic that is systematically off (a missing rank, a wrong H), which moves most elements. + """ + assert out.dtype == ref.dtype == torch.bfloat16 and out.shape == ref.shape, where + o, r = out.float(), ref.float() + _, exponent = torch.frexp(torch.maximum(o.abs(), r.abs())) + step = torch.ldexp(torch.ones_like(o), exponent - 8) + diff = (o - r).abs() + assert bool((diff <= step).all()), f"{where}: max diff {diff.max().item()} beyond one BF16 step" + frac = (diff != 0).float().mean().item() + NORMED_STEP_SEEN.append(frac) + assert frac <= NORMED_STEP_FRACTION, f"{where}: {frac:.4f} of normed elements off by one step" + + +def _assert_quant(outs, ref, m: int, where: str) -> None: + """The MXFP8 copy: bit-exact against the recipe applied to the op's own normed rows (a pure function of + them), and equal to the reference's wherever the reference's normed 32-column block is the op's.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + mxfp8_rows, + sf_written, + ) + + q, sf = outs[2], outs[3] + own_q, own_sf = mxfp8_rows(outs[0]) + written = sf_written(m, HIDDEN // 32, sf.device) + _assert_bitwise(q, own_q, f"{where} e4m3 rows") + _assert_bitwise(sf[written], own_sf[written], f"{where} UE8M0 scales") + same_block = (outs[0].view(m, -1, 32) == ref[0].view(m, -1, 32)).all(-1) + qb = q.view(torch.uint8).view(m, -1, 32) + rqb = ref[2].view(torch.uint8).view(m, -1, 32) + assert bool((qb == rqb).all(-1)[same_block].all()), ( + f"{where}: e4m3 differs on a block whose normed rows agree" + ) + + +def _check_outputs(outs, ref, m: int, quant: bool, where: str) -> None: + assert len(outs) == (4 if quant else 2), (where, len(outs)) + _assert_bitwise(outs[1], ref[1], f"{where} updated residual") + _assert_normed(outs[0], ref[0], f"{where} normed") + if quant: + _assert_quant(outs, ref, m, where) + + +# ---------------------------------------------------------------- calls + + +def _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, opts, own_op=None): + own_op = OWN_OPS[-1] if own_op is None else own_op + return allreduce_entry.m3_mc_allreduce_norm( + x, + x2, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + ws.max_tokens, + quant, + prepushed, + opts, + own_op=own_op, + ) + + +def _plain_case(spec: Dict[str, Any], seed: int, where: str, own_op: bool) -> None: + m, maxt, quant, prepushed = spec["m"], spec["maxt"], spec["quant"], spec["prepushed"] + ws = WORKSPACES[maxt] + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] if spec["x2"] else None + residual, norm_w = _residual(seed, m), _norm_w(seed) + x, x2 = xs[RANK], (None if x2s is None else x2s[RANK]) + if prepushed >= 2: + _push(ws, x, 0) + if prepushed == 3: + _push(ws, x2, 1) + with validating(allreduce_entry.m3_mc_allreduce_norm): + outs = _call_plain(ws, x, x2, residual, norm_w, quant, prepushed, spec["opts"], own_op) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, residual, norm_w, EPS, quant, spec.get("drafter", False) + ) + _check_outputs(outs, ref, m, quant, where) + dist.barrier() + + +# ---------------------------------------------------------------- checks + + +def _op_name(own_op: bool) -> str: + return "own op" if own_op else "image op" + + +def check_plain_cells() -> None: + """Every cell of ``m3_mc_allreduce_norm.CELLS`` through ``__call__`` under ``validating``, gated on its + ``reference``, with each op of ``OWN_OPS``.""" + for own_op in OWN_OPS: + for i, cell in enumerate(allreduce_entry.m3_mc_allreduce_norm.CELLS): + _plain_case(cell.spec, 1000 + 17 * i, f"{cell.why} ({_op_name(own_op)})", own_op) + + +def check_finalize_cells() -> None: + """Every cell of ``m3_mc_finalize_allreduce_norm.CELLS``, gated on its ``reference``, with each op of + ``OWN_OPS``.""" + for own_op in OWN_OPS: + _finalize_cells(own_op) + + +def _finalize_cells(own_op: bool) -> None: + op = finalize_entry.m3_mc_finalize_allreduce_norm + for i, cell in enumerate(op.CELLS): + spec = cell.spec + m, maxt, quant, topk = spec["m"], spec["maxt"], spec["quant"], spec["topk"] + seed = 2000 + 17 * i + ws = WORKSPACES[maxt] + per_rank = [_finalize_inputs(seed, r, m, topk, spec["dropped"]) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + with validating(op): + outs = op( + fc2, + perm, + scales, + shared, + residual, + norm_w, + EPS, + ws.uc, + ws.mc.data_ptr(), + ws.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + maxt, + quant, + spec["opts"], + own_op=own_op, + ) + torch.cuda.synchronize() + ref = op.reference( + [p[0] for p in per_rank], + [p[1] for p in per_rank], + [p[2] for p in per_rank], + [p[3] for p in per_rank], + residual, + norm_w, + EPS, + quant, + ) + _check_outputs(outs, ref, m, quant, f"{cell.why} ({_op_name(own_op)})") + dist.barrier() + + +def _capture_reference( + updated: torch.Tensor, weight: torch.Tensor, eps: float, out: torch.Tensor +) -> None: + """The capture the boundary replaces: m3_spec_tail.capture_norm of the updated residual (no second addend).""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_spec_tail, + ) + + m3_spec_tail.capture_norm(updated, None, weight, eps, out) + + +def check_capture_norm_in_boundary() -> None: + """CAPTURE_NORM_IN_BOUNDARY: with cap_out the boundary also writes the capture, BITWISE equal to capture_norm of + its own updated residual, into a strided [M, H] block of a [rows, 3H] capture buffer (other bytes untouched); its + normed / updated / MXFP8 outputs are bitwise the ones of the same call without cap_out. The row kernel (the plain + entry: x, x + x2, prepushed) and the cluster kernel (MoE finalize), 1 / 4 / 8 rows, residual scales from 1e-3 to + 1e3 (the statistic's range), the fc_norm eps 1e-6 and 1e-5; on the op the target's boundaries call.""" + cap_eps_values = (1e-6, 1e-5) + own_op = OWN_OPS[-1] + for case, (m, scale, kind) in enumerate( + ( + (1, 1.0, "plain"), + (4, 1e3, "plain_x2"), + (8, 1e-3, "prepushed"), + (4, 1.0, "prepushed"), + (1, 4.0, "finalize"), + (4, 1e3, "finalize"), + (8, 1e-3, "finalize"), + ) + ): + seed = 5000 + 31 * case + cap_eps = cap_eps_values[case % 2] + ws = WORKSPACES[8] + residual = (_residual(seed, m).float() * scale).to(torch.bfloat16) + norm_w = _norm_w(seed) + cap_w = _norm_w(seed + 101) + quant = case % 3 == 1 + outs_by = [] + for with_cap in (False, True): + buf = torch.full((m + 3, 3 * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + cap = buf[:m, HIDDEN : 2 * HIDDEN] if with_cap else None + if kind == "finalize": + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, m, 4, 1) + fc2 = (fc2.float() * scale).to(torch.bfloat16) + outs = finalize_entry.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), + ws.flags_for(CLUSTER), RANK, WORLD, CLUSTER, False, 8, quant, 7, + cap, cap_w if with_cap else None, cap_eps, own_op=own_op, + ) # fmt: skip + else: + x = (_rows(seed, RANK, m).float() * scale).to(torch.bfloat16) + x2 = ( + (_rows(seed + 3, RANK, m).float() * scale).to(torch.bfloat16) + if kind == "plain_x2" + else None + ) + prepushed = 2 if kind == "prepushed" else 0 + if prepushed: + _push(ws, x, 0) + outs = allreduce_entry.m3_mc_allreduce_norm( + x, x2, residual, norm_w, EPS, ws.uc, ws.mc.data_ptr(), ws.flags_for(CLUSTER), RANK, WORLD, + CLUSTER, False, 8, quant, prepushed, (7, 7, 15, 15, 15, 15, 7, 7)[m - 1], + cap, cap_w if with_cap else None, cap_eps, own_op=own_op, + ) # fmt: skip + torch.cuda.synchronize() + outs_by.append((outs, buf)) + dist.barrier() + (plain, _), (capped, buf) = outs_by + where = f"capture in boundary, {kind}, {m} rows, scale {scale}" + assert len(plain) == len(capped), where + for i, (a, b) in enumerate(zip(plain, capped)): + if i == 3: + # The scale buffer is torch.empty: compare the bytes the kernel writes (the valid rows'). + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written, + ) + + written = sf_written(m, HIDDEN // 32, a.device) + a, b = a[written], b[written] + _assert_bitwise(b, a, f"{where}: output {i} with vs without cap_out") + ref = torch.empty((m, HIDDEN), dtype=torch.bfloat16, device="cuda") + _capture_reference(capped[1], cap_w, cap_eps, ref) + torch.cuda.synchronize() + _assert_bitwise(buf[:m, HIDDEN : 2 * HIDDEN], ref, f"{where}: capture vs capture_norm") + untouched = buf.clone() + untouched[:m, HIDDEN : 2 * HIDDEN] = 7.0 + assert bool((untouched == 7.0).all()), f"{where}: bytes outside the capture block changed" + dist.barrier() + + +def check_the_target_wrappers_select_bucket_and_options() -> None: + """The target's own wrappers (bucket by token count, option table) on top of the entries, every token count of + the 8-token bucket plus one of each larger bucket, and the finalize wrapper; the MXFP8 publication is the + prequant slot's business and is left off.""" + buckets = ar_module.McBuckets([WORKSPACES[b] for b in BUCKETS]) + for m in (*range(1, 9), 11, 40): + seed = 3000 + m + xs = [_rows(seed, r, m) for r in range(WORLD)] + x2s = [_rows(seed + 3, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + normed, updated = ar_module.m3_mc_allreduce_norm( + xs[RANK], residual, norm_w, EPS, buckets, x2=x2s[RANK] + ) + torch.cuda.synchronize() + ref = allreduce_entry.m3_mc_allreduce_norm.reference(xs, x2s, residual, norm_w, EPS, False) + _check_outputs([normed, updated], ref, m, False, f"wrapper m={m}") + dist.barrier() + for m in (9, 33): + seed = 3100 + m + per_rank = [_finalize_inputs(seed, r, m, 4, 2) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + fc2, perm, scales, shared = per_rank[RANK] + normed, updated = ar_module.m3_mc_finalize_allreduce_norm( + fc2, perm, scales, shared, residual, norm_w, EPS, buckets + ) + torch.cuda.synchronize() + ref = finalize_entry.m3_mc_finalize_allreduce_norm.reference( + *[[p[j] for p in per_rank] for j in range(4)], residual, norm_w, EPS, False + ) + _check_outputs([normed, updated], ref, m, False, f"finalize wrapper m={m}") + dist.barrier() + + +def check_cuda_graph_replays() -> None: + """Two boundaries captured in one graph (8- and 16-token buckets), replayed with fresh inputs. + + The decode graph holds every boundary call; each replay rotates the same per-CTA parities the eager calls do, + so a replay that reused a stale half or skipped a re-arm reads the previous replay's rows and fails the + bitwise residual gate (every replay gets new inputs). + """ + shapes = ((3, 8, 15), (12, 16, 15)) + norm_w = _norm_w(4000) + bufs = [] + for m, maxt, opts in shapes: + bufs.append( + dict( + m=m, + maxt=maxt, + opts=opts, + x=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + x2=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + residual=torch.empty(m, HIDDEN, dtype=torch.bfloat16, device="cuda"), + ) + ) + + def body(): + return [ + _call_plain( + WORKSPACES[b["maxt"]], b["x"], b["x2"], b["residual"], norm_w, False, 0, b["opts"] + ) + for b in bufs + ] + + def fill(seed: int): + refs = [] + for j, b in enumerate(bufs): + s = seed + 31 * j + xs = [_rows(s, r, b["m"]) for r in range(WORLD)] + x2s = [_rows(s + 3, r, b["m"]) for r in range(WORLD)] + b["x"].copy_(xs[RANK]) + b["x2"].copy_(x2s[RANK]) + b["residual"].copy_(_residual(s, b["m"])) + refs.append( + allreduce_entry.m3_mc_allreduce_norm.reference( + xs, x2s, b["residual"].clone(), norm_w, EPS, False + ) + ) + torch.cuda.synchronize() + dist.barrier() + return refs + + refs = fill(4100) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + warm = body() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(warm[j], refs[j], b["m"], False, f"graph warm-up m={b['m']}") + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = body() + for rep in range(4): + refs = fill(4200 + 100 * rep) + graph.replay() + torch.cuda.synchronize() + for j, b in enumerate(bufs): + _check_outputs(captured[j], refs[j], b["m"], False, f"replay {rep} m={b['m']}") + dist.barrier() + del graph + dist.barrier() + + +def check_the_gate_sees_a_missing_rank() -> None: + """The residual gate is not blind to the reduction: dropping any one rank's partial from the reference + moves most elements of the updated residual (each rank's rows are distinct by construction).""" + m, seed = 4, 5000 + xs = [_rows(seed, r, m) for r in range(WORLD)] + residual, norm_w = _residual(seed, m), _norm_w(seed) + full = allreduce_entry.m3_mc_allreduce_norm.reference(xs, None, residual, norm_w, EPS, False) + for r in range(WORLD): + partial = [x if i != r else torch.zeros_like(x) for i, x in enumerate(xs)] + wrong = allreduce_entry.m3_mc_allreduce_norm.reference( + partial, None, residual, norm_w, EPS, False + ) + frac = (wrong[1] != full[1]).float().mean().item() + assert frac > 0.9, f"dropping rank {r} changes only {frac:.3f} of the residual" + dist.barrier() + + +def _expect_rejected(fn, match: str, where: str) -> None: + try: + fn() + except AssertionError as exc: + assert match in str(exc), f"{where}: rejected for the wrong reason: {exc}" + else: + raise AssertionError(f"{where}: the guard admitted it") + + +def check_guards() -> None: + """``is_valid`` refuses what the op accepts and answers wrongly. Nothing here reaches the kernel: the guard + raises first, on every rank alike, so no rank is left polling.""" + m, seed = 4, 6000 + ws8, ws16 = WORKSPACES[8], WORKSPACES[16] + x, residual, norm_w = _rows(seed, RANK, m), _residual(seed, m), _norm_w(seed) + op = allreduce_entry.m3_mc_allreduce_norm + + def call(ws_uc, flags, norm): + with validating(op): + op( + x, + None, + residual, + norm, + EPS, + ws_uc, + ws8.mc.data_ptr(), + flags, + RANK, + WORLD, + CLUSTER, + False, + 8, + False, + 0, + 15, + ) + + strided = torch.empty(2 * HIDDEN, dtype=torch.bfloat16, device="cuda")[::2] + strided.copy_(norm_w) + _expect_rejected( + lambda: call(ws8.uc, ws8.flags_for(CLUSTER), strided), "contiguous", "strided norm_w" + ) + _expect_rejected( + lambda: call(ws8.uc[: ws8.uc.numel() // 2], ws8.flags_for(CLUSTER), norm_w), + "words", + "short workspace", + ) + _expect_rejected( + lambda: call(ws8.uc, ws16.flags_for(CLUSTER), norm_w), "counters", "another bucket's flags" + ) + + fin = finalize_entry.m3_mc_finalize_allreduce_norm + fc2, perm, scales, shared = _finalize_inputs(seed, RANK, 9, 4, 0) + bad = perm.clone() + bad[0, 0] = fc2.shape[0] + + def call_fin(): + with validating(fin): + fin( + fc2, + bad, + scales, + shared, + _residual(seed, 9), + norm_w, + EPS, + ws16.uc, + ws16.mc.data_ptr(), + ws16.flags_for(CLUSTER), + RANK, + WORLD, + CLUSTER, + False, + 16, + False, + 15, + ) + + _expect_rejected(call_fin, "perm spans", "perm past the FC2 rows") + dist.barrier() + + +CHECKS = ( + check_plain_cells, + check_finalize_cells, + check_capture_norm_in_boundary, + check_the_target_wrappers_select_bucket_and_options, + check_cuda_graph_replays, + check_the_gate_sees_a_missing_rank, + check_guards, +) + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, OWN_OPS, allreduce_entry, finalize_entry, ar_module, validating + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + torch.cuda.set_device(rank) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", rank) + ) + RANK, WORLD = rank, world + assert WORLD in WORLD_SIZES, f"the matrix runs at {WORLD_SIZES} ranks on one node, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_mc_allreduce as _ar # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm as _plain, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm import ( # noqa: E501 + m3_mc_finalize_allreduce_norm as _finalize, + ) + + allreduce_entry, finalize_entry, ar_module = _plain, _finalize, _ar + # The same library is loaded on every rank, so the ranks agree on the ops they drive. + OWN_OPS = (False, True) if _ar.own_op() else (False,) + print(f"[rank {RANK}] ops: {', '.join(_op_name(o) for o in OWN_OPS)}", flush=True) + # Spawned by file path, not collected by pytest: reach the sibling helper the plain way (confined to the + # rank processes). /comm -> tests/unittest/_torch/modeling_v2. + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + from _validating import validating as _validating_cm + + validating = _validating_cm + for maxt in BUCKETS: + WORKSPACES[maxt] = _make_workspace(maxt) + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peers' polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print( + f"[rank {RANK}] {len(CHECKS)} checks passed; normed one-step fraction max " + f"{max(NORMED_STEP_SEEN):.5f} over {len(NORMED_STEP_SEEN)} calls", + flush=True, + ) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the devices this run owns (one rank each)" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size in WORLD_SIZES, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at {WORLD_SIZES}" + ) + 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", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + 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), LOCAL_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: Optional[str] = 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: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py new file mode 100644 index 000000000000..95374cf56848 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py @@ -0,0 +1,403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU certification of the MoE engine's all-reduce push (``m3_moe_engine`` with ``arp``), 8 ranks, two trays (SM107). + +The decode forward's MoE boundary (MOE_PUSH rows): the engine pushes its bf16 sum (routed + shared) into set A of +the boundary's multicast exchange while it writes it, and the boundary all-reduce runs with ``prepushed`` 2, reading +the pushed words instead of loading and pushing the rows itself. Certified against the unpushed boundary the target +runs otherwise: the same engine call without ``arp``, then ``m3_mc_allreduce_norm`` with ``prepushed`` 0. Every +output must be bit for bit the same: the engine's own output tensor, the updated residual, the normed rows and, with +``quant``, the MXFP8 hand-off it publishes. Both read the same bf16 words, so the reduction and the norm see equal +inputs; any difference is a push that stored the wrong word, slot or parity. + +Each rank holds its own expert shard (operand seed per rank) and the same hidden rows, residual and norm weight, as +the TP2 decode step does. Eager calls alternate pushed and unpushed in both orders so both call parities of the +exchange are exercised; a CUDA graph of the pushed engine + boundary is replayed on fresh rows against eager +unpushed references. + +Its own launcher, as ``_m3_mc_allreduce_norm_op_matrix.py`` is: run plainly it spawns one rank process per device +in CUDA_VISIBLE_DEVICES and kills them all on the first failure or on the deadline (a rank whose peer never pushes +polls forever). The collected entry point is ``test_modeling_v2_m3_moe_push_op_matrix.py``. + + CUDA_VISIBLE_DEVICES=0,1 python _m3_moe_push_op_matrix.py +""" + +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Optional + +import torch +import torch.distributed as dist + +assert torch.cuda.is_available(), "m3_moe_engine requires CUDA devices" + +_WORKER_FLAG = "--rank-worker" +DEADLINE_S = 1200 + +HIDDEN = 6144 +EPS = 1e-6 +CLUSTER = 8 +MAX_TOKENS = 8 +EMPTY_WORD = -2147483648 +RESERVED_SMS = 24 +OPERAND_SEED = 20261013 + +# Bound inside the rank body: importing the entries pulls in tensorrt_llm. +ar = None +sf_written: Any = None +engine_op: Any = None +ar_push_desc: Any = None +moe_state: Any = None +operands: Any = None +prequant: Any = None +RANK = 0 +WORLD = 1 +WS: Any = None + + +def _gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def _residual(seed: int, m: int) -> torch.Tensor: + return (torch.randn(m, HIDDEN, generator=_gen(seed + 7), device="cuda") * 4.0).to( + torch.bfloat16 + ) + + +def _norm_w(seed: int) -> torch.Tensor: + g = _gen(seed + 11) + return (1.0 + 0.25 * torch.randn(HIDDEN, generator=g, device="cuda")).to(torch.bfloat16) + + +def _make_workspace(): + """The 8-token bucket's exchange, as ``McAllReduceWorkspace.create`` builds it, over the torch.distributed + group.""" + from tensorrt_llm._torch.distributed.ops import _make_mnnvl_mcast_buffer + from tensorrt_llm.mapping import Mapping + + mapping = Mapping(world_size=WORLD, tp_size=WORLD, rank=RANK, gpus_per_node=4) + words = 4 * MAX_TOKENS * WORLD * (HIDDEN // 2) + group = dist.group.WORLD + handle = _make_mnnvl_mcast_buffer(group, words * 4, mapping, False) + uc = handle.get_uc_buffer(RANK, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + dist.barrier() + ws = ar.McAllReduceWorkspace( + uc=uc, + mc=mc, + flags={}, + rank=RANK, + world_size=WORLD, + hidden=HIDDEN, + handle=handle, + comm=group, + max_tokens=MAX_TOKENS, + ) + ws.flags_for(CLUSTER) + return ws + + +def _assert_bitwise(out: torch.Tensor, ref: torch.Tensor, where: str) -> None: + assert out.dtype == ref.dtype and out.shape == ref.shape, ( + where, + out.dtype, + ref.dtype, + out.shape, + ) + a, b = out.contiguous().view(torch.uint8), ref.contiguous().view(torch.uint8) + bad = (a != b).sum().item() + assert bad == 0, f"{where}: {bad} of {a.numel()} bytes differ" + + +_STATES = {} +_LAYERS = {} + + +def _layer(m: int, case: str): + """This rank's engine layer for ``m`` tokens: its own expert shard; ``concentrated`` routes every token to the + same four experts.""" + if m not in _STATES: + sms = torch.cuda.get_device_properties(0).multi_processor_count + _STATES[m] = moe_state( + torch.device("cuda"), num_tokens=m, num_ctas=sms - RESERVED_SMS, shared=True + ) + key = (m, case) + if key not in _LAYERS: + ops = operands.experts(OPERAND_SEED + RANK) + bias = ops.bias.clone() + if case == "concentrated": + bias[list(operands.CONCENTRATED)] += 1.0 + _LAYERS[key] = _STATES[m].layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=ops.shared, + ) + return _LAYERS[key] + + +def _desc(): + flags = WS.flags_for(CLUSTER) + return ar_push_desc( + WS.mc.data_ptr(), flags.data_ptr(), RANK, WORLD, WS.max_tokens, flags.device + ) + + +def _boundary(x, residual, norm_w, quant, prepushed): + prequant.SLOT.clear() + outs = ar.m3_mc_allreduce_norm( + x, residual, norm_w, EPS, WS, use_pdl=True, quant=quant, prepushed=prepushed + ) + assert outs is not None, "the boundary all-reduce declined its operands" + pub = prequant.SLOT.pending + prequant.SLOT.clear() + return outs, (None if pub is None else (pub[3].clone(), pub[4].clone())) + + +def _unpushed(hidden, layer, residual, norm_w, quant): + out = engine_op(hidden, layer) + (normed, res), pub = _boundary(out, residual.clone(), norm_w, quant, 0) + return out.clone(), normed.clone(), res.clone(), pub + + +def _pushed(hidden, layer, residual, norm_w, quant): + out = engine_op(hidden, layer, arp=_desc()) + (normed, res), pub = _boundary(out, residual.clone(), norm_w, quant, 2) + return out.clone(), normed.clone(), res.clone(), pub + + +def _compare(got, ref, quant: bool, where: str) -> None: + _assert_bitwise(got[0], ref[0], f"{where} engine output") + _assert_bitwise(got[2], ref[2], f"{where} updated residual") + _assert_bitwise(got[1], ref[1], f"{where} normed") + if quant: + assert got[3] is not None and ref[3] is not None, f"{where}: no MXFP8 hand-off published" + _assert_bitwise(got[3][0], ref[3][0], f"{where} e4m3 rows") + # Only the scale atoms of the m rows are written; the swizzled buffer's padding is not. + written = sf_written(got[0].shape[0], HIDDEN // 32, got[3][1].device) + _assert_bitwise(got[3][1][written], ref[3][1][written], f"{where} UE8M0 scales") + + +# ---------------------------------------------------------------- checks + + +def check_pushed_boundary_is_the_unpushed_one() -> None: + """Every token count, both routing cases, with and without the MXFP8 hand-off, both call orders.""" + for m in range(1, MAX_TOKENS + 1): + for case in ("random", "concentrated"): + layer = _layer(m, case) + for quant in (False, True): + for seed in (1, 2): + s = ( + 1000 * m + + 10 * seed + + (5 if quant else 0) + + (500 if case == "concentrated" else 0) + ) + hidden = operands.hidden_states(m, s) + residual, norm_w = _residual(s, m), _norm_w(s) + where = f"m={m} {case} quant={quant} seed={seed}" + if seed == 1: + ref = _unpushed(hidden, layer, residual, norm_w, quant) + got = _pushed(hidden, layer, residual, norm_w, quant) + else: + got = _pushed(hidden, layer, residual, norm_w, quant) + ref = _unpushed(hidden, layer, residual, norm_w, quant) + torch.cuda.synchronize() + _compare(got, ref, quant, where) + + +def check_graph_replays() -> None: + """A captured pushed engine + prepushed boundary, replayed on fresh rows, equals the eager unpushed path.""" + for m in (1, 4, 8): + layer = _layer(m, "random") + hidden = operands.hidden_states(m, 90000 + m).clone() + residual = _residual(90000 + m, m).clone() + norm_w = _norm_w(90000 + m) + desc = _desc() + for _ in range( + 2 + ): # warm both parities eagerly; every push is consumed by its prepushed boundary + _pushed(hidden, layer, residual, norm_w, False) + _unpushed(hidden, layer, residual, norm_w, False) + torch.cuda.synchronize() + dist.barrier() + res_in = residual.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + g_out = engine_op(hidden, layer, arp=desc) + g_normed, g_res = ar.m3_mc_allreduce_norm( + g_out, res_in, norm_w, EPS, WS, use_pdl=True, prepushed=2 + ) + torch.cuda.synchronize() + dist.barrier() + for rep in range(6): + s = 91000 + 100 * m + rep + fresh_h, fresh_r = operands.hidden_states(m, s), _residual(s, m) + hidden.copy_(fresh_h) + res_in.copy_(fresh_r) + torch.cuda.synchronize() + dist.barrier() + graph.replay() + torch.cuda.synchronize() + got = (g_out.clone(), g_normed.clone(), g_res.clone(), None) + ref = _unpushed(fresh_h, layer, fresh_r, norm_w, False) + torch.cuda.synchronize() + _compare(got, ref, False, f"graph m={m} replay {rep}") + dist.barrier() + + +CHECKS = [check_pushed_boundary_is_the_unpushed_one, check_graph_replays] + + +def _run_one_rank() -> int: + """Body of one spawned rank: run every check; a failing rank exits nonzero and the launcher kills its + siblings.""" + global RANK, WORLD, ar, sf_written, engine_op, ar_push_desc, moe_state, operands, prequant, WS + os.environ.setdefault("TLLM_DISABLE_MPI", "1") + rank = int(os.environ["RANK"]) + world = int(os.environ["WORLD_SIZE"]) + local = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local) + dist.init_process_group( + "cpu:gloo,cuda:nccl", init_method="env://", device_id=torch.device("cuda", local) + ) + RANK, WORLD = rank, world + assert WORLD == 8, f"the TP8 push is certified at 8 ranks, got {WORLD}" + + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_mc_allreduce as _ar # noqa: E501 + import tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_prequant as _pq # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_allreduce_norm import ( # noqa: E501 + sf_written as _sf_written, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + M3MoeState, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + ar_push_desc as _desc_fn, + ) + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as _op, + ) + + ar, prequant, moe_state, ar_push_desc, engine_op = _ar, _pq, M3MoeState, _desc_fn, _op + sf_written = _sf_written + # Spawned by file path, not collected by pytest: the sibling moe/ operands the plain way. + sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "moe")) + import _m3_moe_operands as _operands + + operands = _operands + WS = _make_workspace() + + for check in CHECKS: + try: + check() + except BaseException: + import traceback # only the rank subprocess's failure path needs it + + print(f"[rank {RANK}] FAILED {check.__name__}", flush=True) + traceback.print_exc() + sys.stdout.flush() + sys.stderr.flush() + # A rank that leaves early wedges its peer's polls; the launcher kills them on this exit. + os._exit(13) + print(f"[rank {RANK}] {check.__name__} passed", flush=True) + dist.barrier() + torch.cuda.synchronize() + print(f"[rank {RANK}] {len(CHECKS)} checks passed", flush=True) + sys.stdout.flush() + os._exit(0) + + +def _spawn_ranks() -> None: + """One worker process per claimed device, each in its own session; on the first nonzero exit or on the + deadline every session is killed.""" + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + world_size = len([d for d in visible.split(",") if d.strip()]) + assert world_size == 2, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + ) + 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", + ) + # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed + # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. + 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), LOCAL_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: Optional[str] = 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}") + + +_SRUN_FLAG = "--srun-rank" + + +def _srun_rank() -> int: + """One rank per Slurm task over two trays (srun -N2 --ntasks-per-node 4): RANK / WORLD_SIZE / LOCAL_RANK from + Slurm; the rendezvous is MASTER_ADDR / MASTER_PORT exported by the job (the first node).""" + os.environ["RANK"] = os.environ["SLURM_PROCID"] + os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] + os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"] + return _run_one_rank() + + +if __name__ == "__main__": + if _SRUN_FLAG in sys.argv: + sys.exit(_srun_rank()) + if _WORKER_FLAG in sys.argv: + sys.exit(_run_one_rank()) + _spawn_ranks() + print("OK") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_rank_job.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_rank_job.py new file mode 100644 index 000000000000..dedb967cd195 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_rank_job.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run this target's collective rank job from pytest, and report what it did. + +The same contract as ``tests/unittest/_torch/modeling_v2/comm/_rank_job.py``: the +launcher beside this file owns the per-device rank processes and the deadline +that kills a wedged collective; this only selects the devices, starts it in a +fresh interpreter (by file path: this tree is not a package) and turns its exit +code into an assertion. + +World size. The tp8 target runs 8 ranks over two GB300 trays of 4 GPUs (one +NVLink domain); a single-node unit-test job cannot form that group. The +collective entries take the world size as a runtime argument and their cells' +specs do not name one (token counts, buckets, push modes, epilogue options), so +the jobs here certify every cell at each world size in ``WORLD_SIZES`` the +node's visible devices can form: 2 and 4 on a GB300 tray, 8 only on a node with +8 devices. A world size the node cannot form is skipped with that reason, never +emulated; the 8-rank, two-tray group itself is exercised only by the target's +multi-node runs. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +TARGET_WORLD_SIZE = 8 +"""tp8's world size: 8 ranks over two GB300 trays, which one node's job cannot form.""" + +WORLD_SIZES = (2, 4, TARGET_WORLD_SIZE) +"""The world sizes the jobs certify the (world-independent) cells at, where the node has that many devices.""" + +_LAUNCHER_GRACE_S = 300 +"""Headroom over the launcher's own deadline, so its message wins the race.""" + + +def visible_devices() -> list: + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible: + return [d for d in visible.split(",") if d.strip()] + import torch + + return [str(i) for i in range(torch.cuda.device_count())] + + +def require_world(world_size: int) -> str: + """The first ``world_size`` visible devices; skips (with the reason) when the node has fewer.""" + devices = visible_devices() + if len(devices) < world_size: + where = ( + "the tp8 target's own 8-rank group spans two trays and is exercised only by its multi-node runs" + if world_size == TARGET_WORLD_SIZE + else "run this job on a node with more devices" + ) + pytest.skip( + f"{world_size} ranks need {world_size} devices; {len(devices)} visible on this node ({where})" + ) + return ",".join(devices[:world_size]) + + +def _load_launcher_constants(launcher: Path): + """Read the launcher's module-level budget without importing it by name (it imports only torch at + module scope).""" + spec = importlib.util.spec_from_file_location(launcher.stem, launcher) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def run(entry: str, world_size: int) -> str: + """Run ``__op_matrix``'s launcher over ``world_size`` devices (skipped when the node has fewer); returns + its stdout.""" + assert world_size in WORLD_SIZES, world_size + launcher = Path(__file__).resolve().with_name(f"_{entry}_op_matrix.py") + env = dict(os.environ, CUDA_VISIBLE_DEVICES=require_world(world_size)) + + # The ranks import tensorrt_llm absolutely. From a source checkout that is not installed, put the repo root + # (tests/unittest/_torch/modeling_v2/models/minimax_m3//comm -> repo root) on their path; from a tests + # tree run against an installed package, they import the installed one. + repo_root = Path(__file__).resolve().parents[8] + if (repo_root / "tensorrt_llm").is_dir(): + env["PYTHONPATH"] = os.pathsep.join( + p for p in (str(repo_root), env.get("PYTHONPATH", "")) if p + ) + else: + assert importlib.util.find_spec("tensorrt_llm") is not None, ( + f"no tensorrt_llm/ at the repo root {repo_root} and no installed tensorrt_llm" + ) + + constants = _load_launcher_constants(launcher) + timeout = constants.DEADLINE_S + _LAUNCHER_GRACE_S + + completed = subprocess.run( + [sys.executable, str(launcher)], + env=env, + capture_output=True, + text=True, + timeout=timeout, + ) + if completed.returncode != 0: + raise AssertionError( + f"{launcher.name} exited {completed.returncode} at world size {world_size}\n" + f"--- stdout ---\n{completed.stdout}\n" + f"--- stderr ---\n{completed.stderr}" + ) + return completed.stdout diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_allreduce.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_allreduce.py new file mode 100644 index 000000000000..c19a7461a368 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_allreduce.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the ``m3_allreduce`` certification matrix. + +The matrix is ``_m3_allreduce_op_matrix.py`` beside this file, its own launcher; see ``_rank_job``. It drives every +cell of the entry against its reference, then requires the fused draft loop's boundaries built on it (plain +all-reduce, then the fused add + RMSNorm) to equal stock Eagle3's bit for bit under every single-node strategy, and +reports both against an fp64 reference as a warning, one line per strategy and case. + +The tp8 drafter runs at 8 ranks over two trays; a single-node job cannot form that group, so the matrix runs at each +world size of ``_rank_job.WORLD_SIZES`` the node's devices can form (the claim, draft boundary = stock boundary, does +not depend on the rank count) and skips the others. MNNVL, the strategy the two-tray group may take, is multi-node +only and not driven here. +""" + +import warnings + +import _rank_job +import pytest +import torch + +# Per world size, _rank_job.require_world skips the sizes this node's devices cannot form. +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +_RESULT_TAG = "[ulp]" + + +# Spawns one rank process per device of the world size; under xdist workers would fight for them. +@pytest.mark.no_xdist +@pytest.mark.parametrize("world_size", _rank_job.WORLD_SIZES, ids=lambda w: f"world{w}") +def test_m3_allreduce_op_matrix(world_size) -> None: + out = _rank_job.run("m3_allreduce", world_size) + lines = [line for line in out.splitlines() if line.startswith(_RESULT_TAG)] + assert lines, f"the launcher reported no measurements:\n{out}" + warnings.warn( + f"draft boundary (= stock Eagle3) vs fp64 at {world_size} ranks, max BF16 ulps\n" + + "\n".join(lines), + stacklevel=1, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py new file mode 100644 index 000000000000..0a9a7d39c177 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collected entry point for the multicast all-reduce + RMSNorm entries' certification matrix. + +Covers both ``m3_mc_allreduce_norm`` and ``m3_mc_finalize_allreduce_norm`` (one +kernel, one exchange protocol). The matrix is ``_m3_mc_allreduce_norm_op_matrix.py`` +beside this file, its own launcher; it drives every cell of each entry's +``CELLS`` against the entry's ``reference``. See ``_rank_job`` for why. + +The tp8 target runs these entries at 8 ranks over two trays, which a single-node +job cannot form. The world size is a runtime argument of the kernel and no cell's +spec names one, so the matrix certifies every cell at each world size of +``_rank_job.WORLD_SIZES`` the node's devices can form (2 and 4 on a 4-device tray, +8 on an 8-device node) and skips the others; it does not emulate the two-tray group. +""" + +import _rank_job +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_mc_finalize_allreduce_norm import ( # noqa: E501 + finalize_partial, +) + +assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" + + +# Spawns one rank process per device of the world size; under xdist workers would fight for them. +@pytest.mark.no_xdist +@pytest.mark.parametrize("world_size", _rank_job.WORLD_SIZES, ids=lambda w: f"world{w}") +def test_m3_mc_allreduce_norm_op_matrix(world_size) -> None: + _rank_job.run("m3_mc_allreduce_norm", world_size) + + +def test_the_cells_are_world_independent() -> None: + """What lets a single-node job certify the tp8 entries: no cell's spec names a world size.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm import ( # noqa: E501 + m3_mc_allreduce_norm, + m3_mc_finalize_allreduce_norm, + ) + + for op in ( + m3_mc_allreduce_norm.m3_mc_allreduce_norm, + m3_mc_finalize_allreduce_norm.m3_mc_finalize_allreduce_norm, + ): + assert op.CELLS + for cell in op.CELLS: + assert not {"world", "world_size", "rank"} & set(cell.spec), cell.why + + +def test_finalize_partial_rounds_the_routed_sum_before_the_shared_expert() -> None: + """As stock: the finalized routed output is BF16, and the shared expert is added to it (shared.add_(routed)). + Four routed rows of 1.0 at FP32 routing weight 0.2509765625 sum to 1.00390625, which rounds to BF16 1.0; plus the + shared -1 that is 0, where adding before the rounding would leave 0.00390625.""" + fc2 = torch.ones((4, 8), dtype=torch.bfloat16) + perm = torch.arange(4, dtype=torch.int32).view(1, 4) + scales = torch.full((1, 4), 0.2509765625, dtype=torch.float32) + shared = torch.full((1, 8), -1.0, dtype=torch.bfloat16) + out = finalize_partial(fc2, perm, scales, shared) + routed = finalize_partial(fc2, perm, scales, None) + assert torch.equal(out, shared.clone().add_(routed)) + assert torch.equal(out, torch.zeros(1, 8, dtype=torch.bfloat16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py new file mode 100644 index 000000000000..546ad3e11ef7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_embedding entry of the MiniMax-M3 tp8 target: an exact gather with zero rows outside the +table.""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.embedding.m3_embedding import ( # noqa: E501 + m3_embedding as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" + +VOCAB, HIDDEN = 4096, 6144 + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + table = torch.randn(VOCAB, HIDDEN, generator=g, device="cuda").to(torch.bfloat16) + ids = torch.randint( + 0, VOCAB, (spec["rows"],), generator=g, device="cuda", dtype=spec["id_dtype"] + ) + if spec["pad"]: + ids[1::3] = -1 + with validating(op): + out = op(ids, table) + assert torch.equal(out, op.reference(ids, table)) + if spec["pad"]: + assert not out[1::3].any(), "padded ids must embed to zero rows" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py new file mode 100644 index 000000000000..e44ad56d743d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" + +import torch + + +def swizzle(byte: torch.Tensor) -> torch.Tensor: + """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" + rows, cols = byte.shape + padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + + (c // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (r // 128) * (128 * padded_cols) + ) + out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): + """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" + v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) + sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) + return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py new file mode 100644 index 000000000000..509c1c09cebd --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_bf16_linear entry of the MiniMax-M3 tp8 target (the drafter's projections).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_bf16_linear import ( # noqa: E501 + m3_bf16_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the + row, RMS well under one; measured on GB300: worst 0.5 unit.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w = (torch.randn(spec["n"], spec["k"], generator=g, device="cuda") * 0.02).to(torch.bfloat16) + with validating(op): + out = op(x, w) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(x, w), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py new file mode 100644 index 000000000000..67796dc0302f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise oracle of the dense MLP's swigluoai C++ op (ops/gemm/m3_dense_swiglu, +trtllm::minimax_m3_sm107_tp8_swiglu_oai) against the Triton silu_and_mul kernel the dense MLP runs today +(modules/swiglu.py at limit 7, alpha 1.702, beta 1): every bf16 gate bit pattern against special and random up values, +every up bit pattern against special and random gates, random pairs, and the dense MLP's shapes at every decode row +count, eager and replayed from a CUDA graph, with and without the programmatic launch.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.silu_and_mul +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_dense_swiglu import ( # noqa: E501 + ALPHA, + BETA, + LIMIT, + m3_dense_swiglu, +) + +assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_swiglu_oai"): + pytest.skip( + "needs the TP8 op library with minimax_m3_sm107_tp8_swiglu_oai", allow_module_level=True + ) + +# The TP8 dense MLP's activation width (gate_up 3072 -> d 1536). +DENSE_D = modeling.DENSE_DOWN[1] + + +def _triton(x: torch.Tensor) -> torch.Tensor: + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=LIMIT, swiglu_alpha=ALPHA, swiglu_beta=BETA + ) + + +def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" + return out + + +def _row(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + return torch.cat([gate.view(1, -1), up.view(1, -1)], dim=1).contiguous() + + +def _all_bf16() -> torch.Tensor: + return ( + torch.arange(1 << 16, dtype=torch.int32, device="cuda").to(torch.int16).view(torch.bfloat16) + ) + + +def _specials(n: int, seed: int) -> list: + g = torch.Generator(device="cuda").manual_seed(seed) + vals = [0.0, -0.0, 1.0, -1.0, 6.99, 7.0, 7.01, -7.0, -7.01, 1e-38, -1e-38, 1e-40, 88.0, -88.0, 3e38, -3e38, + float("inf"), float("-inf"), float("nan")] # fmt: skip + rows = [torch.full((n,), v, dtype=torch.bfloat16, device="cuda") for v in vals] + for scale in (0.05, 1.0, 4.0, 30.0): + rows.append((scale * torch.randn(n, generator=g, device="cuda")).to(torch.bfloat16)) + rows.append( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + return rows + + +def _assert_bitwise(got: torch.Tensor, ref: torch.Tensor, what: str) -> None: + bad = (got.view(torch.int16) != ref.view(torch.int16)).nonzero().flatten() + assert bad.numel() == 0, f"{what}: {bad.numel()} mismatches, first at {bad[:4].tolist()}" + + +def test_every_gate_pattern_is_bitwise_the_triton_kernel() -> None: + gate = _all_bf16() + for j, up in enumerate(_specials(gate.numel(), 1)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"up set {j}") + + +def test_every_up_pattern_is_bitwise_the_triton_kernel() -> None: + up = _all_bf16() + for j, gate in enumerate(_specials(up.numel(), 2)): + x = _row(gate, up) + _assert_bitwise(_op(x), _triton(x), f"gate set {j}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: + g = torch.Generator(device="cuda").manual_seed(100 + seed) + n = 1 << 22 + + def pattern() -> torch.Tensor: + bits = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + return bits.to(torch.int16).view(torch.bfloat16) + + x = _row(pattern(), pattern()) + _assert_bitwise(_op(x), _triton(x), f"seed {seed}") + + +@pytest.mark.parametrize("pdl", [True, False], ids=["pdl", "plain"]) +@pytest.mark.parametrize("rows", range(1, modeling.DECODE_MAX_TOKENS + 1)) +def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pdl: bool) -> None: + """The TP8 dense MLP's gate_up rows (a GEMV output's range plus values past the clamps), eager and replayed from a + CUDA graph after fresh rows are written in place, against the Triton kernel on the same rows.""" + g = torch.Generator(device="cuda").manual_seed(2000 + rows) + x = (3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, DENSE_D : DENSE_D + 8] = x[:, :8] + _assert_bitwise(_op(x, pdl), _triton(x), f"eager rows {rows}") + out = torch.empty(rows, DENSE_D, dtype=torch.bfloat16, device="cuda") + _op(x, pdl) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out.copy_(_op(x, pdl)) + x.copy_((3.0 * torch.randn(rows, 2 * DENSE_D, generator=g, device="cuda")).to(torch.bfloat16)) + graph.replay() + torch.cuda.synchronize() + _assert_bitwise(out, _triton(x), f"graph rows {rows}") + + +def test_the_op_declines_other_parameters_and_layouts() -> None: + x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + + +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: + """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton + kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" + rows, d = cell.spec["rows"], cell.spec["d"] + g = torch.Generator(device="cuda").manual_seed(3000 + rows) + x = (3.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :8] = torch.tensor([8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16) + x[:, d : d + 8] = x[:, :8] + _assert_bitwise(_op(x), _triton(x), cell.why) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py new file mode 100644 index 000000000000..e63d7f3f11ed --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_draft_lowm entry: the Eagle3 drafter's BF16 linears on FlashInfer's fixed-tactic low-M +kernels (direct SIMT and split-K tensor-core), at the TP8 drafter shards, at the rows its table certifies. + +Every cell is held to the fp64 reference within the derived band. The drafter's projection call +(``M3Eagle3Draft.linear``) at a covered row count is the entry's call bit for bit, at an uncovered one the in-tree +GEMM's. +""" + +import zlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + ROWS, + TACTICS, + tactic_for, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + m3_draft_lowm as op, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" + +# Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). +# +# Kernel and reference multiply the same BF16 operands (every product exact in fp32); the kernel sums in fp32 (the +# direct kernel one dot product in K order, the split-K kernel tensor-core blocks and then its split partials), the +# reference in fp64, and both round once to BF16. With w ~ N(0, 1 / K) and x ~ N(0, 1) the outputs are O(1) and the +# fp32 order difference is ~sqrt(K) * 2^-24 of the summed magnitudes, a few 1e-6 at K = 18432: far below half a BF16 +# ulp, so the two roundings disagree by at most one ulp of the element -- up to two units of the row scale for an +# element in its row maximum's binade -- and only where the sum sits that close to a rounding boundary, a few in ten +# thousand elements, hence 0.1 ulp RMS. The same derivation as the decode GEMV's band (test_modeling_v2_m3_mx_gemv), +# with no slack under it: a second ulp would mean a wrong sum. +ELEMENT_ULP = 2.0 +RMS_ULP = 0.1 + +#: The drafter's projections (N, K) and their weight keys on ``M3Eagle3Draft.w``. +DRAFTER_LINEARS = { + (1280, 12288): "qkv", + (6144, 1024): "o", + (4608, 6144): "gu", + (6144, 2304): "dn", + (6144, 18432): "fc", +} + + +@pytest.fixture(autouse=True) +def _policy(): + modeling.set_require(False) + modeling.reset_logged() + yield + modeling.set_require(False) + modeling.reset_logged() + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + return x, w + + +def test_the_table_is_the_tp8_drafter() -> None: + assert set(TACTICS) == set(DRAFTER_LINEARS) + assert {c.spec["nk"] for c in op.CELLS} == set(TACTICS) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + kind, tactic = tactic_for(n, k, m) + out = torch.full((m, n), float("nan"), dtype=torch.bfloat16, device="cuda") + args = ( + x, + w, + out, + kind, + tactic, + True, + ) # the target launches with PDL when the environment enables it + certifying(op) + assert op(*args) is out + ref = op.reference(*args) + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=RMS_ULP) + # PDL only changes the launch attribute. + plain = torch.empty_like(out) + op(x, w, plain, kind, tactic, False) + assert torch.equal(plain.view(torch.int16), out.view(torch.int16)) + + +def _draft(linears=DRAFTER_LINEARS): + """A drafter holding only its projection weights and the low-M bookkeeping (the module is not built).""" + draft = object.__new__(modeling.M3Eagle3Draft) + draft.w = {key: build(*nk, 1, seed=nk[0] + nk[1])[1] for nk, key in linears.items()} + draft._lowm_built, draft._lowm_engaged = set(), set() + draft._push_arm, draft._pushed = None, 0 + return draft + + +@pytest.fixture(scope="module") +def prebuilt(): + draft = _draft() + assert draft.prebuild_linears() == len(op.CELLS) + assert draft.prebuild_linears() == 0, "a second prebuild finds nothing left to compile" + return draft + + +@pytest.mark.parametrize("nk", sorted(DRAFTER_LINEARS), ids=lambda nk: DRAFTER_LINEARS[nk]) +def test_the_projection_is_the_entry_and_falls_back_past_its_rows(prebuilt, nk) -> None: + n, k = nk + key = DRAFTER_LINEARS[nk] + w = prebuilt.w[key] + use_pdl = get_env_enable_pdl() + modeling.set_require(True) + for m in ROWS[nk]: + assert modeling.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt.linear(x, key) + kind, tactic = tactic_for(n, k, m) + want = op(x, w, torch.empty(m, n, dtype=x.dtype, device="cuda"), kind, tactic, use_pdl) + assert torch.equal(got.view(torch.int16), want.view(torch.int16)), m + # Rows the table does not certify (the b2 rows, and past it): the in-tree GEMM, silently, even under require. + for m in sorted(set(range(1, modeling.DECODE_MAX_TOKENS + 2)) - set(ROWS[nk])): + assert not modeling.DRAFT_LOWM.covers((nk, m)) + x, _ = build(n, k, m, seed=m) + got = prebuilt.linear(x, key) + assert torch.equal(got.view(torch.int16), F.linear(x, w).view(torch.int16)), m + + +def test_a_projection_outside_the_table_is_the_in_tree_gemm() -> None: + draft = _draft({(256, 512): "o"}) + x, _ = build(256, 512, 1, seed=3) + assert draft._lowm_linear(draft.w["o"], x) is None + got = draft.linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_a_covered_call_not_prebuilt_raises_under_require() -> None: + draft = _draft({(6144, 1024): "o"}) + x, _ = build(6144, 1024, 1, seed=5) + modeling.set_require(True) + with pytest.raises(modeling.FastPathUnavailable, match="not compiled before warmup"): + draft.linear(x, "o") + # Outside require it falls back to the in-tree GEMM. + modeling.set_require(False) + got = draft.linear(x, "o") + assert torch.equal(got.view(torch.int16), F.linear(x, draft.w["o"]).view(torch.int16)) + + +def test_guard_refuses_a_strided_activation() -> None: + x, w = build(6144, 1024, 2, seed=6) + wide = torch.zeros(2, 2048, dtype=torch.bfloat16, device="cuda") + wide[:, :1024] = x + kind, tactic = tactic_for(6144, 1024, 2) + with pytest.raises(AssertionError, match="contiguous"): + with validating(op): + op( + wide[:, :1024], + w, + torch.empty(2, 6144, dtype=torch.bfloat16, device="cuda"), + kind, + tactic, + False, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py new file mode 100644 index 000000000000..149d230a929c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_fp32_linear entry of the MiniMax-M3 tp8 target (the router beyond the decode routers' rows).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_fp32_linear import ( # noqa: E501 + m3_fp32_linear as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """torch's default fp32 matmul precision on this hardware may be TF32 (~2^-11 relative per product): the band is + stated in fp32 ulps of each row's largest logit and covers it (K = 6144 TF32 products: ~2^13 ulps worst element); + measured on GB300 with the engine's defaults: worst element ~1700 ulps, RMS ~300.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = (torch.randn(spec["m"], spec["k"], generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(spec["e"], spec["k"], generator=g, device="cuda") * 0.02 + with validating(op): + out = op(x, w) + assert out.dtype == torch.float32 and out.shape == (spec["m"], spec["e"]) + assert_within_ulp(out, op.reference(x, w), element_ulp=8192.0, rms_ulp=2048.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py new file mode 100644 index 000000000000..918aa9e5ffab --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv entry: the decode MXFP8 GEMV on a producer-quantized activation. + +Operands as the target holds them: the weight through mxfp8_quantize with the swizzled scale layout (the MXFP8 linear +method's storage), the activation through mxfp8_quantize with the swizzled layout (what the boundary all-reduce and +the attention merge publish). +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" + + +# Element band, in ULP units of the row scale: see compare. +ELEMENT_ULP = 1.5 + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """Element ELEMENT_ULP of the row scale, relative RMS 0.1 ulp. + + Kernel and reference multiply bit-identical dequantized operands; the kernel sums in fp32 (tensor-core blocks, + then the split-K partials), the reference in fp64. Over K <= 6144 that order difference is ~1e-6 relative, so the + two bf16 roundings disagree only where the sum sits within that of a rounding boundary: one bf16 ulp of that + element, which for an element in its row maximum's binade is up to two of ``assert_within_ulp``'s units of the row + scale, and a few in ten thousand elements in aggregate. Any fp32 sum shows it: torch's IEEE fp32 matmul of the + same dequantized operands measured at most 1.22 units against the fp64 reference over every row count 1..48 on + the o_proj and the QKV (6 seeds each, GB300), the kernel at most 1.25, both at most 0.009 ulp RMS. The band + leaves the margin between that and a full extra ulp. The kernel is the same on SM107 (the tp2 table only picks + split and ring), and its TP2 projections' K is 1536..6144. + """ + assert_within_ulp(out, ref, element_ulp=ELEMENT_ULP, rms_ulp=0.1) + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + return x, xq.view(m, k), xsf, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + _, xq, xsf, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(xq, xsf, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + compare(out, op.reference(xq, xsf, w, wsf)) + + +def test_an_unserved_shape_returns_none() -> None: + """Shapes outside the tuned table fall back at the call site; the op says so rather than computing.""" + _, xq, xsf, w, wsf = build(256, 512, 1, seed=1) + assert op(xq, xsf, w, wsf) is None + + +def test_guard_refuses_linear_activation_scales() -> None: + """A linear [M, K / 32] scale buffer is read as the swizzled one.""" + x, _, _, w, wsf = build(6144, 4096, 4, seed=2) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, False) + with pytest.raises(AssertionError, match="swizzled"): + with validating(op): + op(xq.view(4, 4096), xsf, w, wsf) + + +def test_guard_refuses_a_strided_weight() -> None: + _, xq, xsf, w, wsf = build(6144, 4096, 1, seed=3) + wide = torch.zeros(6144, 8192, dtype=w.dtype, device="cuda") + wide[:, :4096] = w + with pytest.raises(AssertionError, match="contiguous e4m3"): + with validating(op): + op(xq, xsf, wide[:, :4096], wsf) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py new file mode 100644 index 000000000000..c8c9461b44da --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the gemm/m3_mx_gemv_bf16 entry: the decode MXFP8 GEMV quantizing a bf16 activation in its prologue. + +The weight as the MXFP8 linear method stores it (mxfp8_quantize, swizzled scales). Besides the reference band, the +prologue's quantization is held to mxfp8_quantize itself: the GEMV on a bf16 x must equal, bit for bit, the +prequantized GEMV on mxfp8_quantize(x) -- the claim the target relies on when it mixes the two paths. +""" + +import zlib + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv_bf16 import ( # noqa: E501 + m3_mx_gemv_bf16 as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" + +# The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. +ELEMENT_ULP = 1.5 + + +def build(n: int, k: int, m: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = torch.randn(m, k, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + return x, wq.view(n, k), wsf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + (n, k), m = cell.spec["nk"], cell.spec["m"] + x, w, wsf = build(n, k, m, seed=zlib.crc32(cell.why.encode()) % 2**31) + with validating(op): + out = op(x, w, wsf, use_pdl=True) # the target launches with PDL + assert out is not None, "a certified shape must be served" + assert out.shape == (m, n) and out.dtype is torch.bfloat16 + # The band of test_modeling_v2_m3_mx_gemv.compare: the quantized operands are the reference's exactly (checked + # below), the sums differ only in order. + assert_within_ulp(out, op.reference(x, w, wsf), element_ulp=ELEMENT_ULP, rms_ulp=0.1) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + pre = m3_mx_gemv(xq.view(m, k), xsf, w, wsf, use_pdl=True) + assert pre is not None and torch.equal(out.view(torch.int16), pre.view(torch.int16)), ( + "the prologue's MXFP8 quantization differs from mxfp8_quantize" + ) + + +def test_a_strided_activation_is_declined() -> None: + """The prologue reads x densely; the op declines a strided x (None) rather than computing on it.""" + x, w, wsf = build(6144, 4096, 2, seed=1) + wide = torch.zeros(2, 8192, dtype=torch.bfloat16, device="cuda") + wide[:, :4096] = x + assert op(wide[:, :4096], w, wsf) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py new file mode 100644 index 000000000000..659b302d73ba --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_producer.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Oracle for #117, the sparse QKV + index GEMV with the producer in its owner epilogue (mx_gemv +``m3_mx_gemv_producer``): q, index-q and every cache byte (NVFP4 K / V data and scales, FP8 index-K) bit for bit equal +to today's chain, the GEMV storing the packed BF16 row and then the producer kernel on it. Both producer builds are +references: the library's ``minimax_m3_producer_pw`` and, when the target's extension is loaded, its own +``minimax_m3_sm107_tp8_producer_pw`` (the copy the target dispatches, with its pre-wait step-input reads).""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv import ( # noqa: E501 + op as gemv, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +N, K = 12 * 128, 6144 +HEADS_Q = 8 +PAGE = 128 +PAGES = 512 +MAX_POSITIONS = 262144 +EPS = 1e-6 +KV_QUANT = (1.0, 0.37, 0.61) + + +def _case(tokens, seed, scale=1.0, bad_slots=False): + g = torch.Generator(device="cuda").manual_seed(seed) + w = (torch.randn(N, K, generator=g, device="cuda") * (scale / K**0.5)).to(torch.bfloat16) + x = torch.randn(tokens, K, generator=g, device="cuda").to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + xq, xsf = torch.ops.trtllm.mxfp8_quantize(x, True) + rotary = torch.randn(MAX_POSITIONS, 2, 32, device="cuda", generator=g) + norms = [(torch.randn(128, device="cuda", generator=g) * 0.3).bfloat16() for _ in range(4)] + slots = torch.randperm(PAGES * PAGE, device="cuda", generator=g)[:tokens].int() + if bad_slots and tokens >= 3: + slots[1] = -1 + slots[2] = PAGES * PAGE + 5 + positions = torch.randint( + 0, MAX_POSITIONS, (tokens,), device="cuda", dtype=torch.int32, generator=g + ) + quant = torch.tensor(KV_QUANT, device="cuda") + return xq.view(tokens, K), xsf, wq.view(N, K), wsf, slots, quant, norms, rotary, positions + + +def _caches(): + return ( + torch.zeros(PAGES, 2, 1, PAGE, 64, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 2, 1, PAGE, 8, dtype=torch.uint8, device="cuda"), + torch.zeros(PAGES, 1, PAGE, 128, dtype=torch.float8_e4m3fn, device="cuda"), + ) + + +def _chain(case, caches, producer): + xq, xsf, wq, wsf, slots, quant, norms, rotary, positions = case + packed = gemv.m3_mx_gemv(xq, xsf, wq, wsf, use_pdl=True) + assert packed is not None, "the sparse QKV shape must be served" + if producer == "lib": + q, iq = torch.ops.trtllm.minimax_m3_producer_pw( + packed, *caches, slots, quant, EPS, *norms, rotary, positions, HEADS_Q, True, True + ) + else: + q, iq = torch.ops.trtllm.minimax_m3_sm107_tp8_producer_pw( + packed, *caches, slots, quant, EPS, *norms, rotary, positions, True, True, True + ) + torch.cuda.synchronize() + return q, iq + + +def _fused(case, caches): + xq, xsf, wq, wsf, slots, quant, norms, rotary, positions = case + tokens = xq.shape[0] + q = torch.zeros(tokens, HEADS_Q, 128, dtype=torch.float8_e4m3fn, device="cuda") + iq = torch.zeros(tokens, 1, 128, dtype=torch.float8_e4m3fn, device="cuda") + kv_data, kv_scale, index_k = caches + ok = gemv.m3_mx_gemv_producer( + xq, xsf, wq, wsf, q, iq, kv_data, kv_scale, index_k, slots, quant, EPS, *norms, rotary, positions, + use_pdl=True, + ) # fmt: skip + assert ok, "the producer GEMV must serve the sparse QKV shape at <= 8 tokens" + torch.cuda.synchronize() + return q, iq + + +def _producers(): + out = ["lib"] + if hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_producer_pw"): + out.append("tp8") + return out + + +@pytest.mark.parametrize( + "tokens, seed, scale, bad_slots", + [ + (1, 21, 1.0, False), + (2, 22, 1.0, False), + (4, 23, 1.0, False), + (4, 24, 40.0, False), + (8, 25, 1.0, False), + (8, 26, 0.01, False), + (8, 27, 1.0, True), + ], + ids=["t1", "t2", "t4", "t4-large", "t8", "t8-tiny", "t8-bad-slots"], +) +@pytest.mark.parametrize("producer", _producers()) +def test_producer_gemv_is_bitwise_the_chain(tokens, seed, scale, bad_slots, producer) -> None: + case = _case(tokens, seed, scale, bad_slots) + ref_caches = _caches() + rq, riq = _chain(case, ref_caches, producer) + caches = _caches() + q, iq = _fused(case, caches) + names = ("q", "index_q", "kv_data", "kv_scale", "index_k") + for name, a in zip(names, (rq, riq, *ref_caches)): + assert bool(a.contiguous().view(torch.uint8).any()), f"reference {name} is all zero" + for name, a, b in zip(names, (rq, riq, *ref_caches), (q, iq, *caches)): + a8, b8 = a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8) + diff = (a8 != b8).nonzero() + assert diff.numel() == 0, ( + f"{name}: {diff.shape[0]} bytes differ, first at {tuple(diff[0].tolist())}" + ) + + +def test_producer_gemv_is_deterministic_and_graph_replays() -> None: + """Three eager runs and a captured replay write the same bytes.""" + case = _case(4, 31) + outs = [] + for _ in range(3): + caches = _caches() + q, iq = _fused(case, caches) + outs.append([t.contiguous().view(torch.uint8).clone() for t in (q, iq, *caches)]) + xq, xsf, wq, wsf, slots, quant, norms, rotary, positions = case + q = torch.zeros(4, HEADS_Q, 128, dtype=torch.float8_e4m3fn, device="cuda") + iq = torch.zeros(4, 1, 128, dtype=torch.float8_e4m3fn, device="cuda") + caches = _caches() + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + gemv.m3_mx_gemv_producer( + xq, xsf, wq, wsf, q, iq, *caches, slots, quant, EPS, *norms, rotary, positions, use_pdl=True + ) # fmt: skip + graph.replay() + torch.cuda.synchronize() + outs.append([t.contiguous().view(torch.uint8).clone() for t in (q, iq, *caches)]) + for run in outs[1:]: + for a, b in zip(outs[0], run): + assert torch.equal(a, b) + + +def test_declines_outside_its_shape() -> None: + """More than 8 tokens, or a projection other than the 12-head packed one, launches nothing.""" + case = _case(8, 41) + xq, xsf, wq, wsf, slots, quant, norms, rotary, positions = case + wide = torch.cat([xq, xq], dim=0) + q = torch.zeros(16, HEADS_Q, 128, dtype=torch.float8_e4m3fn, device="cuda") + iq = torch.zeros(16, 1, 128, dtype=torch.float8_e4m3fn, device="cuda") + xsf16 = torch.ops.trtllm.mxfp8_quantize( + torch.zeros(16, K, dtype=torch.bfloat16, device="cuda"), True + )[1] + caches = _caches() + assert not gemv.m3_mx_gemv_producer( + wide, xsf16, wq, wsf, q, iq, *caches, slots, quant, EPS, *norms, rotary, positions + ) # fmt: skip diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py new file mode 100644 index 000000000000..4374a72aa8a0 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_gemm entry of the MiniMax-M3 tp8 target (the decode-graph GEMM on a prequantized +activation).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 + m3_mxfp8_gemm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most + 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + act, act_sf = mxfp8(spec["m"], spec["k"], g) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + with validating(op): + out = op(act, act_sf, w, w_sf, torch.bfloat16, False) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp( + out, + op.reference(act, act_sf, w, w_sf, torch.bfloat16, False), + element_ulp=2.0, + rms_ulp=0.25, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py new file mode 100644 index 000000000000..dcff36cd9b93 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_mxfp8_linear entry of the MiniMax-M3 tp8 target: MXFP8 quantize + GEMM on each backend the +engine selects (native CUTLASS, FlashInfer, the per-bucket tuned quantize + GEMM).""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 + m3_mxfp8_linear as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_mxfp8_operands import mxfp8 # noqa: E402 +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" + +_FLAGS = { + "native": (False, False, False, False), + "flashinfer": (False, False, True, False), + "tuned": (True, False, False, False), +} + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and + one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" + spec = cell.spec + if ( + spec["path"] == "flashinfer" + and getattr(torch.ops.trtllm, "flashinfer_mm_mxfp8", None) is None + ): + pytest.skip("this build has no trtllm::flashinfer_mm_mxfp8") + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + x = torch.randn(spec["m"], spec["k"], generator=g, device="cuda").to(torch.bfloat16) + w, w_sf = mxfp8(spec["n"], spec["k"], g) + alpha = torch.ones(1, dtype=torch.float32, device="cuda") + args = (x, w, w_sf, alpha, torch.bfloat16, *_FLAGS[spec["path"]]) + with validating(op): + out = op(*args) + assert out.shape == (spec["m"], spec["n"]) and out.dtype == torch.bfloat16 + assert_within_ulp(out, op.reference(*args), element_ulp=2.0, rms_ulp=0.25) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py new file mode 100644 index 000000000000..457d56dd8698 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for MX-TAC: the pinned decode-graph MXFP8 decision (CuTe DSL quantize + CuTe DSL GEMM) gives the same bits +as every other candidate stock's AutoTuner can pick (TRT-LLM quantize, CUTLASS GEMM), the same on every call.""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.mxfp8_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.modeling import ( # noqa: E501 + _pinned_mxfp8_gemm, + _pinned_mxfp8_quantize, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner + +assert torch.cuda.is_available(), "MX-TAC requires a CUDA device" + +# The TP8 per-rank projections (ops/gemm/m3_mx_gemv) at decode rows the GEMVs do not serve (9..16) and a few beyond. +CASES = [ + ((1536, 6144), 16), + ((1280, 6144), 12), + ((6144, 1024), 9), + ((768, 6144), 16), + ((6144, 384), 12), + ((3072, 6144), 24), + ((6144, 1536), 32), + ((1536, 6144), 48), +] + + +@pytest.mark.parametrize("nk,m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +def test_pinned_decision_matches_every_stock_candidate(nk, m, scale) -> None: + n, k = nk + g = torch.Generator(device="cuda").manual_seed(n + k + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + x = (torch.randn(m, k, generator=g, device="cuda") * scale).to(torch.bfloat16) + x[0, ::32] *= 1e3 # one outlier per 32-element block on the first row + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + wq = wq.view(n, k) + xq, xsf = _pinned_mxfp8_quantize(x) + tq, tsf = torch.ops.trtllm.mxfp8_quantize(x, True) + assert torch.equal(xq.view(torch.uint8).view(m, k), tq.view(torch.uint8).view(m, k)) + assert torch.equal(xsf.view(torch.uint8).flatten(), tsf.view(torch.uint8).flatten()) + xq = xq.view(m, k) + out = _pinned_mxfp8_gemm(xq, xsf, wq, wsf, torch.bfloat16) + runner = FlashInferMXFP8GemmRunner(torch.bfloat16) + for tactic in (FlashInferMXFP8GemmRunner.CUTE_DSL, FlashInferMXFP8GemmRunner.CUTLASS): + ref = runner.forward([xq, xsf, wq, wsf], tactic=tactic) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)), tactic + again = _pinned_mxfp8_gemm(xq, xsf, wq, wsf, torch.bfloat16) + assert torch.equal(out.view(torch.int16), again.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py new file mode 100644 index 000000000000..1826394596eb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Bitwise oracle of the mx_gemv swigluoai prologue arithmetic (ops/mx_gemv/_swiglu_oai_bf16x2) against the +Triton silu_and_mul kernel the dense MLP runs today (modules/swiglu.py, alpha 1.702, limit 7, beta 1): every bf16 gate +bit pattern against special and random up values, and every up bit pattern against special and random gates.""" + +import cutlass +import cutlass.cute as cute +import pytest +import torch +from cutlass.cute.runtime import from_dlpack + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv.m3_mx_gemv_kernel import ( # noqa: E501 + SWIGLU_OAI_ALPHA, + SWIGLU_OAI_LIMIT, + _swiglu_oai_bf16x2, +) + +assert torch.cuda.is_available(), "the swigluoai prologue oracle requires a CUDA device" + +THREADS = 256 + + +@cute.kernel +def _swiglu_words_kernel(gate: cute.Tensor, up: cute.Tensor, out: cute.Tensor, n: cutlass.Int32): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + i = bidx * cutlass.Int32(THREADS) + tidx + if i < n: + out[i] = _swiglu_oai_bf16x2(gate[i], up[i]) + + +@cute.jit +def _swiglu_words(gate: cute.Tensor, up: cute.Tensor, out: cute.Tensor, n: cutlass.Int32): + _swiglu_words_kernel(gate, up, out, n).launch( + grid=((n + THREADS - 1) // THREADS, 1, 1), block=(THREADS, 1, 1) + ) + + +_compiled = {} + + +def _prologue(gate_bf16: torch.Tensor, up_bf16: torch.Tensor) -> torch.Tensor: + """The prologue helper over flat bf16 tensors of even length (pairs packed into int32 words).""" + g = gate_bf16.contiguous().view(torch.int32) + u = up_bf16.contiguous().view(torch.int32) + out = torch.empty_like(g) + args = [from_dlpack(t, assumed_align=4) for t in (g, u, out)] + fn = _compiled.get("f") + if fn is None: + fn = _compiled["f"] = cute.compile(_swiglu_words, *args, cutlass.Int32(g.numel())) + fn(*args, cutlass.Int32(g.numel())) + return out.view(torch.bfloat16) + + +def _triton(gate_bf16: torch.Tensor, up_bf16: torch.Tensor) -> torch.Tensor: + """The dense MLP's activation today: torch.ops.trtllm.silu_and_mul on one [1, 2d] row per call site layout.""" + x = torch.cat([gate_bf16.view(1, -1), up_bf16.view(1, -1)], dim=1) + return torch.ops.trtllm.silu_and_mul( + x, swiglu_limit=SWIGLU_OAI_LIMIT, swiglu_alpha=SWIGLU_OAI_ALPHA, swiglu_beta=1.0 + ).view(-1) + + +def _all_bf16() -> torch.Tensor: + return ( + torch.arange(1 << 16, dtype=torch.int32, device="cuda").to(torch.int16).view(torch.bfloat16) + ) + + +def _specials(n: int, seed: int) -> list: + g = torch.Generator(device="cuda").manual_seed(seed) + vals = [0.0, -0.0, 1.0, -1.0, 6.99, 7.0, 7.01, -7.0, -7.01, 1e-38, -1e-38, 1e-40, 88.0, -88.0, 3e38, -3e38, + float("inf"), float("-inf"), float("nan")] # fmt: skip + rows = [torch.full((n,), v, dtype=torch.bfloat16, device="cuda") for v in vals] + for scale in (0.05, 1.0, 4.0, 30.0): + rows.append((scale * torch.randn(n, generator=g, device="cuda")).to(torch.bfloat16)) + rows.append( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + return rows + + +def _assert_bitwise(got: torch.Tensor, ref: torch.Tensor, what: str) -> None: + bad = (got.view(torch.int16) != ref.view(torch.int16)).nonzero().flatten() + assert bad.numel() == 0, f"{what}: {bad.numel()} mismatches, first at {bad[:4].tolist()}" + + +def test_every_gate_pattern_is_bitwise_the_triton_kernel() -> None: + gate = _all_bf16() + for j, up in enumerate(_specials(gate.numel(), 1)): + _assert_bitwise(_prologue(gate, up), _triton(gate, up), f"up set {j}") + + +def test_every_up_pattern_is_bitwise_the_triton_kernel() -> None: + up = _all_bf16() + for j, gate in enumerate(_specials(up.numel(), 2)): + _assert_bitwise(_prologue(gate, up), _triton(gate, up), f"gate set {j}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed: int) -> None: + g = torch.Generator(device="cuda").manual_seed(100 + seed) + n = 1 << 22 + gate = ( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + up = ( + torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda") + .to(torch.int16) + .view(torch.bfloat16) + ) + _assert_bitwise(_prologue(gate, up), _triton(gate, up), f"seed {seed}") + + +# The GEMV with the activation in its prologue against today's two launches (Triton silu_and_mul, then the bf16 GEMV): +# bit for bit on the dense MLP's down projection at every decode row count. +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E402 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv import ( # noqa: E402, E501 + op as gemv, +) + + +@pytest.mark.parametrize("m", range(1, modeling.DECODE_MAX_TOKENS + 1)) +def test_down_gemv_with_the_activation_prologue_is_bitwise_two_launches(m: int) -> None: + n, k = modeling.DENSE_DOWN + g = torch.Generator(device="cuda").manual_seed(1000 + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + wq, wsf = torch.ops.trtllm.mxfp8_quantize(w, True) + wq = wq.view(n, k) + # A gate_up GEMV output's range, plus values past the clamps. + gate_up = (3.0 * torch.randn(m, 2 * k, generator=g, device="cuda")).to(torch.bfloat16) + gate_up[:, :8] = torch.tensor( + [8.0, -8.0, 7.0, -7.0, 0.0, -0.0, 50.0, -50.0], dtype=torch.bfloat16 + ) + act = torch.ops.trtllm.silu_and_mul( + gate_up, swiglu_limit=SWIGLU_OAI_LIMIT, swiglu_alpha=SWIGLU_OAI_ALPHA, swiglu_beta=1.0 + ) + ref = gemv.m3_mx_gemv_bf16(act.contiguous(), wq, wsf, use_pdl=True) + out = gemv.m3_mx_gemv_swiglu(gate_up, wq, wsf, use_pdl=True) + assert ref is not None and out is not None, ( + "the dense down projection must be served at every decode row count" + ) + assert torch.equal(out.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py new file mode 100644 index 000000000000..5a4dd5928876 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_r1b1_gemv_order_gate_a.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rule 1b gate a for the TP8 decode GEMV order override (ops/mx_gemv/op.STRIDED_K_TILES: strided k-tiles): +for every overridden (projection, rows) the output is the same on 3 eager calls and on a CUDA graph replay +(deterministic), and its max abs / max rel error vs the fp64 product of the dequantized operands stays within the +shipped order's band (the shared table's config on the same operands, over several seeds).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mx_gemv import ( # noqa: E501 + DENSE_DOWN, + DENSE_GATE_UP, + O_PROJ, + QKV, + QKV_SPARSE, + m3_mx_gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + mxfp8_dequant, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv import ( # noqa: E501 + op, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +CASES = [ + (nk, m) for nk in (QKV_SPARSE, QKV, O_PROJ, DENSE_GATE_UP, DENSE_DOWN) for m in (1, 2, 4, 8) +] + + +def _err(y, ref): + d = (y.double() - ref).abs() + return float(d.max()), float((d / ref.abs().amax(dim=1, keepdim=True).clamp_min(1e-30)).max()) + + +@pytest.mark.parametrize("nk, m", CASES, ids=[f"{n}x{k}-m{m}" for (n, k), m in CASES]) +def test_gate_a(nk, m) -> None: + n, k = nk + cfg = op.get_config(n, k, m) + assert cfg is not None and len(cfg) > 2 and cfg[2] == 0, ( + "the override is off or does not cover this call" + ) + shipped = cfg[:2] # the contiguous order + worst_new, worst_ship = [0.0, 0.0], [0.0, 0.0] + for seed in range(4): + g = torch.Generator(device="cuda").manual_seed(1000 * seed + m) + w = (torch.randn(n, k, device="cuda", generator=g) * 2).to(torch.float8_e4m3fn) + ws = torch.randint(118, 124, (n * k // 32,), dtype=torch.uint8, device="cuda", generator=g) + x = (torch.randn(m, k, device="cuda", generator=g) * (0.5 + seed)).to(torch.bfloat16) + q, s = torch.ops.trtllm.mxfp8_quantize(x, True) + ref = ( + mxfp8_dequant(q.view(torch.uint8), s).double() + @ mxfp8_dequant(w.view(torch.uint8), ws).double().t() + ) + ys = [m3_mx_gemv(q, s, w, ws, True, cfg, True, None) for _ in range(3)] + assert all(y is not None for y in ys), "the GEMV declined the override" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + out = [] + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out.append(m3_mx_gemv(q, s, w, ws, True, cfg, True, None)) + graph.replay() + torch.cuda.synchronize() + for y in ys[1:] + [out[0]]: + assert torch.equal(ys[0].view(torch.int16), y.view(torch.int16)), ( + "not deterministic (eager / graph)" + ) + ship = m3_mx_gemv(q, s, w, ws, True, shipped, True, None) + torch.cuda.synchronize() + for acc, y in ((worst_new, ys[0]), (worst_ship, ship)): + e = _err(y, ref) + acc[0], acc[1] = max(acc[0], e[0]), max(acc[1], e[1]) + print( + f"\nR1B1 tp8 {n}x{k} m={m} cfg {cfg}: max abs / rel vs fp64 {worst_new[0]:.3e} / {worst_new[1]:.3e} | shipped " + f"order {worst_ship[0]:.3e} / {worst_ship[1]:.3e}", + flush=True, + ) + # The band: the shipped order's own error (bf16 output rounding dominates), with 25 % headroom. + assert worst_new[0] <= 1.25 * worst_ship[0] and worst_new[1] <= 1.25 * worst_ship[1] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_moe_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_moe_operands.py new file mode 100644 index 000000000000..b3b710ec185f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_moe_operands.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""One tp8 rank's MiniMax-M3 MoE layer operands for the decode-engine entries' tests. + +The routed experts are random checkpoint-format NVFP4 tensors ([w3 (up); w1 (gate)] rows, e4m3 block scales) put +through TRT-LLM's own TRTLLM-Gen NVFP4 loader shuffle (``NVFP4TRTLLMGenFusedMoEMethod._shuffle_all_experts``), which +writes the buffers the engines read in place. The per-expert scalars follow the loader's conventions (act_beta and +act_clamp divided by fc31_alpha, fc31_scale_c = fc2_input_scale x fc31_alpha) and are sized so the FC1 sums land +around the SwiGLU-OAI clamp and the outputs near unit scale. The shared expert's MXFP8 weights come from +mxfp8_quantize with the swizzled scale layout, as the MXFP8 linear method stores them. +""" + +import functools +from types import SimpleNamespace + +import torch + +# Hidden size, the routed experts' intermediate per rank at TP8 (1536 / 8 x 2: every expert on every rank), experts, +# top-k, and the shared expert's intermediate per rank. +H, I_TP, NUM_EXPERTS, TOP_K = 6144, 384, 128, 4 +SHARED_I = 384 +SWIGLU_ALPHA, SWIGLU_LIMIT = 1.702, 7.0 +ROUTED_SCALING = 2.0 +MINIMAX2_ROUTING = 5 # RoutingMethodType.MiniMax2 +# NVFP4 global scales: 448 x 6 over the amax each side is sized for. +FC31_INPUT_SCALE = 448.0 * 6.0 / 8.0 +FC2_INPUT_SCALE = 448.0 * 6.0 / 24.0 + + +def _rand_e4m3_scales(shape, gen) -> torch.Tensor: + """E4M3 block scales in [0.5, 2): biased exponents 6 and 7, any mantissa.""" + return torch.randint(0x30, 0x40, shape, dtype=torch.uint8, device="cuda", generator=gen).view( + torch.float8_e4m3fn + ) + + +@functools.lru_cache(maxsize=None) +def experts(seed: int = 20261013): + """The routed experts and the router of one MoE layer, in the layout and conventions of the TRTLLM-Gen loader. + + Returns a SimpleNamespace: ``tensors`` (the ten expert operands in ``M3MoeState.layer``'s order), ``bias``, + ``router_weight``, ``fc31_input_scale`` and ``shared`` (gate_up, its scales, down, its scales). + """ + from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4TRTLLMGenFusedMoEMethod + + gen = torch.Generator(device="cuda").manual_seed(seed) + e = NUM_EXPERTS + w31 = torch.randint( + 0, 256, (e, 2 * I_TP, H // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + w2 = torch.randint(0, 256, (e, H, I_TP // 2), dtype=torch.uint8, device="cuda", generator=gen) + w31s = _rand_e4m3_scales((e, 2 * I_TP, H // 16), gen) + w2s = _rand_e4m3_scales((e, H, I_TP // 16), gen) + module = SimpleNamespace( + w3_w1_weight=w31, + w2_weight=w2, + w3_w1_weight_scale=w31s, + w2_weight_scale=w2s, + is_gated_activation=True, + bias=False, + ) + NVFP4TRTLLMGenFusedMoEMethod()._shuffle_all_experts(module) + torch.cuda.synchronize() + + # Raw FC1 sums: rms(x) x FC31_INPUT_SCALE x rms(E2M1 value) x rms(block scale) x sqrt(H) ~ 1.0e5 at rms(x) = 1; + # fc31_alpha brings them to ~1.5 in the dequantized domain, the FC2 alpha the outputs to ~1. + base = 1.5 / (FC31_INPUT_SCALE * 2.93 * 1.3 * H**0.5) + fc31_alpha = base * (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) + fc2_alpha = (0.8 + 0.4 * torch.rand(e, device="cuda", generator=gen)) / ( + FC2_INPUT_SCALE * 2.93 * 1.3 * I_TP**0.5 + ) + tensors = ( + w31, + w31s, + w2, + w2s, + fc31_alpha, + FC2_INPUT_SCALE * fc31_alpha, # fc31_scale_c + torch.full((e,), SWIGLU_ALPHA, device="cuda"), # act_alpha + 1.0 / fc31_alpha, # act_beta: the (up + 1) offset over fc31_alpha + SWIGLU_LIMIT / fc31_alpha, # act_clamp over fc31_alpha + fc2_alpha, + ) + bias = (0.02 * torch.randn(e, device="cuda", generator=gen)).float() + router_weight = (torch.randn(e, H, device="cuda", generator=gen) / H**0.5).float().contiguous() + fc31_input_scale = torch.tensor([FC31_INPUT_SCALE], dtype=torch.float32, device="cuda") + + gu = (0.02 * torch.randn(2 * SHARED_I, H, device="cuda", generator=gen)).to(torch.bfloat16) + dn = (0.05 * torch.randn(H, SHARED_I, device="cuda", generator=gen)).to(torch.bfloat16) + gu_q, gu_sf = torch.ops.trtllm.mxfp8_quantize(gu, True) + dn_q, dn_sf = torch.ops.trtllm.mxfp8_quantize(dn, True) + shared = (gu_q.view(2 * SHARED_I, H), gu_sf, dn_q.view(H, SHARED_I), dn_sf) + return SimpleNamespace( + tensors=tensors, + bias=bias, + router_weight=router_weight, + fc31_input_scale=fc31_input_scale, + shared=shared, + ) + + +# Routing cases: random (up to 4 M distinct experts), and every token on the same four experts (the bias outweighs +# any sigmoid difference), which leaves later slot groups empty. +CONCENTRATED = (5, 40, 77, 126) + + +def layer_for(state, case: str, shared: bool = True): + """One layer on ``state`` over :func:`experts` for a routing case. ``shared=False``: the shared expert's weights + are zero, so the engine returns the routed part alone (bf16(0 + bf16(routed))).""" + ops = experts() + bias = ops.bias.clone() + if case == "concentrated": + bias[list(CONCENTRATED)] += 1.0 + shared_weights = ops.shared + if not shared: + gu, gu_sf, dn, dn_sf = ops.shared + shared_weights = (torch.zeros_like(gu), gu_sf, torch.zeros_like(dn), dn_sf) + return state.layer( + *ops.tensors, + bias, + router_weight=ops.router_weight, + fc31_input_scale=ops.fc31_input_scale, + shared_weights=shared_weights, + ) + + +def stock_routed( + hidden: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor +) -> torch.Tensor: + """The stock TRTLLM-Gen NVFP4 MoE (fp4_quantize, then fp4_block_scale_moe_runner with SwiGLU-OAI) on the given + routing, over :func:`experts`: what the built-in MiniMaxM3MoE computes for the routed part.""" + import tensorrt_llm._torch.custom_ops # noqa: F401 -- registers the stock ops + + ops = experts() + w31, w31s, w2, w2s, fc31_alpha, scale_c, act_alpha, act_beta, act_clamp, fc2_alpha = ops.tensors + xq, xsf = torch.ops.trtllm.fp4_quantize(hidden, ops.fc31_input_scale, 16, False, False) + out = torch.ops.trtllm.fp4_block_scale_moe_runner( + None, None, xq, xsf.view(torch.float8_e4m3fn).flatten(), + w31, w31s, None, act_alpha, act_beta, act_clamp, w2, w2s, None, + scale_c, fc31_alpha, fc2_alpha, + NUM_EXPERTS, TOP_K, None, None, I_TP, 0, NUM_EXPERTS, ROUTED_SCALING, MINIMAX2_ROUTING, True, 0, + topk_weights=topk_weights.to(torch.bfloat16), topk_ids=topk_ids.to(torch.int32), + ) # fmt: skip + return out[0] if isinstance(out, (list, tuple)) else out + + +def hidden_states(m: int, seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(m, H, device="cuda", generator=gen).to(torch.bfloat16) + + +def check_routed_part(op, state, m: int, seed: int) -> None: + """Exact, independent of the reference: the engine's routing (its route_dbg record: ids, then the bf16 weights' + fp32 bits) is the reference routing bit for bit, and its routed part is the stock TRTLLM-Gen MoE's output on that + routing bit for bit -- the claim the engines are built on.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + moe_route, + ) + + layer = layer_for(state, "random", shared=False) + hidden = hidden_states(m, seed=seed) + out = op(hidden, layer) + ids = state.route_dbg[:, :4].long() + weights = state.route_dbg[:, 4:].contiguous().view(torch.float32) + ref_ids, ref_weights = moe_route(hidden.float(), layer.router_weight, layer.bias) + assert torch.equal(ids, ref_ids) and torch.equal(weights, ref_weights), ( + "routing differs from the reference" + ) + stock = stock_routed(hidden, ids, weights) + assert torch.equal(out.view(torch.int16), stock.view(torch.int16)), ( + "routed part differs from TRTLLM-Gen's" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py new file mode 100644 index 000000000000..7a88ec1adca9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py @@ -0,0 +1,86 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared operands and gates of the router entries' tests (m3_router_gemv, m3_router_gemv_fp4, m3_router_tc: one +kernel source).""" + +import zlib + +import torch + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp # noqa: E402 — needs the path declared above + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv import ( # noqa: E501 + m3_router_gemv as gemv, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv import ( # noqa: E501 + near_tie, + nvfp4_quantize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 + m3_router_gemv_fp4 as gemv_fp4, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_tc import ( # noqa: E501 + m3_router_tc as tc, +) + +__all__ = [ + "SF_SCALE", + "_build", + "_compare_fp4", + "_compare_logits", + "_seed", + "gemv", + "gemv_fp4", + "tc", +] + +#: The experts' NVFP4 input scale the tc cells use: 448 * 6 / 8, block scales well inside E4M3. +SF_SCALE = 336.0 + + +def _build(m: int, e: int, k: int, seed: int): + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(m, k, generator=g, device="cuda") * 2.0).to(torch.bfloat16) + w = torch.randn(e, k, generator=g, device="cuda") * 0.02 + return x, w + + +def _seed(cell) -> int: + return zlib.crc32(cell.why.encode()) % 2**31 + + +def _compare_logits(out: torch.Tensor, ref: torch.Tensor) -> None: + """fp32 logits that differ from a true fp32 product only in summation order. + + Over K = 6144 products a reordered fp32 sum moves by a few ulps of the sum of magnitudes, which near + cancellation is ~sqrt(K) times the row's largest logit's ulp at most. Gated in fp32 ulps of each row's largest + logit: 64 element-wise, 8 in relative RMS. Measured on GB300: GEMV 1.6 / 1.4 ulp, tensor-core router 3.8 / 2.0 + ulp at 2..48 tokens; the tensor-core logits against a true fp32 product (the TF32 rounding the band must see) + sit at 2200+ / 1700 ulp. + """ + assert out.dtype == ref.dtype == torch.float32 and out.shape == ref.shape + assert_within_ulp(out, ref, element_ulp=64.0, rms_ulp=8.0) + + +def _codes(packed: torch.Tensor) -> torch.Tensor: + """[M, K / 2] packed e2m1 bytes -> [M, K] 4-bit codes, low nibble first.""" + lo, hi = packed & 0xF, packed >> 4 + return torch.stack((lo, hi), dim=-1).reshape(packed.shape[0], -1) + + +def _compare_fp4( + q: torch.Tensor, sf: torch.Tensor, x: torch.Tensor, sf_scale: torch.Tensor +) -> None: + """Scales bit-exact; codes bit-exact except within ``TIE_WINDOW`` of an e2m1 midpoint, where the kernel's + approximate reciprocals choose the side (and then only by one code).""" + ref_q, ref_sf, scaled = nvfp4_quantize(x, sf_scale) + assert q.dtype == torch.uint8 and q.shape == ref_q.shape + assert torch.equal(sf, ref_sf), f"{(sf != ref_sf).sum().item()} block scales differ" + got, want = _codes(q), _codes(ref_q) + differ = got != want + assert not bool((differ & ~near_tie(scaled)).any()), ( + f"{(differ & ~near_tie(scaled)).sum().item()} e2m1 codes differ away from a rounding tie" + ) + magnitude = ((got & 7).int() - (want & 7).int()).abs() + assert bool((magnitude[differ] <= 1).all()), "a tie resolved more than one code away" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py new file mode 100644 index 000000000000..8836a2e7726d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the moe/m3_moe_engine entry: the 1..8-token MoE engine on its caller-owned state. + +One Rubin (SM107) rank of the tp8 deployment: 128 experts (all local, routed and shared intermediate 384 per +rank), top 4, hidden 6144, the shared expert folded in, ``num_ctas`` = SMs - 24 as the target builds it. Operands: +``_m3_moe_operands``. +""" + +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + HIDDEN, + INTER, + SHARED_INTER, + M3MoeState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_moe_engine import ( # noqa: E501 + m3_moe_engine as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_moe_operands import ( # noqa: E402 + I_TP, + SHARED_I, + H, + check_routed_part, + hidden_states, + layer_for, +) +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + +_STATES = {} + + +def test_the_operands_are_the_tp8_engine_geometry() -> None: + assert (H, I_TP, SHARED_I) == (HIDDEN, INTER, SHARED_INTER) == (6144, 384, 384) + + +def test_the_engine_refuses_the_tp4_width() -> None: + with pytest.raises(ValueError, match="384"): + M3MoeState(torch.device("cuda"), num_tokens=1, inter=768) + + +def _state(m: int) -> M3MoeState: + if m not in _STATES: + dev = torch.device("cuda") + sms = torch.cuda.get_device_properties(dev).multi_processor_count + _STATES[m] = M3MoeState(dev, num_tokens=m, num_ctas=sms - 24, shared=True) + return _STATES[m] + + +def compare(out: torch.Tensor, ref: torch.Tensor) -> None: + """The engine against the plain-torch MoE: BAND_ELEMENT bf16 ulp of the row's largest magnitude per element, + BAND_RMS ulp relative RMS. + + Both sides quantize the same operands with the same recipes (the activation's and the intermediate's NVFP4 codes + and block scales, the shared expert's MXFP8); they differ in summation order (tensor-core fp32 against fp64) and in + the PTX approximations (rcp / ex2 / tanh). Such a difference changes the output only where it moves a value across + a rounding boundary: one bf16 ulp of a slot output or of the sum, or, rarely, one intermediate NVFP4 code, which + reaches the output as one weight column's share. Measured over every cell, both routing cases and 3 seeds each + on the tp4 engine (GB300, intermediate 768): at most 1.48 ulp element, 0.18 ulp RMS; the tp8 engine runs the + same recipes over half the intermediate. The band leaves ~3x on both; a misread recipe (the activation's + 1 / (sf x (1 / sc)) spelling of the reciprocal on ~0.06% of its codes) measured 12 ulp element, 7 ulp RMS. + """ + assert_within_ulp(out, ref, element_ulp=BAND_ELEMENT, rms_ulp=BAND_RMS) + + +BAND_ELEMENT = 4.0 +BAND_RMS = 0.5 + + +@pytest.mark.parametrize("case", ("random", "concentrated")) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, case) -> None: + m = cell.spec["tokens"] + layer = layer_for(_state(m), case) + hidden = hidden_states(m, seed=zlib.crc32((cell.why + case).encode()) % 2**31) + with validating(op): + out = op(hidden, layer) + again = op(hidden, layer) + assert out.shape == (m, H) and out.dtype is torch.bfloat16 + # The workspace is tagged per call and never reset: a repeated call is the same call. + assert torch.equal(out.view(torch.int16), again.view(torch.int16)), ( + "a repeated call changed the output" + ) + compare(out, op.reference(hidden, layer)) + + +@pytest.mark.parametrize("m", (1, 2, 8)) +def test_routed_part_is_trtllm_gen_bit_for_bit(m) -> None: + check_routed_part(op, _state(m), m, seed=100 + m) + + +def test_layers_on_one_state_do_not_leak_into_each_other() -> None: + """Two layers share one state's workspace, in stream order: A, B, A gives A's output twice, bit for bit, and the + caller's ``out`` receives the same rows as a fresh tensor.""" + state = _state(4) + a, b = layer_for(state, "random"), layer_for(state, "concentrated") + hidden = hidden_states(4, seed=7) + first = op(hidden, a) + op(hidden, b) + out = torch.empty(4, H, dtype=torch.bfloat16, device="cuda") + second = op(hidden, a, out=out) + assert second.data_ptr() == out.data_ptr() + assert torch.equal(first.view(torch.int16), out.view(torch.int16)) + + +@pytest.mark.parametrize( + "bad", + ("misaligned hidden", "fp32 out", "short out", "long out", "int32 arp"), +) +def test_malformed_buffers_are_rejected_by_the_call(bad) -> None: + """Why the entry has no is_valid: the compiled call checks each buffer against its signature and raises.""" + state = _state(2) + layer = layer_for(state, "random") + hidden, kw = hidden_states(2, seed=3), {} + if bad == "misaligned hidden": + hidden = torch.randn(2 * H + 8, device="cuda").to(torch.bfloat16)[2 : 2 + 2 * H].view(2, H) + elif bad == "fp32 out": + kw["out"] = torch.empty(2, H, dtype=torch.float32, device="cuda") + elif bad == "short out": + kw["out"] = torch.empty(1, H, dtype=torch.bfloat16, device="cuda") + elif bad == "long out": + kw["out"] = torch.empty(3, H, dtype=torch.bfloat16, device="cuda") + else: + kw["arp"] = torch.zeros(8, dtype=torch.int32, device="cuda") + with pytest.raises(ValueError): + op(hidden, layer, **kw) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py new file mode 100644 index 000000000000..ca6b047d862d --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_router_tc entry of the MiniMax-M3 tp8 target.""" + +import pytest +import torch + +__extra_import_path__ = ["../../../..", "."] +from _m3_router import ( # noqa: E402, F401 + SF_SCALE, + _build, + _compare_fp4, + _compare_logits, + _seed, + gemv, + gemv_fp4, + tc, +) +from _validating import validating # noqa: E402 — needs the path declared above + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + + +@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +def test_tc_certified_cells(cell) -> None: + spec = cell.spec + x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) + sf_scale = torch.tensor([SF_SCALE], dtype=torch.float32, device="cuda") if spec["fp4"] else None + with validating(tc): + outs = tc(x, w, sf_scale, False) + ref = tc.reference(x, w, sf_scale, False) + assert len(outs) == len(ref) == (3 if spec["fp4"] else 1) + _compare_logits(outs[0], ref[0]) + if spec["fp4"]: + _compare_fp4(outs[1], outs[2], x, sf_scale) + + +def test_tc_reference_is_tf32_and_the_gate_sees_it() -> None: + """The tensor-core router rounds the weight to TF32: against a true fp32 product its logits sit far outside + the band, so the TF32 rounding in ``reference`` is load-bearing rather than decoration.""" + x, w = _build(16, 128, 6144, 5) + out = tc(x, w, None, False)[0] + with pytest.raises(AssertionError): + _compare_logits(out, gemv.reference(x, w, 1, False)) + + +def test_guard_refuses_a_misaligned_activation() -> None: + """A row view starting 8 bytes into the buffer keeps a legal row stride but not a legal bulk-copy source.""" + x, w = _build(2, 128, 6144, 3) + flat = torch.empty(2 * 6144 + 4, dtype=torch.bfloat16, device="cuda") + shifted = flat[4:].view(2, 6144) + shifted.copy_(x) + with pytest.raises(AssertionError, match="16-byte aligned"): + with validating(tc): + tc(shifted, w, None, False) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py new file mode 100644 index 000000000000..05f9e80edbd9 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the norm/m3_embedding_firstnorm entry: the decode embedding gather + the first layer's RMSNorm. + +The table is the model's full bf16 [200064, 6144] embedding shape (random values at the checkpoint's scale), eps the +checkpoint's 1e-6. +""" + +import functools +import zlib + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + VOCAB, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.norm.m3_embedding_firstnorm import ( # noqa: E501 + m3_embedding_firstnorm as op, +) + +__extra_import_path__ = ["../../../.."] +from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" + +H = 6144 +EPS = 1e-6 + + +@functools.lru_cache(maxsize=None) +def _operands(): + g = torch.Generator(device="cuda").manual_seed(20261013) + table = (0.02 * torch.randn(VOCAB, H, generator=g, device="cuda")).to(torch.bfloat16) + weight = (1.0 + 0.1 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + return table, weight + + +def _buffers(rows): + return (torch.empty(rows, H, dtype=torch.bfloat16, device="cuda"), + torch.empty(rows, H, dtype=torch.bfloat16, device="cuda")) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + rows, id_dtype = cell.spec["rows"], cell.spec["id_dtype"] + table, weight = _operands() + g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) + ids = torch.randint(0, VOCAB, (rows,), generator=g, device="cuda").to(id_dtype) + raw, out = _buffers(rows) + with validating(op): + normed, residual = op(ids, table, weight, EPS, raw, out) + assert normed.data_ptr() == out.data_ptr() and residual.data_ptr() == raw.data_ptr() + ref_normed, ref_raw = op.reference(ids, table, weight, EPS, raw, out) + # The gather is a copy: exact. + assert torch.equal(raw.view(torch.int16), ref_raw.view(torch.int16)) + # The norm: the kernel's fp32 sum of squares (warp butterflies, then across 4 warps) and fast rsqrt against an + # fp64 reference, then x * rstd * w rounded once to bf16. The fp32 side is ~1e-6 relative off, so the two roundings + # disagree only next to a bf16 rounding boundary: one ulp at most per element, a small fraction in aggregate. + # Measured over these cells (3 seeds each, GB300): at most 0.51 ulp element, 0.026 ulp RMS. + assert_within_ulp(out, ref_normed, element_ulp=1.0, rms_ulp=0.1) + + +def test_ids_outside_the_table_give_zero_rows() -> None: + """Documented, not refused: an id outside [0, V) gathers a zero row whose norm is zero.""" + table, weight = _operands() + ids = torch.tensor([-1, 5, VOCAB], dtype=torch.int64, device="cuda") + raw, out = _buffers(3) + op(ids, table, weight, EPS, raw, out) + assert not raw[0].any() and not raw[2].any() and not out[0].any() and not out[2].any() + assert torch.equal(raw[1], table[5]) + + +@pytest.mark.parametrize("rows", range(1, 9)) +def test_against_the_stock_embedding_and_rmsnorm(rows) -> None: + """What this target certifies on SM107 against the built-in path it replaces (the model's Embedding gather, then + the first layer's RMSNorm module), including a constant row, large and tiny rows and both vocab ends: the gather + bit for bit, the norm within one bf16 ulp (the sum-of-squares tree differs from flashinfer's: measured ~5e-6 of + the elements one ulp off over 12000 random rows).""" + from tensorrt_llm._torch.modules.rms_norm import RMSNorm + + table, weight = _operands() + table = table.clone() + g = torch.Generator(device="cuda").manual_seed(rows) + table[1] = 3.0 + table[2] = (300.0 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + table[3] = (1e-4 * torch.randn(H, generator=g, device="cuda")).to(torch.bfloat16) + norm = RMSNorm(hidden_size=H, eps=EPS, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + norm.weight.copy_(weight) + for ids in ( + torch.randint(0, VOCAB, (rows,), generator=g, device="cuda"), + torch.tensor([1, 2, 3, 0, VOCAB - 1, 1, 2, 3][:rows], device="cuda"), + ): + raw, out = _buffers(rows) + normed, residual = op(ids, table, weight, EPS, raw, out) + x = torch.nn.functional.embedding(ids, table) + assert torch.equal(residual.view(torch.int16), x.view(torch.int16)) + stock = norm(x) + off = normed != stock + assert int(off.sum()) <= 4, f"{int(off.sum())} normed elements differ from stock" + assert_within_ulp(normed, stock, element_ulp=1.0, rms_ulp=0.05) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py new file mode 100644 index 000000000000..15f7afe40867 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Inputs shared by the tp4 target's spec entry tests: TP4 dimensions, single-device stand-ins for the cross-rank +exchanges, and the comparison bands.""" + +from __future__ import annotations + +import zlib + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + make_keys, +) + +#: MiniMax-M3 at TP8: hidden size, vocabulary, the target head's shard and RMSNorm epsilon. +HIDDEN = 6144 +VOCAB = 200064 +WORLD = 8 +TARGET_SHARD = VOCAB // WORLD +EPS = 1e-6 + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): + """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" + uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") + keys = torch.zeros(slot_rows, dtype=torch.int64, device="cuda") + counters = torch.tensor([0, call], dtype=torch.int32, device="cuda") + return uc, keys, counters + + +def plant_peers( + uc: torch.Tensor, + counters: torch.Tensor, + own_best: torch.Tensor, + rank: int, + world: int, + shard: int, + g: torch.Generator, + slot_rows: int = SHARD_ROWS, + ties: bool = False, +) -> None: + """Write the other ranks' keys for each row into this rank's slots, as their multicast pushes would. + + Each peer's maximum is the own row maximum shifted by +-0.25 (well clear of the accumulation-order band), so + some rows are won by a peer and some by this rank; with ``ties`` a peer on an even row reports exactly the own + maximum, and the lower global index must win. + """ + if world == 1: + return + m = own_best.shape[0] + half = int(counters[1].item()) & 1 + for r in range(world): + if r == rank: + continue + sign = torch.where(torch.rand(m, generator=g, device="cuda") < 0.5, -1.0, 1.0) + value = own_best.float() + 0.25 * sign + if ties: + value = torch.where(torch.arange(m, device="cuda") % 2 == 0, own_best.float(), value) + index = r * shard + torch.randint(0, shard, (m,), generator=g, device="cuda") + keys = make_keys(value, index) + for t in range(m): + uc[(half * slot_rows + t) * world + r] = keys[t] + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py new file mode 100644 index 000000000000..605c96c773cc --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_capture_norm entry (Eagle3 capture with fc_norm folded in).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_capture_norm import ( # noqa: E501 + m3_capture_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" + +LAYERS = 3 + + +def _case(spec, seed): + g = gen(seed) + m, slot = spec["m"], spec["slot"] + h = randn((m, HIDDEN), g, 2.0) + r = randn((m, HIDDEN), g, 2.0) if spec["residual"] else None + w = (1.0 + randn((HIDDEN,), g, 0.1).float()).to(torch.bfloat16) + buf = torch.full((m, LAYERS * HIDDEN), 7.0, dtype=torch.bfloat16, device="cuda") + dst = buf[:, slot * HIDDEN : (slot + 1) * HIDDEN] + return (h, r, w, EPS, dst), buf + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer + untouched.""" + args, buf = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + slot = cell.spec["slot"] + others = torch.cat([buf[:, : slot * HIDDEN], buf[:, (slot + 1) * HIDDEN :]], dim=1) + assert bool((others == 7.0).all()) + + +def test_guard_refuses_a_strided_weight() -> None: + (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + w2 = torch.stack([w, w], dim=1)[:, 0] + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(h, r, w2, eps, dst) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py new file mode 100644 index 000000000000..28d5cbefc927 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU test: a first-draft-inputs reader that skips the drafter's fc_norm (the capture norm is armed) gets the captured +rows unchanged when every captured layer was written with the step's norm, and otherwise a copy with only the raw +layers normalized by their own fc_norm (M3Eagle3Worker._m3_renorm_raw_captures); the per-forward reset clears the +record and the unconsumed fc fork result.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) + +M3Eagle3Worker = modeling.M3Eagle3Worker + +H = 4 +LAYERS = (1, 29, 56) + + +def _worker(normed): + w = object.__new__(M3Eagle3Worker) + w._m3_capture_normed = set(normed) + return w + + +@pytest.fixture(autouse=True) +def _scaling_norm(monkeypatch): + """The RMSNorm the worker applies, as ``x * w`` (CPU): enough to tell which blocks were normalized here.""" + monkeypatch.setattr(modeling, "flashinfer_rmsnorm", lambda x, w, eps: x * w) + modeling.set_require(False) + + +def _drafter(): + # fc_norm i scales by (i + 2). + return SimpleNamespace(w={f"fc_norm{i}": i + 2 for i in range(len(LAYERS))}, rms_eps=1e-6) + + +def test_all_captured_layers_normed_returns_the_rows_as_they_lie() -> None: + hs = torch.randn(4, H * len(LAYERS)) + out = _worker(set(LAYERS))._m3_renorm_raw_captures( + hs, SimpleNamespace(layers_to_capture=LAYERS, hidden_size=H), _drafter() + ) + assert out is hs + + +@pytest.mark.parametrize("normed", [set(), {1, 29}, {29, 56}, {1, 56}]) +def test_raw_layers_get_their_own_fc_norm(normed) -> None: + hs = torch.randn(4, H * len(LAYERS)) + ref = hs.clone() + out = _worker(normed)._m3_renorm_raw_captures( + hs, SimpleNamespace(layers_to_capture=LAYERS, hidden_size=H), _drafter() + ) + assert torch.equal(hs, ref) # the capture buffer itself is not modified + for i, layer_id in enumerate(LAYERS): + block, src = out[:, i * H : (i + 1) * H], ref[:, i * H : (i + 1) * H] + assert torch.equal(block, src if layer_id in normed else src * (i + 2)), layer_id + + +def test_arm_resets_the_record_and_the_fork_result() -> None: + w = _worker(set(LAYERS)) + w._m3_fc_hidden = object() + w._m3_fc_early_stream = None # no fork pending + w._m3_join_pending = lambda: None # (tp8) nothing else pending + w.arm_capture_norm( + None, None, None, None + ) # no drafter: resets the per-forward state and returns + assert w._m3_capture_normed == set() + assert w._m3_fc_hidden is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py new file mode 100644 index 000000000000..7afbbaa79e26 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_target.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The worker's capture target (the capture norm in the boundary all-reduce): the capture buffer's column block of a +captured layer with its fc_norm weight and eps; None without an armed capture norm, for a layer that is not captured, +past the buffer's rows, or when the loaded all-reduce op has no capture epilogue (then the capture_norm kernel runs).""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_mc_allreduce, +) + +H = 64 +LAYERS = (2, 30, 58) + + +def _setup(armed=True): + hidden = torch.arange(8 * 3 * H, dtype=torch.float32).reshape(8, 3 * H).to(torch.bfloat16) + weights = torch.randn(3, H).to(torch.bfloat16) + meta = SimpleNamespace(hidden_states=hidden, hidden_size=H, layers_to_capture=LAYERS) + worker = SimpleNamespace(_m3_capture_norm=(weights, 1e-5) if armed else None) + return worker, meta, hidden, weights + + +def _target(worker, meta, layer_id, num_tokens): + return modeling.M3Eagle3Worker.capture_target(worker, meta, layer_id, num_tokens) + + +@pytest.fixture +def supported(monkeypatch): + monkeypatch.setattr(m3_mc_allreduce, "capture_supported", lambda: True) + + +def test_column_block_of_each_captured_layer(supported) -> None: + worker, meta, hidden, weights = _setup() + for i, layer_id in enumerate(LAYERS): + out, w, eps = _target(worker, meta, layer_id, 4) + assert out.shape == (4, H) + assert out.data_ptr() == hidden[:4, i * H : (i + 1) * H].data_ptr() + assert out.stride() == hidden.stride() + assert torch.equal(w, weights[i]) + assert eps == 1e-5 + + +def test_none_when_not_capturable(supported) -> None: + worker, meta, _, _ = _setup() + assert _target(worker, meta, 3, 4) is None # not a captured layer + assert _target(worker, meta, LAYERS[0], 9) is None # past the buffer's rows + unarmed, meta2, _, _ = _setup(armed=False) + assert _target(unarmed, meta2, LAYERS[0], 4) is None + + +def test_none_without_the_capture_epilogue(monkeypatch) -> None: + monkeypatch.setattr(m3_mc_allreduce, "capture_supported", lambda: False) + worker, meta, _, _ = _setup() + assert _target(worker, meta, LAYERS[0], 4) is None + + +def test_capture_supported_reads_the_op_schema() -> None: + # The schema of the op the boundaries call: this target's copy (m3_mc_allreduce.OWN_OP), else the image's. + name = ( + "minimax_m3_sm107_tp8_mc_allreduce_norm" + if m3_mc_allreduce.own_op() + else "minimax_m3_sm107_tp2_mc_allreduce_norm" + ) + op = getattr(torch.ops.trtllm, name, None) + want = op is not None and any(a.name == "cap_out" for a in op.default._schema.arguments) + assert m3_mc_allreduce.capture_supported() is want + + +def test_own_op_follows_the_kernel_constant(monkeypatch) -> None: + monkeypatch.setattr(m3_mc_allreduce, "OWN_OP", False) + assert m3_mc_allreduce.own_op() is False + monkeypatch.setattr(m3_mc_allreduce, "OWN_OP", True) + loaded = getattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_mc_allreduce_norm", None) is not None + assert m3_mc_allreduce.own_op() is loaded diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py new file mode 100644 index 000000000000..5a406341fa57 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_workspaces.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The lazily created exchange workspaces refuse to be created inside a CUDA-graph capture (K3 audit #150, check #10). + +The draft-step and sampled-accept workspaces are created collectively (both ranks, a multicast buffer) on the first +eager call. Creating one inside a capture would allocate and synchronize mid-capture, so a capture that finds it missing +reports a miss (``ops/m3_report.capture_missed``) and gets None (the caller then takes its stock path) instead. +""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( + m3_draft_step, + m3_report, + m3_sampled_accept, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU" +) + + +def _refuse(*args, **kwargs): + raise AssertionError("a workspace must not be created inside a CUDA-graph capture") + + +@pytest.mark.parametrize( + "module, creator", + [ + (m3_draft_step, "create_workspace"), + (m3_sampled_accept, "SampledAcceptWorkspace.create"), + ], + ids=["draft_step", "sampled_accept"], +) +def test_a_capture_without_the_workspace_misses_instead_of_creating_it( + monkeypatch, module, creator +): + monkeypatch.setattr(module, "_WS", None) + owner, _, attr = creator.rpartition(".") + monkeypatch.setattr(getattr(module, owner) if owner else module, attr, _refuse) + misses = [] + monkeypatch.setattr(m3_report, "capture_missed", lambda *args: misses.append(args)) + + device = torch.device("cuda", torch.cuda.current_device()) + found = [] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + found.append(module.workspace(None, device)) + torch.cuda.current_stream().wait_stream(stream) + + assert found == [None] + assert module._WS is None + assert len(misses) == 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py new file mode 100644 index 000000000000..8e8982524c8b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py @@ -0,0 +1,50 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_chunk_norm entry (the Eagle3 fc_norm over the captured layers).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_chunk_norm import ( # noqa: E501 + m3_chunk_norm as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, c, h = spec["m"], spec["c"], spec["h"] + buf = randn((m, c * h + spec["extra_cols"]), g, 2.0) + src = buf[:, : c * h + spec["extra_cols"]] + w = (1.0 + randn((c, h), g, 0.1).float()).to(torch.bfloat16) + dst = torch.full((m, c * h), 7.0, dtype=torch.bfloat16, device="cuda") + return (src, w, EPS, dst, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out = op(*args) + assert_bf16_close(out, op.reference(*args)) + + +def test_guard_refuses_a_strided_weight() -> None: + src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() + wt.copy_(w) + with pytest.raises(AssertionError, match="dense"): + with validating(op): + op(src, wt, eps, dst, pdl) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py new file mode 100644 index 000000000000..f07eecb1f0c2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the fc fork's counter wait (modeling's DI_FC_FLAG, ops/m3_fc_wait): the one-thread wait kernel +queued on the main stream releases its successor only after the fork's bump on another stream (so the successor reads +what the fork wrote, with no stream join before it), eagerly and on CUDA-graph replays; it counts an overrun instead +of hanging when no bump comes.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_fc_wait, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + +N = 1 << 16 + + +def _wait(counts, spins=m3_fc_wait.MAX_SPINS): + m3_fc_wait.fc_wait_kernel[(1,)](counts, MAX_SPINS=spins, num_warps=1, launch_pdl=True) + + +def _step(counts, aux, src, dst, value, delay_cycles): + """One fork / wait / consume: the fork (after a delay) writes ``value`` into src and bumps the counter; the main + stream waits on the counter, copies src to dst, then joins the fork (deferred, as the worker does).""" + main = torch.cuda.current_stream() + aux.wait_stream(main) + with torch.cuda.stream(aux): + torch.cuda._sleep(delay_cycles) + src.fill_(value) + counts[:1].add_(1) + _wait(counts) + dst.copy_(src) + main.wait_stream(aux) + + +@pytest.mark.parametrize("delay_cycles", [0, 200_000, 2_000_000]) +def test_wait_releases_after_the_fork(delay_cycles) -> None: + counts = torch.zeros(3, dtype=torch.int64, device="cuda") + src = torch.zeros(N, dtype=torch.float32, device="cuda") + dst = torch.empty_like(src) + aux = torch.cuda.Stream() + for it in range(1, 9): + _step(counts, aux, src, dst, float(it), delay_cycles) + torch.cuda.synchronize() + assert bool((dst == float(it)).all()), f"iteration {it}: read before the fork finished" + assert counts.tolist() == [8, 8, 0] + + +def test_wait_releases_on_graph_replays() -> None: + counts = torch.zeros(3, dtype=torch.int64, device="cuda") + src = torch.zeros(N, dtype=torch.float32, device="cuda") + dst = torch.empty_like(src) + val = torch.zeros(1, dtype=torch.float32, device="cuda") + aux = torch.cuda.Stream() + cap = torch.cuda.Stream() + + def body(): + main = torch.cuda.current_stream() + aux.wait_stream(main) + with torch.cuda.stream(aux): + torch.cuda._sleep(1_000_000) + src.copy_(val.expand(N)) + counts[:1].add_(1) + _wait(counts) + dst.copy_(src) + main.wait_stream(aux) + + with torch.cuda.stream(cap): + body() # eager warm-up (compiles the wait kernel outside capture) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(cap): + with torch.cuda.graph(graph, stream=cap): + body() + torch.cuda.synchronize() + base = counts.tolist() + for it in range(1, 11): + val.fill_(float(it)) + graph.replay() + torch.cuda.synchronize() + assert bool((dst == float(it)).all()), f"replay {it}: read before the fork finished" + assert counts.tolist() == [base[0] + 10, base[1] + 10, 0] + + +def test_wait_without_a_fork_counts_an_overrun() -> None: + counts = torch.zeros(3, dtype=torch.int64, device="cuda") + _wait(counts, spins=1000) + torch.cuda.synchronize() + assert counts.tolist() == [0, 1, 1] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py new file mode 100644 index 000000000000..f1f5f6160e32 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_input entry (embedding + both RMSNorms + concat in one launch), against its reference and +against the in-tree drafter's modules.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_spec_tail, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_input import ( # noqa: E501 + m3_draft_input as op, +) +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.rms_norm import RMSNorm + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + EPS, + HIDDEN, + VOCAB, + assert_bf16_close, + gen, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" + + +def _case(spec, seed): + g = gen(seed) + m, h = spec["m"], spec["h"] + table = randn((VOCAB, h), g) + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + if spec["invalid_id"]: + ids[0] = -1 + hidden = randn((4 * m if spec["gather"] else m, h), g) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 2 if spec["gather"] else None + w_e = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + w_h = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + out = torch.full((m, 2 * h), 7.0, dtype=torch.bfloat16, device="cuda") + residual = torch.full((m, h), 7.0, dtype=torch.bfloat16, device="cuda") + return (ids, table, hidden, rows, w_e, w_h, EPS, EPS, out, residual, True) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Normalized halves within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic in a + different summation order, rsqrtf); the gathered residual and the zero embedding of an invalid id exact.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + out, residual = op(*args) + ref_out, ref_res = op.reference(*args) + torch.cuda.synchronize() + assert_bf16_close(out, ref_out) + assert torch.equal(residual, ref_res) + if cell.spec["invalid_id"]: + assert not bool(out[0, : cell.spec["h"]].float().any()) + + +def test_guard_refuses_rows_past_hidden() -> None: + gather = next(c for c in op.CELLS if c.spec["gather"]) + args = list(_case(gather.spec, 1)) + args[3] = args[3] + 100 + with pytest.raises(AssertionError, match="rows must index hidden"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize("m,gather", [(4, False), (32, False), (8, True)], ids=str) +def test_the_in_tree_draft_input(m, gather) -> None: + """Against the in-tree drafter's own modules: the Embedding lookup, then the layer's input_layernorm and + hidden_norm (in-tree RMSNorm) and the concat of Eagle3DecoderLayer.forward. Within one bf16 ulp on at most 0.1% of + the elements (the RMS statistic summed in another order); the residual (the carried hidden rows) exact.""" + g = gen(seed_of(f"in-tree draft input {m} {gather}")) + h = HIDDEN + embed = Embedding(VOCAB, h, dtype=torch.bfloat16).cuda() + embed.weight.data = randn((VOCAB, h), g) + norms = [] + for _ in range(2): + norm = RMSNorm(hidden_size=h, eps=EPS, dtype=torch.bfloat16).cuda() + norm.weight.data = (1.0 + randn((h,), g, 0.1).float()).to(torch.bfloat16) + norms.append(norm) + input_layernorm, hidden_norm = norms + ids = torch.randint(0, VOCAB, (m,), generator=g, device="cuda", dtype=torch.int32) + hidden = randn((4 * m if gather else m, h), g, 2.0) + rows = torch.arange(m, device="cuda", dtype=torch.int64) * 4 + 3 if gather else None + carried = hidden[rows] if gather else hidden + stock = torch.cat([input_layernorm(embed(ids)), hidden_norm(carried)], dim=-1) + out, residual = m3_spec_tail.draft_input( + ids, embed.weight, hidden, rows, input_layernorm.weight, EPS, hidden_norm.weight, EPS + ) + torch.cuda.synchronize() + assert_bf16_close(out, stock) + assert torch.equal(residual, carried) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_push_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_push_take.py new file mode 100644 index 000000000000..8f0896d6fca6 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_push_take.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The drafter boundary push bookkeeping (``M3Eagle3Draft.arm_push`` / ``take_pushed``, ``DRAFT_AR_PUSH``): the +boundary all-reduce runs prepushed only when the armed call pushed, and only on that call's own output; any other +input raises instead of reducing rows the exchange does not hold. Host state only (no GPU).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) + + +@pytest.fixture +def draft(): + """A drafter with only its push bookkeeping (the module and its weights are not built).""" + d = object.__new__(modeling.M3Eagle3Draft) + d._push_arm, d._pushed = None, 0 + return d + + +def test_not_pushed_is_false_and_disarms(draft) -> None: + draft._push_arm = (1, 2, 3, 0, 8, 8) + assert draft.take_pushed(torch.zeros(4)) is False + assert draft._push_arm is None + + +def test_pushed_output_is_true_once(draft) -> None: + out = torch.zeros(4, 8) + draft._pushed = out.data_ptr() + assert draft.take_pushed(out) is True + assert draft._pushed == 0 + assert draft.take_pushed(out) is False + + +def test_other_input_raises(draft) -> None: + out = torch.zeros(4, 8) + other = out.clone() + draft._pushed = out.data_ptr() + with pytest.raises(RuntimeError, match="not the boundary all-reduce's input"): + draft.take_pushed(other) + assert draft._pushed == 0 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py new file mode 100644 index 000000000000..5ae74445fe33 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_silu entry (the drafter MLP's plain SwiGLU as one C++ launch): bit for bit the Triton +silu_and_mul at its defaults, over every bf16 gate and every bf16 up pattern.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_silu import ( # noqa: E501 + m3_draft_silu as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" + +ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) + + +def _rows(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """[rows, 2d] gate | up from equally long flat gate and up values, d = 2304 (the TP8 drafter shard).""" + d = 2304 + n = gate.numel() + rows = -(-n // d) + pad = rows * d - n + g = torch.cat([gate, gate[:pad]]).view(rows, d) + u = torch.cat([up, up[:pad]]).view(rows, d) + return torch.cat([g, u], dim=1).contiguous() + + +def _assert_bitwise(x: torch.Tensor, what: str) -> None: + got = op(x) + assert got is not None, what + ref = op.reference(x) + torch.cuda.synchronize() + diff = got.view(torch.int16) != ref.view(torch.int16) + assert not bool(diff.any()), f"{what}: {int(diff.sum())} of {diff.numel()} differ" + + +@pytest.mark.parametrize("up", [0.0, 1.0, -1.0, 3.5, -0.0, 1e-3], ids=lambda v: f"up={v}") +def test_every_gate_pattern_is_bitwise_the_triton_kernel(up) -> None: + gate = ALL.cuda() + _assert_bitwise(_rows(gate, torch.full_like(gate, up)), f"gate sweep, up {up}") + + +@pytest.mark.parametrize("gate", [1.0, -1.0, 0.5, -3.0, 20.0, -0.0], ids=lambda v: f"gate={v}") +def test_every_up_pattern_is_bitwise_the_triton_kernel(gate) -> None: + up = ALL.cuda() + _assert_bitwise(_rows(torch.full_like(up, gate), up), f"up sweep, gate {gate}") + + +@pytest.mark.parametrize("seed", range(4)) +def test_random_pairs_are_bitwise_the_triton_kernel(seed) -> None: + g = gen(100 + seed) + n = 1 << 22 + gate = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + up = torch.randint(-(1 << 15), 1 << 15, (n,), generator=g, device="cuda").to(torch.int16) + _assert_bitwise(_rows(gate.view(torch.bfloat16), up.view(torch.bfloat16)), f"seed {seed}") + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The drafter's gate_up output range (and past it), in a CUDA graph as the drafter runs it.""" + g = gen(seed_of(cell.why)) + rows, d = cell.spec["rows"], cell.spec["d"] + x = (4.0 * torch.randn(rows, 2 * d, generator=g, device="cuda")).to(torch.bfloat16) + x[:, :6] = torch.tensor([60.0, -60.0, 0.0, -0.0, 88.0, -88.0], dtype=torch.bfloat16) + certifying(op) + ref = op.reference(x) + eager = op(x) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(x) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(eager.view(torch.int16), ref.view(torch.int16)) + assert torch.equal(captured.view(torch.int16), ref.view(torch.int16)) + + +def test_declines_what_it_does_not_cover() -> None: + x = torch.zeros(2, 2 * 2304, dtype=torch.bfloat16, device="cuda") + assert op(x[:, 8:]) is None # strided rows + assert ( + op(torch.zeros(2, 2 * 12, dtype=torch.bfloat16, device="cuda")) is None + ) # d not a multiple of 8 + assert op(x.float()) is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py new file mode 100644 index 000000000000..877e52628100 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -0,0 +1,747 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_draft_step entry (one fused greedy draft pass: the BF16 draft head, streamed from its lossless +12-bit coding or as BF16 fragments, the cross-rank argmax and the draft loop's bookkeeping), and its tokens against the +in-tree drafter's head and greedy draft sampling. + +The in-tree path: each rank's LM head shard GEMM (BF16 out), ``LogitsProcessor``'s float cast, then either the argmax +of the gathered logits (``torch.argmax``) or ``greedy_sample_draft_with_tp_gather``'s per-rank ``torch.max`` and the +pick among the ranks' (index, value) pairs. The edge cases use integer data, whose fp32 sums are exact in any order, +so the kernel and the in-tree path see the same BF16 logits and must agree bit for bit; every in-tree comparison runs +with the coded head (the default) and with the BF16 fragments, and on integer data the coded head's winning keys are +the BF16 fragments' bit for bit. +""" + +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_draft_step as draft_step_kernels, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_LEN, + DRAFT_SHARD, + HIDDEN, + MAX_TOKENS, + bf16_fragments, + from_fragments, + head_weight, + max_tile_pieces, + shard_logits, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_step import ( # noqa: E501 + m3_draft_step as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 — needs the path declared above + WORLD, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" + +T = DRAFT_LEN + 1 +_LOW32 = 0xFFFFFFFF + +# ------------------------------------------------------------------ the per-row rule against the in-tree token +# +# The kernel and the in-tree head (a cuBLAS BF16 GEMM) both round an FP32 sum of the same exact BF16 products to BF16, +# in different summation orders. A row's token can differ from the in-tree one only where some candidate's FP32 sum +# can land on either side of a BF16 rounding midpoint. Each row is held to: +# +# kernel token == in-tree token, or the row is provably undecidable: with the fp64 reference logit L_i (exact BF16 +# products, fp64 sum) and the FP32-order bound E_i = gamma_K * sum_k |w_ik * x_k|, the BF16 value of any FP32 sum of +# candidate i lies in [lo_i, hi_i] = [bf16(L_i - E_i), bf16(L_i + E_i)]; the kernel's token must win for some +# choice of values within those intervals (larger value, then the lower global id), and a BF16 rounding midpoint +# must lie within E of the kernel's or the in-tree token's L. +# +# gamma_K = K * 2^-22 is la-xh2's bound (4x the round-to-nearest unit roundoff 2^-24 per addition: any summation +# order, and a truncating tensor-core accumulator); K counts the row's non-zero products (adding an exact zero does +# not round). lo / hi take one more FP32 step outward before the BF16 rounding, which covers rounding L -+ E to FP32. +# Each undecidable row is reported as a warning with its midpoint distances, and each test reports its row count; any +# other mismatch fails. +GAMMA_UNIT = 2.0**-22 + + +def _bf16_outward(v: torch.Tensor, direction: float) -> torch.Tensor: + f = v.float() + f = torch.nextafter(f, torch.full_like(f, direction)) + b = f.to(torch.bfloat16).float() + return torch.where(b == 0, torch.zeros_like(b), b) + + +def bounds(h: torch.Tensor, w: torch.Tensor): + """(fp64 logits L, bound E, lo, hi) [rows, N] of the head input rows h against the head rows w.""" + cols = (h != 0).any(dim=0) # zero columns add exact zeros + hd, wd = h[:, cols].double(), w[:, cols].double() + logits = hd @ wd.t() + bound = (h != 0).sum(dim=1, keepdim=True).double() * GAMMA_UNIT * (hd.abs() @ wd.abs().t()) + del wd + inf = float("inf") + return logits, bound, _bf16_outward(logits - bound, -inf), _bf16_outward(logits + bound, inf) + + +def midpoint_distance(v: float) -> float: + """Distance of v to the nearest BF16 rounding midpoint.""" + b = torch.tensor(v, dtype=torch.float64).float().to(torch.bfloat16).float() + if float(b) == 0.0: + return abs(v) + ulp = 2.0 ** (int(torch.frexp(b).exponent) - 8) + return min(abs(v - (float(b) + ulp / 2)), abs(v - (float(b) - ulp / 2))) + + +def _winnable(tok: int, ids: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor) -> bool: + """Whether candidate ``tok`` wins the row for some values within the candidates' [lo, hi].""" + k = int((ids == tok).nonzero()[0, 0]) + others = ids != tok + beat = (hi[k] > lo[others]) | ((hi[k] == lo[others]) & (tok < ids[others])) + return bool(beat.all()) + + +def check_row(name: str, kernel_tok: int, stock_tok: int, ids, logits, bound, lo, hi) -> bool: + """The per-row rule for one row over the candidates ``ids`` [C] (logits / bound / lo / hi [C]); True when the + row is undecidable (reported as a warning), False when the tokens match; a mismatch outside the bound fails.""" + if kernel_tok == stock_tok: + return False + assert _winnable(kernel_tok, ids, lo, hi), ( + f"{name}: token {kernel_tok} differs from the in-tree {stock_tok} outside the FP32-order bound" + ) + ik = int((ids == kernel_tok).nonzero()[0, 0]) + i_stock = int((ids == stock_tok).nonzero()[0, 0]) + lk, ls, ek, es = ( + float(logits[ik]), + float(logits[i_stock]), + float(bound[ik]), + float(bound[i_stock]), + ) + dk, ds = midpoint_distance(lk), midpoint_distance(ls) + assert dk <= ek or ds <= es, f"{name}: neither logit is within its bound of a BF16 midpoint" + warnings.warn( + f"[draft-step undecidable] {name}: kernel {kernel_tok} (L {lk:.9g}, midpoint distance {dk:.3e}, bound " + f"{ek:.3e}) vs in-tree {stock_tok} (L {ls:.9g}, midpoint distance {ds:.3e}, bound {es:.3e})" + ) + return True + + +def report(name: str, rows: int, undecidable: int) -> None: + warnings.warn(f"[draft-step rows] {name}: {rows} rows, {undecidable} undecidable") + + +def key_value(keys: torch.Tensor) -> torch.Tensor: + """The fp32 value a key orders by (the inverse of the kernel's order-preserving image).""" + u = (keys >> 32) & _LOW32 + bits = torch.where(u >= 0x80000000, u - 0x80000000, _LOW32 - u) + bits = torch.where(bits >= 2**31, bits - 2**32, bits) # as a signed 32-bit pattern + return bits.to(torch.int32).view(torch.float32) + + +def key_index(keys: torch.Tensor) -> torch.Tensor: + return _LOW32 - (keys & _LOW32) + + +def _bookkeeping(nreq: int, step: int, n_x_rows: int, g: torch.Generator): + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if step == 0 + else None + ) + seq_lens = torch.full((nreq,), T, dtype=torch.int32, device="cuda") if step == 0 else None + next_new = torch.full((nreq, T), -5, dtype=torch.int32, device="cuda") + pos_src = torch.randint(100, 3000, (n_x_rows,), generator=g, device="cuda", dtype=torch.int32) + kv_lens = torch.randint(200, 3000, (nreq,), generator=g, device="cuda", dtype=torch.int32) + ids_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + pos_next = torch.full((nreq,), -9, dtype=torch.int32, device="cuda") + return num_acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next + + +def _head(w: torch.Tensor, coded: bool): + """(wf, (exc_off, exc_pos, exc_w, exp_base)) of a BF16 shard: its 12-bit coding (checked to round-trip) or its + BF16 fragments.""" + if not coded: + return bf16_fragments(w), (None, None, None, 0) + packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) + assert round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w) + return packed, (exc_off, exc_pos, exc_w.view(-1), base) + + +def _args( + x, + x_rows, + nreq, + wf, + rank, + world, + step, + g, + uc=None, + keys=None, + counters=None, + num_acc=None, + coded=(None, None, None, 0), +): + """The op's argument tuple for one rank (fresh workspace and bookkeeping buffers unless given); ``coded`` = the + coded head's (exc_off, exc_pos, exc_w, exp_base).""" + acc, seq_lens, next_new, pos_src, kv_lens, ids_next, pos_next = _bookkeeping( + nreq, step, x.shape[0], g + ) + if num_acc is None: + num_acc = acc + if uc is None: + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + out_keys = torch.zeros(nreq, dtype=torch.int64, device="cuda") + tiles = wf.shape[0] if coded[0] is None else coded[0].numel() - 1 + return (x, x_rows, nreq, wf, rank * tiles * 16, keys, counters, uc, uc.data_ptr(), rank, world, + next_new[:, 1:], step, ids_next, pos_src, pos_next, kv_lens, num_acc, seq_lens, DRAFT_LEN, num_sms(), + True, out_keys, *coded) # fmt: skip + + +def _fixed(args) -> bool: + """fixed_order for these op arguments as dispatch picks it (ops/m3_draft_step._fixed_order): the warp-order + sum where the launch grid lets a tile take more than 2 warp pieces (D-FIX; the TP8 shard on 212 SMs does).""" + x, wf, num_ctas, exc_off = args[0], args[3], args[20], args[23] + tiles = wf.shape[0] if exc_off is None else exc_off.numel() - 1 + return max_tile_pieces(tiles, num_ctas, x.shape[1] // (16 if exc_off is None else 64)) > 2 + + +def _case(spec, seed): + g = gen(seed) + nreq, step, rank, world, vs, k = ( + spec["nreq"], + spec["step"], + spec["rank"], + spec["world"], + spec["vs"], + spec["k"], + ) + pass0 = step == 0 + x = randn((nreq * T if pass0 else nreq, k), g) + wf, coded = _head(randn((vs, k), g, 0.02), spec["coded"]) + num_acc = ( + torch.randint(1, T + 1, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if pass0 + else None + ) + x_rows = (torch.arange(nreq, device="cuda") * T + num_acc.long() - 1) if pass0 else None + uc, keys, counters = head_workspace(world, slot_rows=MAX_TOKENS) + h = x[x_rows] if pass0 else x + own = shard_logits(h, wf, *coded).max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=MAX_TOKENS, ties=spec.get("ties", False) + ) + return _args(x, x_rows, nreq, wf, rank, world, step, g, uc, keys, counters, num_acc, coded) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """The bookkeeping is exact on every row (next positions, KV and sequence lengths). Each row's token is the + in-tree one (this rank's BF16 shard logits from the head's GEMM, their first maximum, against the peers' keys) or + the row is provably undecidable (check_row); the winning key's index is the token, and its value lies in the + token's [lo, hi] (or is the peer's pushed key when a peer wins).""" + args = _case(cell.spec, seed_of(cell.why)) + x, x_rows, nreq, wf, offset, _, counters, uc = args[:8] + rank, world = args[9], args[10] + half = int(counters[1].item()) & 1 + pushed = ( + uc[half * MAX_TOKENS * world : (half * MAX_TOKENS + nreq) * world].view(nreq, world).clone() + ) + h = x[x_rows] if x_rows is not None else x[:nreq] + w = head_weight(wf, x.shape[1], *args[23:27]) + best, idx = local_argmax(torch.nn.functional.linear(h, w).float(), offset) + stock = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, MAX_TOKENS) + certifying(op) + ref = op.reference(*args) + op(*args, fixed_order=_fixed(args)) + torch.cuda.synchronize() + assert args[6].tolist() == [0, 1] and not bool(args[5].any()), "counters / keys not re-armed" + step = args[12] + others = [c for c in range(args[11].shape[1]) if c != step] + assert torch.equal(args[11][:, others], ref["draft_out"][:, others]) + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + assert torch.equal(args[15], ref["pos_next"]) + assert torch.equal(args[16], ref["kv_lens"]) + if args[18] is not None: + assert torch.equal(args[18], ref["seq_lens"]) + # Candidates per row: this rank's head rows (fp64 bounds) and each peer's pushed key (an exact value). + logits, bound, lo, hi = bounds(h, w) + own_ids = offset + torch.arange(w.shape[0], device="cuda") + peers = [p for p in range(world) if p != rank] + undecidable = 0 + for r in range(nreq): + pv = key_value(pushed[r, peers]).double() + ids = torch.cat([own_ids, key_index(pushed[r, peers])]) + undecidable += check_row( + f"{cell.why[:40]} row {r}", + int(args[13][r]), + int(stock[r]), + ids, + torch.cat([logits[r], pv]), + torch.cat([bound[r], torch.zeros_like(pv)]), + torch.cat([lo[r], pv.float()]), + torch.cat([hi[r], pv.float()]), + ) + report(cell.why[:40], nreq, undecidable) + out_keys = args[22] + assert torch.equal(key_index(out_keys), args[13].long()), "the winning key's index is the token" + got = key_value(out_keys) + for r in range(nreq): + c = int(args[13][r]) - offset + if 0 <= c < w.shape[0]: + assert float(lo[r, c]) <= float(got[r]) <= float(hi[r, c]), ( + f"row {r}: winning value outside [lo, hi]" + ) + else: + assert int(out_keys[r]) in pushed[r].tolist(), ( + f"row {r}: a peer won with a key it did not push" + ) + + +def test_fragments_round_trip() -> None: + g = gen(seed_of("fragments")) + w = randn((64, 96), g) + wf = bf16_fragments(w) + assert wf.shape == (4, 6, 256) + assert torch.equal(from_fragments(wf), w) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 3) # K 96: no coding, the BF16 fragments + assert not prepared.coded + assert torch.equal(prepared.wf, wf) and prepared.vocab_offset == 3 * 64 and prepared.weight is w + + +# ------------------------------------------------------------------ against the in-tree drafter, bit for bit + + +def _stock_tokens(shards: list, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """The in-tree greedy draft tokens of the rows x over the vocab-sharded BF16 head: (argmax of the gathered float + logits, greedy_sample_draft_with_tp_gather's per-rank (index, value) pick).""" + logits = [torch.nn.functional.linear(x, w).float() for w in shards] + gathered = torch.argmax(torch.cat(logits, dim=-1), dim=-1).to(torch.int32) + combined = [] + for rank, shard_logit in enumerate(logits): + worker = SimpleNamespace(mapping=SimpleNamespace(tp_rank=rank)) + combined.append(Eagle3OneModelWorker._get_local_max_and_combined(worker, shard_logit)) + pick = Eagle3OneModelWorker._get_draft_tokens_from_gathered + pick = getattr(pick, "_torchdynamo_orig_callable", pick) + tp_gather = pick(None, torch.cat(combined, dim=-1)).to(torch.int32) + return gathered, tp_gather + + +def _fused_tokens(shards: list, x: torch.Tensor, coded: bool, step: int = 1) -> list: + """Each rank's fused draft tokens: every rank's own key from a one-rank run, then each rank's four-rank pass with + its peers' keys in its slots, as their multicast pushes would put them.""" + g = gen(seed_of("fused tokens")) + nreq = x.shape[0] + heads = [_head(w, coded) for w in shards] + own_keys = [] + for rank, (wf, cd) in enumerate(heads): + args = list(_args(x, None, nreq, wf, rank, 1, step, g, coded=cd)) + args[9] = 0 # one-rank run: this rank's slot only, global ids of `rank` + op(*args, fixed_order=_fixed(args)) + own_keys.append(args[22].clone()) + tokens = [] + for rank, (wf, cd) in enumerate(heads): + uc, keys, counters = head_workspace(WORLD, slot_rows=MAX_TOKENS) + for peer in range(WORLD): + if peer != rank: + for t in range(nreq): + uc[t * WORLD + peer] = own_keys[peer][t] + args = _args(x, None, nreq, wf, rank, WORLD, step, g, uc, keys, counters, coded=cd) + with validating(op): + op(*args, fixed_order=_fixed(args)) + torch.cuda.synchronize() + assert torch.equal(args[11][:, step], args[13]), "draft_out and ids_next disagree" + tokens.append(args[13].clone()) + return tokens + + +def _integer_head(g: torch.Generator, rows: int = WORLD * DRAFT_SHARD) -> torch.Tensor: + """A [rows, 6144] head of small integers over the first 4 columns (zero elsewhere): its logits against integer + rows are exact integers in any summation order, many of which share a BF16 value.""" + w = torch.zeros((rows, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, :4] = torch.randint(-60, 61, (rows, 4), generator=g, device="cuda").to(torch.bfloat16) + return w + + +def _shards(w: torch.Tensor) -> list: + return [w[r * DRAFT_SHARD : (r + 1) * DRAFT_SHARD].contiguous() for r in range(WORLD)] + + +def _assert_stock(w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None) -> None: + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + if expect is not None: + assert gathered.tolist() == expect, (gathered.tolist(), expect) + for rank, tokens in enumerate(_fused_tokens(shards, x, coded)): + assert torch.equal(tokens, gathered), ( + f"rank {rank}: {tokens.tolist()} vs stock {gathered.tolist()}" + ) + + +def _ones_rows(nreq: int) -> torch.Tensor: + x = torch.zeros((nreq, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[:, :4] = 1.0 + return x + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_forced_bf16_ties_take_the_lowest_global_id(coded) -> None: + """A head row copied onto every rank that wins every row of ones (logit 244, every other at most 240): the lowest + global id must win, inside a shard and across ranks. Then integer rows in [-2, 2], whose logits (up to 480 in + magnitude, BF16 spacing 2 above 256) collide often, exactly or after BF16 rounding.""" + g = gen(seed_of("forced ties")) + w = _integer_head(g) + win = torch.tensor([61.0, 61.0, 61.0, 61.0], device="cuda").to(torch.bfloat16) + for gid in (DRAFT_SHARD + 5000, DRAFT_SHARD + 77, DRAFT_SHARD + 9000, 23750): + w[gid, :4] = win + _assert_stock(w, _ones_rows(5), coded, expect=[23750] * 5) + # Random integer rows: ties after BF16 rounding happen naturally. + x = torch.randint(-2, 3, (8, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + x[:, 4:] = 0 + _assert_stock(w, x, coded) + + +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_signed_zero_maxima_take_the_lowest_global_id(coded) -> None: + """Rows whose maximum is zero: an all-zero head input (every logit +-0) and logits <= 0 with zeros by cancellation + on several ranks. -0 and +0 compare equal, so the lowest global id reaching zero wins.""" + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 0] = -1.0 + w[:, 1] = -1.0 + for gid in (DRAFT_SHARD + 5, DRAFT_SHARD + 40000, DRAFT_SHARD + 70007, 24999): + w[gid, 0] = 1.0 + w[gid, 1] = -1.0 + x = torch.zeros((3, HIDDEN), dtype=torch.bfloat16, device="cuda") + x[1, :2] = 1.0 + x[2, :2] = 2.0 + x[2, 2] = -1.0 # w[:, 2] is zero: a -0 product on every row + _assert_stock(w, x, coded, expect=[0, 24999, 24999]) + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + (1000, 1001, "lower"), # 1001 rounds to 1000 in BF16: a tie, the lower id wins + (1001, 1002, "lower"), # 1001 and the tie-to-even midpoint 1002 both round to 1000 + (1000, 1004, "higher"), # adjacent BF16 values: the larger wins + (1003, 1000, "lower"), # 1003 rounds up to 1004 > 1000 + ], + ids=["one_fp32_step", "midpoint", "adjacent_bf16", "rounds_up"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_near_equal_top2_follow_the_bf16_logits(lower, higher, winner, where, coded) -> None: + """Top-2 logits one to four fp32 steps apart: the winner is the one the BF16 logits decide (the in-tree logits' + dtype), not the fp32 sum. The pair sits inside rank 1's shard or on ranks 0 and 1; every other logit is at most + 240 in magnitude.""" + g = gen(seed_of(f"near top2 {lower} {higher} {where}")) + w = _integer_head(g) + if where == "in_shard": + lo_id, hi_id = DRAFT_SHARD + 100, DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD - 1000, DRAFT_SHARD + 3 + for gid, value in ((lo_id, lower), (hi_id, higher)): + q, r = divmod(value, 4) # four integers <= 251, exact in BF16 + w[gid, :4] = torch.tensor([q + (i < r) for i in range(4)], device="cuda").to(torch.bfloat16) + expect = lo_id if winner == "lower" else hi_id + _assert_stock(w, _ones_rows(2), coded, expect=[expect, expect]) + + +# ------------------------------------------------------------------ non-integer rows at a BF16 rounding midpoint + +#: A BF16 rounding midpoint (between 6.0 and 6.03125; a tie there rounds to the even 6.0) and the FP32 spacing there. +MIDPOINT = 6.0 + 2.0**-6 +FP32_ULP = 2.0**-21 + + +def _near_midpoint_head(g: torch.Generator, targets: dict): + """A full-vocabulary head [8 * 25008, 6144] and one head-input row x with 65 non-zero entries: 1, 2^-12, 2^-20 + and 2^-28 on columns 0..3 and 61 random non-integer BF16 values. Every head row carries random non-integer BF16 + values on those 61 columns (logits about 0.2 in magnitude); each row of ``targets`` (global id -> fp64 value) also + gets columns 0..3 set so its fp64 logit is the target to within 2^-30 (a coarse value, then three finer + corrections, each BF16-rounded).""" + x = torch.zeros((1, HIDDEN), dtype=torch.bfloat16, device="cuda") + scales = (1.0, 2.0**-12, 2.0**-20, 2.0**-28) + for col, scale in enumerate(scales): + x[0, col] = scale + x[0, 4:65] = randn((61,), g, 0.5) + w = torch.zeros((WORLD * DRAFT_SHARD, HIDDEN), dtype=torch.bfloat16, device="cuda") + w[:, 4:65] = randn((WORLD * DRAFT_SHARD, 61), g, 0.05) + xd = x[0].double() + for gid, target in targets.items(): + rest = float(target - (w[gid].double() * xd).sum()) + for col, scale in enumerate(scales): + v = torch.tensor(rest / scale, dtype=torch.float64).to(torch.bfloat16) + w[gid, col] = v + rest -= float(v) * scale + assert abs(rest) <= 2.0**-30, (gid, rest) + return w, x + + +def _check_fused_against_stock( + name: str, w: torch.Tensor, x: torch.Tensor, coded: bool, expect=None +) -> int: + """Each rank's fused tokens against the in-tree tokens under the per-row rule (full-vocabulary candidates); + ``expect`` (a token per row) additionally pins rows the bound decides. Returns the undecidable rows.""" + shards = _shards(w) + gathered, tp_gather = _stock_tokens(shards, x) + assert torch.equal(gathered, tp_gather), "the two in-tree pickers disagree" + fused = _fused_tokens(shards, x, coded) + for tokens in fused[1:]: + assert torch.equal(tokens, fused[0]), "the ranks disagree on the global token" + logits, bound, lo, hi = bounds(x, w) + ids = torch.arange(w.shape[0], device="cuda") + undecidable = 0 + for r in range(x.shape[0]): + undecidable += check_row( + f"{name} row {r}", + int(fused[0][r]), + int(gathered[r]), + ids, + logits[r], + bound[r], + lo[r], + hi[r], + ) + if expect is not None: + assert int(gathered[r]) == expect[r] and int(fused[0][r]) == expect[r], ( + name, + int(gathered[r]), + int(fused[0][r]), + expect[r], + ) + others = ids != expect[r] + assert _winnable(expect[r], ids, lo[r], hi[r]), ( + f"{name}: the bound does not let {expect[r]} win" + ) + rivals = [int(i) for i in ids[others][hi[r][others] >= lo[r][expect[r]]]] + assert not any(_winnable(i, ids, lo[r], hi[r]) for i in rivals), ( + f"{name}: the row is not decided" + ) + report(name, x.shape[0], undecidable) + return undecidable + + +@pytest.mark.parametrize( + "lower,higher,winner", + [ + ( + -2, + 2, + None, + ), # straddling the midpoint by 2 FP32 ulps: either side is in-tree up to the order + (2, -2, None), + (0, -2, None), # exactly on the midpoint (ties to the even 6.0) against just below it + ( + -2048, + 2048, + "higher", + ), # ~1e-3 either side, well outside the bound: decided by BF16 rounding + ( + -4096, + -2048, + "lower", + ), # both below the midpoint, ~1e-3 apart: a BF16 tie the lower id wins + ], + ids=["straddle", "straddle_reversed", "on_midpoint", "decided_up", "decided_tie"], +) +@pytest.mark.parametrize("where", ["in_shard", "across_ranks"]) +@pytest.mark.parametrize("coded", [True, False], ids=["coded", "bf16"]) +def test_top2_at_a_bf16_midpoint(lower, higher, winner, where, coded) -> None: + """Non-integer rows whose top-2 fp64 logits sit the given number of FP32 ulps from a BF16 rounding midpoint, inside + rank 1's shard or on ranks 0 and 1, every other logit far below. The per-row rule holds on every row; outside the + bound the token is the one the BF16 rounding decides, bit for bit.""" + g = gen(seed_of(f"midpoint {lower} {higher} {where}")) + if where == "in_shard": + lo_id, hi_id = DRAFT_SHARD + 100, DRAFT_SHARD + 200 + else: + lo_id, hi_id = DRAFT_SHARD - 1000, DRAFT_SHARD + 3 + w, x = _near_midpoint_head( + g, {lo_id: MIDPOINT + lower * FP32_ULP, hi_id: MIDPOINT + higher * FP32_ULP} + ) + expect = None if winner is None else [lo_id if winner == "lower" else hi_id] + _check_fused_against_stock(f"midpoint {where} {lower:+d}/{higher:+d} ulp", w, x, coded, expect) + + +# ------------------------------------------------------------------ the coded head against the BF16 fragments + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_coded_head_is_the_bf16_head_bit_for_bit(nreq) -> None: + """Integer data (fp32 sums exact in any order) with exceptions on both sides of the exponent window: the coded + head's winning keys (value and global id) and its bookkeeping are the BF16 fragments' bit for bit, on one rank and + across the ranks.""" + g = gen(seed_of(f"coded vs bf16 {nreq}")) + sign = torch.randint(0, 2, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * 2 - 1 + w = (torch.randint(1, 41, (DRAFT_SHARD, HIDDEN), generator=g, device="cuda") * sign).to( + torch.bfloat16 + ) + # Exceptions on both sides of the window of 1..40: exact zeros, and 2^15 (fp32 sums stay below 2^24). + w[::5, ::512] = 0.0 + w[::13, 7::600] = 32768.0 + x = torch.randint(-2, 3, (nreq, HIDDEN), generator=g, device="cuda").to(torch.bfloat16) + wf_c, cd = _head(w, True) + assert cd[1].numel() > 0, "the shard must carry exceptions" + wf_b, cb = _head(w, False) + outs = [] + for wf, c in ((wf_c, cd), (wf_b, cb)): + args = _args(x, None, nreq, wf, 1, 1, 1, gen(seed_of("bookkeeping")), coded=c) + args = list(args) + args[9] = 0 + with validating(op): + op(*args, fixed_order=_fixed(args)) + torch.cuda.synchronize() + outs.append(args) + coded_args, bf16_args = outs + for i in (11, 13, 15, 16, 22): # draft_out, ids_next, pos_next, kv_lens, out_keys + assert torch.equal(coded_args[i], bf16_args[i]), i + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), DRAFT_SHARD) + assert torch.equal(coded_args[22], make_keys(best, idx)), ( + "the winning keys are not the in-tree ones" + ) + + +@pytest.mark.parametrize("nreq", [1, 2, 8]) +def test_tiny_exception_products_are_exact(nreq) -> None: + """Weights of 1 (exponent window base 113) with column 0 zero except 2^-30 / 2^-29 on rows 0 / 1, and a one-hot + activation on column 0: the exceptions' exact products are the only non-zero logits, so the coded head must pick + row 1 with the BF16 head's key bit for bit (a correction against a non-zero placeholder cancels them to 0).""" + n, k = 256, 1024 + w = torch.ones((n, k), dtype=torch.bfloat16, device="cuda") + w[:, 0] = 0.0 + w[0, 0] = 2.0**-30 + w[1, 0] = 2.0**-29 + x = torch.zeros((nreq, k), dtype=torch.bfloat16, device="cuda") + x[:, 0] = 1.0 + wf_c, cd = _head(w, True) + assert cd[3] == 113 + keys = [] + for wf, c in ((wf_c, cd), _head(w, False)): + args = _args(x, None, nreq, wf, 0, 1, 1, gen(seed_of("tiny exceptions")), coded=c) + with validating(op): + # 16 tiles on an all-SM grid: every tile is streamed by all 16 warps, which only the fixed-order sum + # (D-FIX) adds deterministically (is_valid rejects the atomics there). + op(*args, fixed_order=True) + torch.cuda.synchronize() + assert args[13].tolist() == [1] * nreq + keys.append(args[22]) + best, idx = local_argmax(torch.nn.functional.linear(x, w).float(), 0) + assert torch.equal(keys[0], keys[1]) and torch.equal(keys[0], make_keys(best, idx)) + + +def test_prepare_head_codes_the_shard_and_fails_closed(monkeypatch) -> None: + """With the coding on, prepare_head keeps the coded shard when it decodes to the weight bit for bit; a coding + that does not is dropped for the BF16 fragments. (The tp8 target ships HEAD_CODED = False.)""" + monkeypatch.setattr(draft_step_kernels, "HEAD_CODED", True) + g = gen(seed_of("prepare coded")) + w = randn((4096, 1024), g, 0.02) + head = SimpleNamespace(weight=w, bias=None) + prepared = draft_step_kernels.prepare_head(head, 2) + assert prepared.coded and prepared.vocab_offset == 2 * 4096 and prepared.weight is w + assert torch.equal(head_weight(prepared.wf, 1024, prepared.exc_off, prepared.exc_pos, prepared.exc_w, + prepared.exp_base), w) # fmt: skip + + def corrupt(weight, *args, **kwargs): + packed, base, exc_off, exc_pos, exc_w = encode_chunked(weight) + packed = packed.clone() + packed[11] ^= 0x10 + return packed, base, exc_off, exc_pos, exc_w + + monkeypatch.setattr(draft_step_kernels, "encode_chunked", corrupt) + dropped = draft_step_kernels.prepare_head(head, 2) + assert not dropped.coded and dropped.exc_off is None + assert torch.equal(dropped.wf, bf16_fragments(w)) + + +def test_prepare_head_keeps_bf16_when_coding_is_off(monkeypatch) -> None: + """HEAD_CODED = False keeps the BF16 fragments even for a shard whose coding would decode bit for bit.""" + g = gen(seed_of("prepare bf16")) + w = randn((4096, 1024), g, 0.02) + monkeypatch.setattr(draft_step_kernels, "HEAD_CODED", False) + prepared = draft_step_kernels.prepare_head(SimpleNamespace(weight=w, bias=None), 2) + assert not prepared.coded and prepared.exc_off is None + assert torch.equal(prepared.wf, bf16_fragments(w)) + + +def test_guard_refuses_head_rows_past_x() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[1] = args[1] + 1000 + with pytest.raises(AssertionError, match="x_rows must index x"): + with validating(op): + op(*args, fixed_order=_fixed(args)) + + +def test_guard_refuses_too_few_ctas() -> None: + args = list(_case(op.CELLS[1].spec, 2)) + args[20] = 8 + with pytest.raises(AssertionError, match="tiles per CTA"): + with validating(op): + op(*args, fixed_order=_fixed(args)) + + +# ------------------------------------------------- one pass as two launches (uGPU, ops/m3_draft_step.UGPU) + + +@pytest.mark.parametrize( + "cell", + [c for c in op.CELLS if not c.spec["coded"]], + ids=[c.why[:44] for c in op.CELLS if not c.spec["coded"]], +) +def test_two_launches_per_domain_are_the_one_launch_pass(cell) -> None: + """A pass split into two launches (tiles [0, split) and [split, n), one per locality domain stream, the exchange + and bookkeeping by the last of both launches' CTAs) writes exactly what the one-launch pass writes: tokens, + winning keys, draft columns, next positions, KV and sequence lengths, and re-arms the counters and keys.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_locality import ( # noqa: E501 + LocalityMap, + ) + + one = _case(cell.spec, seed_of(cell.why)) + two = _case(cell.spec, seed_of(cell.why)) + op(*one) + lm = LocalityMap.get(torch.cuda.current_device()) + streams = lm.domain_streams() + n_tiles = two[3].shape[0] + split = max(1, n_tiles // 2) + per = min(lm.sms_per_domain, split, n_tiles - split) + main = torch.cuda.current_stream() + fork = torch.cuda.Event() + fork.record(main) + for st, lo, hi in ((streams[0], 0, split), (streams[1], split, n_tiles)): + st.wait_event(fork) + with torch.cuda.stream(st): + op(*two[:20], per, False, *two[22:], tile_lo=lo, tile_hi=hi, total_ctas=2 * per) + for st in streams: + done = torch.cuda.Event() + done.record(st) + main.wait_event(done) + torch.cuda.synchronize() + for i, name in ( + (11, "draft_out"), + (13, "ids_next"), + (15, "pos_next"), + (16, "kv_lens"), + (22, "out_keys"), + ): + assert torch.equal(one[i], two[i]), name + if one[18] is not None: + assert torch.equal(one[18], two[18]), "seq_lens" + assert two[6].tolist() == [0, 1] and not bool(two[5].any()), "counters / keys not re-armed" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py new file mode 100644 index 000000000000..ec29cf2a88e7 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step_fixed_order.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""D-FIX: the draft step's tile pieces and their summation order (ops/m3_draft_step.FIXED_ORDER). + +A CTA's (tile, chunk) range is cut into 16 equal warp segments; the shared-memory atomics add a tile's pieces in +arrival order, which is deterministic only while no tile takes more than 2 pieces. ``fixed_order`` sums them in warp +order instead. The geometry tests enumerate the real grid (the shard on 212 SMs) so a grid or shard change cannot +silently bring back an order-dependent sum; the GPU tests hold the fixed order to the atomics on exact data and to the +warp-order oracle on a planted tile whose sum depends on the order. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( + m3_draft_step as kds, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_step import ( # noqa: E501 + DRAFT_SHARD, + MAX_TOKENS, + bf16_fragments, + max_tile_pieces, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, head_workspace # noqa: E402 + +OP = "minimax_m3_sm107tp8_draft_step" +SMS = 212 # Rubin's SMs: the grid the draft step launches on +K = 6144 +CHUNKS = K // 16 +# This target's shard on Rubin: 7-8 tiles per CTA; a 7-tile CTA's tiles take up to 4 pieces (the atomics' sum is +# order-dependent, so the fixed order is on). +ORDER_DEPENDENT = True + +gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), + reason="the sm_107 TP8 target runs on SM107 only", +) + + +@pytest.mark.parametrize("grid", [SMS, kds.GRID]) +def test_real_grid_pieces(grid) -> None: + """One CTA per SM and the DS-GRID launch (ops/m3_draft_step.GRID).""" + pieces = max_tile_pieces(DRAFT_SHARD // 16, grid, CHUNKS) + assert (pieces > 2) == ORDER_DEPENDENT, (grid, pieces) + + +def test_ds_grid_launches_the_configured_grid(monkeypatch) -> None: + monkeypatch.setattr(kds, "_num_sms", lambda: SMS) + assert kds._grid() == min(kds.GRID, SMS) + monkeypatch.setattr(kds, "GRID", 0) + assert kds._grid() == SMS + + +def test_max_tile_pieces_counts_warp_segments() -> None: + # One tile per CTA: all 16 warps stream it. 7 tiles per CTA: a 168-chunk segment can sit inside a tile (TP8). + assert max_tile_pieces(16, 16, CHUNKS) == 16 + assert max_tile_pieces(7, 1, CHUNKS) == 4 + assert max_tile_pieces(8, 1, CHUNKS) == 2 + assert max_tile_pieces(15, 1, CHUNKS) == 2 + + +@pytest.mark.parametrize("flag", [True, False]) +def test_the_kernel_engages_the_fixed_order_where_the_geometry_needs_it(flag, monkeypatch) -> None: + monkeypatch.setattr(kds, "FIXED_ORDER", flag) + x = torch.empty(1, K) + + def head(n_tiles): + return types.SimpleNamespace(wf=torch.empty(n_tiles, 1), exc_off=None, bnd=None) + + assert kds._fixed_order(1, head(DRAFT_SHARD // 16), x, SMS) == (flag and ORDER_DEPENDENT) + assert kds._fixed_order(1, head(DRAFT_SHARD // 16), x, kds.GRID) == (flag and ORDER_DEPENDENT) + assert kds._fixed_order(1, head(16), x, 16) == flag # 16 pieces per tile + + +def _launcher(wf, x, fixed, nctas): + """(launch, ids, keys): one pass's op launch on buffers made here (outside any capture), world 1.""" + uc, keys, counters = head_workspace(1, slot_rows=MAX_TOKENS) + m = x.shape[0] + i32 = dict(dtype=torch.int32, device="cuda") + draft_out = torch.zeros(m, 4, **i32) + ids = torch.zeros(m, **i32) + pos_src, pos_next, kv_lens = ( + torch.full((m,), 100, **i32), + torch.zeros(m, **i32), + torch.full((m,), 500, **i32), + ) + out = torch.zeros(m, dtype=torch.int64, device="cuda") + + def launch(): + getattr(torch.ops.trtllm, OP)(x, None, m, wf, 0, keys, counters, uc, uc.data_ptr(), 0, 1, draft_out, 1, ids, + pos_src, pos_next, kv_lens, None, None, 3, nctas, False, out, + fixed_order=fixed) # fmt: skip + + return launch, ids, out + + +@gpu +@pytest.mark.parametrize("nctas,vs", [(SMS, DRAFT_SHARD), (16, 256)]) +@pytest.mark.parametrize("m", [1, 4, 8]) +def test_fixed_order_equals_the_atomics_on_exact_data(nctas, vs, m) -> None: + """Integer data: every fp32 sum is exact in any order, so both summations give the same bits.""" + g = gen(17 + m) + w = torch.randint(-1, 2, (vs, K), generator=g, device="cuda").to(torch.bfloat16) + x = torch.randint(-1, 2, (m, K), generator=g, device="cuda").to(torch.bfloat16) + wf = bf16_fragments(w) + out = {} + for fixed in (False, True): + launch, ids, keys = _launcher(wf, x, fixed, nctas) + launch() + torch.cuda.synchronize() + out[fixed] = (keys, ids) + assert torch.equal(out[False][0], out[True][0]) + ref = (x.double() @ w.double().t()).float().bfloat16().float().argmax(dim=1).int() + assert torch.equal(out[True][1], ref) + + +@gpu +@pytest.mark.parametrize("graph", [False, True]) +def test_fixed_order_is_the_warp_order_sum_on_a_planted_tile(graph) -> None: + """One tile per CTA (16 pieces): row 0's pieces are +2^25 (warp 0), +1 (warps 1..14), -2^25 (warp 15). In warp + order the +1s are absorbed by 2^25 and the row sums to 0, so the competitor (row 255, logit 0.5) wins; arrival + order could give up to 14. Every launch must give the warp-order result.""" + w = torch.zeros(256, K, device="cuda") + seg = CHUNKS // 16 * 16 # one warp segment in K + w[0, 0] = 2.0**25 + for s in range(1, 15): + w[0, s * seg] = 1.0 + w[0, 15 * seg] = -(2.0**25) + w[255, 0] = 0.5 + wf = bf16_fragments(w.to(torch.bfloat16)) + x = torch.ones(1, K, device="cuda", dtype=torch.bfloat16) + launch, ids, _ = _launcher(wf, x, True, 16) + run = launch + if graph: + launch() + torch.cuda.synchronize() + gr = torch.cuda.CUDAGraph() + cs = torch.cuda.Stream() + with torch.cuda.stream(cs), torch.cuda.graph(gr, stream=cs): + launch() + run = gr.replay + for _ in range(16): + ids.zero_() + run() + torch.cuda.synchronize() + assert int(ids[0]) == 255 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py new file mode 100644 index 000000000000..5e6ae369f8fa --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py @@ -0,0 +1,376 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the forward-tail spec fold (ops/m3_fold_tail and M3SpecStepFold.in_forward). + +The tail must write exactly what the sampler's store (slot_scatter) and then the engine's next-step gathers +(step_gather, rows = the stored rows in order, first previous row 0) write, for the rows the step armed, and nothing +when the step was not armed. +""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( + m3_fold_tail, + m3_spec_tail, +) + +assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" + +VOCAB = 200064 +NUM_SLOTS = 9 +MAX_ROWS = 8 +STALE = -4 + + +def _store(widths, g): + wn, wx, wd = widths + i32 = dict(dtype=torch.int32, device="cuda") + return types.SimpleNamespace( + new_tokens=torch.randint( + 0, VOCAB, (wn, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_new_tokens=torch.randint( + 0, VOCAB, (wx, NUM_SLOTS, 1), generator=g, device="cuda", dtype=torch.int32 + ), + next_draft_tokens=torch.randint( + 0, VOCAB, (NUM_SLOTS, wd), generator=g, device="cuda", dtype=torch.int32 + ), + new_tokens_lens=torch.randint(1, 5, (NUM_SLOTS,), generator=g, **i32), + ) + + +def _clone(ns): + return types.SimpleNamespace(**{k: v.clone() for k, v in vars(ns).items()}) + + +def _outputs(rows, t, g): + d = t - 1 + return { + "new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_new_tokens": torch.randint( + 0, VOCAB, (rows, t), generator=g, device="cuda", dtype=torch.int32 + ), + "next_draft_tokens": torch.randint( + 0, VOCAB, (rows, d), generator=g, device="cuda", dtype=torch.int32 + ), + "new_tokens_lens": torch.randint( + 1, t + 1, (rows,), generator=g, device="cuda", dtype=torch.int32 + ), + } + + +def _engine(t): + i32 = dict(dtype=torch.int32, device="cuda") + return ( + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS * (t - 1),), STALE, **i32), + torch.full((MAX_ROWS * t,), STALE, **i32), + torch.full((MAX_ROWS,), STALE, **i32), + ) + + +def _reference(store, engine, slots, out, t): + """slot_scatter into the store, then step_gather of the next step's inputs from it.""" + assert m3_spec_tail.slot_scatter( + slots, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["new_tokens_lens"], + store.new_tokens_lens, + out["next_draft_tokens"], + store.next_draft_tokens, + ) + p = slots.numel() + ids, draft, pos_off, kv_off = engine + slots32 = slots.to(torch.int32) + assert m3_spec_tail.step_gather( + store.next_new_tokens, + store.next_draft_tokens, + store.new_tokens_lens, + slots32, + slots32.repeat_interleave(t), + p, + t, + t - 1, + ids, + 0, + draft, + 0, + pos_off, + kv_off, + 0, + ) + + +def _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=False, one_warp=False): + ids, draft, pos_off, kv_off = engine + assert m3_fold_tail.fold_tail( + ctrl, + slots_buf, + rows, + out["new_tokens"], + store.new_tokens.squeeze(-1).T, + out["next_new_tokens"], + store.next_new_tokens.squeeze(-1).T, + out["next_draft_tokens"], + store.next_draft_tokens, + out["new_tokens_lens"], + store.new_tokens_lens, + ids, + draft, + pos_off, + kv_off, + t, + cpp=cpp, + one_warp=one_warp, + ) + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("widths", [(4, 4, 3), (8, 8, 7)], ids=["store=step", "store wider"]) +def test_tail_matches_scatter_then_gather(rows, widths, cpp) -> None: + """Store bit for bit; engine rows [0, P*T) / [0, P) bit for bit (the forward reads only those rows).""" + t = 4 + g = torch.Generator(device="cuda").manual_seed(1000 * rows + widths[0]) + store = _store(widths, g) + out = _outputs(rows, t, g) + slots = torch.randperm(NUM_SLOTS, generator=g, device="cuda")[:rows] + ref_store, tail_store = _clone(store), _clone(store) + ref_engine, tail_engine = _engine(t), _engine(t) + _reference(ref_store, ref_engine, slots, out, t) + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + slots_buf = torch.zeros(MAX_ROWS, dtype=torch.long, device="cuda") + slots_buf[:rows] = slots + _tail(tail_store, tail_engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(ref_store, name), getattr(tail_store, name)), name + spans = (rows * t, rows * (t - 1), rows * t, rows) + for k, (ref, got, n) in enumerate(zip(ref_engine, tail_engine, spans)): + assert torch.equal(ref[:n], got[:n]), k + assert torch.all(got[n:] == STALE), k + + +@pytest.mark.parametrize("cpp", [False, True], ids=["triton", "cpp"]) +def test_unarmed_tail_writes_nothing(cpp) -> None: + """ctrl[0] = 0 (the step was not armed): the store and the engine buffers are untouched.""" + t, rows = 4, 2 + g = torch.Generator(device="cuda").manual_seed(7) + store = _store((4, 4, 3), g) + out = _outputs(rows, t, g) + before, engine = _clone(store), _engine(t) + ctrl = torch.zeros(1, dtype=torch.int32, device="cuda") + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp) + torch.cuda.synchronize() + for name in vars(store): + assert torch.equal(getattr(before, name), getattr(store, name)), name + for buf in engine: + assert torch.all(buf == STALE) + + +def _fold(): + fold = modeling.M3SpecStepFold() + t = 4 + engine = _engine(t) + fold.bind_engine_buffers(*engine, max_batch_size=MAX_ROWS) + g = torch.Generator(device="cuda").manual_seed(11) + fold.bind_sampler_store(_store((4, 4, 3), g)) + return fold, g, t + + +def test_fold_handshake_stages_only_confirmed_rows() -> None: + """begin_step arms; the sampler confirms only the armed rows from the tail's outputs; the next step is staged only + for exactly those rows in that order.""" + fold, g, t = _fold() + assert fold.begin_step([3, 5], 2) is False # nothing staged before the first forward + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [3, 5]) is True + assert fold.begin_step([3, 5], 2) is True + out = _outputs(2, t, g) + fold.store_in_forward(out) + assert ( + fold.sampler_rows_stored(out, 0, [5, 3]) is False + ) # other order: the sampler stores itself + assert fold.begin_step([5, 3], 2) is False + + +def test_fold_handshake_needs_the_tail_outputs() -> None: + """Outputs the tail did not run on (another dict, a mixed step, an unarmed step) are never reported stored.""" + fold, g, t = _fold() + fold.begin_step([1], 1) + fold.store_in_forward(_outputs(1, t, g)) + assert fold.sampler_rows_stored(_outputs(1, t, g), 0, [1]) is False + fold.begin_step(None, 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 0, [1]) is False + fold.begin_step([1], 1) + out = _outputs(1, t, g) + fold.store_in_forward(out) + assert fold.sampler_rows_stored(out, 1, [1]) is False # rows were skipped (chunked context) + + +def test_a_padded_step_after_an_armed_one_writes_nothing() -> None: + """K3 #150 check #5: CUDA-graph padding rows carry no seq slot, so the engine passes rows=None for a padded step. + Even right after an armed step, that step uploads ctrl 0 and its forward tail leaves the store and the engine + buffers untouched (a dummy row never writes a live slot).""" + fold, g, t = _fold() + fold.begin_step([2, 4], 2) + fold.store_in_forward(_outputs(2, t, g)) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 2 + + fold.begin_step(None, 4) + torch.cuda.synchronize() + assert int(fold._ctrl.item()) == 0 + store_before = _clone(fold.store) + engine_before = [buf.clone() for buf in fold.step_inputs] + fold.store_in_forward(_outputs(4, t, g)) + torch.cuda.synchronize() + for name in vars(store_before): + assert torch.equal(getattr(store_before, name), getattr(fold.store, name)), name + for before, buf in zip(engine_before, fold.step_inputs): + assert torch.equal(before, buf) + + +def test_the_fold_stores_in_the_forward_once_bound() -> None: + """The fold always stores and stages at the forward's tail; before the engine binds its buffers nothing is armed.""" + fold = modeling.M3SpecStepFold() + assert fold.in_forward + assert fold.begin_step([0], 1) is False + + +def test_the_store_copy_event_is_recorded_at_binding(): + """An external event the sampler records after its side-stream store copies and the forward's tail waits on, + recorded once at binding so a captured wait has a completed record to wait on.""" + fold = modeling.M3SpecStepFold() + device = torch.device("cuda") + fold.bind_engine_buffers( + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + torch.zeros(8, dtype=torch.int32, device=device), + 4, + ) + assert fold.store_copies_done is not None + fold.store_copies_done.synchronize() + + +def test_the_cpp_tail_launches_plainly() -> None: + """The C++ tail writes the next step's inputs, which pre-wait readers take before their grid wait (MSA contract): + its kernel node carries no programmatic (PDL) serialization.""" + from cuda.bindings import runtime as rt + + t, rows = 4, 2 + g = torch.Generator(device="cuda").manual_seed(13) + store, out, engine = _store((4, 4, 3), g), _outputs(rows, t, g), _engine(t) + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=True) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + _tail(store, engine, ctrl, slots_buf, out, t, rows, cpp=True) + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + # cuda-python changed the signature: (graph[, numNodes]) on older bindings, (graph, nodes, numNodes) on newer ones. + try: + err, _, n = rt.cudaGraphGetNodes(raw) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + except TypeError: + err, _, n = rt.cudaGraphGetNodes(raw, None, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, None, n) + flags = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind == rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + # Older bindings return no attribute for a kernel launched without one (a plain launch). + flags.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + assert flags == [0], flags + + +@pytest.mark.parametrize("rows", [1, 2, 4, 8]) +@pytest.mark.parametrize("widths", [(4, 4, 3), (8, 8, 7)], ids=["store=step", "store wider"]) +def test_one_warp_form_is_bitwise_the_kernel(rows, widths) -> None: + """The in-kernel form (one warp over the rows: what the drafter's last-pass draft_step runs in its last CTA, H-R2-3) + writes exactly what the one-CTA-per-row kernel writes, armed and unarmed, eager and from a graph.""" + t = 4 + for trial in range(3): + g = torch.Generator(device="cuda").manual_seed(100 * rows + trial) + store, out, engine = _store(widths, g), _outputs(rows, t, g), _engine(t) + slots_buf = torch.randperm(NUM_SLOTS, generator=g, device="cuda")[:MAX_ROWS].to(torch.long) + for armed in (rows, 0): + ctrl = torch.tensor([armed], dtype=torch.int32, device="cuda") + a_store, a_engine = _clone(store), tuple(e.clone() for e in engine) + b_store, b_engine = _clone(store), tuple(e.clone() for e in engine) + _tail(a_store, a_engine, ctrl, slots_buf, out, t, rows, cpp=True) + _tail(b_store, b_engine, ctrl, slots_buf, out, t, rows, cpp=True, one_warp=True) + torch.cuda.synchronize() + for k in vars(a_store): + assert torch.equal(getattr(a_store, k), getattr(b_store, k)), (k, armed, trial) + for x, y in zip(a_engine, b_engine): + assert torch.equal(x, y), (armed, trial) + # From a captured graph, replayed on fresh outputs. + g = torch.Generator(device="cuda").manual_seed(7) + store, out, engine = _store(widths, g), _outputs(rows, t, g), _engine(t) + slots_buf = torch.arange(MAX_ROWS, dtype=torch.long, device="cuda") + ctrl = torch.tensor([rows], dtype=torch.int32, device="cuda") + b_store, b_engine = _clone(store), tuple(e.clone() for e in engine) + _tail(b_store, b_engine, ctrl, slots_buf, out, t, rows, cpp=True, one_warp=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _tail(b_store, b_engine, ctrl, slots_buf, out, t, rows, cpp=True, one_warp=True) + for step in range(3): + fresh = _outputs(rows, t, torch.Generator(device="cuda").manual_seed(20 + step)) + for k in out: + out[k].copy_(fresh[k]) + a_store, a_engine = _clone(b_store), tuple(e.clone() for e in b_engine) + _tail(a_store, a_engine, ctrl, slots_buf, out, t, rows, cpp=True) + graph.replay() + torch.cuda.synchronize() + for k in vars(a_store): + assert torch.equal(getattr(a_store, k), getattr(b_store, k)), (k, step) + for x, y in zip(a_engine, b_engine): + assert torch.equal(x, y), step + + +def test_the_store_wait_precedes_only_a_plain_launch_under_capture(): + """The store wait is an external-event wait node under graph capture, which drops the next kernel's PDL edge: + it may precede the fold tail (a plain launch); any other launch fails the capture (eager steps are not checked).""" + fold = modeling.M3SpecStepFold() + device = torch.device("cuda") + fold.bind_engine_buffers( + *(torch.zeros(8, dtype=torch.int32, device=device) for _ in range(4)), 4 + ) + assert fold.STORE_WAIT_BEFORE == frozenset({"fold_tail"}) + fold._wait_store_copies(before="draft_step") # eager: allowed + x = torch.zeros(4, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + fold._wait_store_copies(before="fold_tail") + x.add_(1) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(x, torch.ones(4, device=device)) + bad = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="PDL edge"): + with torch.cuda.graph(bad, stream=stream): + fold._wait_store_copies(before="draft_step") diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py new file mode 100644 index 000000000000..0f6b4f255102 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test of the per-request forced-acceptance draw (``FORCE_AL_PER_REQUEST``, the ``m3_force_vec_req`` entry, +synth only): each request's forced count comes from its own key (its KV length at the verify), so it does not depend +on the call history, on the request's slot in the batch, or on the other requests.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec_req, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SEED, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _pool(): + g = torch.Generator(device="cpu") + g.manual_seed(_FORCE_ACCEPT_RNG_SEED) + return torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, dtype=torch.float32, generator=g).cuda() + + +def _draw(pool, keys, base_total=2, max_total=4, frac=0.6): + out = torch.empty(keys.numel(), dtype=torch.int32, device="cuda") + m3_force_vec_req( + pool, + keys, + out, + keys.numel(), + base_total, + max_total, + frac, + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + max(16, 1 << (keys.numel() - 1).bit_length()), + ) + return out + + +def test_a_request_draws_the_same_whatever_ran_before_or_beside_it() -> None: + pool = _pool() + keys = torch.tensor([165_521, 40_003, 323_007, 7], dtype=torch.int32, device="cuda") + first = _draw(pool, keys) + for _ in range(5): # no shared counter: repeated calls draw the same + assert torch.equal(_draw(pool, keys), first) + perm = torch.tensor([2, 0, 3, 1], device="cuda") + assert torch.equal(_draw(pool, keys[perm]), first[perm]) # slot-independent + alone = torch.cat([_draw(pool, keys[i : i + 1]) for i in range(4)]) + assert torch.equal(alone, first) # independent of the other requests + + +def test_the_counts_follow_the_forced_fraction() -> None: + pool = _pool() + keys = torch.arange(100_000, 100_000 + 64 * 400, dtype=torch.int32, device="cuda") + counts = torch.cat([_draw(pool, keys[i : i + 64]) for i in range(0, keys.numel(), 64)]) + assert set(counts.tolist()) <= {2, 3} + extra = (counts == 3).float().mean().item() + assert abs(extra - 0.6) < 0.02, extra diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py new file mode 100644 index 000000000000..6b63ba265112 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec entry of the MiniMax-M3 tp8 target: the fractional forced-acceptance draw.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(5) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + counter = torch.tensor([41], dtype=torch.int64, device="cuda") + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + counter, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_SLOT_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert int(counter.item()) == 42, "the draw bumps the step counter once" + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py new file mode 100644 index 000000000000..871dfa22d617 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_force_vec_req entry of the MiniMax-M3 tp8 target: the per-request fractional forced-acceptance +draw (one draw per request at its own key, no shared counter).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 + m3_force_vec_req as op, +) +from tensorrt_llm._torch.speculative.interface import ( + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE, +) + +__extra_import_path__ = ["../../../.."] +from _validating import validating # noqa: E402 — needs the path declared above + +assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +def test_certified_cells(cell) -> None: + spec = cell.spec + g = torch.Generator(device="cuda").manual_seed(7) + pool = torch.rand(_FORCE_ACCEPT_RNG_POOL_SIZE, generator=g, device="cuda") + # Each request's key is its KV length at the verify: long, distinct, arbitrary. + keys = torch.randint( + 1, 1 << 20, (spec["batch"],), generator=g, device="cuda", dtype=torch.int64 + ).to(torch.int32) + out = torch.empty(spec["batch"], dtype=torch.int32, device="cuda") + args = ( + pool, + keys, + out, + spec["batch"], + 2, + 4, + spec["frac"], + _FORCE_ACCEPT_RNG_COUNTER_STRIDE, + _FORCE_ACCEPT_RNG_POOL_SIZE - 1, + spec["block"], + ) + ref = op.reference(*args) + with validating(op): + op(*args) + assert torch.equal(out, ref) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py new file mode 100644 index 000000000000..821c6eeb08ae --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the target's copy of FlashInfer's split-K low-M GEMM (ops/m3_lowm_splitk, +modeling.DRAFT_LOWM_OWN_SPLITK): bit for bit FlashInfer's kernel at every certified split-K cell of the drafter's +linears, eager and in a CUDA graph, with and without PDL.""" + +import pytest +import torch +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic +from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + TACTICS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, + run_splitk_dense_swiglu, +) + +assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" + +# Every split-K cell of the TP8 drafter table: (N, K, max rows, tactic) for kind "s". +CELLS = [ + (nk[0], nk[1], rows, tactic) + for nk, entries in sorted(TACTICS.items()) + for rows, kind, tactic in entries + if kind == "s" +] + + +def test_the_table_has_the_four_certified_split_k_cells() -> None: + assert sorted((n, k, m) for n, k, m, _ in CELLS) == [ + (6144, 1024, 1), + (6144, 1024, 4), + (6144, 2304, 1), + (6144, 2304, 4), + ] + + +@pytest.mark.parametrize("use_pdl", [False, True], ids=["plain", "pdl"]) +@pytest.mark.parametrize("cell", CELLS, ids=[f"{n}x{k}@{m}" for n, k, m, _ in CELLS]) +def test_bitwise_flashinfer(cell, use_pdl) -> None: + n, k, m, tactic = cell + g = torch.Generator(device="cuda").manual_seed(n + 7 * k + m) + for trial in range(3): + x = (torch.randn(m, k, generator=g, device="cuda") * (1.0 + trial)).to(torch.bfloat16) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + ref = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + got = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + fi_run(x, w.t(), None, ref, use_pdl, FiTactic(*tactic)) + run_splitk_dense(x, w.t(), None, got, use_pdl, SplitKTactic(*tactic)) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), f"trial {trial}" + # In a CUDA graph, as the drafter runs it. + cap = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_splitk_dense(x, w.t(), None, cap, use_pdl, SplitKTactic(*tactic)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(cap.view(torch.int16), ref.view(torch.int16)) + + +# The drafter MLP's down projection: the SwiGLU folded into the GEMM's B load (c2). +DOWN_CELLS = [c for c in CELLS if (c[0], c[1]) == (6144, 2304)] + + +@pytest.mark.parametrize("use_pdl", [False, True], ids=["plain", "pdl"]) +@pytest.mark.parametrize("cell", DOWN_CELLS, ids=[f"{n}x{k}@{m}" for n, k, m, _ in DOWN_CELLS]) +def test_swiglu_prologue_is_bitwise_silu_then_split_k(cell, use_pdl) -> None: + n, k, m, tactic = cell + g = torch.Generator(device="cuda").manual_seed(31 * n + m) + w = (torch.randn(n, k, generator=g, device="cuda") / k**0.5).to(torch.bfloat16) + scratch = torch.full((8, k), 7, dtype=torch.bfloat16, device="cuda") + for trial in range(3): + gate_up = (torch.randn(m, 2 * k, generator=g, device="cuda") * (2.0 + trial)).to( + torch.bfloat16 + ) + gate_up[:, :6] = torch.tensor([60.0, -60.0, 0.0, -0.0, 88.0, -88.0], dtype=torch.bfloat16) + act = torch.ops.trtllm.silu_and_mul(gate_up) + ref = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + fi_run(act, w.t(), None, ref, use_pdl, FiTactic(*tactic)) + got = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + run_splitk_dense_swiglu(gate_up, w.t(), got, use_pdl, SplitKTactic(*tactic), scratch) + torch.cuda.synchronize() + assert torch.equal(got.view(torch.int16), ref.view(torch.int16)), f"trial {trial}" + assert torch.equal(scratch[:m].view(torch.int16), act.view(torch.int16)), ( + "the B scratch holds the activation" + ) + cap = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_splitk_dense_swiglu(gate_up, w.t(), cap, use_pdl, SplitKTactic(*tactic), scratch) + scratch.fill_(3) # a replay rewrites what it reads + graph.replay() + torch.cuda.synchronize() + assert torch.equal(cap.view(torch.int16), ref.view(torch.int16)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py new file mode 100644 index 000000000000..6066e23c92a3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the tp8 m3_sampled_accept entry (fused temperature + top-p verify over the 25008-column TP8 shard) +and its workspace size.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_sampled_accept import ( # noqa: E501 + m3_sampled_accept as op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_sampled_accept_ws_words import ( # noqa: E501 + m3_sampled_accept_ws_words as ws_words, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" + +DRAFT_LEN = 3 +#: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). +HEAD = 12 +#: Draws per distribution cell, and the total-variation band against the reference distribution. With 12 +#: categories and >= 4000 draws the expected TV of an exact sampler is ~0.02 (sd ~0.006); 0.06 is ~6 sd above it +#: while a wrong nucleus cut (one category too many or too few, >= 1 / 12 of the mass moved) lands at >= 0.08. +DRAWS = 4000 +TV_BAND = 0.06 + + +def _logits(spec, g, rows): + vs = spec["vs"] + if spec["check"] == "flat": + return (torch.randn((rows, vs), generator=g, device="cuda") * 0.3).to(torch.bfloat16) + if spec["check"] == "argmax": + return (torch.randn((rows, vs), generator=g, device="cuda") * 3).to(torch.bfloat16) + row = torch.full((vs,), -20.0, device="cuda") + hot = torch.randperm(vs, generator=g, device="cuda")[:HEAD] + row[hot] = torch.linspace(2.0, 0.2, HEAD, device="cuda") + return row.to(torch.bfloat16).expand(rows, vs).contiguous() + + +def _case(spec, seed): + g = gen(seed) + nreq = spec["nreq"] + rows = nreq * (DRAFT_LEN + 1) + logits = _logits(spec, g, rows) + words = ws_words(1) + uc = torch.zeros(words, dtype=torch.int64, device="cuda") + counters = torch.zeros(2, dtype=torch.int32, device="cuda") + tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + extra = torch.zeros(1, dtype=torch.int32, device="cuda") + temperatures = torch.full((rows,), spec["temperature"], dtype=torch.float32, device="cuda") + top_ps = torch.full((rows,), spec["top_p"], dtype=torch.float32, device="cuda") + seeds = torch.randint(0, 2**62, (rows,), generator=g, device="cuda", dtype=torch.int64) + offsets = torch.zeros(rows, dtype=torch.int64, device="cuda") + # Drafts drawn from the likeliest tokens, so some prefixes match. + top = logits[0].float().topk(3).indices.to(torch.int32) + spec["vocab_offset"] + draft = top[torch.randint(0, 3, (nreq, DRAFT_LEN), generator=g, device="cuda")].contiguous() + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, DRAFT_LEN + 1), -7, dtype=torch.int32, device="cuda") + force = spec.get("force", "none") + force_total = 3 if force == "total" else 0 + force_vec = (torch.randint(1, DRAFT_LEN + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if force == "vec" else None) # fmt: skip + return [logits, spec["vocab_offset"], temperatures, top_ps, seeds, offsets, uc, uc.data_ptr(), counters, tokens, + 0, 1, draft, batch_idx, acc, num, gather, force_total, next_new, force_vec, extra, True] # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Statistical for the draw, exact for everything derived from it. + + Every drawn token has nonzero reference probability (lies in the nucleus); the acceptance outputs equal the + reference's strict acceptance of the drawn tokens, every call; the call counter advances once per call; and, + for the distribution cells, the draws over many calls (fresh Philox offsets) are within TV_BAND of the + reference distribution (see DRAWS / TV_BAND). + """ + spec = cell.spec + args = _case(spec, seed_of(cell.why)) + logits, offset, rows = args[0], spec["vocab_offset"], args[0].shape[0] + certifying(op) + ref = op.reference(*args) + probs = ref["probs"] + calls = 1 if spec["check"] == "argmax" else -(-DRAWS // rows) + counts = torch.zeros(logits.shape[1], dtype=torch.float64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + drawn = args[9][:rows].clone() + local = drawn.long() - offset + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + assert bool((probs.gather(1, local[:, None]) > 0).all()), ( + "a drawn token outside the nucleus" + ) + acc = ref["accept"](drawn) + assert torch.equal(args[14], acc["accepted"]) + assert torch.equal(args[15], acc["num_accepted"]) + assert torch.equal(args[16], acc["gather_ids"]) + assert torch.equal(args[18], acc["next_new"]) + assert args[8].tolist() == [0, c + 1] + counts.index_add_(0, local, torch.ones(rows, dtype=torch.float64, device="cuda")) + if spec["check"] == "distribution": + emp = counts / counts.sum() + tv = 0.5 * (emp - probs[0].double()).abs().sum().item() + assert tv <= TV_BAND, f"total variation {tv:.3f} against the reference nucleus distribution" + + +def _walk_outputs(args, one_round_k, ilp, calls=3, offset_step=1000): + """The draws and acceptance outputs of `calls` calls with the owner walk one token (ilp False) or 8 tokens at a + time (ilp True), at SA-1R variant `one_round_k`, each call at its own Philox offsets.""" + if one_round_k: + words = torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words(1, one_round_k) + args[6] = torch.zeros(words, dtype=torch.int64, device="cuda") + args[7] = args[6].data_ptr() + rows = args[0].shape[0] + got = [] + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * offset_step) + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept( + *args, True, False, one_round_k, None, ilp + ) + got.append([args[i].clone() for i in (9, 14, 15, 16, 18)]) + return got + + +@pytest.mark.parametrize("one_round_k", [0, 82]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_the_ilp_walk_changes_nothing(cell, one_round_k) -> None: + """ops/m3_sampled_accept.ILP_WALK (SA-PW) only computes the owner's token weights 8 at a time; the walk adds + them in the same order with the same early exit, so the draws and acceptance outputs are the one-token walk's, + bit for bit, call after call, with SA-1R off and on.""" + one = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, False) + ilp = _walk_outputs(_case(cell.spec, seed_of(cell.why)), one_round_k, True) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.95]) +def test_the_ilp_walk_changes_nothing_over_many_rounds(top_p) -> None: + """A flat tail (many tokens of equal weight) under a tiny and a wide nucleus: rows take many exchange rounds and + each round's walk exits at a different token; the 8-token walk draws the same tokens as the one-token walk.""" + spec = dict(nreq=1, vs=25008, vocab_offset=25008, check="flat", top_p=top_p, temperature=1.0) + one = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), 0, False, calls=24) + ilp = _walk_outputs(_case(spec, seed_of(f"ilp walk flat {top_p}")), 0, True, calls=24) + for call, (a, b) in enumerate(zip(one, ilp)): + for x, y in zip(a, b): + assert torch.equal(x, y), ( + f"call {call}: the 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +def test_the_ilp_walk_is_the_same_under_graph_replay() -> None: + """The 8-token walk captured in a CUDA graph and replayed draws what the one-token walk draws eagerly.""" + cell = op.CELLS[0] + eager = _walk_outputs(_case(cell.spec, seed_of(cell.why)), 0, False) + _walk_outputs( + _case(cell.spec, seed_of("ilp walk warm-up")), 0, True, calls=1 + ) # first launch outside the capture + args = _case(cell.spec, seed_of(cell.why)) + rows = args[0].shape[0] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept(*args, True, False, 0, None, True) + torch.cuda.current_stream().wait_stream(stream) + for c in range(3): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + graph.replay() + torch.cuda.synchronize() + for x, y in zip(eager[c], [args[i] for i in (9, 14, 15, 16, 18)]): + assert torch.equal(x, y), ( + f"replay {c}: the graphed 8-token walk changes {x.tolist()} -> {y.tolist()}" + ) + + +@pytest.mark.parametrize("top_p", [0.001, 0.0002]) +def test_long_rejection_chains_stay_in_the_nucleus(top_p) -> None: + """A tiny nucleus over a flat tail: rows that need more exchange rounds than the workspace has slots. + + Each round's candidate is drawn over the tokens heavier than the last one, so a row takes ~ln(1 / top_p) + 1 + rounds on average, with a tail past the 16 exchange slots (about 1% of the rows at top_p 0.001). Rounds reuse the + slots cyclically, told apart by a lap-parity bit in the wire tags, and no round accepts unconditionally: every + draw, however many rounds it took, lies in the reference nucleus. World 1 (one rank's 8 CTAs per row exchange + through the same slot ring); the cross-rank path needs the multi-device job. + """ + spec = dict(nreq=1, vs=25008, vocab_offset=25008, check="flat", top_p=top_p, temperature=1.0) + args = _case(spec, seed_of(f"flat tail top_p {top_p}")) + logits, rows = args[0], args[0].shape[0] + probs = op.reference(*args)["probs"] + calls = -(-DRAWS // rows) + outside = torch.zeros((), dtype=torch.int64, device="cuda") + for c in range(calls): + args[5].copy_(torch.arange(rows, device="cuda", dtype=torch.int64) + c * 1000) + op(*args) + local = args[9][:rows].long() - spec["vocab_offset"] + assert bool(((local >= 0) & (local < logits.shape[1])).all()), "token outside the shard" + outside += (probs.gather(1, local[:, None]) == 0).sum() + assert args[8].tolist() == [0, c + 1] + draws = calls * rows + rounds = 2.0 + args[20].item() / draws + kept = int((probs > 0).sum(1).min().item()) + assert outside.item() == 0, ( + f"{outside.item()} of {draws} draws outside the nucleus (kept >= {kept} of {logits.shape[1]}, " + f"{rounds:.2f} rounds per row on average)" + ) + assert rounds > 4.0, ( + f"only {rounds:.2f} rounds per row: the rows no longer stress the slot ring" + ) + + +@pytest.mark.parametrize("cell", ws_words.CELLS, ids=[c.why[:44] for c in ws_words.CELLS]) +def test_workspace_words(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + certifying(ws_words) + world = cell.spec["world"] + assert ws_words(world) == ws_words.reference(world) + + +def test_guard_refuses_top_k_rows() -> None: + args = _case(op.CELLS[0].spec, 1) + top_ks = torch.full((args[0].shape[0],), 20, dtype=torch.int32, device="cuda") + with pytest.raises(AssertionError, match="top-k rows"): + with validating(op): + op(*args, top_ks=top_ks) + + +def test_guard_refuses_fp64_temperatures() -> None: + args = _case(op.CELLS[0].spec, 2) + args[2] = args[2].double() + with pytest.raises(AssertionError, match="fp32 sampling parameters"): + with validating(op): + op(*args) + + +@pytest.mark.parametrize( + "check,top_p", [("flat", 0.95), ("argmax", 0.95), ("flat", 0.001), ("distribution", 0.5)] +) +def test_the_copy_is_the_tp2_kernel_bit_for_bit(check, top_p) -> None: + """This target's own op (trtllm::minimax_m3_sm107tp8_sampled_accept) is the TP2 kernel it copies, bit for bit: + the same tokens, acceptance outputs, call counter and extra-round count from the same seeds and offsets.""" + tp2 = getattr(torch.ops.trtllm, "minimax_m3_sm107tp2_sampled_accept", None) + if tp2 is None: + pytest.skip("the TP2 op is not loaded") + spec = dict(nreq=1, vs=25008, vocab_offset=25008, check=check, top_p=top_p, temperature=1.0) + for call in range(8): + outs = [] + for which in ("tp8", "tp2"): + args = _case(spec, seed_of(f"copy {check} {top_p}")) + args[5].copy_( + torch.arange(args[0].shape[0], device="cuda", dtype=torch.int64) + call * 1000 + ) + if which == "tp8": + op(*args) + else: + tp2(*args[:21], args[21]) + torch.cuda.synchronize() + outs.append([args[i].clone() for i in (9, 14, 15, 16, 18, 20)] + [args[8].clone()]) + for a, b in zip(*outs): + assert torch.equal(a, b) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py new file mode 100644 index 000000000000..3925faf6082a --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept_1r.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""SA-1R (ops/m3_sampled_accept.ONE_ROUND_K) on trtllm::minimax_m3_sm107tp8_sampled_accept, world 1, the TP8 shard. + +The one-round variants (one_round_k 9: exact top-9 list; 82: pooled top-8 + bound) must return the two-round +kernel's (one_round_k 0) sampled tokens, accepted tokens, accepted counts, gather ids and extra-attempt count BIT FOR +BIT on the same Philox stream, over random AgentX-like rows (several temperatures / top-p, broad and heavy-tailed +heads) and over rows built so the decision sits within ulps of top_p * Z (the margin fallback) or a slice holds more +heavy tokens than its list (the incomplete-list fallback), eager and under CUDA-graph replay. The round-0 decided +counter must engage on the AgentX-like rows and stay below the row count on the boundary rows (fallbacks fire). +Ported from MoE's prototype test (moe/sa1r/6c753632ec, exact on 39,968 rows, 788162). +""" + +import pytest +import torch + +import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* + +VS, D = 25008, 3 +KX = (9, 82) + +if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + pytest.skip("the sm_107 tp8 sampled accept runs on SM 10.7", allow_module_level=True) + + +class Runner: + """One variant with its own workspace / counters (call tag sequence) and outputs.""" + + def __init__(self, k, rows): + self.k, self.rows = k, rows + words = int(torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept_ws_words(1, k)) + self.ws = torch.zeros(words, dtype=torch.int64, device="cuda") + self.counters = torch.zeros(2, dtype=torch.int32, device="cuda") + self.tokens = torch.zeros(48, dtype=torch.int32, device="cuda") + self.extra = torch.zeros(1, dtype=torch.int32, device="cuda") + self.decided = torch.zeros(1, dtype=torch.int32, device="cuda") + nreq = rows // (D + 1) + self.acc = torch.zeros((nreq, D + 1), dtype=torch.int32, device="cuda") + self.nacc = torch.zeros(nreq, dtype=torch.int32, device="cuda") + self.gids = torch.zeros(nreq, dtype=torch.int64, device="cuda") + + def __call__(self, lg, temps, tps, seeds, offs, draft, bidx): + torch.ops.trtllm.minimax_m3_sm107tp8_sampled_accept( + lg, 0, temps, tps, seeds, offs, self.ws, self.ws.data_ptr(), self.counters, self.tokens, 0, 1, draft, + bidx, self.acc, self.nacc, self.gids, 0, None, None, self.extra, False, True, False, self.k, + self.decided) # fmt: skip + + def out(self): + return (self.tokens[: self.rows].clone(), self.acc.clone(), self.nacc.clone(), self.gids.clone(), + self.extra.clone()) # fmt: skip + + +def _params(rows, temp, top_p, seed0): + temps = torch.full((rows,), temp, dtype=torch.float32, device="cuda") + tps = ( + (top_p if torch.is_tensor(top_p) else torch.full((rows,), top_p)) + .float() + .cuda() + .contiguous() + ) + seeds = torch.arange(rows, dtype=torch.int64, device="cuda") * 7919 + seed0 + offs = torch.zeros(rows, dtype=torch.int64, device="cuda") + return temps, tps, seeds, offs + + +def _sets(n, rows, g, kind): + out = [] + for _ in range(n): + lg = torch.randn(rows, VS, device="cuda", generator=g) * (1.0 if kind == "flat" else 2.0) + if kind == "lm": # a few heavy tokens over a broad tail + hot = torch.randint(0, VS, (rows, 8), device="cuda", generator=g) + lg.scatter_(1, hot, 12.0 + 3.0 * torch.rand(rows, 8, device="cuda", generator=g)) + elif kind == "tail": # 64 near-equal top tokens: lists overflow + hot = torch.randint(0, VS, (rows, 64), device="cuda", generator=g) + lg.scatter_(1, hot, 9.0 + 0.5 * torch.rand(rows, 64, device="cuda", generator=g)) + out.append(lg.to(torch.bfloat16).contiguous()) + return out + + +def _draft_from(lg, g): + nreq = lg.shape[0] // (D + 1) + am = lg.float().argmax(1).view(nreq, D + 1)[:, :D].int() + rnd = torch.randint(0, VS, (nreq, D), device="cuda", generator=g, dtype=torch.int32) + return torch.where(torch.rand(nreq, D, device="cuda", generator=g) < 0.7, am, rnd).contiguous() + + +def _compare(ref, got, what): + for n, x, y in zip( + ("tokens", "acc_tokens", "num_acc", "gather_ids", "extra_attempts"), ref, got + ): + assert torch.equal(x, y), f"{what}: {n} differs" + + +CASES = [(1.0, 0.95, "lm"), (0.7, 0.8, "lm"), (1.3, 0.5, "lm"), (1.0, 1.0, "lm"), (1.0, 0.95, "flat"), + (1.0, 0.95, "tail"), (0.6, 0.99, "tail")] # fmt: skip + + +def test_one_round_is_two_round_bitwise(): + g = torch.Generator(device="cuda").manual_seed(2) + total = 0 + decided = {k: 0 for k in KX} + for rows, calls in ((32, 20), (4, 60)): + for ci, (t, p, kind) in enumerate(CASES): + runners = {k: Runner(k, rows) for k in (0, *KX)} + temps, tps, seeds, offs = _params(rows, t, p, 1000 + ci) + sets = _sets(4, rows, g, kind) + bidx = torch.arange(rows // (D + 1), dtype=torch.int32, device="cuda") + for i in range(calls): + lg = sets[i % 4] + draft = _draft_from(lg, g) + offs.fill_(i) + for r in runners.values(): + r(lg, temps, tps, seeds, offs, draft, bidx) + for k in KX: + _compare( + runners[0].out(), runners[k].out(), f"K={k} rows={rows} case={ci} call={i}" + ) + total += rows + if kind == "lm" and p < 1.0: + for k in KX: + decided[k] += int(runners[k].decided.item()) + assert int(runners[0].decided.item()) == 0 + assert total >= 5000 + for k in KX: + assert decided[k] > 0, ( + f"K={k}: no row decided in round 0 on the AgentX-like rows (SA-1R not engaged)" + ) + + +def _boundary_rows(rows, n_heavy_same_slice): + """Heavy tokens heavier than a 50-token candidate group; top_p swept across ulps of S / Z (margin fallback), or + n heavy tokens in one slice (more than its list: incomplete fallback).""" + lg = torch.full((rows, VS), -40.0, device="cuda") + sl = ((VS + 7) // 8 + 7) // 8 * 8 + heavy = list(range(5, 5 + n_heavy_same_slice)) if n_heavy_same_slice else [5, 3 * sl + 11] + grp = [2 * sl + 3 + 37 * j for j in range(25)] + [5 * sl + 7 + 41 * j for j in range(25)] + lg[:, heavy] = 5.0 + lg[:, grp] = 4.0 + lg = lg.to(torch.bfloat16).contiguous() + w = torch.exp(lg[0].double() - lg[0].double().max()) + crit = float(w[heavy].sum()) / float(w.sum()) + ks = torch.arange(rows, dtype=torch.float64) - rows // 2 + return lg, (crit * (1.0 + ks * 2.0**-24)).float() + + +@pytest.mark.parametrize("graph", [False, True]) +def test_one_round_boundary_and_incomplete(graph): + rows = 32 + g = torch.Generator(device="cuda").manual_seed(3) + seen = {k: [0, 0] for k in KX} # decided rows, rows + for n_same in (0, 12): + lg, tps_rows = _boundary_rows(rows, n_same) + for shift in range(-4, 4): + tps = (tps_rows.double() * (1.0 + shift * 32 * 2.0**-24)).float() + runners = {k: Runner(k, rows) for k in (0, *KX)} + temps, tp_t, seeds, offs = _params(rows, 1.0, tps, 5000 + shift) + bidx = torch.arange(rows // (D + 1), dtype=torch.int32, device="cuda") + draft = _draft_from(lg, g) + graphs = {} + if graph: + for r in runners.values(): + r(lg, temps, tp_t, seeds, offs, draft, bidx) + torch.cuda.synchronize() + for k, r in runners.items(): + graphs[k] = torch.cuda.CUDAGraph() + with torch.cuda.graph(graphs[k]): + r(lg, temps, tp_t, seeds, offs, draft, bidx) + for r in runners.values(): + r.decided.zero_() + for i in range(12): + offs.fill_(i) + for k, r in runners.items(): + if graph: + graphs[k].replay() + else: + r(lg, temps, tp_t, seeds, offs, draft, bidx) + for k in KX: + _compare(runners[0].out(), runners[k].out(), + f"boundary n_same={n_same} shift={shift} call={i} K={k}") # fmt: skip + for k in KX: + seen[k][0] += int(runners[k].decided.item()) + seen[k][1] += 12 * rows + for k in KX: + assert 0 < seen[k][0] < seen[k][1], ( + f"K={k}: decided {seen[k]}: the boundary sweep missed both sides" + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py new file mode 100644 index 000000000000..ca795ef4f3ca --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_shard_accept entry (greedy verify from this rank's logits shard).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + SHARD_ROWS, + exchange_tokens, + local_argmax, + make_keys, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 + m3_shard_accept as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import ( # noqa: E402 + craft_drafts, + gen, + head_workspace, + num_sms, + plant_peers, + randn, + seed_of, +) +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" + + +def _case(spec, seed, edit_logits=None): + g = gen(seed) + nreq, d, vs, world, rank = ( + spec["nreq"], + spec["draft_len"], + spec["vs"], + spec["world"], + spec["rank"], + ) + m = nreq * (d + 1) + if spec["strided"]: + logits = randn((m, vs + 72), g, 3.0)[:, :vs] + else: + logits = randn((m, vs), g, 3.0) + if edit_logits is not None: + edit_logits(logits) + uc, keys, counters = head_workspace(world, slot_rows=SHARD_ROWS) + own = logits.float().max(dim=1).values + plant_peers( + uc, counters, own, rank, world, vs, g, slot_rows=SHARD_ROWS, ties=spec.get("ties", False) + ) + offset = rank * vs + best, idx = local_argmax(logits.float(), offset) + tokens = exchange_tokens(make_keys(best, idx), uc, counters, rank, world, SHARD_ROWS) + draft = craft_drafts(tokens, nreq, d) + batch_idx = torch.arange(nreq, dtype=torch.int32, device="cuda") + acc = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + num = torch.full((nreq,), -7, dtype=torch.int32, device="cuda") + gather = torch.full((nreq,), -7, dtype=torch.int64, device="cuda") + next_new = torch.full((nreq, d + 1), -7, dtype=torch.int32, device="cuda") + force_total = 3 if spec["force"] == "total" else 0 + force_vec = (torch.randint(1, d + 2, (nreq,), generator=g, device="cuda", dtype=torch.int32) + if spec["force"] == "vec" else None) # fmt: skip + return (logits, offset, keys, counters, uc, uc.data_ptr(), rank, world, draft, batch_idx, acc, num, gather, + force_total, next_new, num_sms(), force_vec) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Bit-exact: the op reduces BF16 logits without arithmetic, ties (within the shard and across ranks) go to + the lowest global index on both sides.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) + assert args[3].tolist() == [0, 1] and not bool(args[2].any()) + + +def test_guard_refuses_nan_logits() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0][1, 7] = float("nan") + with pytest.raises(AssertionError, match="NaN"): + with validating(op): + op(*args) + + +def test_guard_refuses_a_batch_index_past_the_requests() -> None: + args = list(_case(op.CELLS[0].spec, 2)) + args[9] = args[9] + 1 + with pytest.raises(AssertionError, match="batch_idx"): + with validating(op): + op(*args) + + +def _signed_zero_maxima(logits: torch.Tensor) -> None: + """Every row's maximum is zero, reached as -0.0 at a lower index and +0.0 at a higher one.""" + logits.fill_(-1.0) + logits[:, 3] = -0.0 + logits[:, 9] = 0.0 + + +@pytest.mark.parametrize("ties", [False, True], ids=["peer_rows", "peer_ties"]) +def test_signed_zero_maxima_take_the_lower_index(ties) -> None: + """-0.0 and +0.0 compare equal, so the lower index wins as with torch.argmax (also across ranks).""" + spec = dict(op.CELLS[1].spec, ties=ties) + args = _case(spec, seed_of(f"signed-zero-{ties}"), edit_logits=_signed_zero_maxima) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + assert torch.equal(args[10], ref["accepted"]) + assert torch.equal(args[11], ref["num_accepted"]) + assert torch.equal(args[12], ref["gather_ids"]) + assert torch.equal(args[14], ref["next_new"]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py new file mode 100644 index 000000000000..fab3b7fd7853 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_slot_scatter entry (the spec sampler's per-slot stores in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_slot_scatter import ( # noqa: E501 + m3_slot_scatter as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" + +STALE = -4 + + +def _case(spec, seed): + """Worker outputs and the store's transposed [slots, width] views of its [width, slots, 1] buffers.""" + g = gen(seed) + n, nslots = spec["n"], spec["slots"] + + def rnd(shape): + return torch.randint(0, VOCAB, shape, generator=g, device="cuda", dtype=torch.int32) + + store_new = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_next = torch.full((spec["store_w"], nslots, 1), STALE, dtype=torch.int32, device="cuda") + store_lens = torch.full((nslots,), STALE, dtype=torch.int32, device="cuda") + store_draft = torch.full((nslots, spec["store_d"]), STALE, dtype=torch.int32, device="cuda") + slots = torch.randperm(nslots, generator=g, device="cuda")[:n] + lens = torch.randint(1, 5, (n,), generator=g, device="cuda", dtype=torch.int32) + return (slots, rnd((n, spec["src_w"])), store_new.squeeze(-1).T, rnd((n, spec["src_w"])), + store_next.squeeze(-1).T, lens, store_lens, rnd((n, spec["src_d"])), store_draft) # fmt: skip + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, zero padding included, other slots untouched.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in (("da", args[2]), ("db", args[4]), ("dl", args[6]), ("dd", args[8])): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_repeated_slots() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[0] = args[0].clone() + args[0][1] = args[0][0] + with pytest.raises(AssertionError, match="distinct"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py new file mode 100644 index 000000000000..c4ed9c08c2c3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for m3_spec_offsets (RuntimeHooks.spec_offsets_op): the overlap scheduler's two speculative step corrections +in one plain launch, bit for bit the in-place ATen adds, eager and in a CUDA graph.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_spec_offsets import ( # noqa: E501 + m3_spec_offsets as op, +) + +assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" + +I32 = torch.iinfo(torch.int32) + + +def _inputs(tokens, seqs, seed, wrap=False): + g = torch.Generator(device="cuda").manual_seed(seed) + lo, hi = (I32.min, I32.max) if wrap else (0, 1 << 20) + pos = torch.randint(lo, hi, (tokens,), generator=g, device="cuda", dtype=torch.int32) + kv = torch.randint(lo, hi, (seqs,), generator=g, device="cuda", dtype=torch.int32) + pos_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (tokens,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + kv_off = torch.randint( + -8 if not wrap else lo, + 8 if not wrap else hi, + (seqs,), + generator=g, + device="cuda", + dtype=torch.int32, + ) + if wrap: + pos[:2], pos_off[:2] = I32.max, 1 # overflow wraps, as the ATen add does + kv[:1], kv_off[:1] = I32.min, -1 + return pos, pos_off, kv, kv_off + + +CELLS = [c.spec for c in op.CELLS] + [ + dict(tokens=1, seqs=1), + dict(tokens=300, seqs=75), + dict(tokens=4, seqs=0), +] + + +@pytest.mark.parametrize("spec", CELLS, ids=lambda s: "-".join(f"{k}{v}" for k, v in s.items())) +def test_bitwise_aten_adds(spec) -> None: + for trial in range(3): + pos, pos_off, kv, kv_off = _inputs( + spec["tokens"], spec["seqs"], 97 * trial + spec["tokens"], spec.get("wrap", False) + ) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + assert op(pos, pos_off, kv, kv_off) is True + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"trial {trial}" + + +def test_graph_replays_read_the_live_offsets() -> None: + pos, pos_off, kv, kv_off = _inputs(32, 8, 5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + for step in range(3): + fresh = _inputs(32, 8, 11 + step) + for dst, src in zip((pos, pos_off, kv, kv_off), fresh): + dst.copy_(src) + ref_pos, ref_kv = pos.clone(), kv.clone() + op.reference(ref_pos, pos_off, ref_kv, kv_off) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(pos, ref_pos) and torch.equal(kv, ref_kv), f"step {step}" + + +def test_runner_views_are_accepted() -> None: + """The runner passes position_ids[0, ctx:] and kv_lens_cuda[ctx_reqs:seqs]: 1-D slices of contiguous buffers.""" + position_ids = torch.arange(40, dtype=torch.int32, device="cuda").view(1, 40) + kv_lens = torch.arange(12, dtype=torch.int32, device="cuda") + pos_off = torch.full((40,), 2, dtype=torch.int32, device="cuda") + kv_off = torch.full((12,), 3, dtype=torch.int32, device="cuda") + assert op(position_ids[0, 8:], pos_off[:32], kv_lens[2:10], kv_off[:8]) is True + torch.cuda.synchronize() + assert torch.equal( + position_ids[0], torch.cat([torch.arange(8), torch.arange(8, 40) + 2]).int().cuda() + ) + assert torch.equal( + kv_lens, + torch.cat([torch.arange(2), torch.arange(2, 10) + 3, torch.arange(10, 12)]).int().cuda(), + ) + + +@pytest.mark.parametrize( + "bad", + ["int64", "strided", "length", "2d"], +) +def test_declines_other_layouts_without_writing(bad) -> None: + pos, pos_off, kv, kv_off = _inputs(8, 2, 3) + if bad == "int64": + kv = kv.long() + elif bad == "strided": + pos = torch.zeros(16, dtype=torch.int32, device="cuda")[::2] + elif bad == "length": + pos_off = pos_off[:7] + else: + pos = pos.view(2, 4) + before = (pos.clone(), kv.clone()) + assert op(pos, pos_off, kv, kv_off) is False + torch.cuda.synchronize() + assert torch.equal(pos, before[0]) and torch.equal(kv, before[1]) + + +def test_launches_plainly() -> None: + """It writes step inputs that pre-wait readers take before their grid wait (MSA contract, dispatch + PRODUCER_PREWAIT_STEP_INPUTS): its kernel node carries no programmatic (PDL) serialization.""" + from cuda.bindings import runtime as rt + + pos, pos_off, kv, kv_off = _inputs(32, 8, 7) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + op(pos, pos_off, kv, kv_off) + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + err, _, n = rt.cudaGraphGetNodes(raw, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + flags = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind == rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + flags.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + assert flags == [0], flags diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py new file mode 100644 index 000000000000..1c23a9357876 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_gather entry (the overlap scheduler's step-input gathers in one launch).""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_step_gather import ( # noqa: E501 + m3_step_gather as op, +) + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" + +#: Rows of the engine's flat buffers (the max batch), and their stale fill. +MAX_ROWS = 8 +STALE = -4 + + +def _case(spec, seed): + g = gen(seed) + p, t, d, nslots, b0 = spec["p"], spec["t"], spec["d"], spec["slots"], spec["b0"] + i32 = dict(dtype=torch.int32, device="cuda") + new_tok = torch.randint(0, VOCAB, (t, nslots, 1), generator=g, device="cuda", dtype=torch.int32) + if spec.get("two_d"): + new_tok = new_tok.squeeze(-1) + next_draft = torch.randint(0, VOCAB, (nslots, d), generator=g, device="cuda", dtype=torch.int32) + lens = torch.randint(1, t + 1, (nslots,), generator=g, device="cuda", dtype=torch.int32) + slots = torch.randperm(nslots, generator=g, device="cuda")[:p].to(torch.int32) + ppi = slots.repeat_interleave(t) + ids = torch.full((MAX_ROWS * t,), STALE, **i32) + draft = torch.full((MAX_ROWS * d,), STALE, **i32) + pos_off = torch.full((MAX_ROWS * t,), STALE, **i32) + kv_off = torch.full((MAX_ROWS,), STALE, **i32) + return ( + new_tok, + next_draft, + lens, + slots, + ppi, + p, + t, + d, + ids, + spec["ids0"], + draft, + spec["d0"], + pos_off, + kv_off, + b0, + ) + + +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" + args = _case(cell.spec, seed_of(cell.why)) + certifying(op) + ref = op.reference(*args) + op(*args) + torch.cuda.synchronize() + for name, got in ( + ("ids", args[8]), + ("draft", args[10]), + ("pos_off", args[12]), + ("kv_off", args[13]), + ): + assert torch.equal(got, ref[name]), name + + +def test_guard_refuses_a_slot_past_the_store() -> None: + args = list(_case(op.CELLS[0].spec, 1)) + args[3] = args[3].clone() + args[3][0] = 8 + with pytest.raises(AssertionError, match="slots must index the store"): + with validating(op): + op(*args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py new file mode 100644 index 000000000000..a4076e46f9d1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py @@ -0,0 +1,228 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GPU test for the m3_step_prologue entry (the decode graph's step-input prologue in one launch).""" + +import random + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_step_prologue import ( # noqa: E501 + m3_step_prologue as op, +) +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +__extra_import_path__ = ["../../../..", "."] +from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above +from _validating import certifying, validating # noqa: E402, F401 + +assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" +# The reference and the pack tests run this branch's step_inputs (the step epoch, prologue_op): a run that mounts only +# the target's files gets the image's older module, whose graph scatter has no epoch argument. +assert "epoch" in step_inputs._graph_scatter_kernel.arg_names, ( + "tensorrt_llm/_torch/pyexecutor/step_inputs.py is not this branch's (no step epoch): mount the shared pyexecutor " + "files (step_inputs.py, runtime_hooks.py, engine/runners/decoder/runner.py, h2d_cache.py) with the target" +) + +MAX_DESC = StepInputPack.MAX_DESC +WINDOW = StepInputPack.GRAPH_WINDOW +DATA_START = StepInputPack.HEADER_BYTES + StepInputPack.DESC_BYTES * MAX_DESC +ARENA = 1 << 20 + + +def _layout(spec, seed): + """Copies as (dst, src, bytes) arena offsets (a host copy's src: a window offset) and the host copies' bytes.""" + rng = random.Random(seed) + h2d, d2d, data = [], [], [] + used = DATA_START + pos = 0 + for _ in range(spec["h2d"]): + n = rng.choice([1, 3, 4, 7, 16, 24, 32, 48, 64, 5, 12]) + pos += rng.randint(0, 40) # gaps and every alignment + h2d.append((pos, used, n)) + data.append(bytes(rng.getrandbits(8) for _ in range(n))) + pos += n + used += -(-n // 16) * 16 + assert pos < ARENA // 2 and used <= WINDOW + # Device copies write the last quarter; unchained ones read the third. + dpos, spos = 3 * ARENA // 4, ARENA // 2 + for i in range(spec["d2d"]): + if spec.get("chained") and h2d: + dst_h, _, n_h = h2d[i % len(h2d)] + off = rng.randint(0, n_h - 1) + src, n = dst_h + off, n_h - off + else: + n = rng.choice([4, 8, 16, 32, 6, 64]) + spos += rng.randint(0, 24) + src = spos + spos += n + dpos += rng.randint(0, 24) + d2d.append((dpos, src, n)) + dpos += n + assert dpos < ARENA and spos < 3 * ARENA // 4 + return h2d, d2d, data + + +def _window(layout, base: int) -> torch.Tensor: + """The window of ``layout`` for an arena at device address ``base``.""" + h2d, d2d, data = layout + w = torch.zeros(WINDOW, dtype=torch.uint8) + words = [len(h2d), len(d2d)] + for dst, src, n in h2d: + words += [base + dst, src, n] + for dst, src, n in d2d: + words += [base + dst, base + src, n] + w[: 8 * len(words)] = torch.tensor(words, dtype=torch.int64).view(torch.uint8) + for (_, off, n), b in zip(h2d, data): + w[off : off + n] = torch.frombuffer(bytearray(b), dtype=torch.uint8) + return w.cuda() + + +def _arenas(seed): + a = torch.randint(0, 256, (ARENA,), generator=gen(seed), device="cuda", dtype=torch.int32) + a = a.to(torch.uint8) + return a, a.clone() + + +@pytest.mark.parametrize("with_epoch", [True, False], ids=["epoch", "no_epoch"]) +@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +def test_certified_cells(cell, with_epoch, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name + """Byte for byte the default prologue's destinations, header clear and epoch.""" + seed = seed_of(cell.why) + layout = _layout(cell.spec, seed) + arena, ref_arena = _arenas(seed) + window, ref_window = _window(layout, arena.data_ptr()), _window(layout, ref_arena.data_ptr()) + epoch = torch.full((1,), 41, dtype=torch.int64, device="cuda") if with_epoch else None + ref_epoch = epoch.clone() if with_epoch else None + certifying(op) + op(window, epoch, MAX_DESC) + op.reference(ref_window, ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena) + # Each window holds its own arena's addresses: past the consumed header, both are left as staged. + for w, a in ((window, arena), (ref_window, ref_arena)): + assert not w[:16].any(), "the header is consumed" + assert torch.equal(w[16:], _window(layout, a.data_ptr())[16:]) + if with_epoch: + assert int(epoch.item()) == int(ref_epoch.item()) == 42 + + +def test_graph_replays_scatter_each_new_window() -> None: + """Captured once on a static window: every replay writes that replay's window and bumps the epoch once.""" + arena, ref_arena = _arenas(7) + window = torch.zeros(WINDOW, dtype=torch.uint8, device="cuda") + epoch = torch.zeros(1, dtype=torch.int64, device="cuda") + ref_epoch = epoch.clone() + op(window, epoch, MAX_DESC) # the empty window: writes nothing + op.reference(window.clone(), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + op(window, epoch, MAX_DESC) + for step, spec in enumerate( + [dict(h2d=24, d2d=8, chained=True), dict(h2d=0, d2d=0), dict(h2d=40, d2d=4)] + ): + layout = _layout(spec, 100 + step) + window.copy_(_window(layout, arena.data_ptr())) + graph.replay() + op.reference(_window(layout, ref_arena.data_ptr()), ref_epoch, MAX_DESC) + torch.cuda.synchronize() + assert torch.equal(arena, ref_arena), step + assert not window[:16].any(), step + assert int(epoch.item()) == int(ref_epoch.item()) == step + 2, step + + +def _pack_case(prologue_op): + pack = StepInputPack( + torch.device("cuda", torch.cuda.current_device()), + parallel_scatter=True, + in_graph=True, + early_copy=True, + step_epoch=True, + prologue_op=prologue_op, + ) + a, b, c = (torch.zeros(16, dtype=torch.int32, device="cuda") for _ in range(3)) + out = torch.zeros(48, dtype=torch.int32, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _pack_step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + + def host(v): + return torch.tensor(v, dtype=torch.int32).pin_memory() + + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, host(values[0])) + step_inputs.copy_(b, host(values[1])) + if overlap: + step_inputs.copy_(a[:4], host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if not flush: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_pack_with_the_one_launch_prologue_matches_the_default() -> None: + """The same pack steps through the default prologue and the one-launch one: the same outputs and epochs, and the + one-launch prologue is engaged (the default's warm-up compile is not what bumped its epoch).""" + packs = [_pack_case(None), _pack_case(op)] + assert packs[0][0].prologue_op is None and packs[1][0].prologue_op is op + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + outs = [] + for pack, graph, dsts, out in packs: + # A replay with a new pack, one with an eagerly issued pack, and one with nothing to scatter. + _pack_step(pack, graph, dsts, (va, vb), overlap=step == 2, flush=step != 3) + torch.cuda.synchronize() + outs.append((out.tolist(), int(pack.epoch.item()))) + assert outs[0] == outs[1], step + assert outs[1][1] == step + 2, step # one bump at warm-up, then one per replay + + +def _kernel_nodes_launch_plainly(capture) -> list: + """The PDL attributes of every kernel node of a graph ``capture(stream)`` records: the step-input writers must be + plain launches (MSA pre-wait contract: no programmatic serialization or event on them).""" + from cuda.bindings import runtime as rt + + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + capture() + raw = rt.cudaGraph_t(init_value=graph.raw_cuda_graph()) + # cuda-python changed the signature: (graph[, numNodes]) on older bindings, (graph, nodes, numNodes) on newer ones. + try: + err, _, n = rt.cudaGraphGetNodes(raw) + err, nodes, n = rt.cudaGraphGetNodes(raw, n) + except TypeError: + err, _, n = rt.cudaGraphGetNodes(raw, None, 0) + err, nodes, n = rt.cudaGraphGetNodes(raw, None, n) + found = [] + for node in nodes[:n]: + err, kind = rt.cudaGraphNodeGetType(node) + if kind != rt.cudaGraphNodeType.cudaGraphNodeTypeKernel: + continue + err, attr = rt.cudaGraphKernelNodeGetAttribute( + node, rt.cudaLaunchAttributeID.cudaLaunchAttributeProgrammaticStreamSerialization + ) + # Older bindings return no attribute for a kernel launched without one (a plain launch). + found.append(0 if attr is None else int(attr.programmaticStreamSerializationAllowed)) + return found + + +def test_the_prologue_launches_plainly() -> None: + window = torch.zeros(WINDOW, dtype=torch.uint8, device="cuda") + epoch = torch.zeros(1, dtype=torch.int64, device="cuda") + op(window, epoch, MAX_DESC) + torch.cuda.synchronize() + flags = _kernel_nodes_launch_plainly(lambda: op(window, epoch, MAX_DESC)) + assert flags == [0], flags diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py new file mode 100644 index 000000000000..e578e1f8e55b --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_capture_boundary.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The generic layer's feed-forward exit on Eagle3 capture layers, in both targets' bodies (``PrefillTarget._layer``, +``DecodeTarget._generic_layer``). + +Every layer's MoE / MLP output goes unreduced to the fused boundary (all-reduce + residual + the next layer's input +norm, the final norm after the last layer), and a capture layer captures the reduced residual sum that boundary +writes, with no residual of its own. CPU only: the attention, the feed-forward and the boundaries are stand-ins. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) + +LAYERS = 4 + + +class _Spec: + def __init__(self, layers): + self.layers = set(layers) + + def is_layer_capture(self, i): + return i in self.layers + + +@pytest.fixture(params=["prefill", "decode"]) +def layer(request, monkeypatch): + """(core, the target's generic layer bound to a target over ``core``).""" + calls = [] + w = {f"l{i}_{n}": f"{n}{i}" for i in range(LAYERS) for n in ("in_norm", "post_norm")} + w["final_norm"] = "final" + c = SimpleNamespace( + w=w, num_layers=LAYERS, moe_layers=frozenset({2, 3}), rms_eps=1e-6, calls=calls + ) + monkeypatch.setattr(modeling, "_can_mc", lambda x, folded=False: False) + monkeypatch.setattr( + modeling, + "_native_allreduce_norm", + lambda core, x, r, norm: (("normed", x, norm), ("updated", x, r)), + ) + + def boundary(self, x, x2, residual, norm_w): + calls.append((x, x2, norm_w)) + return ("next", norm_w), ("sum", x, x2) + + if request.param == "prefill": + cls = modeling.PrefillTarget + monkeypatch.setattr(cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i))) + monkeypatch.setattr(cls, "_swiglu_mlp", lambda self, h, prefix, *a, **k: ("mlp", prefix)) + else: + cls = modeling.DecodeTarget + monkeypatch.setattr( + cls, "_moe", lambda self, i, h, **k: (("moe", i), ("routed", i), False, False) + ) + monkeypatch.setattr( + cls, "_swiglu_mlp", lambda self, h, prefix, *a, **k: (("mlp", prefix), False) + ) + monkeypatch.setattr(cls, "_attention", lambda self, i, *a, **k: (("attn", i), False)) + monkeypatch.setattr(cls, "_generic_boundary", boundary) + target = cls(c) + return c, target._layer if request.param == "prefill" else target._generic_layer + + +@pytest.mark.parametrize("i", range(LAYERS)) +def test_a_capture_layer_captures_the_boundary_residual_sum(layer, i) -> None: + core, run = layer + captured = [] + out = run( + i, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([i]), + lambda spec, layer, h, r: captured.append((layer, h, r)), + ) + norm = core.w[f"l{i + 1}_in_norm"] if i + 1 < LAYERS else "final" + assert core.calls[-1][2] == norm, ( + "the boundary applies the next layer's norm (the final norm after the last)" + ) + assert captured == [(i, out[1], None)] + if i in core.moe_layers: + assert core.calls[-1][1] == ("routed", i), "the routed part folds into the boundary" + + +def test_a_layer_without_capture_captures_nothing(layer) -> None: + _, run = layer + captured = [] + run( + 1, + None, + torch.ones(2), + torch.zeros(2), + None, + _Spec([]), + lambda *a: captured.append(a), + ) + assert not captured diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py new file mode 100644 index 000000000000..fe4a08348fb1 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_dispatch.py @@ -0,0 +1,479 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The tp8 (Rubin, two trays) decode target's fast-path table (``modeling.py``: the row ranges and the selection +functions): its row ranges against the catalog cells that certify them, and the selection for every decode row count. + +Every range's covered shapes are listed in ``ENABLED``; a range not listed covers nothing. + +CPU only: nothing here launches a kernel. +""" + +import inspect +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling as dispatch, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_draft_step as draft_step_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + triton_sparse_decode as tsd, +) + +ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 2) # every decode row count and the first one past it + + +@pytest.fixture(autouse=True) +def _policy(): + """Every test starts outside require mode with nothing logged, and leaves it so.""" + dispatch.set_require(False) + dispatch.reset_logged() + yield + dispatch.set_require(False) + dispatch.reset_logged() + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_every_covered_shape_is_a_certified_cell(fast_path) -> None: + uncertified = fast_path.keys - fast_path.certified() + assert not uncertified, ( + f"{fast_path.path} covers {sorted(uncertified, key=str)} but " + f"{type(fast_path.entry).__name__}.CELLS certifies none of them: add a cell or narrow the range" + ) + + +def test_the_fast_path_names_are_unique() -> None: + names = [r.path for r in dispatch.FAST_PATHS] + assert len(names) == len(set(names)) + + +def test_the_decode_target_takes_the_gen_workers_batch() -> None: + """Two requests verifying 1 + 3 drafts; a wider step is the prefill target's.""" + assert dispatch.DECODE_MAX_TOKENS == 8 + assert [dispatch.is_decode_step(n) for n in ROWS] == [True] * 8 + [False] + + +# ------------------------------------------------------------------ the cells turned on + + +def test_the_moe_engine_takes_every_decode_step() -> None: + """The MoE lane's m3_moe engine (TP8 width 384, shared expert folded in) at 1..8 rows, unless killed; its sum + pushes into the multicast boundary all-reduce at the engine's rows while that all-reduce is on.""" + expected = set() if dispatch.KILL_MOE_ENGINE else set(range(1, 9)) + assert set(dispatch.MOE_ENGINE.keys) == expected + for rows in ROWS: + want = dispatch.MOE_ENGINE.path if rows in expected else None + assert dispatch.moe_path(rows) == want, rows + push_rows = dispatch.MOE_ENGINE_MAX_TOKENS if dispatch.MC_AR_TP2 else 0 + assert dispatch.MOE_PUSH_MAX_TOKENS == push_rows + assert set(dispatch.MOE_PUSH.keys) == set(range(1, push_rows + 1)) + + +# ------------------------------------------------------------------ on: the sparse layers' block selection + + +def test_the_index_scorer_takes_one_and_two_requests_with_their_schedules() -> None: + assert dispatch.INDEX_SCORE.keys == frozenset({1, 2}) + assert dispatch.index_score_schedule(1) == (1272, True) + assert dispatch.index_score_schedule(2) == (768, True) + # every covered request count has a split (the graph captures it) + assert set(dispatch.INDEX_SCORE.keys) <= set(dispatch.INDEX_SCORE_SPLIT_K) + + +def test_the_long_select_takes_every_key_width_up_to_its_maximum() -> None: + assert dispatch.LONG_SELECT.keys == frozenset(dispatch.LONG_SELECT_KEY_WIDTHS) + assert (dispatch.LONG_SELECT_MIN_BLOCKS, dispatch.LONG_SELECT_MAX_BLOCKS) == (129, 9216) + + +def test_the_384_thread_select_takes_every_key_width_with_its_launch_key() -> None: + assert dispatch.SELECT_LONG_DSL is True + assert dispatch.LONG_SELECT_DSL.keys == frozenset(dispatch.LONG_SELECT_DSL_KEY_WIDTHS) + keys = dispatch.long_select_dsl_keys() + assert [key[1:] for key in keys] == [(False, dispatch.SELECT_LONG_TRIGGER_PHASE)] + + +@pytest.mark.parametrize("num_blocks", [129, 1294, 5888, 9216]) +def test_a_long_select_decline_inside_its_range_is_a_miss(num_blocks) -> None: + """The kernel declining operands its cells certify must not fall back silently under require.""" + assert dispatch.long_select_covers(num_blocks) + # Outside require: logged once, the in-tree selector runs. + dispatch.long_select_fallback(num_blocks, 8) + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="long_select"): + dispatch.long_select_fallback(num_blocks, 8) + + +@pytest.mark.parametrize("num_blocks", [1, 16, 128, 9217, 20000]) +def test_rows_outside_the_long_select_range_are_uncovered_even_under_require(num_blocks) -> None: + """Short rows and rows past LONG_SELECT_MAX_BLOCKS take the in-tree selector in every mode (long contexts must not + fail under require).""" + assert not dispatch.long_select_covers(num_blocks) + dispatch.set_require(True) + dispatch.long_select_fallback(num_blocks, 8) + + +def test_taken_logs_without_raising_under_require() -> None: + dispatch.set_require(True) + dispatch.taken(dispatch.INDEX_SCORE.path, 1) + dispatch.taken(dispatch.LONG_SELECT.path, 8) + + +def test_an_index_score_cell_of_another_schedule_certifies_nothing() -> None: + spec = dict(seq_lens=(165521,), qlen=4, heads=1, split_k=512, early=True) + assert tuple(dispatch.INDEX_SCORE.certifies(spec)) == () + assert tuple(dispatch.INDEX_SCORE.certifies(dict(spec, split_k=1272, heads=2))) == () + assert tuple(dispatch.INDEX_SCORE.certifies(dict(spec, split_k=1272))) == (1,) + + +#: Every TP8 per-rank MXFP8 projection (N, K): sparse-layer QKV + index, dense QKV, o_proj, shared gate_up / down, +#: dense gate_up / down. +TP8_PROJECTIONS = ( + (1536, 6144), + (1280, 6144), + (6144, 1024), + (768, 6144), + (6144, 384), + (3072, 6144), + (6144, 1536), +) +_DECODE_ROWS = range(1, dispatch.DECODE_MAX_TOKENS + 1) + +#: Fast paths a lane has turned on, each with the shapes it covers (the rest of the skeleton stays off). +ENABLED = { + # SpecTail: sampled (temperature + top-p) verify of 1-2 requests at draft length 3 (the shard kernel). + "m3_sampled_accept": frozenset({(3, 1), (3, 2)}), + # SpecTail: greedy verify (shard argmax + accept) and the fused draft loop of 1-2 requests at draft length 3, + # and that loop's draft-input rows (1 / 2 later-pass rows, 4 / 8 first-pass rows). + "m3_shard_accept": frozenset({(3, 1), (3, 2)}), + "m3_draft_step": frozenset({(3, 1), (3, 2)}), + "m3_draft_input": frozenset({1, 2, 4, 8}), + # SpecTail: the drafter's BF16 linears on FlashInfer's low-M kernels (TP8 port) at the certified b1 rows. + "m3_draft_lowm": frozenset( + { + ((1280, 12288), 1), + ((1280, 12288), 4), + ((6144, 1024), 1), + ((6144, 1024), 4), + ((4608, 6144), 1), + ((4608, 6144), 4), + ((6144, 2304), 1), + ((6144, 2304), 4), + ((6144, 18432), 1), + } + ), + # SpecTail: the fused first drafter inputs of 1-2 fused-verify requests, the capture norm and the chunk norm at + # their 4 / 8 verify rows (own ranges and kills since the TP2 fast-path audit). + "m3_fused_first_inputs": frozenset({1, 2}), + "m3_capture_norm": frozenset({4, 8}), + "m3_chunk_norm": frozenset({4, 8}), + # MSA: the one-launch long-context sparse decode at the TP8 heads (pre-wait up to its rows, plain above), the + # identity-selection producer, and the sparse layers' NVFP4 QKV + indexer producer at the TP8 heads. + "m3_sparse_decode_lc8": frozenset({"plain", "prewait"}), + "m3_producer": frozenset({"nvfp4_tp8"}), + "sparse_decode_producer": frozenset({"multi_query", "per_row"}), + # MoE: the m3_moe engine (TP8 width 384, shared expert folded in) at 1..8 rows, unless killed. + "m3_moe_engine": frozenset() if dispatch.KILL_MOE_ENGINE else frozenset(range(1, 9)), + # The TF32 tensor-core router of 9..16-row pure generation steps on the generic layer (above the engine's rows). + "m3_router_tc": frozenset(range(9, 17)), + # Indexer: the index scorer at 1-2 requests, the long-row select at every key width (129..9216 blocks). + "m3_index_score": frozenset({1, 2}), + "long_select": frozenset({2, 4, 9}), + "long_select_dsl": frozenset({4, 11, 24}), + # LongAttn: the dense layers' full-context decode of 1-2 requests, 1 + 3 verify rows or one draft-pass row. + "m3_fullctx_decode": frozenset({(4, 1), (1, 1), (4, 2), (1, 2)}), + # GEMM/AR: the decode MXFP8 GEMVs at every TP8 projection and decode row, on an MXFP8 and a bf16 activation. + "m3_mx_gemv": frozenset((nk, m) for nk in TP8_PROJECTIONS for m in _DECODE_ROWS) + if dispatch.MX_GEMV_TP2 + else frozenset(), + "m3_mx_gemv_bf16": frozenset((nk, m) for nk in TP8_PROJECTIONS for m in _DECODE_ROWS) + if dispatch.MX_GEMV_TP2 + else frozenset(), + # GEMM/AR: the multicast boundary all-reduce, the o_proj's push into it, the fused embedding + first RMSNorm. + "mc_allreduce_norm": frozenset(_DECODE_ROWS) if dispatch.MC_AR_TP2 else frozenset(), + "attn_push": frozenset(_DECODE_ROWS) + if (dispatch.MC_AR_TP2 and dispatch.ATTN_PUSH_TP2) + else frozenset(), + "dense_push": frozenset(_DECODE_ROWS) + if (dispatch.MC_AR_TP2 and dispatch.DENSE_PUSH_TP2) + else frozenset(), + # MoE: the engine's push into the multicast boundary all-reduce, at the engine's rows. + "moe_push": frozenset(range(1, dispatch.MOE_PUSH_MAX_TOKENS + 1)), + "embed_firstnorm": frozenset(_DECODE_ROWS) if dispatch.EMBED_FIRSTNORM_TP2 else frozenset(), + # SpecTail: the fused draft loop's all-reduce boundaries (1..8 rows) on the multicast all-reduce (opts 16). + "m3_draft_mc_allreduce_norm": frozenset(range(1, dispatch.DRAFT_MC_AR_MAX_ROWS + 1)), + # LongAttn: the Eagle3 drafter layer's generation passes of 1-2 requests on the same recipes. + "m3_drafter_fullctx": frozenset({(4, 1), (1, 1), (4, 2), (1, 2)}), +} + + +@pytest.mark.parametrize("fast_path", dispatch.FAST_PATHS, ids=lambda r: r.path) +def test_skeleton_no_fast_path_covers_a_shape(fast_path) -> None: + expected = ENABLED.get(fast_path.path, frozenset()) + assert fast_path.keys == expected, f"{fast_path.path} covers {sorted(fast_path.keys, key=str)}" + + +def test_the_sparse_decode_writes_the_mxfp8_handoff() -> None: + """The sparse decode (LC8, the Triton merge) writes the MXFP8 o_proj hand-off; the range keys the target's + SPARSE_DECODE_LC8 / SPARSE_DECODE_PRODUCER cover are the kernel module's.""" + assert tsd.MERGE_MXFP8_HANDOFF is True + assert dispatch.SPARSE_DECODE_LC8.keys is tsd.LC8_KEYS + assert dispatch.SPARSE_DECODE_PRODUCER.keys is tsd.PRODUCER_KEYS + + +def test_sparse_decode_lc8_covers_the_decode_steps_with_the_handoff() -> None: + """Every decode step (at most DECODE_MAX_TOKENS rows) writes the MXFP8 o_proj hand-off; the one-launch decode + stages its pre-wait reads up to m3_attn_lc8.PREWAIT_MAX_ROWS rows and merges once.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_attn_lc8, + ) + + assert dispatch.SPARSE_DECODE_LC8.keys == frozenset({"plain", "prewait"}) + assert tsd.MERGE_MXFP8_MIN_TOKENS <= 1 and dispatch.DECODE_MAX_TOKENS <= tsd.MergeMxfp8.MAX_ROWS + assert ( + m3_attn_lc8.PREWAIT_MAX_ROWS == 8 and m3_attn_lc8.MERGE_ONCE and m3_attn_lc8.ONE_DIV_MERGE + ) + + +def test_the_dense_layers_take_the_fullctx_decode_at_the_gen_workers_shapes() -> None: + """One or two requests, 1 + 3 verify rows or one draft-pass row, each at its own pinned recipe.""" + assert dispatch.FULLCTX_DECODE_ENABLED + assert dispatch.FULLCTX_DECODE.keys == {(4, 1), (1, 1), (4, 2), (1, 2)} + assert set(dispatch.FULLCTX_DECODE_RECIPES) == set(dispatch.FULLCTX_DECODE.keys) + for recipe in dispatch.FULLCTX_DECODE_RECIPES.values(): + assert 2 <= recipe["splits"] <= 128 # the standalone reducer's range + + +def test_the_fullctx_entry_cells_run_at_the_forward_plan_capacity() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention import ( # noqa: E501 + m3_fullctx_decode as entry, + ) + + assert entry.MAX_KV_LEN == dispatch.FULLCTX_DECODE_MAX_KV_LEN + + +def test_the_fullctx_capacity_covers_the_agentx_cache_bound() -> None: + """The forward compares the cache's length bound (max_seq_len plus the speculative rows it reserves) with the plan + capacity; at the AgentX max_seq_len of 1 << 20 with up to 8 draft rows the dense layers must still be taken.""" + assert dispatch.FULLCTX_DECODE_MAX_KV_LEN >= (1 << 20) + 8 + + +def test_the_drafter_takes_the_fullctx_decode_at_the_same_shapes() -> None: + assert dispatch.DRAFTER_FULLCTX_ENABLED + assert dispatch.DRAFTER_FULLCTX.keys == dispatch.FULLCTX_DECODE.keys + + +# ------------------------------------------------------------------ landed: the decode MXFP8 GEMVs (GEMM/AR lane) + + +def test_the_gemv_rows_follow_the_kill_constant() -> None: + expected = set(TP8_PROJECTIONS) if dispatch.MX_GEMV_TP2 else set() + assert set(dispatch.MX_GEMV_ROWS) == expected + assert set(dispatch.MX_GEMV_BF16_MAX_TOKENS) == expected + + +@pytest.mark.parametrize("rows", range(1, dispatch.DECODE_MAX_TOKENS + 1)) +@pytest.mark.parametrize("nk", TP8_PROJECTIONS, ids=str) +def test_every_projection_takes_the_gemv_at_every_decode_row(nk, rows) -> None: + """MXFP8 activation (a producer quantized it) -> the GEMV; bf16 -> the GEMV quantizing in its prologue (staged: + the stock decode-graph GEMMs while MX_GEMV_TP2 is off).""" + on = dispatch.MX_GEMV_TP2 + assert dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) == ( + dispatch.MX_GEMV_PATH if on else dispatch.PREQUANT_GEMM_PATH + ) + assert dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) == ( + dispatch.MX_GEMV_BF16_PATH if on else dispatch.GENERIC_GEMM_PATH + ) + + +@pytest.mark.parametrize("nk", TP8_PROJECTIONS, ids=str) +def test_wider_than_the_decode_step_is_the_generic_gemm(nk) -> None: + rows = dispatch.DECODE_MAX_TOKENS + 1 + assert ( + dispatch.linear_path(*nk, rows, prequant=True, low_m_gemm=True) + == dispatch.PREQUANT_GEMM_PATH + ) + assert ( + dispatch.linear_path(*nk, rows, prequant=False, low_m_gemm=True) + == dispatch.GENERIC_GEMM_PATH + ) + + +def test_an_unknown_projection_is_the_generic_gemm() -> None: + assert ( + dispatch.linear_path(256, 512, 1, prequant=True, low_m_gemm=True) + == dispatch.PREQUANT_GEMM_PATH + ) + assert ( + dispatch.linear_path(256, 512, 1, prequant=False, low_m_gemm=True) + == dispatch.GENERIC_GEMM_PATH + ) + + +def test_taken_logs_once_per_path_and_rows(monkeypatch) -> None: + lines = [] + monkeypatch.setattr(dispatch.logger, "info", lines.append) + for _ in range(3): + dispatch.taken(dispatch.MX_GEMV_PATH, 4, "[6144, 4096] mxfp8 in") + dispatch.taken(dispatch.MX_GEMV_PATH, 8, "[6144, 4096] mxfp8 in") + assert len(lines) == 2 and all("m3_mx_gemv taken at" in line for line in lines) + + +def test_skeleton_the_dense_decode_split_is_the_built_in_one() -> None: + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {} + + +def test_skeleton_speculation_runs_the_in_tree_loop() -> None: + # SpecTail: fused greedy / sampled verify and the fused draft loop at b1/b2 (BF16 head; no coded heads). + assert dispatch.SHARD_ACCEPT_MAX_REQUESTS == 2 + assert dispatch.SAMPLED_VERIFY_MAX_REQUESTS == 2 + assert dispatch.DRAFT_STEP_MAX_REQUESTS == 2 + assert dispatch.FC_EARLY_MAX_ROWS == 0 + assert not draft_step_kernel.HEAD_CODED + + +def test_a_miss_raises_only_under_require() -> None: + dispatch.missed("some_path", 1, "a precondition") + dispatch.set_require(True) + with pytest.raises(dispatch.FastPathUnavailable, match="some_path"): + dispatch.missed("some_path", 1, "a precondition") + + +def test_skeleton_first_drafter_inputs_are_in_tree() -> None: + # SpecTail turned the sampled verify on (SAMPLED_VERIFY_MAX_REQUESTS 2, test_skeleton_speculation_...). + assert dispatch.FIRST_DRAFTER_INPUTS_MAX_REQUESTS == 0 + + +@pytest.mark.parametrize( + "rows, capture_normed, taken", + [ + (4, False, True), + (8, False, True), + (8, True, True), + (12, False, False), + (12, True, False), + (16, True, False), + ], +) +def test_the_fused_first_inputs_take_only_their_covered_rows(rows, capture_normed, taken) -> None: + """The fused first drafter inputs (and the fc fork) take a fused-verify step only at the requests + FUSED_FIRST_INPUTS covers and, unless the captures applied the fc_norm, the rows CHUNK_NORM covers; anything else + runs the in-tree first drafter inputs.""" + + worker = object.__new__(dispatch.M3Eagle3Worker) + worker._m3_capture_norm = (None, 1e-6) if capture_normed else None + meta = SimpleNamespace(runtime_draft_len=dispatch.SHARD_ACCEPT_DRAFT_LEN) + assert worker._m3_first_inputs_covered(rows, meta) is taken + + +# ------------------------------------------------------- landed: the multicast boundary all-reduce (GEMM/AR lane) + + +def test_both_boundaries_take_the_multicast_all_reduce() -> None: + """One 8-token bucket for the whole decode step; the MoE boundary publishes the next QKV's MXFP8 rows (killed: + every boundary on the built-in all-reduce while MC_AR_TP2 is off).""" + on = dispatch.MC_AR_TP2 + assert tuple(p.max_tokens for p in dispatch.DECODE_PLANS) == (dispatch.DECODE_MAX_TOKENS,) + for plan in dispatch.DECODE_PLANS: + assert (plan.attn_mc_ar, plan.moe_mc_ar, plan.qkv_prequant) == (on, on, on), plan + assert dispatch.ATTN_PUSH_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if on and dispatch.ATTN_PUSH_TP2 else 0 + ) + assert dispatch.DENSE_PUSH_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if on and dispatch.DENSE_PUSH_TP2 else 0 + ) + for rows in _DECODE_ROWS: + assert dispatch.plan_for(rows) is dispatch.DECODE_PLANS[0] + # The MoE pushes only from its engine, at the engine's rows. + assert dispatch.MOE_PUSH_MAX_TOKENS == (dispatch.MOE_ENGINE_MAX_TOKENS if on else 0) + + +def test_the_generic_boundaries_take_the_multicast_all_reduce() -> None: + """Prefill, mixed steps and the Eagle3 capture layers: the same certified rows, alone and folded.""" + rows = dispatch.DECODE_MAX_TOKENS if dispatch.MC_AR_TP2 else 0 + assert dispatch.GENERIC_MC_AR_MAX_TOKENS == rows + assert dispatch.GENERIC_MC_AR_FOLDED_MAX_TOKENS == rows + + +# ------------------------------------------------------------------ landed: the fused embedding + first RMSNorm + + +def test_the_fused_embedding_rows_follow_the_kill_constant() -> None: + assert dispatch.EMBED_FIRSTNORM_MAX_TOKENS == ( + dispatch.DECODE_MAX_TOKENS if dispatch.EMBED_FIRSTNORM_TP2 else 0 + ) + + +# ------------------------------------------------------- landed: MX-TAC (pinned decode-graph MXFP8 GEMM tactic) + + +def test_the_decode_mxfp8_decision_is_pinned_to_cute_dsl() -> None: + from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + FlashInferMXFP8GemmRunner, + MXFP8QuantizeRunner, + ) + + assert dispatch.MXFP8_DECODE_QUANT_TACTIC == MXFP8QuantizeRunner.CUTE_DSL + assert dispatch.MXFP8_DECODE_GEMM_TACTIC == FlashInferMXFP8GemmRunner.CUTE_DSL + + +@pytest.mark.parametrize("target", ["DecodeTarget", "PrefillTarget"]) +def test_the_moe_forward_reports_the_engine_taken(target) -> None: + """The engagement evidence comes from the forward (not the engine build's warm call).""" + + src = inspect.getsource(getattr(dispatch, target)._moe) + assert src.index("taken(MOE_ENGINE.path") < src.index("m3_moe_engine(h.contiguous()") + + +# ------------------------------------------------------------------ the drafter's low-M linears (TP8 port) + + +def test_draft_lowm_covers_exactly_the_certified_rows_of_its_tactic_table() -> None: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm import ( # noqa: E501 + m3_draft_lowm as entry, + ) + + assert set(entry.TACTICS) == set(entry.ROWS) == set(dispatch.DRAFT_LOWM_ROWS) + assert dispatch.DRAFT_LOWM.keys == frozenset( + (nk, m) + for nk, rows in entry.ROWS.items() + for m in rows + if m <= dispatch.DRAFT_LOWM_MAX_ROWS + ) + for nk, table in entry.TACTICS.items(): + rows = [row[0] for row in table] + assert rows == sorted(rows) and len(set(rows)) == len(rows), nk + for m in entry.ROWS[nk]: + assert entry.tactic_for(*nk, m) is not None, (nk, m) + + +def test_prebuild_draft_head_is_idempotent_and_keeps_replaced_copies(monkeypatch) -> None: + """A repeated cache_derived_state with an unchanged head weight keeps the built (placed) copy; a changed weight + rebuilds it and keeps the old copy alive, so a CUDA graph captured earlier never reads unmapped memory.""" + + import torch + + built = [] + + def fake_prepare(lm_head, tp_rank): + head = SimpleNamespace(weight=lm_head.weight) + built.append(head) + return head + + monkeypatch.setattr(draft_step_kernel, "prepare_head", fake_prepare) + worker = object.__new__(dispatch.M3Eagle3Worker) + worker._m3_draft_head = None + worker._m3_draft_head_version = None + worker._m3_draft_head_graveyard = [] + worker.mapping = SimpleNamespace(tp_rank=0) + monkeypatch.setattr(worker, "_m3_drafter_supported", lambda dm: True, raising=False) + model = SimpleNamespace(lm_head=SimpleNamespace(weight=torch.zeros(4, 4))) + + assert worker.prebuild_draft_head(model) and len(built) == 1 + first = worker._m3_draft_head + assert worker.prebuild_draft_head(model) and len(built) == 1 and worker._m3_draft_head is first + model.lm_head.weight.add_(1.0) # new values: version bump + assert worker.prebuild_draft_head(model) and len(built) == 2 + assert worker._m3_draft_head is built[1] and worker._m3_draft_head_graveyard == [first] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py new file mode 100644 index 000000000000..8f28a67e722f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decode robustness without a GPU kernel: the multicast all-reduce epilogue options keep the cluster-barrier bit, +an engine build that fails on one TP rank drops that engine on every rank (two ranks, and the tp8 target's eight, +simulated with threads), and the Eagle3 fc fork's output is recorded on the stream that reads it.""" + +import threading +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 + modeling, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 + m3_mc_allreduce as mc, +) + + +@pytest.fixture(autouse=True) +def _restore_modes(): + require = modeling.require_mode() + yield + modeling.set_require(require) + modeling.reset_logged() + + +@pytest.mark.parametrize("finalize", [False, True]) +@pytest.mark.parametrize("bucket", [8, 16, 32, 48]) +def test_ar_opts_keep_cluster_barrier_bit(monkeypatch, bucket, finalize): + # Tables without bit 2 and with buckets missing: every lookup, fallback included, keeps the barrier. + monkeypatch.setattr(mc, "AR_OPTS_BY_M", (0, 1, 4, 5, 8, 9, 12, 13)) + monkeypatch.setattr(mc, "AR_OPTS_BY_BUCKET", {}) + monkeypatch.setattr(mc, "AR_FIN_OPTS_BY_BUCKET", {}) + for m in range(1, min(bucket, 8) + 1): + opts = mc._ar_opts(m, bucket, finalize=finalize) + assert opts & mc._OPT_CLUSTER_BARRIER and 0 <= opts <= 15, (m, bucket, finalize, opts) + + +def test_ar_opts_shipped_tables_unchanged(): + for m in range(1, 9): + assert mc._ar_opts(m, 8) == mc.AR_OPTS_BY_M[m - 1] + for b, o in mc.AR_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b) == o + for b, o in mc.AR_FIN_OPTS_BY_BUCKET.items(): + assert mc._ar_opts(b, b, finalize=True) == o + + +class _FakeComm: + """An MPI-like TP communicator for ``world`` threads: allreduce sums the ranks' ints.""" + + def __init__(self, world): + self.world = world + self.barrier = threading.Barrier(world) + self.lock = threading.Lock() + self.vals = [] + self.calls = 0 + + def allreduce(self, v): + with self.lock: + self.vals.append(v) + self.calls += 1 + self.barrier.wait() + total = sum(self.vals) + self.barrier.wait() + with self.lock: + self.vals.clear() + self.barrier.wait() + return total + + def Get_size(self): # noqa: N802 (MPI name) + return self.world + + +def _mapping(tp_size=2, adp=False): + return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) + + +def _run_ranks(fns): + out, errs = [None] * len(fns), [] + + def run(i): + try: + out[i] = fns[i]() + except BaseException as e: # noqa: BLE001 + errs.append(e) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not errs, errs + return out + + +def _patch_comm(monkeypatch, world: int) -> _FakeComm: + from tensorrt_llm._torch.distributed import ops + + comm = _FakeComm(world) + monkeypatch.setattr(mc, "_tp_comm", lambda mapping: comm) + monkeypatch.setattr(ops, "mpi_disabled", lambda: False) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return comm + + +@pytest.fixture +def two_rank_comm(monkeypatch): + return _patch_comm(monkeypatch, 2) + + +def test_tp_agree_single_rank_and_adp_take_no_collective(monkeypatch): + monkeypatch.setattr(mc, "_tp_comm", mock.Mock(side_effect=AssertionError("no collective"))) + assert mc.tp_agree(_mapping(1), True) is True + assert mc.tp_agree(_mapping(1), False) is False + assert mc.tp_agree(_mapping(8, adp=True), False) is False + + +def test_tp_agree_refuses_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError): + mc.tp_agree(_mapping(2), True) + + +@pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) +def test_tp_agree_built_two_ranks(two_rank_comm, oks): + built = [object() if ok else None for ok in oks] + out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + assert out == (built if all(oks) else [None, None]) + + +def _core(fail: bool, world: int = 2) -> SimpleNamespace: + """The engine-building half of the core (``MiniMaxM3Tp8ModelingV2``) on one of ``world`` simulated TP ranks, + its engine builds stubbed (and failing where ``fail``).""" + core = SimpleNamespace(mapping=_mapping(world), moe_engines={}) + cls = modeling.MiniMaxM3Tp8ModelingV2 + + def build(i, num_tokens, shared=True): + if fail: + raise RuntimeError("compile failed on this rank") + return ("engine", num_tokens) + + core._build_engine = build + for name in ("_prebuild_engines", "_try_build", "engine", "_lookup"): + setattr(core, name, getattr(cls, name).__get__(core)) + return core + + +@pytest.mark.parametrize( + "world,failing", + [(2, ()), (2, (1,)), (8, ()), (8, (5,)), (8, (0, 7))], + ids=["w2_ok", "w2_rank1_fails", "w8_ok", "w8_rank5_fails", "w8_ranks0_7_fail"], +) +def test_engine_build_agreed_across_ranks(monkeypatch, world, failing): + comm = _patch_comm(monkeypatch, world) + cores = [_core(rank in failing, world) for rank in range(world)] + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + keys = sorted(cores[0].moe_engines[3]) + assert keys and comm.calls == world * len(keys) + assert not any(hasattr(c, "moe16_engines") for c in cores) + for c in cores[1:]: + assert c.moe_engines[3].keys() == cores[0].moe_engines[3].keys() + for key in keys: + for c in cores: + assert (c.moe_engines[3][key] is None) == bool(failing), key + # Every rank then takes the same path in the forward: the engine, or (after a miss) the generic MoE. + modeling.set_require(False) + for n in sorted(modeling.MOE_ENGINE.keys): + got = [c.engine(3, n, shared=True) is None for c in cores] + assert len(set(got)) == 1, (n, got) + + +@pytest.mark.parametrize("capturing", [False, True]) +def test_fc_join_records_output_on_main(monkeypatch, capturing): + M3Eagle3Worker = modeling.M3Eagle3Worker + main, fork = mock.Mock(name="main"), mock.Mock(name="fork") + hidden = mock.Mock(name="hidden", is_cuda=True) + monkeypatch.setattr(torch.cuda, "current_stream", lambda *a: main) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + worker = SimpleNamespace(_m3_fc_hidden=hidden, _m3_fc_flagged=False) + M3Eagle3Worker._m3_fc_join(worker, fork) + main.wait_stream.assert_called_once_with(fork) + if capturing: + hidden.record_stream.assert_not_called() + else: + hidden.record_stream.assert_called_once_with(main) + worker._m3_fc_hidden = None + M3Eagle3Worker._m3_fc_join(worker, fork) 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..9f228694696c 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -427,6 +427,15 @@ 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/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm": ( + "2 ranks on sm_107; CI has only the single-GPU GR100 stage (auto:gr100-x1)" + ), + "unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm": ( + "4 ranks on sm_107; CI has only the single-GPU GR100 stage (auto:gr100-x1)" + ), + "unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm": ( + "2, 4 or 8 ranks on sm_107; CI has only the single-GPU GR100 stage (auto:gr100-x1)" + ), "unittest/_torch/modeling_v2/models/glm_moe_dsa/glm_5_2_nvfp4__sm_107__tp16/comm": ( "4 ranks on sm_107; CI has only the single-GPU GR100 stage (auto:gr100-x1)" ), diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py index cb9e43bb2c6b..0c592a14603c 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_routing.py @@ -626,11 +626,12 @@ def test_routes_on_the_kv_cache_algorithm_the_engine_runs( # experts split by expert rather than by tensor: another weight layout, no target (dict(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8), None), (dict(world_size=8, tp_size=8, moe_tp_size=8, enable_attention_dp=True), None), - (dict(world_size=2, tp_size=2, moe_tp_size=2), None), + (dict(world_size=2, tp_size=2, moe_tp_size=2, moe_ep_size=1), "tp2"), + (dict(world_size=2, tp_size=2, moe_tp_size=1, moe_ep_size=2), None), ], ) def test_minimax_m3_parallel_names_the_target_topology(mapping_kwargs, parallel): - """MiniMax-M3 has one target per tensor-parallel width: tp4 and tp8 load different per-rank shapes.""" + """MiniMax-M3 has one target per tensor-parallel width: tp2, tp4 and tp8 load different per-rank shapes.""" from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing assert routing._parallel(Mapping(**mapping_kwargs)) == parallel @@ -661,6 +662,64 @@ def _m3_text_config(**layout): return dict(zip(_M3_SHAPE_FIELDS, routing._SHAPE), **(layout or _M3_LAYOUT)) +_M3_TARGETS = "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3" +_M3_TP2 = dict(world_size=2, tp_size=2, moe_tp_size=2, moe_ep_size=1) +_M3_TP4 = dict(world_size=4, tp_size=4, moe_tp_size=4, moe_ep_size=1) +_M3_TP8 = dict(world_size=8, tp_size=8, moe_tp_size=8, moe_ep_size=1) + + +@pytest.mark.parametrize( + "capability, mapping_kwargs, target", + [ + ((10, 3), _M3_TP4, "minimax_m3_nvfp4__sm_103__tp4.ModelingV2MiniMaxM3Nvfp4Sm103Tp4"), + ((10, 3), _M3_TP8, "minimax_m3_nvfp4__sm_103__tp8.ModelingV2MiniMaxM3Nvfp4Sm103Tp8"), + ((10, 7), _M3_TP2, "minimax_m3_nvfp4__sm_107__tp2.ModelingV2MiniMaxM3Nvfp4Sm107Tp2"), + ((10, 7), _M3_TP4, "minimax_m3_nvfp4__sm_107__tp4.ModelingV2MiniMaxM3Nvfp4Sm107Tp4"), + ((10, 7), _M3_TP8, "minimax_m3_nvfp4__sm_107__tp8.ModelingV2MiniMaxM3Nvfp4Sm107Tp8"), + # no GB300 tp2 target + ((10, 3), _M3_TP2, None), + ], +) +def test_minimax_m3_route_picks_the_target_by_architecture_and_width( + monkeypatch, capability, mapping_kwargs, target +): + """Each (SM, tensor-parallel width) has its own target directory. The target modules are stubbed: this checks the + decision, not that the modeling code imports on this host.""" + import sys + import types + + from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + from tensorrt_llm.quantization.mode import QuantAlgo + + for directory, name in ( + ("minimax_m3_nvfp4__sm_103__tp4", "ModelingV2MiniMaxM3Nvfp4Sm103Tp4"), + ("minimax_m3_nvfp4__sm_103__tp8", "ModelingV2MiniMaxM3Nvfp4Sm103Tp8"), + ("minimax_m3_nvfp4__sm_107__tp2", "ModelingV2MiniMaxM3Nvfp4Sm107Tp2"), + ("minimax_m3_nvfp4__sm_107__tp4", "ModelingV2MiniMaxM3Nvfp4Sm107Tp4"), + ("minimax_m3_nvfp4__sm_107__tp8", "ModelingV2MiniMaxM3Nvfp4Sm107Tp8"), + ): + stub = types.ModuleType(f"{_M3_TARGETS}.{directory}.modeling") + setattr(stub, name, type(name, (), {"__module__": stub.__name__})) + monkeypatch.setitem(sys.modules, stub.__name__, stub) + monkeypatch.setattr(routing, "_quant_algo", lambda _config: QuantAlgo.MIXED_PRECISION) + + model_config = types.SimpleNamespace( + pretrained_config=types.SimpleNamespace(text_config=_m3_text_config()), + mapping=Mapping(**mapping_kwargs), + spec_config=None, + ) + llm_args = _m3_llm_args("auto", "nvfp4") + if target is None: + with pytest.raises(NoTarget, match="sm107 \\(Rubin\\) with 2, 4 or 8 ranks"): + routing.route(model_config, llm_args, capability) + return + resolved = routing.route(model_config, llm_args, capability) + directory, name = target.split(".") + assert resolved.__module__ == f"{_M3_TARGETS}.{directory}.modeling" + assert resolved.__name__ == name + + @pytest.mark.parametrize("spec", [None, "eagle3"], ids=["no_speculation", "one_model_eagle3"]) @pytest.mark.parametrize("tp", [4, 8]) def test_minimax_m3_route_takes_the_released_checkpoint(monkeypatch, spec, tp): @@ -736,6 +795,41 @@ def test_minimax_m3_route_declines_another_rope_base_or_sparse_layout(monkeypatc ) +@pytest.mark.parametrize("mapping_kwargs", [_M3_TP2, _M3_TP4, _M3_TP8], ids=["tp2", "tp4", "tp8"]) +@pytest.mark.parametrize( + "layout", + [ + _OTHER_LAYOUT, + dict( + _M3_LAYOUT, + sparse_attention_config=dict( + _M3_LAYOUT["sparse_attention_config"], use_sparse_attention=False + ), + ), + ], + ids=["other_revision", "sparse_disabled"], +) +def test_minimax_m3_rubin_route_declines_another_layout(monkeypatch, mapping_kwargs, layout): + """The sm107 targets build the same RoPE table and sparse layout as the sm103 ones, so they decline the same + checkpoints.""" + from types import SimpleNamespace + + from tensorrt_llm._torch._experimental.modeling_v2._router_index import NoTarget + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing + + monkeypatch.setattr(routing, "_quant_algo", lambda cfg: QuantAlgo.MIXED_PRECISION) + with pytest.raises(NoTarget, match="sparse attention layers other than 3-59|rope_theta"): + routing.route( + SimpleNamespace( + pretrained_config=SimpleNamespace(text_config=_m3_text_config(**layout)), + mapping=Mapping(**mapping_kwargs), + spec_config=None, + ), + _m3_llm_args("auto", "nvfp4"), + routing._SM_RUBIN, + ) + + def test_minimax_m3_layout_reads_rope_theta_from_rope_parameters(): """A config that keeps the RoPE base under ``rope_parameters`` (newer transformers) reads the same.""" from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import routing diff --git a/tests/unittest/_torch/pyexecutor/engine/test_decoder_runtime_hooks.py b/tests/unittest/_torch/pyexecutor/engine/test_decoder_runtime_hooks.py index 3bfbaee8c913..f6ccb731c0cf 100644 --- a/tests/unittest/_torch/pyexecutor/engine/test_decoder_runtime_hooks.py +++ b/tests/unittest/_torch/pyexecutor/engine/test_decoder_runtime_hooks.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """DecoderRunner and DecoderMetadata wiring of the step-loop hooks a model declares.""" +from types import SimpleNamespace + import pytest from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import ( @@ -30,12 +32,16 @@ def _runner( - sparse_attention_config: BaseSparseAttentionConfig | None, hooks: RuntimeHooks + sparse_attention_config: BaseSparseAttentionConfig | None, + hooks: RuntimeHooks, + max_draft_len: int = 0, ) -> DecoderRunner: - # Only the two attributes the threshold reads. + # Only the attributes the threshold reads. runner = object.__new__(DecoderRunner) runner._sparse_attention_config = sparse_attention_config runner.runtime_hooks = hooks + runner._config = SimpleNamespace(max_draft_len=max_draft_len) + runner._get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1 return runner @@ -64,10 +70,14 @@ def test_a_sparse_config_that_asks_for_the_split_sets_the_threshold(): assert _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() == 512 -def test_short_identity_threshold_is_the_topk_span(): +@pytest.mark.parametrize("max_draft_len", [0, 3]) +def test_short_identity_threshold_keeps_one_step_of_headroom_in_the_topk_span(max_draft_len): + """The short graph family runs the identity selection: one generation step's tokens below the top-k span, the most + the graph key's host length estimate can lag the attended length under the overlap scheduler.""" config = MiniMaxM3SparseAttentionConfig(sparse_topk_blocks=16, sparse_block_size=128) - assert _runner(config, _SHORT_IDENTITY)._short_long_seq_len_threshold() == 16 * 128 + threshold = _runner(config, _SHORT_IDENTITY, max_draft_len)._short_long_seq_len_threshold() + assert threshold == 16 * 128 - (max_draft_len + 1) @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py b/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py index c054adb09421..374a712d3325 100644 --- a/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py +++ b/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py @@ -28,6 +28,7 @@ ) from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.spec_sampler_base import SampleStateTensorsSpec from tensorrt_llm.bindings import SamplingConfig from tensorrt_llm.llmapi.llm_args import ( BaseSparseAttentionConfig, @@ -599,3 +600,46 @@ def test_sparse_sequence_mode_uses_promoted_context_cursor(self) -> None: runner._get_seq_len_mode(batch, overlap_state, frozenset({request.py_request_id})) ) self.assertFalse(runner._get_seq_len_mode(batch, overlap_state, frozenset())) + + +def _overlap_spec_batch(max_draft_len: int, *num_tokens: int) -> tuple: + """A decode batch under the overlap scheduler with speculative decoding: each request has a previous batch, and the + in-flight step's sampled tokens and drafts are device tensors the host has not read yet.""" + batch = _generation_batch(*num_tokens) + for request in batch.generation_requests: + request.py_batch_idx = 0 + request.py_is_first_draft = False + tensors = object.__new__(SampleStateTensorsSpec) + tensors.new_tokens = torch.zeros(1, len(num_tokens), 1, dtype=torch.int32) + tensors.next_draft_tokens = torch.zeros(len(num_tokens), max_draft_len, dtype=torch.int32) + return batch, tensors + + +@pytest.mark.parametrize("max_draft_len", [1, 3]) +def test_identity_short_mode_keeps_every_attended_token_within_the_topk_span(max_draft_len) -> None: + """With the identity threshold (top-k span minus one generation step's tokens), a decode step the runner keys short + attends at most top-k blocks for every verify token, whatever the in-flight step accepted. The host knows a + request's tokens up to the step before the in-flight one; the in-flight step accepts a in [1, draft + 1] tokens and + this step adds draft more, so its last verify token attends T + a + draft tokens. The bare top-k span let a + request attend one block past it (bug B: identity steps at 2049 tokens).""" + span = 16 * 128 + tokens_per_step = max_draft_len + 1 + spec_config = SimpleNamespace(max_draft_len=max_draft_len, _use_shared_kv_cache=True) + + def worst_attended(host_tokens: int) -> int: + return host_tokens + tokens_per_step + max_draft_len + + def short(threshold: int, host_tokens: int) -> bool: + runner = _decoder_runner(short_long_seq_len_threshold=threshold) + runner.spec_config = spec_config + batch, tensors = _overlap_spec_batch(max_draft_len, host_tokens) + return runner._get_seq_len_mode(batch, tensors) + + for host_tokens in range(span - 4 * tokens_per_step, span + 1): + if short(span - tokens_per_step, host_tokens): + assert worst_attended(host_tokens) <= span, host_tokens + # The unpadded span keys a request short whose worst case attends past it. + assert any( + short(span, t) and worst_attended(t) > span + for t in range(span - 4 * tokens_per_step, span + 1) + ) diff --git a/tests/unittest/_torch/pyexecutor/test_first_content_emission.py b/tests/unittest/_torch/pyexecutor/test_first_content_emission.py new file mode 100644 index 000000000000..932a5c130b1f --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_first_content_emission.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the streaming emission schedule in ``PyExecutor``. + +``_should_emit_response`` decides on which decoding iterations a streaming +request gets a response. With ``stream_interval = S`` it emits on iteration 1, +on finish and every S iterations; it also emits on iterations +``2..first_content_emit_iters`` (the model's ``RuntimeHooks.first_content_emit_iters``; 1, the +default, keeps the stock schedule; the MiniMax-M3 sm_107 targets set 2), +so text held back by a reasoning parser after the first token reaches the +client one step later instead of one stream interval later. These tests call +the real method unbound against a minimal fake ``self`` and check that: + +* with ``first_content_emit_iters = 1`` (the default) the schedule is the stock + one: iteration 1, every S iterations, and finish; +* ``first_content_emit_iters = 2`` adds iteration 2 only; +* larger values add iterations ``2..N``; +* a finished request always emits, and ``emit_first_iter=False`` (early + first-token emission) still suppresses only a non-terminal iteration 1; +* ``stream_interval = 1`` emits every iteration whatever the setting. +""" + +from __future__ import annotations + +import types + +import pytest + +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks, runtime_hooks_of + +pytestmark = pytest.mark.cpu_only + + +def _fake_self(stream_interval: int, first_content_emit_iters: int): + return types.SimpleNamespace( + stream_interval=stream_interval, first_content_emit_iters=first_content_emit_iters + ) + + +def _request(decoding_iter: int, finished: bool = False): + return types.SimpleNamespace(py_decoding_iter=decoding_iter, is_finished=finished) + + +def _emitted_iters(stream_interval, first_content_emit_iters, last=60, emit_first_iter=True): + fake = _fake_self(stream_interval, first_content_emit_iters) + return [ + it + for it in range(1, last + 1) + if PyExecutor._should_emit_response(fake, _request(it), emit_first_iter) + ] + + +def _stock_iters(stream_interval, last=60): + return [it for it in range(1, last + 1) if it == 1 or it % stream_interval == 0] + + +@pytest.mark.parametrize("stream_interval", [2, 3, 20]) +def test_kill_matches_stock_schedule(stream_interval): + assert _emitted_iters(stream_interval, 1) == _stock_iters(stream_interval) + + +@pytest.mark.parametrize("stream_interval", [3, 20]) +def test_default_adds_iteration_two_only(stream_interval): + expected = sorted(set(_stock_iters(stream_interval)) | {2}) + assert _emitted_iters(stream_interval, 2) == expected + + +def test_larger_window_adds_iterations_up_to_n(): + expected = sorted(set(_stock_iters(20)) | {2, 3, 4, 5}) + assert _emitted_iters(20, 5) == expected + + +@pytest.mark.parametrize("first_content_emit_iters", [1, 2, 5]) +@pytest.mark.parametrize("decoding_iter", [1, 2, 7, 19]) +def test_finished_request_always_emits(first_content_emit_iters, decoding_iter): + fake = _fake_self(20, first_content_emit_iters) + for emit_first_iter in (True, False): + assert PyExecutor._should_emit_response( + fake, _request(decoding_iter, finished=True), emit_first_iter + ) + + +@pytest.mark.parametrize("first_content_emit_iters", [1, 2, 5]) +def test_early_first_token_suppresses_only_iteration_one(first_content_emit_iters): + with_first = _emitted_iters(20, first_content_emit_iters) + without_first = _emitted_iters(20, first_content_emit_iters, emit_first_iter=False) + assert without_first == [it for it in with_first if it != 1] + + +@pytest.mark.parametrize("first_content_emit_iters", [1, 2, 5]) +def test_stream_interval_one_emits_every_iteration(first_content_emit_iters): + assert _emitted_iters(1, first_content_emit_iters) == list(range(1, 61)) + + +class _PlainModel: + """A model that declares no runtime hooks (every model but the MiniMax-M3 sm_107 targets).""" + + +class _M3Model: + """A model whose runtime hooks opt into the iteration-2 emission, as the MiniMax-M3 sm_107 targets do.""" + + runtime_hooks = RuntimeHooks(first_content_emit_iters=2) + + +def test_runtime_hooks_choose_the_schedule(): + """PyExecutor reads the schedule from the model's runtime hooks: a model without hooks keeps the stock schedule + (1, 20, 40, ...), one that opts in also emits on iteration 2.""" + plain = runtime_hooks_of(_PlainModel()).first_content_emit_iters + m3 = runtime_hooks_of(_M3Model()).first_content_emit_iters + + assert _emitted_iters(20, plain) == [1, 20, 40, 60] + assert _emitted_iters(20, m3) == [1, 2, 20, 40, 60] diff --git a/tests/unittest/_torch/pyexecutor/test_h2d_cache.py b/tests/unittest/_torch/pyexecutor/test_h2d_cache.py index bd86fc1d1fe1..2c168e1cb32d 100644 --- a/tests/unittest/_torch/pyexecutor/test_h2d_cache.py +++ b/tests/unittest/_torch/pyexecutor/test_h2d_cache.py @@ -7,11 +7,14 @@ import pytest import torch +from tensorrt_llm._torch.memory_buffer_utils import Buffers from tensorrt_llm._torch.pyexecutor import h2d_cache from tensorrt_llm._torch.pyexecutor.h2d_cache import ( H2DSkipCache, copy_h2d_if_changed, + copy_h2d_pool_shared, fill_tail_if_changed, + fill_tail_pool_shared, ) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -316,6 +319,105 @@ def test_a_tail_record_keeps_its_storage_so_a_new_buffer_cannot_inherit_it(cache assert fresh[3:].eq(-1).all() +# --------------------------------------------------------------------------- +# CUDA graph metadata pool: one block shared by every graph key +# --------------------------------------------------------------------------- + +NUM_SLOTS = 8 + + +def _graph_key_views(name): + """Two graph keys' metadata views from one pool, as create_cuda_graph_metadata hands them out.""" + pool = Buffers() + large = pool.get_buffer([NUM_SLOTS], torch.int32, name, reserve_buffer=True) + small = pool.get_buffer([NUM_SLOTS], torch.int32, name, reserve_buffer=True) + assert small.data_ptr() == large.data_ptr() + return large, small + + +def _capture_live_row_write(buf, rows, values): + """A graph whose replay writes buf[:rows] on the device, as on_update_kv_lens writes the live slots.""" + slots = torch.tensor(values, dtype=torch.int32, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + buf[:rows].copy_(slots) + torch.cuda.current_stream().wait_stream(stream) + return graph + + +def _key_b_step(buf, live): + """Key B's prepare: its live slots, then -1 over the padded rows its captured producers also run.""" + copy_h2d_pool_shared(buf[: len(live)], _host(live)) + fill_tail_pool_shared(buf, len(live), -1) + + +@pytest.mark.parametrize("skip_pool_shared", [False, True], ids=["fixed", "kill_switch"]) +def test_two_graph_keys_over_one_block_keep_the_padded_rows_unused( + cache, monkeypatch, skip_pool_shared +): + monkeypatch.setattr(h2d_cache, "SKIP_POOL_SHARED_UPLOADS", skip_pool_shared) + key_a, key_b = _graph_key_views("msa_out_cache_loc") + replay_a = _capture_live_row_write(key_a, 6, [10, 11, 12, 13, 14, 15]) + + _key_b_step(key_b, [20, 21, 22]) + torch.cuda.synchronize() + assert key_b.tolist() == [20, 21, 22, -1, -1, -1, -1, -1] + # Key A's step: its replay writes six live slots into the shared block. + replay_a.replay() + # Key B's next step, with the same live slots. + _key_b_step(key_b, [20, 21, 22]) + torch.cuda.synchronize() + if skip_pool_shared: + # The old behaviour: both writes were skipped, so key B's live and padded rows hold key A's live slots. + assert key_b.tolist() == [10, 11, 12, 13, 14, 15, -1, -1] + else: + assert key_b.tolist() == [20, 21, 22, -1, -1, -1, -1, -1] + + +@pytest.mark.parametrize("skip_pool_shared", [False, True], ids=["fixed", "kill_switch"]) +def test_two_graph_keys_over_one_block_keep_their_uploads(cache, monkeypatch, skip_pool_shared): + monkeypatch.setattr(h2d_cache, "SKIP_POOL_SHARED_UPLOADS", skip_pool_shared) + key_a, key_b = _graph_key_views("msa_block_table") + replay_a = _capture_live_row_write(key_a, 4, [7, 7, 7, 7]) + table = [1, 2, 3, 4] + + assert copy_h2d_pool_shared(key_b[:4], _host(table)) is True + replay_a.replay() + copied = copy_h2d_pool_shared(key_b[:4], _host(table)) + torch.cuda.synchronize() + if skip_pool_shared: + assert copied is False + assert key_b[:4].tolist() == [7, 7, 7, 7] + else: + assert copied is True + assert key_b[:4].tolist() == table + + +def test_pool_shared_writes_drop_the_shadow_and_tail_record(cache): + buf = _device_zeros(6) + values = _host([1, 2]) + assert copy_h2d_if_changed(buf[:2], values) is True + fill_tail_if_changed(buf, 2, -1) + copy_h2d_pool_shared(buf[:2], _host([5, 6])) + buf[3] = 9 + # Neither record vouches for the buffer after a write outside the skip cache. + assert copy_h2d_if_changed(buf[:2], values) is True + fill_tail_if_changed(buf, 2, -1) + torch.cuda.synchronize() + assert buf.tolist() == [1, 2, -1, -1, -1, -1] + + +def test_pool_shared_writes_without_an_active_cache_are_plain_writes(): + assert not h2d_cache.enabled() + buf = _device_zeros(4) + assert copy_h2d_pool_shared(buf[:2], _host([3, 4]), pack=True) is True + fill_tail_pool_shared(buf, 2, -1) + torch.cuda.synchronize() + assert buf.tolist() == [3, 4, -1, -1] + + def test_failed_copy_through_another_view_invalidates_the_old_record(cache, monkeypatch): dst = _device_zeros(4) host = _host([1, 2, 3, 4]) diff --git a/tests/unittest/_torch/pyexecutor/test_h2d_changed_span.py b/tests/unittest/_torch/pyexecutor/test_h2d_changed_span.py new file mode 100644 index 000000000000..2a987a260072 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_h2d_changed_span.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""H2DSkipCache(changed_span=True): a changed packed 2-D upload stages only each row's changed span.""" + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs +from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache, copy_h2d_if_changed +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + +ROWS, WIDTH = 4, 1024 + + +@pytest.fixture +def pack(): + return StepInputPack(torch.device("cuda"), parallel_scatter=True) + + +def _table(seed, rows=ROWS, width=WIDTH): + generator = torch.Generator().manual_seed(seed) + return torch.randint( + 0, 1 << 20, (rows, width), generator=generator, dtype=torch.int32 + ).pin_memory() + + +def _upload(pack, cache, dst, src): + """One step: open the pack, route the copy through the cache, issue the pack; returns the staged copies.""" + with h2d_cache.active(cache), step_inputs.active(pack): + step_inputs.begin() + copied = copy_h2d_if_changed(dst, src, pack=True) + staged = list(pack.h2d) + step_inputs.flush() + torch.cuda.synchronize() + return copied, staged + + +@pytest.mark.parametrize("changed_span", [False, True]) +def test_device_rows_equal_the_source_after_every_step(pack, changed_span): + """Bitwise: whatever is staged, the device view holds the source after each step.""" + cache = H2DSkipCache(changed_span=changed_span) + buf = torch.full((8, 2 * WIDTH), -7, dtype=torch.int32, device="cuda") + dst = buf[:ROWS, :WIDTH] + src = _table(0) + generator = torch.Generator().manual_seed(1) + for step in range(40): + src = src.clone().pin_memory() + # A request crosses into a new page: one entry of one row changes (sometimes several, in several rows). + for _ in range(1 + step % 3): + row = int(torch.randint(0, ROWS, (1,), generator=generator)) + col = int(torch.randint(0, WIDTH, (1,), generator=generator)) + src[row, col] = step + _upload(pack, cache, dst, src) + assert torch.equal(dst.cpu(), src) + # Rows and columns outside the view were never written. + assert bool((buf[ROWS:] == -7).all()) and bool((buf[:, WIDTH:] == -7).all()) + + +def test_one_changed_entry_stages_one_element(pack): + cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(2) + _upload(pack, cache, dst, src) + + src = src.clone().pin_memory() + src[2, 517] += 1 + copied, staged = _upload(pack, cache, dst, src) + + assert copied is True + assert staged == [(dst[2, 517].data_ptr(), staged[0][1], 4)] + assert torch.equal(dst.cpu(), src) + + +def test_each_changed_row_stages_its_first_to_last_changed_column(pack): + cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(3) + _upload(pack, cache, dst, src) + + src = src.clone().pin_memory() + src[0, 10] += 1 + src[0, 20] += 1 + src[3, 1000] += 1 + _, staged = _upload(pack, cache, dst, src) + + assert [(ptr, nbytes) for ptr, _, nbytes in staged] == [ + (dst[0, 10].data_ptr(), 11 * 4), + (dst[3, 1000].data_ptr(), 4), + ] + assert torch.equal(dst.cpu(), src) + + +def test_wider_view_uploads_the_columns_the_shadow_does_not_cover(pack): + cache = H2DSkipCache(changed_span=True) + buf = torch.full((ROWS, WIDTH), -1, dtype=torch.int32, device="cuda") + src = _table(4) + _upload(pack, cache, buf[:, : WIDTH // 2], src[:, : WIDTH // 2].contiguous().pin_memory()) + + _, staged = _upload( + pack, cache, buf[:, : WIDTH // 2 + 3], src[:, : WIDTH // 2 + 3].contiguous().pin_memory() + ) + + assert [nbytes for _, _, nbytes in staged] == [3 * 4] * ROWS + assert torch.equal(buf[:, : WIDTH // 2 + 3].cpu(), src[:, : WIDTH // 2 + 3]) + + +def test_a_mostly_changed_table_is_staged_whole(pack): + cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + _upload(pack, cache, dst, _table(5)) + + src = _table(6) + _, staged = _upload(pack, cache, dst, src) + + assert [nbytes for _, _, nbytes in staged] == [ROWS * WIDTH * 4] + assert torch.equal(dst.cpu(), src) + + +def test_off_by_default_the_whole_view_is_staged(pack): + cache = H2DSkipCache() + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(7) + _upload(pack, cache, dst, src) + + src = src.clone().pin_memory() + src[1, 1] += 1 + _, staged = _upload(pack, cache, dst, src) + + assert [nbytes for _, _, nbytes in staged] == [ROWS * WIDTH * 4] + + +def test_an_unpacked_copy_copies_the_whole_view(pack): + cache = H2DSkipCache(changed_span=True) + dst = torch.zeros((ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(8) + with h2d_cache.active(cache): + copy_h2d_if_changed(dst, src) + src = src.clone().pin_memory() + src[1, 1] += 1 + assert copy_h2d_if_changed(dst, src) is True + torch.cuda.synchronize() + assert torch.equal(dst.cpu(), src) + + +def test_a_view_with_another_row_stride_is_staged_whole(pack): + """The shadow describes other device bytes: nothing of it may be trusted.""" + cache = H2DSkipCache(changed_span=True) + buf = torch.zeros((2 * ROWS, WIDTH), dtype=torch.int32, device="cuda") + src = _table(9) + _upload(pack, cache, buf[:ROWS], src) + + _, staged = _upload(pack, cache, buf[::2], src) + + assert [nbytes for _, _, nbytes in staged] == [WIDTH * 4] * ROWS + assert torch.equal(buf[::2].cpu(), src) diff --git a/tests/unittest/_torch/pyexecutor/test_kv_connector_executor_lifetime.py b/tests/unittest/_torch/pyexecutor/test_kv_connector_executor_lifetime.py index b056d62e2493..1a8913059208 100644 --- a/tests/unittest/_torch/pyexecutor/test_kv_connector_executor_lifetime.py +++ b/tests/unittest/_torch/pyexecutor/test_kv_connector_executor_lifetime.py @@ -525,6 +525,7 @@ def _preemption_executor() -> PyExecutor: # Wide enough that a mid-stream request emits nothing of its own, so any # response seen below came from the finish path rather than from streaming. executor.stream_interval = 100 + executor.first_content_emit_iters = 1 return executor diff --git a/tests/unittest/_torch/pyexecutor/test_kv_lens_update_once.py b/tests/unittest/_torch/pyexecutor/test_kv_lens_update_once.py new file mode 100644 index 000000000000..f404a80713e1 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_kv_lens_update_once.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""RuntimeHooks.kv_lens_update_once: DecoderInputPreparer.preprocess_inputs re-derives the KV-length metadata once, +after the overlap scheduler's KV-length correction, on exactly the steps that correct it, and leaves the inputs as +before.""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.inputs import DecoderInputPreparer +from tensorrt_llm._torch.pyexecutor.runtime_hooks import RuntimeHooks + + +class _Metadata: + """The fields preprocess_inputs reads; on_update_kv_lens records the kv_lens it saw.""" + + def __init__(self, gen: int, tokens_per_gen: int, with_kv_cache: bool = True): + self.kv_cache_manager = object() if with_kv_cache else None + self.num_seqs = gen + self.num_contexts = 0 + self.num_generations = gen + self.num_ctx_tokens = 0 + self.num_chunked_ctx_requests = 0 + self.kv_lens_cuda = torch.arange(100, 100 + gen, dtype=torch.int32) + self.seen: list[list[int]] = [] + + def on_update_kv_lens(self) -> None: + self.seen.append(self.kv_lens_cuda.tolist()) + + +def _preparer(once: bool, disable_overlap: bool = False) -> DecoderInputPreparer: + r = object.__new__(DecoderInputPreparer) + r.runtime_hooks = RuntimeHooks(kv_lens_update_once=once) + r._config = types.SimpleNamespace(disable_overlap_scheduler=disable_overlap) + r._buffers = types.SimpleNamespace( + previous_pos_id_offsets_cuda=torch.full((64,), -2, dtype=torch.int32), + previous_kv_lens_offsets_cuda=torch.full((16,), -3, dtype=torch.int32), + ) + r.mapping = types.SimpleNamespace(has_cp_helix=lambda: False) + r.guided_decoder = None + return r + + +def _inputs(md: _Metadata, tokens: int) -> dict: + return { + "attn_metadata": md, + "input_ids": torch.zeros(tokens, dtype=torch.int32), + "position_ids": torch.arange(tokens, dtype=torch.int32).view(1, -1), + } + + +@pytest.mark.parametrize("gen", [1, 4]) +def test_corrected_step_updates_once_with_the_same_inputs(gen) -> None: + t = 4 + outs = {} + for once in (False, True): + md = _Metadata(gen, t) + inputs = _preparer(once).preprocess_inputs( + _inputs(md, gen * t), enable_spec_decode=True, runtime_draft_len=t - 1 + ) + outs[once] = (md.seen, inputs["position_ids"].tolist(), md.kv_lens_cuda.tolist()) + corrected = [100 + i - 3 for i in range(gen)] + # Without the hook: once before the correction, once after it. + assert outs[False][0] == [[100 + i for i in range(gen)], corrected] + # With it: only the call after the correction, which sees the same lengths. + assert outs[True][0] == [corrected] + assert outs[True][1:] == outs[False][1:] + + +@pytest.mark.parametrize( + "spec, disable_overlap, with_kv_cache", + [(False, False, True), (True, True, True), (True, False, False)], + ids=["no spec decode", "no overlap scheduler", "no kv cache manager"], +) +def test_uncorrected_steps_keep_their_update(spec, disable_overlap, with_kv_cache) -> None: + """Steps that do not correct the lengths still re-derive them once (the per-forward invalidation), with the same + inputs as without the hook.""" + outs = {} + for once in (False, True): + md = _Metadata(2, 4, with_kv_cache) + inputs = _preparer(once, disable_overlap).preprocess_inputs( + _inputs(md, 8), enable_spec_decode=spec, runtime_draft_len=3 + ) + outs[once] = (md.seen, inputs["position_ids"].tolist(), md.kv_lens_cuda.tolist()) + assert len(outs[True][0]) == 1 + assert outs[True] == outs[False] + + +def test_off_by_default() -> None: + assert RuntimeHooks().kv_lens_update_once is False diff --git a/tests/unittest/_torch/pyexecutor/test_runtime_opt_ins.py b/tests/unittest/_torch/pyexecutor/test_runtime_opt_ins.py new file mode 100644 index 000000000000..4f3d8b2eec74 --- /dev/null +++ b/tests/unittest/_torch/pyexecutor/test_runtime_opt_ins.py @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The process-wide features a model opts into through RuntimeHooks and the executor arms +(``_arm_runtime_opt_ins``): the SM107 MXFP8 GEMM tactic table and the late-compile detector. A model +that declares neither leaves both off, even at a table shape on SM107; the MiniMax-M3 sm_107 targets +declare both.""" + +import importlib + +import pytest +import torch + +from tensorrt_llm._torch import late_compile +from tensorrt_llm._torch.custom_ops import mxfp8_cute_pinned_tactics as pin +from tensorrt_llm._torch.pyexecutor.py_executor import _arm_runtime_opt_ins +from tensorrt_llm._torch.pyexecutor.runtime_hooks import runtime_hooks_of + +pytestmark = pytest.mark.cpu_only + +M3_SHAPE = (64, 6144, 6144) # (m, n, k): a table entry +M3_TARGETS = "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__{}.modeling" + + +class _PlainModel(torch.nn.Module): + """A model that declares no hooks.""" + + +@pytest.fixture(autouse=True) +def _sm107_unarmed(monkeypatch): + monkeypatch.setattr(pin, "_ARMED", [False]) + monkeypatch.setattr(pin, "ENABLED", True) + monkeypatch.setattr(pin, "get_sm_version", lambda: 107) + monkeypatch.setattr(late_compile, "_ready", [False]) + + +def test_a_model_without_the_opt_in_arms_nothing_at_an_m3_shape(): + hooks = runtime_hooks_of(_PlainModel()) + + assert _arm_runtime_opt_ins(hooks) is False + assert pin.tactic_for(*M3_SHAPE) is None + assert not late_compile._ready[0] + + +@pytest.mark.parametrize("tp", ["tp2", "tp4", "tp8"]) +def test_the_m3_sm107_targets_arm_both(tp): + hooks = importlib.import_module(M3_TARGETS.format(tp))._runtime_hooks(None) + + assert hooks.mxfp8_pinned_tactics and hooks.late_compile_detect + assert _arm_runtime_opt_ins(hooks) is True + m, n, k = M3_SHAPE + assert pin.tactic_for(m, n, k) == pin.TACTICS[(n, k, pin._bucket(m))] + + +def test_the_kill_switch_still_wins_once_armed(monkeypatch): + pin.arm() + monkeypatch.setattr(pin, "ENABLED", False) + + assert pin.tactic_for(*M3_SHAPE) is None diff --git a/tests/unittest/_torch/pyexecutor/test_step_inputs.py b/tests/unittest/_torch/pyexecutor/test_step_inputs.py index c1b4b95a02d4..9ca87ba292fd 100644 --- a/tests/unittest/_torch/pyexecutor/test_step_inputs.py +++ b/tests/unittest/_torch/pyexecutor/test_step_inputs.py @@ -602,3 +602,343 @@ def worker(index): assert results == {0: (True, 1), 1: (True, 1)} assert dsts[0].tolist() == [1] * 4 assert dsts[1].tolist() == [2] * 4 + + +def _random_pack_run(parallel: bool, seed: int, overlap: bool): + """Stage a fixed random mix of host copies (2-D page-table rows included) and device copies that read a host + copy's destination; return every destination after the flush.""" + g = torch.Generator().manual_seed(seed) + pack = StepInputPack(_device(), parallel_scatter=parallel) + big = _zeros(4096) + table = torch.zeros((8, 64), dtype=torch.int32, device=_device()) + out = _zeros(256) + with step_inputs.active(pack): + step_inputs.begin() + for k in range(24): + n = int(torch.randint(1, 64, (1,), generator=g)) + step_inputs.copy_( + big[k * 128 : k * 128 + n], + torch.randint(0, 1 << 30, (n,), generator=g, dtype=torch.int32), + ) + step_inputs.copy_( + table[:3, :40], torch.randint(0, 1 << 30, (3, 40), generator=g, dtype=torch.int32) + ) + if overlap: + # A later host copy overwrites part of an earlier one: order matters, so the serial scatter must run. + step_inputs.copy_(big[0:4], torch.full((4,), -7, dtype=torch.int32)) + step_inputs.copy_(out[:64], big[:64]) + step_inputs.copy_(out[64:96], big[128:160]) + step_inputs.flush() + torch.cuda.synchronize() + return big.cpu(), table.cpu(), out.cpu() + + +@pytest.mark.parametrize("overlap", [False, True], ids=["disjoint", "overlapping"]) +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_parallel_scatter_writes_what_the_serial_scatter_writes(seed, overlap): + serial = _random_pack_run(False, seed, overlap) + parallel = _random_pack_run(True, seed, overlap) + for a, b in zip(serial, parallel): + assert torch.equal(a, b) + + +def _compiled_variants(kernel) -> int: + """Compiled variants of a Triton kernel on this device, or -1 where this Triton keeps them elsewhere.""" + caches = getattr(kernel, "device_caches", None) + if not caches: + return -1 + cache = caches.get(torch.cuda.current_device()) + return len(cache[0]) if cache else 0 + + +@pytest.mark.parametrize("count", [1, 15, 16, 17, 120]) # 120: 240 of the 256 descriptors +def test_parallel_scatter_is_bitwise_for_every_copy_count(count): + """``first`` (the device copies' descriptor offset) is the step's host-copy count: one compiled variant serves + every count, so a new count never compiles mid-serving.""" + pack = StepInputPack(_device(), parallel_scatter=True) + width = 3 + dst = _zeros(count * width) + mirror = _zeros(count * width) + g = torch.Generator().manual_seed(count) + values = torch.randint( + -(1 << 31), (1 << 31) - 1, (count, width), generator=g, dtype=torch.int32 + ) + with step_inputs.active(pack): + step_inputs.begin() + for i in range(count): + step_inputs.copy_(dst[i * width : (i + 1) * width], values[i]) + for i in range(count): + step_inputs.copy_(mirror[i * width : (i + 1) * width], dst[i * width : (i + 1) * width]) + assert len(pack.h2d) == count and len(pack.d2d) == count + step_inputs.flush() + torch.cuda.synchronize() + assert torch.equal(dst.cpu(), values.flatten()) + assert torch.equal(mirror.cpu(), values.flatten()) + variants = _compiled_variants(step_inputs._parallel_scatter_kernel) + assert variants in (-1, 1, 2), variants # at most one per DEVICE_SRC + + +def test_parallel_scatter_only_for_order_free_packs(): + assert step_inputs._parallel_ok([(0, 0, 8), (8, 0, 8)], []) + assert not step_inputs._parallel_ok([(0, 0, 8), (4, 0, 8)], []) + # A device copy may read a host copy's destination (the launches run in order) ... + assert step_inputs._parallel_ok([(0, 0, 8)], [(100, 0, 8)]) + # ... but not another device copy's destination, and device destinations must be disjoint. + assert not step_inputs._parallel_ok([], [(100, 0, 8), (200, 100, 8)]) + assert not step_inputs._parallel_ok([], [(100, 0, 8), (104, 300, 8)]) + + +def test_parallel_scatter_is_off_by_default(): + assert not StepInputPack(_device()).parallel_scatter + + +def _graph_case(early=False, epoch=False): + """A pack with the in-graph prologue and a graph that starts with it and then reads the destinations.""" + pack = StepInputPack( + _device(), parallel_scatter=True, in_graph=True, early_copy=early, step_epoch=epoch + ) + a, b, c = _zeros(16), _zeros(16), _zeros(16) + out = _zeros(48) + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + assert step_inputs.capture_prologue() is pack + out.copy_(torch.cat([a, b, c])) + return pack, graph, (a, b, c), out + + +def _step(pack, graph, dsts, values, overlap=False, flush=True): + a, b, c = dsts + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, _host(values[0])) + step_inputs.copy_(b, _host(values[1])) + if overlap: + step_inputs.copy_(a[:4], _host([-1, -2, -3, -4])) + step_inputs.copy_(c, a) # a device copy reading a host copy's destination + if flush: + step_inputs.flush(in_graph=True) + else: + step_inputs.discard() + step_inputs.flush(in_graph=True) + graph.replay() + + +def test_in_graph_pack_is_issued_by_the_replay(): + pack, graph, dsts, out = _graph_case() + for step in range(5): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + _step(pack, graph, dsts, (va, vb)) + torch.cuda.synchronize() + assert out.tolist() == va + vb + va, step + # Nothing was issued outside the graph: the window carried the copies. + assert int(pack.ghost[:8].view(torch.int64)[0]) == 2 + + +def test_in_graph_pack_not_order_free_is_issued_eagerly(): + pack, graph, dsts, out = _graph_case() + va = list(range(16)) + vb = list(range(100, 116)) + _step(pack, graph, dsts, (va, vb), overlap=True) + torch.cuda.synchronize() + expect_a = [-1, -2, -3, -4] + va[4:] + assert out.tolist() == expect_a + vb + expect_a + # The window was zeroed, so the prologue wrote nothing over the eager result. + assert not pack.ghost[:16].any() + + +def test_in_graph_replay_without_a_pack_reapplies_nothing(): + pack, graph, dsts, out = _graph_case() + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + _step(pack, graph, dsts, ([0] * 16, [0] * 16), flush=False) + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_pack_lands_every_step(early): + pack, graph, dsts, out = _graph_case(early) + assert pack.early_copy is early + for step in range(6): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + _step(pack, graph, dsts, (va, vb), overlap=step % 2 == 1) + torch.cuda.synchronize() + head = [-1, -2, -3, -4] if step % 2 == 1 else [] + expect_a = head + va[len(head) :] + assert out.tolist() == expect_a + vb + expect_a, step + + +def _history_run(pack, graph, dsts, out, steps, before_replay=None, before_upload=None): + """Steps back to back without host syncs; each replay's output is copied to a history row behind it (same + stream), so a later upload that overwrote the window too early shows up as the wrong step's values.""" + hist = _zeros(steps, out.numel()) + expect = [] + for step in range(steps): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + expect.append(va + vb + va) + a, b, c = dsts + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, _host(va)) + step_inputs.copy_(b, _host(vb)) + step_inputs.copy_(c, a) + if before_upload is not None: + before_upload(step) + step_inputs.flush(in_graph=True) + if before_replay is not None: + before_replay(step) + graph.replay() + hist[step].copy_(out) + torch.cuda.synchronize() + return hist.tolist(), expect + + +def test_early_copy_waits_for_the_previous_prologue_before_overwriting_the_window(): + """A slow GPU: the replay of step N (and its prologue) is held back behind a long kernel while the host already + uploads step N+1; that upload must wait for step N's prologue to consume the window.""" + pack, graph, dsts, out = _graph_case(early=True) + hist, expect = _history_run( + pack, graph, dsts, out, 4, before_replay=lambda step: torch.cuda._sleep(_SLEEP_CYCLES // 4) + ) + assert hist == expect + + +def test_early_copy_prologue_waits_for_a_slow_upload(): + """A slow copy stream: each step's upload is held back behind a long kernel on the copy stream; the prologue must + wait for it (staged_ready), and the next host fill must wait for the upload to have read the window.""" + pack, graph, dsts, out = _graph_case(early=True) + + def slow_copy(step): + with torch.cuda.stream(pack.copy_stream): + torch.cuda._sleep(_SLEEP_CYCLES // 4) + + hist, expect = _history_run(pack, graph, dsts, out, 4, before_upload=slow_copy) + assert hist == expect + + +def test_early_copy_with_a_slow_host(): + """A slow host: the GPU idles between steps (the window lands long before the replay).""" + import time + + pack, graph, dsts, out = _graph_case(early=True) + hist, expect = _history_run( + pack, graph, dsts, out, 3, before_replay=lambda step: time.sleep(0.05) + ) + assert hist == expect + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_in_graph_replay_after_quiesce_reapplies_nothing(early): + pack, graph, dsts, out = _graph_case(early) + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +def test_early_copy_needs_in_graph(): + assert not StepInputPack(_device(), parallel_scatter=True, early_copy=True).early_copy + + +def test_prologue_op_needs_early_copy_and_is_off_by_default(): + def op(window, epoch, max_desc): + raise AssertionError("not called") + + assert StepInputPack(_device(), parallel_scatter=True, early_copy=True).prologue_op is None + assert ( + StepInputPack(_device(), parallel_scatter=True, in_graph=True, prologue_op=op).prologue_op + is None + ) + + +def test_in_graph_needs_parallel_scatter_and_is_off_by_default(): + assert not StepInputPack(_device()).in_graph + assert not StepInputPack(_device(), in_graph=True).in_graph + + +def test_in_graph_replay_outside_the_pack_context_reapplies_nothing(): + pack, graph, dsts, out = _graph_case() + _step(pack, graph, dsts, (list(range(16)), list(range(16)))) + torch.cuda.synchronize() + dsts[0].fill_(7) + # What the graph runner does for a replay outside the pack's context. + assert not step_inputs.is_active(pack) + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert dsts[0].tolist() == [7] * 16 + + +def test_step_epoch_is_off_by_default_and_needs_in_graph(): + assert StepInputPack(_device(), parallel_scatter=True, in_graph=True).epoch is None + assert StepInputPack(_device(), parallel_scatter=True, step_epoch=True).epoch is None + with step_inputs.active(StepInputPack(_device(), parallel_scatter=True, in_graph=True)): + assert step_inputs.step_epoch() is None + assert step_inputs.step_epoch() is None + + +@pytest.mark.parametrize("early", [False, True], ids=["prologue_copy", "early_copy"]) +def test_step_epoch_advances_once_per_replay_and_never_resets(early): + pack, graph, dsts, out = _graph_case(early, epoch=True) + with step_inputs.active(pack): + assert step_inputs.step_epoch() is pack.epoch + torch.cuda.synchronize() + start = int(pack.epoch.item()) # the warm-up compile launched the prologue once + assert start == 1 + for step in range(5): + va = [step * 100 + i for i in range(16)] + vb = [step * 1000 - i for i in range(16)] + # A replay with a new pack, one with an eagerly issued pack, and one with nothing to scatter. + _step(pack, graph, dsts, (va, vb), overlap=step == 2, flush=step != 3) + torch.cuda.synchronize() + assert int(pack.epoch.item()) == start + step + 1, step + # The copies are unchanged by the epoch. + _step(pack, graph, dsts, (list(range(16)), list(range(16, 32)))) + torch.cuda.synchronize() + assert out.tolist() == list(range(16)) + list(range(16, 32)) + list(range(16)) + assert int(pack.epoch.item()) == start + 6 + + +def test_step_epoch_is_ordered_before_the_graph_body(): + """A kernel after the prologue in the same graph reads the bumped epoch of its own replay.""" + pack = StepInputPack(_device(), parallel_scatter=True, in_graph=True, step_epoch=True) + seen = torch.zeros(4, dtype=torch.int64, device="cuda") + slot = torch.zeros(1, dtype=torch.int64, device="cuda") + graph = torch.cuda.CUDAGraph() + with step_inputs.active(pack): + torch.cuda.synchronize() + with torch.cuda.graph(graph): + step_inputs.capture_prologue() + seen.index_copy_(0, slot, pack.epoch) + slot.add_(1) + torch.cuda.synchronize() + start = int(pack.epoch.item()) + for _ in range(4): + with step_inputs.active(pack): + pack.quiesce_graph_window() + graph.replay() + torch.cuda.synchronize() + assert seen.tolist() == [start + 1, start + 2, start + 3, start + 4] + + +def test_eager_flush_does_not_advance_the_step_epoch(): + pack = StepInputPack(_device(), parallel_scatter=True, in_graph=True, step_epoch=True) + torch.cuda.synchronize() + start = int(pack.epoch.item()) + dst = _zeros(4) + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(dst, _host([1, 2, 3, 4])) + step_inputs.flush() + torch.cuda.synchronize() + assert dst.tolist() == [1, 2, 3, 4] + assert int(pack.epoch.item()) == start diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_one_model_sampling_rejection.py b/tests/unittest/_torch/speculative/hw_agnostic/test_one_model_sampling_rejection.py index a53c2406b55e..c318ae4b7b20 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_one_model_sampling_rejection.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_one_model_sampling_rejection.py @@ -26,6 +26,7 @@ from tensorrt_llm._torch.speculative.spec_sampler_base import ( UNSUPPORTED_BAD_WORDS_MSG, UNSUPPORTED_EMBEDDING_BIAS_MSG, + UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG, UNSUPPORTED_MIN_LENGTH_MSG, UNSUPPORTED_MIN_P_MSG, UNSUPPORTED_NO_REPEAT_NGRAM_MSG, @@ -45,6 +46,11 @@ ("bad", dict(bad="nope"), UNSUPPORTED_BAD_WORDS_MSG), ("no_repeat_ngram_size", dict(no_repeat_ngram_size=3), UNSUPPORTED_NO_REPEAT_NGRAM_MSG), ("top_p_decay", dict(top_p_decay=0.5), UNSUPPORTED_TOP_P_DECAY_MSG), + ( + "logits_processor", + dict(logits_processor=lambda req_id, logits, ids, stream, client_id: None), + UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG, + ), ] # Values a frontend forwards explicitly at their neutral setting. These must not @@ -162,6 +168,11 @@ def test_check_arguments_admits_drafter_modes(): ("min_length", dict(py_min_length=16), UNSUPPORTED_MIN_LENGTH_MSG), ("bad_words", dict(py_bad_words=[[7]]), UNSUPPORTED_BAD_WORDS_MSG), ("no_repeat_ngram_size", dict(py_no_repeat_ngram_size=3), UNSUPPORTED_NO_REPEAT_NGRAM_MSG), + ( + "logits_post_processor", + dict(py_logits_post_processors=[lambda req_id, logits, ids, stream, client_id: None]), + UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG, + ), ] @@ -272,3 +283,21 @@ def test_executor_side_check_still_rejects_top_p_decay(): def test_executor_side_check_admits_a_neutral_request(): """The stub covers every field the path reads; a neutral request passes.""" _validate_on_executor(_executor_request()) + + +@pytest.mark.cpu_only +def test_executor_side_check_admits_an_empty_post_processor_list(): + """An empty list is what a request without callbacks may carry: not rejected.""" + request = _executor_request() + request.py_logits_post_processors = [] + _validate_on_executor(request) + + +@pytest.mark.cpu_only +def test_executor_side_check_rejects_a_post_processor_on_the_full_path(): + """Through the whole validate_request, as the executor admits a request.""" + request = _executor_request() + request.py_logits_post_processors = [lambda req_id, logits, ids, stream, client_id: None] + with pytest.raises(ValueError) as excinfo: + _validate_on_executor(request) + assert str(excinfo.value) == UNSUPPORTED_LOGITS_POST_PROCESSOR_MSG diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py index 9d8d13f6e429..3e88f3051b00 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -612,7 +612,7 @@ def copy_to_host(tensor): transfers.append(("copy", tuple(tensor.shape))) return tensor.clone() - def record_event(): + def record_event(side_stream_event=None): transfers.append(("record", None)) return types.SimpleNamespace(synchronize=lambda: transfers.append(("sync", None))) diff --git a/tests/unittest/_torch/speculative/test_rng_upload_pack.py b/tests/unittest/_torch/speculative/test_rng_upload_pack.py new file mode 100644 index 000000000000..1e740855ecb5 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_rng_upload_pack.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Philox seed/offset uploads of a sampled one-model step, staged into the step-input pack.""" + +import types +from typing import Optional + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor import step_inputs +from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack +from tensorrt_llm._torch.speculative.interface import SpecMetadata + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), + pytest.mark.skipif( + not (torch.cuda.is_available() and hasattr(torch.ops.trtllm, "step_inputs_scatter")), + reason="torch.ops.trtllm.step_inputs_scatter is not registered", + ), +] + +MAX_DRAFT_LEN = 3 +ROWS_PER_REQUEST = MAX_DRAFT_LEN + 1 + + +def _meta(pack_rng_uploads: bool) -> SpecMetadata: + meta = SpecMetadata( + max_num_requests=4, max_draft_len=MAX_DRAFT_LEN, max_total_draft_tokens=MAX_DRAFT_LEN + ) + meta.pack_rng_uploads = pack_rng_uploads + return meta + + +def _request(slot: int, seed: Optional[int]) -> types.SimpleNamespace: + return types.SimpleNamespace( + py_seq_slot=slot, + py_decoding_iter=0, + py_request_id=slot, + sampling_config=types.SimpleNamespace(seed=seed), + ) + + +def _populate(meta: SpecMetadata) -> None: + requests = [_request(2, seed=1234), _request(0, seed=None)] + normalized = [(0.7, 0, 0.9, 0.0, ROWS_PER_REQUEST) for _ in requests] + meta._populate_request_rng_state(requests, normalized) + + +def _rng_state(meta: SpecMetadata) -> dict[str, torch.Tensor]: + torch.cuda.synchronize() + return { + name: getattr(meta, name).clone() + for name in ("request_seeds", "request_offsets", "seeds", "offsets") + } + + +def test_packed_uploads_land_at_the_flush_with_the_plain_values(): + plain = _meta(pack_rng_uploads=False) + _populate(plain) + expected = _rng_state(plain) + + packed = _meta(pack_rng_uploads=True) + pack = StepInputPack(torch.device("cuda", torch.cuda.current_device())) + with step_inputs.active(pack): + step_inputs.begin() + _populate(packed) + before_flush = _rng_state(packed) + step_inputs.flush() + after_flush = _rng_state(packed) + + for name, tensor in expected.items(): + assert torch.equal(after_flush[name], tensor), name + # Staged, not copied: the buffers still hold their zero init until the flush. + assert not torch.equal(before_flush["seeds"], expected["seeds"]) + + +def test_without_an_open_pack_the_uploads_are_plain_copies(): + plain = _meta(pack_rng_uploads=False) + _populate(plain) + expected = _rng_state(plain) + + packed = _meta(pack_rng_uploads=True) + _populate(packed) + + for name, tensor in _rng_state(packed).items(): + assert torch.equal(tensor, expected[name]), name + + +def test_packing_is_off_by_default(): + assert not SpecMetadata( + max_num_requests=1, max_draft_len=MAX_DRAFT_LEN, max_total_draft_tokens=MAX_DRAFT_LEN + ).pack_rng_uploads diff --git a/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py index 1ea8cdaa58f4..5a56b77e1659 100644 --- a/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py +++ b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py @@ -252,3 +252,119 @@ def test_a_fold_that_declines_leaves_the_sampler_store_path(draft_width): assert fold.rows_stored_calls == [(0, SEQ_SLOTS)] assert len(fold.scatter_calls) == 1 _assert_store_equals(sampler, _reference_store(before, outputs)) + + +def _host_outputs(state) -> dict[str, torch.Tensor]: + state.sampler_event.synchronize() + return dict( + new_tokens=state.host.new_tokens.clone(), + new_tokens_lens=state.host.new_tokens_lens.clone(), + next_draft_tokens=state.host.next_draft_tokens.clone(), + ) + + +def test_store_copies_go_on_the_side_stream_when_the_sampler_stored_the_rows(): + sampler = _sampler() + scheduled, outputs = _step(MAX_DRAFT_LEN) + + state = sampler.sample_async(scheduled, outputs, []) + + assert not sampler.d2h_side_stream_copies + assert state.sampler_event.side_stream_event is not None + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_side_stream_store_copies_match_the_main_stream_copies(draft_width): + scheduled, outputs = _step(draft_width) + on_main = _host_outputs(_sampler().sample_async(scheduled, outputs, [])) + sampler = _sampler() + sampler.d2h_side_stream_copies = True + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is not None + on_side = _host_outputs(state) + for name, tensor in on_main.items(): + assert torch.equal(on_side[name], tensor), name + + +def test_a_fold_that_stored_the_rows_keeps_the_copies_on_the_main_stream(): + sampler = _sampler() + sampler.bind_spec_step_fold(_fold(rows_stored=True)) + sampler.d2h_side_stream_copies = True + scheduled, outputs = _step(MAX_DRAFT_LEN) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is None + + +def test_a_fold_that_awaits_the_store_copies_moves_them_to_the_side_stream(): + """A fold with store_copies_done: the copies go on the side stream although its forward stored the rows, and + the event is recorded after them.""" + scheduled, outputs = _step(MAX_DRAFT_LEN) + fold = _fold(rows_stored=True) + fold.store_copies_done = torch.cuda.Event(external=True) + fold.store_copies_done.record() + sampler = _sampler() + sampler.bind_spec_step_fold(fold) + sampler.d2h_side_stream_copies = True + expected = _store_snapshot(sampler) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is not None + state.sampler_event.synchronize() + fold.store_copies_done.synchronize() + host = _host_outputs(state) + for name in ("new_tokens", "new_tokens_lens", "next_draft_tokens"): + assert torch.equal(host[name], expected[name].cpu()), name + + +def test_a_fold_that_awaits_the_store_copies_keeps_them_on_the_main_stream_when_side_copies_are_off(): + scheduled, outputs = _step(MAX_DRAFT_LEN) + fold = _fold(rows_stored=True) + fold.store_copies_done = torch.cuda.Event(external=True) + sampler = _sampler() + sampler.bind_spec_step_fold(fold) + + state = sampler.sample_async(scheduled, outputs, []) + + assert state.sampler_event.side_stream_event is None + + +def test_static_slot_ids_are_off_by_default(): + assert not _sampler().static_slot_ids + + +@pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) +def test_static_slot_ids_store_the_same_rows(draft_width): + scheduled, outputs = _step(draft_width) + plain = _sampler() + expected = _reference_store(_store_snapshot(plain), outputs) + sampler = _sampler() + sampler.static_slot_ids = True + + sampler.sample_async(scheduled, outputs, []) + sampler.sample_async(scheduled, outputs, []) # second step reuses the uploaded ids + + _assert_store_equals(sampler, expected) + + +def test_static_slot_ids_follow_a_changed_batch(): + sampler = _sampler() + sampler.static_slot_ids = True + + first_view = sampler._device_slot_ids([4, 1]) + first = first_view.clone() + # Unchanged ids reuse the sampler's buffer without a new upload. The returned tensor is a view of that buffer, + # which the next call rewrites, so snapshot it before changing the batch. + reused_view = sampler._device_slot_ids([4, 1]) + assert reused_view.data_ptr() == first_view.data_ptr() + reused = reused_view.clone() + changed = sampler._device_slot_ids([2, 5, 0]).clone() + + torch.cuda.synchronize() + assert first.tolist() == [4, 1] + assert reused.tolist() == [4, 1] + assert changed.tolist() == [2, 5, 0] From fc05ff385c9ea13ca2f49f1d05685ef9812e7a48 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:03:41 -0700 Subject: [PATCH 2/8] [None][fix] Align Rubin tests with rebased runtime and CI routing Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b300.yml | 2 +- .../test_fallback_fmha_extra_page_sizes.py | 14 +++----------- .../_torch/pyexecutor/test_kv_lens_update_once.py | 1 + .../_torch/pyexecutor/test_runtime_hooks.py | 4 +++- 4 files changed, 8 insertions(+), 13 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 18d596032c05..6b248ec8f3fe 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -142,7 +142,7 @@ l0_b300: # TIMEOUT measured: 3122 cases in 13m54s on one GB300 (the cpu_only files run on the CPU # stage instead), so 30 is 2.2x headroom. - unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm --ignore=unittest/_torch/modeling_v2/models/kimi_k3_vl --ignore=unittest/_torch/modeling_v2/models/minimax_m3 TIMEOUT (30) - - unittest/_torch/modeling_v2/models/minimax_m3 --import-mode=importlib --ignore-glob=unittest/_torch/modeling_v2/models/minimax_m3/**/comm/* --ignore-glob=unittest/_torch/modeling_v2/models/minimax_m3/targets/sm107* TIMEOUT (30) + - unittest/_torch/modeling_v2/models/minimax_m3 --import-mode=importlib --ignore-glob=unittest/_torch/modeling_v2/models/minimax_m3/**/comm/* --ignore-glob=unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__* TIMEOUT (30) # K3 single-GPU kernel and wrapper checks. - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_cluster_waits.py - unittest/_torch/modeling_v2/models/kimi_k3_vl/test_k3_ctm_gemv.py diff --git a/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py index 479d5f279cfa..221b1a6a8082 100644 --- a/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py +++ b/tests/unittest/_torch/attention/test_fallback_fmha_extra_page_sizes.py @@ -1,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""FallbackFmha's refusal of paged prefill at a KV manager's TRTLLM-Gen extra page sizes: kept on every SM but SM107, -where no TRTLLM-Gen FMHA exists and the thop paged-context kernel serves those pages (its SM107 P128 oracle is -test_modeling_v2_m3_draft_attention_paged_context.py). CPU only: the check runs before any kernel.""" +"""Manager-specific extra page sizes stay on the TRTLLM-Gen route on every SM.""" from types import SimpleNamespace @@ -37,19 +35,13 @@ def _context_batch(tokens_per_block: int = 128, paged: bool = True): ) -@pytest.mark.parametrize("sm", [90, 100, 103, 120]) -def test_extra_page_paged_prefill_is_refused_off_sm107(monkeypatch, sm) -> None: +@pytest.mark.parametrize("sm", [90, 100, 103, 107, 120]) +def test_extra_page_paged_prefill_is_refused(monkeypatch, sm) -> None: monkeypatch.setattr(fallback, "get_sm_version", lambda: sm) with pytest.raises(RuntimeError, match="cannot serve paged prefill with 128-token KV pages"): _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT) -def test_extra_page_paged_prefill_is_served_on_sm107(monkeypatch) -> None: - monkeypatch.setattr(fallback, "get_sm_version", lambda: 107) - with pytest.raises(_PastRefusal): - _fmha()._is_supported(None, None, None, _context_batch(), None, phase=FmhaPhase.CONTEXT) - - @pytest.mark.parametrize("sm", [100, 107]) def test_other_page_sizes_and_unpaged_prefill_pass_the_check(monkeypatch, sm) -> None: monkeypatch.setattr(fallback, "get_sm_version", lambda: sm) diff --git a/tests/unittest/_torch/pyexecutor/test_kv_lens_update_once.py b/tests/unittest/_torch/pyexecutor/test_kv_lens_update_once.py index f404a80713e1..fb40d61f1c27 100644 --- a/tests/unittest/_torch/pyexecutor/test_kv_lens_update_once.py +++ b/tests/unittest/_torch/pyexecutor/test_kv_lens_update_once.py @@ -39,6 +39,7 @@ def _preparer(once: bool, disable_overlap: bool = False) -> DecoderInputPreparer previous_kv_lens_offsets_cuda=torch.full((16,), -3, dtype=torch.int32), ) r.mapping = types.SimpleNamespace(has_cp_helix=lambda: False) + r.decode_prologue = None r.guided_decoder = None return r diff --git a/tests/unittest/_torch/pyexecutor/test_runtime_hooks.py b/tests/unittest/_torch/pyexecutor/test_runtime_hooks.py index 21abc91bb27a..b1e149b92601 100644 --- a/tests/unittest/_torch/pyexecutor/test_runtime_hooks.py +++ b/tests/unittest/_torch/pyexecutor/test_runtime_hooks.py @@ -65,7 +65,9 @@ def test_a_model_without_hooks_gets_the_all_off_default(model): def test_every_field_defaults_to_off(): for field in dataclasses.fields(RuntimeHooks): value = getattr(DEFAULT_RUNTIME_HOOKS, field.name) - if dataclasses.is_dataclass(value): + if field.name == "first_content_emit_iters": + assert value == 1 + elif dataclasses.is_dataclass(value): assert all( getattr(value, item.name) in (False, None) for item in dataclasses.fields(value) ) From 9aeee2c6fc2b9e51d61ded7648d4cd9214441ddd Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:05:37 -0700 Subject: [PATCH 3/8] [None][fix] Drop unrelated GLM import workarounds Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../dsa/test_modeling_v2_hist_topk_op_matrix.py | 1 - .../spec/test_modeling_v2_draft_ehproj_table.py | 1 - .../spec/test_modeling_v2_draft_sample_glue.py | 1 - 3 files changed, 3 deletions(-) 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__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py index 980cd15a09ef..d3ac1acda903 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__tp16/dsa/test_modeling_v2_hist_topk_op_matrix.py @@ -7,7 +7,6 @@ need the 4-node launch its docstring describes. """ -__extra_import_path__ = ["."] import _glm_rank_job import pytest import torch 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__tp16/spec/test_modeling_v2_draft_ehproj_table.py index af00750b1940..b536ec7d560e 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__tp16/spec/test_modeling_v2_draft_ehproj_table.py @@ -7,7 +7,6 @@ 16-rank run over four nodes. """ -__extra_import_path__ = ["."] import _draft_step_rank_job import pytest 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__tp16/spec/test_modeling_v2_draft_sample_glue.py index d784f9480daf..faf00183c753 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__tp16/spec/test_modeling_v2_draft_sample_glue.py @@ -7,7 +7,6 @@ 16-rank run over four nodes. """ -__extra_import_path__ = ["."] import _draft_step_rank_job import pytest From ef86dedbf35bb586fe0c2952489b0a9c2bd95d19 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:39:22 -0700 Subject: [PATCH 4/8] [None][fix] Address Rubin kernel and runtime review findings Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../minimaxM3Sm107Tp2/fullCtxReduce.cu | 16 +- .../minimaxM3Sm107Tp2/sampledAccept.cu | 34 +++- .../minimaxM3Sm107Tp2/sparseDecodeLc2.cu | 20 +- .../kernels/minimaxM3Sm107Tp4/foldTail.cu | 2 +- .../kernels/minimaxM3Sm107Tp4/foldTail.h | 3 +- .../minimaxM3Sm107Tp4/fullCtxReduce.cu | 16 +- .../minimaxM3Sm107Tp4/sampledAccept.cu | 35 +++- .../minimaxM3Sm107Tp4/sparseDecodeLc4.cu | 20 +- .../kernels/minimaxM3Sm107Tp8/foldTail.cu | 2 +- .../kernels/minimaxM3Sm107Tp8/foldTail.h | 3 +- .../minimaxM3Sm107Tp8/fullCtxReduce.cu | 16 +- .../minimaxM3Sm107Tp8/sampledAccept.cu | 35 +++- .../minimaxM3Sm107Tp8/sparseDecodeLc8m.cu | 25 ++- .../minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp | 16 +- .../minimaxM3Sm107Tp2FullCtxReduceOp.cpp | 5 + .../minimaxM3Sm107Tp2SampledAcceptOp.cpp | 2 + .../minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp | 10 +- .../minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp | 16 +- .../minimaxM3Sm107Tp4FoldTailParams.h | 13 +- .../minimaxM3Sm107Tp4FullCtxReduceOp.cpp | 5 + .../minimaxM3Sm107Tp4SampledAcceptOp.cpp | 2 + .../minimaxM3Sm107Tp4SparseDecodeLcOp.cpp | 8 + .../minimaxM3Sm107Tp8DraftStepOp.cpp | 8 +- .../minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp | 3 +- .../minimaxM3Sm107Tp8FoldTailParams.h | 13 +- .../minimaxM3Sm107Tp8FullCtxReduceOp.cpp | 5 + .../minimaxM3Sm107Tp8SampledAcceptOp.cpp | 2 + .../minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp | 13 ++ .../ops/attention/m3_sparse_decode_lc.py | 1 + .../ops/m3_lowm_splitk.py | 10 +- .../ops/m3_moe/op.py | 7 +- .../ops/triton_sparse_decode.py | 6 +- .../ops/triton_sparse_decode_mq.py | 4 +- .../ops/attention/m3_sparse_decode_lc.py | 20 +- .../ops/m3_lowm_splitk.py | 10 +- .../ops/m3_moe/op.py | 7 +- .../ops/m3_moe/op16.py | 4 +- .../ops/triton_sparse_decode.py | 6 +- .../ops/triton_sparse_decode_mq.py | 4 +- .../ops/attention/m3_sparse_decode_lc8.py | 1 + .../ops/m3_lowm_splitk.py | 10 +- .../ops/m3_moe/op.py | 7 +- .../ops/triton_sparse_decode.py | 6 +- .../ops/triton_sparse_decode_mq.py | 4 +- .../sparse/minimax_m3/cache_manager.py | 8 +- .../minimax_m3/embedding_firstnorm/op.py | 8 +- .../cute_dsl_kernels/minimax_m3/moe/op.py | 48 ++++- .../cute_dsl_kernels/minimax_m3/moe/op16.py | 3 + .../cute_dsl_kernels/minimax_m3/mx_gemv/op.py | 36 ++-- .../_torch/pyexecutor/cuda_graph_runner.py | 13 +- tensorrt_llm/_torch/pyexecutor/step_inputs.py | 15 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_lists/test-db/l0_gr100.yml | 1 + .../msa/test_minimax_m3_kv_size_estimate.py | 44 +++++ .../test_modeling_v2_m3_silu_and_mul.py | 3 +- ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 3 +- .../test_modeling_v2_m3_draft_attention.py | 3 +- ...ing_v2_m3_draft_attention_paged_context.py | 24 ++- .../test_modeling_v2_m3_drafter_fullctx.py | 3 +- ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 3 +- ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 3 +- .../test_modeling_v2_m3_fullctx_decode.py | 3 +- .../test_modeling_v2_m3_index_score.py | 5 +- .../test_modeling_v2_m3_msa_attention.py | 3 +- .../attention/test_modeling_v2_m3_producer.py | 3 +- .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 3 +- .../test_modeling_v2_m3_qkv_insert.py | 3 +- .../test_modeling_v2_m3_select_blocks.py | 3 +- .../test_modeling_v2_m3_select_long.py | 3 +- .../test_modeling_v2_m3_select_long_dsl.py | 3 +- .../test_modeling_v2_m3_tokmap_guard.py | 3 +- .../test_modeling_v2_msa_kv_update.py | 3 +- .../comm/_m3_allreduce_op_matrix.py | 4 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 4 +- .../comm/_m3_moe_push_op_matrix.py | 4 +- ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 3 +- .../test_modeling_v2_m3_embedding.py | 3 +- .../gemm/test_modeling_v2_m3_bf16_linear.py | 3 +- .../gemm/test_modeling_v2_m3_dense_seam.py | 3 +- .../gemm/test_modeling_v2_m3_dense_swiglu.py | 3 +- .../gemm/test_modeling_v2_m3_draft_lowm.py | 3 +- .../gemm/test_modeling_v2_m3_fp32_linear.py | 3 +- .../gemm/test_modeling_v2_m3_mx_gemv.py | 3 +- .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 3 +- .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 3 +- .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 3 +- .../gemm/test_modeling_v2_m3_pinned_gemm.py | 3 +- .../gemm/test_modeling_v2_m3_ugpu_head.py | 3 +- .../moe/test_modeling_v2_m3_moe_engine.py | 3 +- .../moe/test_modeling_v2_m3_moe_epoch_wrap.py | 3 +- .../moe/test_modeling_v2_m3_moe_sh_take.py | 3 +- ...test_modeling_v2_m3_embedding_firstnorm.py | 3 +- .../spec/test_modeling_v2_m3_capture_norm.py | 3 +- .../spec/test_modeling_v2_m3_chunk_norm.py | 3 +- ...test_modeling_v2_m3_draft_ar_push_words.py | 9 +- .../spec/test_modeling_v2_m3_draft_input.py | 3 +- .../spec/test_modeling_v2_m3_draft_silu.py | 3 +- .../spec/test_modeling_v2_m3_draft_step.py | 3 +- .../spec/test_modeling_v2_m3_fold_tail.py | 3 +- .../spec/test_modeling_v2_m3_force_vec.py | 3 +- .../spec/test_modeling_v2_m3_force_vec_req.py | 3 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 3 +- .../test_modeling_v2_m3_sampled_accept.py | 3 +- .../spec/test_modeling_v2_m3_shard_accept.py | 3 +- .../spec/test_modeling_v2_m3_slot_scatter.py | 3 +- .../spec/test_modeling_v2_m3_spec_offsets.py | 3 +- .../spec/test_modeling_v2_m3_step_gather.py | 3 +- .../spec/test_modeling_v2_m3_step_prologue.py | 3 +- .../test_modeling_v2_m3_rank_agreement.py | 16 +- .../test_modeling_v2_m3_silu_and_mul.py | 3 +- .../test_modeling_v2_m3_dense_decode.py | 3 +- ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 3 +- .../test_modeling_v2_m3_draft_attention.py | 3 +- .../test_modeling_v2_m3_drafter_fullctx.py | 3 +- ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 3 +- ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 3 +- .../test_modeling_v2_m3_fullctx_decode.py | 3 +- .../test_modeling_v2_m3_index_score.py | 3 +- .../test_modeling_v2_m3_index_score_t2sc.py | 5 +- .../test_modeling_v2_m3_msa_attention.py | 3 +- .../attention/test_modeling_v2_m3_producer.py | 3 +- .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 3 +- .../test_modeling_v2_m3_qkv_insert.py | 3 +- .../test_modeling_v2_m3_select_blocks.py | 3 +- .../test_modeling_v2_m3_select_long.py | 3 +- .../test_modeling_v2_m3_select_long_dsl.py | 3 +- .../test_modeling_v2_m3_sparse_decode.py | 3 +- .../test_modeling_v2_m3_sparse_decode_lc.py | 3 +- .../test_modeling_v2_m3_tokmap_guard.py | 3 +- .../test_modeling_v2_msa_kv_update.py | 3 +- .../comm/_m3_allreduce_op_matrix.py | 4 +- .../comm/_m3_capture_in_boundary_op_matrix.py | 4 +- .../comm/_m3_draft_ar_norm_op_matrix.py | 4 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 4 +- ...test_modeling_v2_m3_capture_in_boundary.py | 3 +- .../comm/test_modeling_v2_m3_draft_ar_norm.py | 3 +- ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 3 +- .../test_modeling_v2_m3_embedding.py | 3 +- .../gemm/test_modeling_v2_m3_bf16_linear.py | 3 +- .../gemm/test_modeling_v2_m3_dense_swiglu.py | 3 +- .../gemm/test_modeling_v2_m3_draft_lowm.py | 3 +- .../gemm/test_modeling_v2_m3_fp32_linear.py | 3 +- .../gemm/test_modeling_v2_m3_mx_gemv.py | 3 +- .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 3 +- .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 3 +- .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 3 +- .../test_modeling_v2_m3_mxfp8_pinned_gemm.py | 3 +- ...st_modeling_v2_m3_mxfp8_pinned_quantize.py | 3 +- .../gemm/test_modeling_v2_m3_target_head.py | 3 +- .../moe/test_modeling_v2_m3_moe16_engine.py | 3 +- .../moe/test_modeling_v2_m3_moe_engine.py | 3 +- ...test_modeling_v2_m3_embedding_firstnorm.py | 3 +- .../spec/test_modeling_v2_m3_capture_norm.py | 3 +- .../spec/test_modeling_v2_m3_chunk_norm.py | 3 +- .../spec/test_modeling_v2_m3_draft_input.py | 3 +- .../spec/test_modeling_v2_m3_draft_silu.py | 3 +- .../spec/test_modeling_v2_m3_draft_step.py | 3 +- .../test_modeling_v2_m3_exact_numerics.py | 5 +- .../spec/test_modeling_v2_m3_fold_tail.py | 3 +- .../spec/test_modeling_v2_m3_force_vec.py | 3 +- .../spec/test_modeling_v2_m3_force_vec_req.py | 3 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 3 +- .../test_modeling_v2_m3_sampled_accept.py | 3 +- .../spec/test_modeling_v2_m3_shard_accept.py | 3 +- .../spec/test_modeling_v2_m3_slot_scatter.py | 3 +- .../spec/test_modeling_v2_m3_spec_offsets.py | 3 +- .../spec/test_modeling_v2_m3_step_gather.py | 3 +- .../spec/test_modeling_v2_m3_step_prologue.py | 3 +- .../test_modeling_v2_m3_prebuild.py | 3 +- .../test_modeling_v2_m3_rank_agreement.py | 16 +- .../test_modeling_v2_m3_silu_and_mul.py | 3 +- ...t_modeling_v2_m3_dense_qk_norm_rope_fp8.py | 3 +- .../test_modeling_v2_m3_draft_attention.py | 3 +- .../test_modeling_v2_m3_drafter_fullctx.py | 3 +- ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 3 +- ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 3 +- .../test_modeling_v2_m3_fullctx_decode.py | 3 +- .../test_modeling_v2_m3_index_score.py | 5 +- .../test_modeling_v2_m3_msa_attention.py | 3 +- .../attention/test_modeling_v2_m3_producer.py | 3 +- .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 3 +- .../test_modeling_v2_m3_qkv_insert.py | 3 +- .../test_modeling_v2_m3_select_blocks.py | 3 +- .../test_modeling_v2_m3_select_long.py | 3 +- .../test_modeling_v2_m3_select_long_dsl.py | 3 +- .../test_modeling_v2_m3_sparse_decode_lc8.py | 5 +- .../test_modeling_v2_m3_tokmap_guard.py | 3 +- .../test_modeling_v2_msa_kv_update.py | 3 +- .../comm/_m3_allreduce_op_matrix.py | 4 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 4 +- .../comm/_m3_moe_push_op_matrix.py | 12 +- ...eling_v2_m3_mc_allreduce_norm_op_matrix.py | 3 +- .../test_modeling_v2_m3_embedding.py | 3 +- .../gemm/test_modeling_v2_m3_bf16_linear.py | 3 +- .../gemm/test_modeling_v2_m3_dense_swiglu.py | 3 +- .../gemm/test_modeling_v2_m3_draft_lowm.py | 3 +- .../gemm/test_modeling_v2_m3_fp32_linear.py | 3 +- .../gemm/test_modeling_v2_m3_mx_gemv.py | 3 +- .../gemm/test_modeling_v2_m3_mx_gemv_bf16.py | 3 +- .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 3 +- .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 3 +- .../gemm/test_modeling_v2_m3_pinned_gemm.py | 3 +- .../test_modeling_v2_m3_swiglu_oai_bf16x2.py | 3 +- .../moe/test_modeling_v2_m3_moe_engine.py | 3 +- ...test_modeling_v2_m3_embedding_firstnorm.py | 3 +- .../spec/test_modeling_v2_m3_capture_norm.py | 3 +- .../test_modeling_v2_m3_capture_normed.py | 6 + .../spec/test_modeling_v2_m3_chunk_norm.py | 3 +- .../spec/test_modeling_v2_m3_draft_input.py | 3 +- .../spec/test_modeling_v2_m3_draft_silu.py | 3 +- .../spec/test_modeling_v2_m3_draft_step.py | 3 +- .../spec/test_modeling_v2_m3_fold_tail.py | 3 +- .../spec/test_modeling_v2_m3_force_vec.py | 3 +- .../spec/test_modeling_v2_m3_force_vec_req.py | 3 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 3 +- .../test_modeling_v2_m3_sampled_accept.py | 3 +- .../spec/test_modeling_v2_m3_shard_accept.py | 3 +- .../spec/test_modeling_v2_m3_slot_scatter.py | 3 +- .../spec/test_modeling_v2_m3_spec_offsets.py | 3 +- .../spec/test_modeling_v2_m3_step_gather.py | 3 +- .../spec/test_modeling_v2_m3_step_prologue.py | 3 +- .../test_modeling_v2_m3_rank_agreement.py | 16 +- .../test_shared_kernel_contracts.py | 180 ++++++++++++++++++ .../pyexecutor/test_cuda_graph_runner.py | 48 +++++ .../_torch/pyexecutor/test_step_inputs.py | 54 ++++++ .../test_spec_sampler_step_fold.py | 9 +- 226 files changed, 1202 insertions(+), 322 deletions(-) create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/test_shared_kernel_contracts.py diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu index 236a13e66d1f..270b6ca53865 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu @@ -29,6 +29,7 @@ #include #include +#include #include TRTLLM_NAMESPACE_BEGIN @@ -56,14 +57,14 @@ __device__ int activeSplits(Params const& p, int b, int row) { int const token = row / p.headsPerKv; int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; - int const end = min(base + p.qTokensPerCta, p.seqLenQ); + int const end = base + min(p.qTokensPerCta, p.seqLenQ - base); visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); } - int const tiles = (visible + p.tileKv - 1) / p.tileKv; - int const tilesPerGroup = p.splits * p.instsKv; - int const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; - int const local = max(groups * p.instsKv, p.instsKv); - return (tiles + local - 1) / local; + int64_t const tiles = (static_cast(visible) + p.tileKv - 1) / p.tileKv; + int64_t const tilesPerGroup = static_cast(p.splits) * p.instsKv; + int64_t const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; + int64_t const local = (groups > 1 ? groups : 1) * p.instsKv; + return static_cast((tiles + local - 1) / local); } template @@ -199,7 +200,8 @@ __global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) { - if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1) + if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1 || params.qTokensPerCta <= 0 + || params.tileKv <= 0 || params.instsKv <= 0 || params.instsKv > std::numeric_limits::max() / kMaxSplits) { return cudaErrorInvalidValue; } diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu index 878282b18839..0e1a4464b5cc 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sampledAccept.cu @@ -19,6 +19,9 @@ #include "tensorrt_llm/common/assert.h" +#include +#include + #include #include @@ -539,7 +542,13 @@ __global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const template cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) { - static int max_clusters = 0; + int device; + cudaError_t const deviceStatus = cudaGetDevice(&device); + if (deviceStatus != cudaSuccess) + return deviceStatus; + static std::mutex configMutex; + static std::unordered_map deviceClusters; + int max_clusters = 0; cudaLaunchConfig_t cfg = {}; cfg.gridDim = dim3(a.rows * kSplits); cfg.blockDim = dim3(kThreads); @@ -554,18 +563,23 @@ cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; cfg.attrs = attrs; cfg.numAttrs = 2; - if (max_clusters == 0) { - cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); - if (err == cudaSuccess) - { - err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); - } - if (err != cudaSuccess) + std::lock_guard lock(configMutex); + max_clusters = deviceClusters[device]; + if (max_clusters == 0) { - return err; + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } } + deviceClusters[device] = max_clusters; } TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu index d96f51d5d150..ca0822613211 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.cu @@ -29,6 +29,9 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h" +#include +#include + #include #include #include @@ -588,10 +591,21 @@ cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t { // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs // every variant eagerly first). - static cudaError_t const configured = configure(); - if (configured != cudaSuccess) + int device; + cudaError_t const deviceStatus = cudaGetDevice(&device); + if (deviceStatus != cudaSuccess) + return deviceStatus; + static std::mutex configMutex; + static std::unordered_set configuredDevices; { - return configured; + std::lock_guard lock(configMutex); + if (configuredDevices.count(device) == 0) + { + cudaError_t const configured = configure(); + if (configured != cudaSuccess) + return configured; + configuredDevices.insert(device); + } } auto kernel = m3_lc_attn_kernel; cudaLaunchConfig_t cfg = {}; diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu index e6a4f0da2039..309db8a47f3d 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu @@ -32,7 +32,7 @@ namespace __global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) { int const i = blockIdx.x; - if (i >= *p.ctrl) + if (i >= *p.ctrl || i >= p.maxRows) { return; } diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h index 54a58b0569e0..92b658cdc193 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h @@ -61,6 +61,7 @@ struct FoldTailParams int32_t* posOff; // [rows * T] int32_t* kvOff; // [rows] int32_t tokensPerStep; + int32_t maxRows; // validated capacity of the source and staging buffers }; //! One CTA of 32 lanes per row, `rows` CTAs; the first ctrl[0] rows write. Widths and tokensPerStep must be <= 32. @@ -115,7 +116,7 @@ __device__ __forceinline__ void foldTailRow(FoldTailParams const& p, int i, int //! The whole tail from one warp of a running kernel (lane = threadIdx.x % 32): rows i < *p.ctrl in order. __device__ __forceinline__ void foldTailWarp(FoldTailParams const& p, int lane) { - int const rows = *p.ctrl; + int const rows = min(*p.ctrl, p.maxRows); for (int i = 0; i < rows; ++i) { foldTailRow(p, i, lane); diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu index 2937d461854c..ffab1d36db50 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu @@ -29,6 +29,7 @@ #include #include +#include #include TRTLLM_NAMESPACE_BEGIN @@ -56,14 +57,14 @@ __device__ int activeSplits(Params const& p, int b, int row) { int const token = row / p.headsPerKv; int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; - int const end = min(base + p.qTokensPerCta, p.seqLenQ); + int const end = base + min(p.qTokensPerCta, p.seqLenQ - base); visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); } - int const tiles = (visible + p.tileKv - 1) / p.tileKv; - int const tilesPerGroup = p.splits * p.instsKv; - int const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; - int const local = max(groups * p.instsKv, p.instsKv); - return (tiles + local - 1) / local; + int64_t const tiles = (static_cast(visible) + p.tileKv - 1) / p.tileKv; + int64_t const tilesPerGroup = static_cast(p.splits) * p.instsKv; + int64_t const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; + int64_t const local = (groups > 1 ? groups : 1) * p.instsKv; + return static_cast((tiles + local - 1) / local); } template @@ -244,7 +245,8 @@ __global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) { - if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1) + if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1 || params.qTokensPerCta <= 0 + || params.tileKv <= 0 || params.instsKv <= 0 || params.instsKv > std::numeric_limits::max() / kMaxSplits) { return cudaErrorInvalidValue; } diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu index 3d041421a92b..ce95ebb5b71e 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu @@ -19,6 +19,9 @@ #include "tensorrt_llm/common/assert.h" +#include +#include + #include #include @@ -872,7 +875,13 @@ __global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const template cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) { - static int max_clusters = 0; + int device; + cudaError_t const deviceStatus = cudaGetDevice(&device); + if (deviceStatus != cudaSuccess) + return deviceStatus; + static std::mutex configMutex; + static std::unordered_map deviceClusters; + int max_clusters = 0; cudaLaunchConfig_t cfg = {}; cfg.gridDim = dim3(a.rows * kSplits); cfg.blockDim = dim3(kThreads); @@ -887,18 +896,24 @@ cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; cfg.attrs = attrs; cfg.numAttrs = 2; - if (max_clusters == 0) { - cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); - if (err == cudaSuccess) - { - err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); - } - if (err != cudaSuccess) + std::lock_guard lock(configMutex); + max_clusters = deviceClusters[device]; + if (max_clusters == 0) { - return err; + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters( + &max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } } + deviceClusters[device] = max_clusters; } TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu index e26804fb47fa..e83ec8131197 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.cu @@ -33,6 +33,9 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h" +#include +#include + #include #include #include @@ -679,10 +682,21 @@ cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t { // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs // every variant eagerly first). - static cudaError_t const configured = configure(); - if (configured != cudaSuccess) + int device; + cudaError_t const deviceStatus = cudaGetDevice(&device); + if (deviceStatus != cudaSuccess) + return deviceStatus; + static std::mutex configMutex; + static std::unordered_set configuredDevices; { - return configured; + std::lock_guard lock(configMutex); + if (configuredDevices.count(device) == 0) + { + cudaError_t const configured = configure(); + if (configured != cudaSuccess) + return configured; + configuredDevices.insert(device); + } } auto kernel = m3_lc_attn_kernel; cudaLaunchConfig_t cfg = {}; diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu index be98efebeb5d..e80a7810b9ba 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu @@ -32,7 +32,7 @@ namespace __global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) { int const i = blockIdx.x; - if (i >= *p.ctrl) + if (i >= *p.ctrl || i >= p.maxRows) { return; } diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h index 86572b786a81..2c565e6c0c7f 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h @@ -61,6 +61,7 @@ struct FoldTailParams int32_t* posOff; // [rows * T] int32_t* kvOff; // [rows] int32_t tokensPerStep; + int32_t maxRows; // validated capacity of the source and staging buffers }; //! One CTA of 32 lanes per row, `rows` CTAs; the first ctrl[0] rows write. Widths and tokensPerStep must be <= 32. @@ -115,7 +116,7 @@ __device__ __forceinline__ void foldTailRow(FoldTailParams const& p, int i, int //! The whole tail from one warp of a running kernel (lane = threadIdx.x % 32): rows i < *p.ctrl in order. __device__ __forceinline__ void foldTailWarp(FoldTailParams const& p, int lane) { - int const rows = *p.ctrl; + int const rows = min(*p.ctrl, p.maxRows); for (int i = 0; i < rows; ++i) { foldTailRow(p, i, lane); diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu index e7505fa48e13..ea33729aed7a 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu @@ -29,6 +29,7 @@ #include #include +#include #include TRTLLM_NAMESPACE_BEGIN @@ -56,14 +57,14 @@ __device__ int activeSplits(Params const& p, int b, int row) { int const token = row / p.headsPerKv; int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; - int const end = min(base + p.qTokensPerCta, p.seqLenQ); + int const end = base + min(p.qTokensPerCta, p.seqLenQ - base); visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); } - int const tiles = (visible + p.tileKv - 1) / p.tileKv; - int const tilesPerGroup = p.splits * p.instsKv; - int const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; - int const local = max(groups * p.instsKv, p.instsKv); - return (tiles + local - 1) / local; + int64_t const tiles = (static_cast(visible) + p.tileKv - 1) / p.tileKv; + int64_t const tilesPerGroup = static_cast(p.splits) * p.instsKv; + int64_t const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; + int64_t const local = (groups > 1 ? groups : 1) * p.instsKv; + return static_cast((tiles + local - 1) / local); } template @@ -244,7 +245,8 @@ __global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) { - if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1) + if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1 || params.qTokensPerCta <= 0 + || params.tileKv <= 0 || params.instsKv <= 0 || params.instsKv > std::numeric_limits::max() / kMaxSplits) { return cudaErrorInvalidValue; } diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu index 6e5f943b2438..1229b05b029f 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu @@ -19,6 +19,9 @@ #include "tensorrt_llm/common/assert.h" +#include +#include + #include #include @@ -874,7 +877,13 @@ __global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const template cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) { - static int max_clusters = 0; + int device; + cudaError_t const deviceStatus = cudaGetDevice(&device); + if (deviceStatus != cudaSuccess) + return deviceStatus; + static std::mutex configMutex; + static std::unordered_map deviceClusters; + int max_clusters = 0; cudaLaunchConfig_t cfg = {}; cfg.gridDim = dim3(a.rows * kSplits); cfg.blockDim = dim3(kThreads); @@ -889,18 +898,24 @@ cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; cfg.attrs = attrs; cfg.numAttrs = 2; - if (max_clusters == 0) { - cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); - if (err == cudaSuccess) - { - err = cudaOccupancyMaxActiveClusters(&max_clusters, sampled_accept_kernel, &cfg); - } - if (err != cudaSuccess) + std::lock_guard lock(configMutex); + max_clusters = deviceClusters[device]; + if (max_clusters == 0) { - return err; + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters( + &max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } } + deviceClusters[device] = max_clusters; } TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu index d51138cc15d2..4b680941fc88 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.cu @@ -39,6 +39,8 @@ #include #include #include +#include +#include TRTLLM_NAMESPACE_BEGIN @@ -689,12 +691,25 @@ cudaError_t configure() template cudaError_t launch(Params const& prm, int num_tokens, bool use_pdl, cudaStream_t stream) { - // Kernel attributes are set once per process, on the first launch (outside CUDA-graph capture: warmup runs - // every variant eagerly first). - static cudaError_t const configured = configure(); - if (configured != cudaSuccess) + int device = 0; + cudaError_t const deviceStatus = cudaGetDevice(&device); + if (deviceStatus != cudaSuccess) { - return configured; + return deviceStatus; + } + static std::mutex configMutex; + static std::unordered_set configuredDevices; + { + std::lock_guard const lock(configMutex); + if (configuredDevices.count(device) == 0) + { + cudaError_t const configured = configure(); + if (configured != cudaSuccess) + { + return configured; + } + configuredDevices.insert(device); + } } auto kernel = m3_tp8_lc8_attn_kernel; cudaLaunchConfig_t cfg = {}; diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp index 99c182799cab..1fbe0b43b07d 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DenseQkNormRopeFp8Op.cpp @@ -39,10 +39,14 @@ torch::Tensor minimaxM3Sm107Tp2DenseQkNormRopeFp8(torch::Tensor const& qkv, int6 TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.device() == qkv.device() && kWeight.device() == qkv.device() && qWeight.is_contiguous() + && kWeight.is_contiguous()); + TORCH_CHECK(positionIds.device() == qkv.device() && positionIds.is_contiguous()); TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); - TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + TORCH_CHECK(!isNeox || (rotaryDim >= 8 && (rotaryDim & (rotaryDim - 1)) == 0), + "unsupported rotary_dim for the NeoX pairing"); at::cuda::CUDAGuard guard(qkv.device()); auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); dqk::launch(qkv.data_ptr(), out.data_ptr(), static_cast(qkv.size(0)), static_cast(numHeadsQ), @@ -64,15 +68,20 @@ torch::Tensor minimaxM3Sm107Tp2DenseQkNormRopeFp8Kv(torch::Tensor const& qkv, in TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); TORCH_CHECK(numHeadsK == numHeadsV, "the pool keeps one head count for K and V"); TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.device() == qkv.device() && kWeight.device() == qkv.device() && qWeight.is_contiguous() + && kWeight.is_contiguous()); + TORCH_CHECK(positionIds.device() == qkv.device() && positionIds.is_contiguous()); TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); - TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + TORCH_CHECK(!isNeox || (rotaryDim >= 8 && (rotaryDim & (rotaryDim - 1)) == 0), + "unsupported rotary_dim for the NeoX pairing"); TORCH_CHECK(kvBuffers.is_cuda() && kvBuffers.scalar_type() == at::kFloat8_e4m3fn && kvBuffers.dim() == 5); TORCH_CHECK(kvBuffers.size(1) == 2 && kvBuffers.size(2) == numHeadsK && kvBuffers.size(4) == dqk::kHeadDim && kvBuffers.stride(4) == 1, "kv_buffers must be the HND pool view [pages, 2, kv heads, tokens per block, 128]"); TORCH_CHECK(outCacheLoc.is_cuda() && outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous()); + TORCH_CHECK(kvBuffers.device() == qkv.device() && outCacheLoc.device() == qkv.device()); TORCH_CHECK(numLiveTokens >= 0 && numLiveTokens <= qkv.size(0) && numLiveTokens <= outCacheLoc.numel()); at::cuda::CUDAGuard guard(qkv.device()); auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); @@ -87,7 +96,8 @@ torch::Tensor minimaxM3Sm107Tp2DenseQkNormRopeFp8Kv(torch::Tensor const& qkv, in view.numLiveTokens = static_cast(numLiveTokens); if (qCompact.has_value()) { - TORCH_CHECK(qCompact->is_cuda() && qCompact->is_contiguous() && qCompact->scalar_type() == at::kFloat8_e4m3fn + TORCH_CHECK(qCompact->device() == qkv.device() && qCompact->is_contiguous() + && qCompact->scalar_type() == at::kFloat8_e4m3fn && qCompact->numel() == qkv.size(0) * numHeadsQ * dqk::kHeadDim, "q_compact must be a contiguous FP8 [tokens, q heads, 128] tensor"); view.qCompact = static_cast<__nv_fp8_e4m3*>(qCompact->data_ptr()); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp index 8db408008dc2..58c03d90d7d8 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2FullCtxReduceOp.cpp @@ -20,6 +20,7 @@ #include #include +#include #include TRTLLM_NAMESPACE_BEGIN @@ -41,6 +42,10 @@ void minimaxM3Sm107Tp2FullCtxReduce(torch::Tensor const& partialO, torch::Tensor TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() && partialStats.sizes() == partialO.sizes().slice(0, 4)); TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + constexpr int64_t kMaxInt = std::numeric_limits::max(); + TORCH_CHECK(qTokensPerCta > 0 && qTokensPerCta <= kMaxInt && tileKv > 0 && tileKv <= kMaxInt && instsKv > 0 + && instsKv <= kMaxInt / fcr::kMaxSplits, + "fullctx_reduce: producer geometry must be positive and fit int32"); int const rows = static_cast(partialO.size(3)); TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); int const headsPerKv = rows / static_cast(seqLenQ); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp index 415042d5c082..acf9e7cfd008 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SampledAcceptOp.cpp @@ -47,6 +47,8 @@ void minimaxM3Sm107Tp2SampledAccept(torch::Tensor const& logits, int64_t vocabOf std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, bool prewaitParams, bool earlyTrigger, bool ilpWalk) { + TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.scalar_type() == at::kBFloat16 && logits.stride(1) == 1, + "sampled_accept: logits must be CUDA BF16 [rows, vocab] with unit column stride"); c10::cuda::CUDAGuard const deviceGuard{logits.device()}; sa::Args a; a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp index fb8fc9eaa8ee..c7cd269a1268 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SparseDecodeLc2Op.cpp @@ -17,6 +17,8 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/sparseDecodeLc2.h" +#include + #include #include #include @@ -71,7 +73,7 @@ void minimaxM3Sm107Tp2SparseDecodeLc(torch::Tensor const& q, torch::Tensor const checkBytes(kBlockScale, "k_block_scale"); checkBytes(vBlockScale, "v_block_scale"); TORCH_CHECK(k.dim() == 4 && k.size(2) == 128 && k.size(3) == kHeadDim / 2 && k.stride(3) == 1 - && v.sizes() == k.sizes() && v.strides() == k.strides(), + && k.stride(2) == kHeadDim / 2 && v.sizes() == k.sizes() && v.strides() == k.strides(), "minimax_m3_sm107_tp2_sparse_decode_lc: k and v must be [pages, H, 128, 64] with matching strides"); int64_t const kvHeads = k.size(1); TORCH_CHECK(kBlockScale.dim() == 4 && kBlockScale.size(1) == kvHeads && kBlockScale.size(3) == 8 @@ -124,6 +126,12 @@ void minimaxM3Sm107Tp2SparseDecodeLc(torch::Tensor const& q, torch::Tensor const prm.block_table = blockTable.data_ptr(); prm.bt_stride = blockTable.stride(0); prm.seq_lens = seqLens.data_ptr(); + TORCH_CHECK( + decodeQueryLen > 0 && decodeQueryLen <= std::numeric_limits::max() && q.size(0) % decodeQueryLen == 0, + "sparse_decode: decode_query_len must be positive and divide the query token count"); + TORCH_CHECK( + seqLens.device() == q.device() && seqLens.is_contiguous() && seqLens.numel() >= q.size(0) / decodeQueryLen, + "sparse_decode: seq_lens must be contiguous and hold every request"); prm.dql = static_cast(decodeQueryLen); prm.sm_scale = static_cast(smScale); prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp index d92f2c1cd29b..f1279a907159 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DenseQkNormRopeFp8Op.cpp @@ -39,10 +39,14 @@ torch::Tensor minimaxM3Sm107Tp4DenseQkNormRopeFp8(torch::Tensor const& qkv, int6 TORCH_CHECK(qkv.is_cuda() && qkv.is_contiguous() && qkv.dim() == 2 && qkv.scalar_type() == at::kBFloat16); TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.device() == qkv.device() && kWeight.device() == qkv.device() && qWeight.is_contiguous() + && kWeight.is_contiguous()); + TORCH_CHECK(positionIds.device() == qkv.device() && positionIds.is_contiguous()); TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); - TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + TORCH_CHECK(!isNeox || (rotaryDim >= 8 && (rotaryDim & (rotaryDim - 1)) == 0), + "unsupported rotary_dim for the NeoX pairing"); at::cuda::CUDAGuard guard(qkv.device()); auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); cudaError_t const err @@ -66,15 +70,20 @@ torch::Tensor minimaxM3Sm107Tp4DenseQkNormRopeFp8Kv(torch::Tensor const& qkv, in TORCH_CHECK(qkv.size(1) == (numHeadsQ + numHeadsK + numHeadsV) * dqk::kHeadDim); TORCH_CHECK(numHeadsK == numHeadsV, "the pool keeps one head count for K and V"); TORCH_CHECK(qWeight.scalar_type() == at::kBFloat16 && kWeight.scalar_type() == at::kBFloat16); + TORCH_CHECK(qWeight.device() == qkv.device() && kWeight.device() == qkv.device() && qWeight.is_contiguous() + && kWeight.is_contiguous()); + TORCH_CHECK(positionIds.device() == qkv.device() && positionIds.is_contiguous()); TORCH_CHECK(qWeight.numel() == dqk::kHeadDim && kWeight.numel() == dqk::kHeadDim); TORCH_CHECK(positionIds.scalar_type() == at::kInt && positionIds.numel() == qkv.size(0)); TORCH_CHECK(rotaryDim > 0 && rotaryDim <= dqk::kHeadDim && rotaryDim % 2 == 0, "unsupported rotary_dim"); - TORCH_CHECK(!isNeox || (rotaryDim * 16) % dqk::kHeadDim == 0, "unsupported rotary_dim for the NeoX pairing"); + TORCH_CHECK(!isNeox || (rotaryDim >= 8 && (rotaryDim & (rotaryDim - 1)) == 0), + "unsupported rotary_dim for the NeoX pairing"); TORCH_CHECK(kvBuffers.is_cuda() && kvBuffers.scalar_type() == at::kFloat8_e4m3fn && kvBuffers.dim() == 5); TORCH_CHECK(kvBuffers.size(1) == 2 && kvBuffers.size(2) == numHeadsK && kvBuffers.size(4) == dqk::kHeadDim && kvBuffers.stride(4) == 1, "kv_buffers must be the HND pool view [pages, 2, kv heads, tokens per block, 128]"); TORCH_CHECK(outCacheLoc.is_cuda() && outCacheLoc.scalar_type() == at::kInt && outCacheLoc.is_contiguous()); + TORCH_CHECK(kvBuffers.device() == qkv.device() && outCacheLoc.device() == qkv.device()); TORCH_CHECK(numLiveTokens >= 0 && numLiveTokens <= qkv.size(0) && numLiveTokens <= outCacheLoc.numel()); at::cuda::CUDAGuard guard(qkv.device()); auto out = torch::empty_like(qkv, qkv.options().dtype(torch::kFloat8_e4m3fn)); @@ -89,7 +98,8 @@ torch::Tensor minimaxM3Sm107Tp4DenseQkNormRopeFp8Kv(torch::Tensor const& qkv, in view.numLiveTokens = static_cast(numLiveTokens); if (qCompact.has_value()) { - TORCH_CHECK(qCompact->is_cuda() && qCompact->is_contiguous() && qCompact->scalar_type() == at::kFloat8_e4m3fn + TORCH_CHECK(qCompact->device() == qkv.device() && qCompact->is_contiguous() + && qCompact->scalar_type() == at::kFloat8_e4m3fn && qCompact->numel() == qkv.size(0) * numHeadsQ * dqk::kHeadDim, "q_compact must be a contiguous FP8 [tokens, q heads, 128] tensor"); view.qCompact = static_cast<__nv_fp8_e4m3*>(qCompact->data_ptr()); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h index b73820b6fc54..7c25f9093405 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h @@ -19,6 +19,7 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" +#include #include TRTLLM_NAMESPACE_BEGIN @@ -54,8 +55,11 @@ inline tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams buildParams(torc "minimax_m3_sm107_tp4_fold_tail: ctrl must be an int32 CUDA tensor"); TORCH_CHECK(slots.is_cuda() && slots.scalar_type() == torch::kInt64 && slots.dim() == 1 && slots.stride(0) == 1, "minimax_m3_sm107_tp4_fold_tail: slots must be a contiguous 1-D int64 CUDA tensor"); - TORCH_CHECK(rows > 0 && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, + TORCH_CHECK( + rows > 0 && rows <= std::numeric_limits::max() && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, "minimax_m3_sm107_tp4_fold_tail: rows must be positive and tokens_per_step in [1, 32]"); + TORCH_CHECK( + rows <= std::numeric_limits::max() / tokensPerStep, "fold_tail: staging indices must fit int32"); check2d(srcNew, "src_new"); check2d(dstNew, "dst_new"); check2d(srcNext, "src_next"); @@ -68,7 +72,14 @@ inline tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams buildParams(torc check1d(draftIds, "draft_ids"); check1d(posOff, "pos_off"); check1d(kvOff, "kv_off"); + TORCH_CHECK(slots.numel() >= rows && srcNew.size(0) >= rows && srcNext.size(0) >= rows && srcDraft.size(0) >= rows + && srcLens.numel() >= rows && kvOff.numel() >= rows, + "minimax_m3_sm107_tp4_fold_tail: source and slot buffers must hold rows entries"); + TORCH_CHECK(ids.numel() >= rows * tokensPerStep && posOff.numel() >= rows * tokensPerStep + && draftIds.numel() >= rows * (tokensPerStep - 1), + "minimax_m3_sm107_tp4_fold_tail: staging outputs are too small for rows * tokens_per_step"); tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams p{}; + p.maxRows = static_cast(rows); p.ctrl = ctrl.data_ptr(); p.slots = slots.data_ptr(); p.srcNew = srcNew.data_ptr(); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp index e2f65f63ae33..6664be44858a 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -44,6 +45,10 @@ void minimaxM3Sm107Tp4FullCtxReduce(torch::Tensor const& partialO, torch::Tensor TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() && partialStats.sizes() == partialO.sizes().slice(0, 4)); TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + constexpr int64_t kMaxInt = std::numeric_limits::max(); + TORCH_CHECK(qTokensPerCta > 0 && qTokensPerCta <= kMaxInt && tileKv > 0 && tileKv <= kMaxInt && instsKv > 0 + && instsKv <= kMaxInt / fcr::kMaxSplits, + "fullctx_reduce: producer geometry must be positive and fit int32"); int const rows = static_cast(partialO.size(3)); TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); int const headsPerKv = rows / static_cast(seqLenQ); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp index 4f7fda5cac90..2c264d043118 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SampledAcceptOp.cpp @@ -47,6 +47,8 @@ void minimaxM3Sm107Tp4SampledAccept(torch::Tensor const& logits, int64_t vocabOf std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, bool prewaitParams, bool earlyTrigger, int64_t oneRoundK, std::optional const& decided, bool ilpWalk) { + TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.scalar_type() == at::kBFloat16 && logits.stride(1) == 1, + "sampled_accept: logits must be CUDA BF16 [rows, vocab] with unit column stride"); c10::cuda::CUDAGuard const deviceGuard{logits.device()}; sa::Args a; a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp index 53121e3b1e7a..8bab6e46800f 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SparseDecodeLcOp.cpp @@ -17,6 +17,8 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/sparseDecodeLc4.h" +#include + #include #include #include @@ -111,6 +113,12 @@ void minimaxM3Sm107Tp4SparseDecodeLc(torch::Tensor const& q, torch::Tensor const prm.block_table = blockTable.data_ptr(); prm.bt_stride = blockTable.stride(0); prm.seq_lens = seqLens.data_ptr(); + TORCH_CHECK( + decodeQueryLen > 0 && decodeQueryLen <= std::numeric_limits::max() && q.size(0) % decodeQueryLen == 0, + "sparse_decode: decode_query_len must be positive and divide the query token count"); + TORCH_CHECK( + seqLens.device() == q.device() && seqLens.is_contiguous() && seqLens.numel() >= q.size(0) / decodeQueryLen, + "sparse_decode: seq_lens must be contiguous and hold every request"); prm.dql = static_cast(decodeQueryLen); prm.sm_scale = static_cast(smScale); prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp index bee82725b41f..0487c5eb0ea8 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DraftStepOp.cpp @@ -145,9 +145,11 @@ void minimaxM3Sm107Tp8DraftStep(torch::Tensor const& x, OptTensor const& xRows, TORCH_CHECK(x.size(0) >= rows, "draft_step: x must hold a row per request"); } TORCH_CHECK(world >= 1 && rank >= 0 && rank < world, "draft_step: rank / world"); - TORCH_CHECK(keys.scalar_type() == at::kLong && keys.numel() >= rows, "draft_step: keys must be int64 [>= rows]"); - TORCH_CHECK(counters.scalar_type() == at::kInt && counters.numel() >= 2, "draft_step: counters must be int32 [2]"); - TORCH_CHECK(wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, + TORCH_CHECK(keys.is_contiguous() && keys.scalar_type() == at::kLong && keys.numel() >= rows, + "draft_step: keys must be int64 [>= rows]"); + TORCH_CHECK(counters.is_contiguous() && counters.scalar_type() == at::kInt && counters.numel() >= 2, + "draft_step: counters must be int32 [2]"); + TORCH_CHECK(wsUc.is_contiguous() && wsUc.scalar_type() == at::kLong && wsUc.numel() >= 2 * ds::kMaxTokens * world, "draft_step: ws_uc must hold [2][8][world] slots"); TORCH_CHECK(draftOut.scalar_type() == at::kInt && draftOut.dim() == 2 && draftOut.stride(1) == 1 && draftOut.size(0) >= rows && step >= 0 && step < draftOut.size(1), diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp index ec93068fc374..7b64d911263f 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp @@ -46,7 +46,8 @@ void minimaxM3Sm107Tp8DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor c && cosSin.numel() % rotaryDim == 0, "cos_sin must be the fp32 [positions, rotary_dim / 2] (cos, sin) table"); int64_t const requests = qkv.size(0) / qLen; - TORCH_CHECK(seqLens.is_cuda() && seqLens.scalar_type() == at::kInt && seqLens.numel() >= requests); + TORCH_CHECK(seqLens.is_cuda() && seqLens.is_contiguous() && seqLens.scalar_type() == at::kInt + && seqLens.numel() >= requests); if (specPositionOffsets.has_value()) { TORCH_CHECK(specPositionOffsets->is_cuda() && specPositionOffsets->is_contiguous() diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h index cc5cde28c9d9..fd8bcc99f10b 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h @@ -19,6 +19,7 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h" +#include #include TRTLLM_NAMESPACE_BEGIN @@ -54,8 +55,11 @@ inline tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams buildParams(torc "minimax_m3_sm107_tp8_fold_tail: ctrl must be an int32 CUDA tensor"); TORCH_CHECK(slots.is_cuda() && slots.scalar_type() == torch::kInt64 && slots.dim() == 1 && slots.stride(0) == 1, "minimax_m3_sm107_tp8_fold_tail: slots must be a contiguous 1-D int64 CUDA tensor"); - TORCH_CHECK(rows > 0 && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, + TORCH_CHECK( + rows > 0 && rows <= std::numeric_limits::max() && tokensPerStep > 0 && tokensPerStep <= kMaxWidth, "minimax_m3_sm107_tp8_fold_tail: rows must be positive and tokens_per_step in [1, 32]"); + TORCH_CHECK( + rows <= std::numeric_limits::max() / tokensPerStep, "fold_tail: staging indices must fit int32"); check2d(srcNew, "src_new"); check2d(dstNew, "dst_new"); check2d(srcNext, "src_next"); @@ -68,7 +72,14 @@ inline tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams buildParams(torc check1d(draftIds, "draft_ids"); check1d(posOff, "pos_off"); check1d(kvOff, "kv_off"); + TORCH_CHECK(slots.numel() >= rows && srcNew.size(0) >= rows && srcNext.size(0) >= rows && srcDraft.size(0) >= rows + && srcLens.numel() >= rows && kvOff.numel() >= rows, + "minimax_m3_sm107_tp8_fold_tail: source and slot buffers must hold rows entries"); + TORCH_CHECK(ids.numel() >= rows * tokensPerStep && posOff.numel() >= rows * tokensPerStep + && draftIds.numel() >= rows * (tokensPerStep - 1), + "minimax_m3_sm107_tp8_fold_tail: staging outputs are too small for rows * tokens_per_step"); tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams p{}; + p.maxRows = static_cast(rows); p.ctrl = ctrl.data_ptr(); p.slots = slots.data_ptr(); p.srcNew = srcNew.data_ptr(); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp index 29b9a0ab391b..b7b06bf2749f 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -44,6 +45,10 @@ void minimaxM3Sm107Tp8FullCtxReduce(torch::Tensor const& partialO, torch::Tensor TORCH_CHECK(partialStats.scalar_type() == at::kFloat && partialStats.is_contiguous() && partialStats.sizes() == partialO.sizes().slice(0, 4)); TORCH_CHECK(seqLensKv.scalar_type() == at::kInt && seqLensKv.numel() >= partialO.size(0)); + constexpr int64_t kMaxInt = std::numeric_limits::max(); + TORCH_CHECK(qTokensPerCta > 0 && qTokensPerCta <= kMaxInt && tileKv > 0 && tileKv <= kMaxInt && instsKv > 0 + && instsKv <= kMaxInt / fcr::kMaxSplits, + "fullctx_reduce: producer geometry must be positive and fit int32"); int const rows = static_cast(partialO.size(3)); TORCH_CHECK(seqLenQ >= 1 && rows % seqLenQ == 0); int const headsPerKv = rows / static_cast(seqLenQ); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp index f6200c9aa906..5df44da43700 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SampledAcceptOp.cpp @@ -47,6 +47,8 @@ void minimaxM3Sm107Tp8SampledAccept(torch::Tensor const& logits, int64_t vocabOf std::optional const& forceVec, std::optional const& extraAttempts, bool usePdl, bool prewaitParams, bool earlyTrigger, int64_t oneRoundK, std::optional const& decided, bool ilpWalk) { + TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.scalar_type() == at::kBFloat16 && logits.stride(1) == 1, + "sampled_accept: logits must be CUDA BF16 [rows, vocab] with unit column stride"); c10::cuda::CUDAGuard const deviceGuard{logits.device()}; sa::Args a; a.logits = reinterpret_cast<__nv_bfloat16 const*>(logits.data_ptr()); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp index dc80380b48bb..2dc706ded147 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SparseDecodeLc8Op.cpp @@ -17,6 +17,8 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/sparseDecodeLc8m.h" +#include + #include #include #include @@ -89,6 +91,11 @@ void minimaxM3Sm107Tp8SparseDecodeLc8(torch::Tensor const& q, torch::Tensor cons { checkBytes(*q8, "q8"); checkBytes(*sf, "sf"); + int64_t const tokens = q.size(0); + TORCH_CHECK(numSfKTiles == 8 && q8->is_contiguous() && q8->numel() >= tokens * 1024, + "sparse_decode_lc8: q8 must hold contiguous T * 1024 bytes with num_sf_k_tiles=8"); + TORCH_CHECK(sf->is_contiguous() && sf->numel() >= ((tokens + 127) / 128) * 128 * 32, + "sparse_decode_lc8: sf must hold the padded 128-row MXFP8 scale tiles"); } tensorrt_llm::kernels::minimaxM3Sm107Tp8::MinimaxM3Tp8SparseDecodeLc8Params prm{}; @@ -111,6 +118,12 @@ void minimaxM3Sm107Tp8SparseDecodeLc8(torch::Tensor const& q, torch::Tensor cons prm.block_table = blockTable.data_ptr(); prm.bt_stride = blockTable.stride(0); prm.seq_lens = seqLens.data_ptr(); + TORCH_CHECK( + decodeQueryLen > 0 && decodeQueryLen <= std::numeric_limits::max() && q.size(0) % decodeQueryLen == 0, + "sparse_decode: decode_query_len must be positive and divide the query token count"); + TORCH_CHECK( + seqLens.device() == q.device() && seqLens.is_contiguous() && seqLens.numel() >= q.size(0) / decodeQueryLen, + "sparse_decode: seq_lens must be contiguous and hold every request"); prm.dql = static_cast(decodeQueryLen); prm.sm_scale = static_cast(smScale); prm.out = reinterpret_cast<__nv_bfloat16*>(out.data_ptr()); diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py index 2733eb1e7524..94b160e3f4d4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py @@ -120,6 +120,7 @@ def sparse_attention_reference( valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) logits = (q64[tok] @ keys.T) * sm_scale logits = logits.masked_fill(~valid[None, :], float("-inf")) + values = values.masked_fill(~valid[:, None], 0) out[tok] = torch.softmax(logits, dim=-1) @ values return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py index 027541ea6990..1cace7ab549a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py @@ -636,6 +636,7 @@ def kernel( split_rank, sD, bidx, + n_idx * self.cta_n, ar_ws, ar_flags, ar_rank, @@ -774,6 +775,7 @@ def epilog_warp( split_rank: cutlass.Int32, sD: cute.Tensor, bidx: cutlass.Int32, + row_offset: cutlass.Int32, ar_ws: cutlass.Int64, ar_flags: cutlass.Int64, ar_rank: cutlass.Int32, @@ -900,6 +902,7 @@ def epilog_warp( sD, epi_tid, bidx, + row_offset, ar_ws, ar_flags, ar_rank, @@ -926,6 +929,7 @@ def _push_tile( sD, epi_tid: cutlass.Int32, bidx: cutlass.Int32, + row_offset: cutlass.Int32, ar_ws: cutlass.Int64, ar_flags: cutlass.Int64, ar_rank: cutlass.Int32, @@ -946,7 +950,8 @@ def _push_tile( half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): word = epi_tid + w0 - t = word // pairs + local_t = word // pairs + t = row_offset + local_t p = word % pairs if t < ar_tokens: n = col0 + 2 * p @@ -961,7 +966,8 @@ def _push_tile( _st_global_b32( ar_ws + off * cutlass.Int64(4), _bf16x2_word( - sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + sD[2 * p, local_t].to(cutlass.Float32), + sD[2 * p + 1, local_t].to(cutlass.Float32), ), ) if cutlass.const_expr(self.use_pdl): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py index 067538849ac8..790194f9a927 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py @@ -444,7 +444,10 @@ def __init__( "M3MoeState allocates and compiles: build it before CUDA-graph capture" ) sms_all = torch.cuda.get_device_properties(device).multi_processor_count - if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + resolved_ctas = ( + (sms_all - SHARED_CTAS if shared else sms_all) if num_ctas is None else num_ctas + ) + if shared and resolved_ctas + SHARED_CTAS > sms_all: raise ValueError( f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " f"(all resident at once)" @@ -456,7 +459,7 @@ def __init__( if not ok: raise RuntimeError(f"m3_moe: {why}") sms = torch.cuda.get_device_properties(self.device).multi_processor_count - self.num_ctas = num_ctas or sms + self.num_ctas = resolved_ctas if not MIN_CTAS <= self.num_ctas <= sms: raise ValueError( f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py index b383cd4036a4..ada666e91599 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py @@ -1213,6 +1213,8 @@ def _check_nvfp4_inputs( ) if head_dim != triton.next_power_of_2(head_dim): raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + if v_paged.shape[:3] != k_paged.shape[:3]: + raise ValueError("NVFP4 K and V caches must have matching page, head and token extents.") for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): if int(data.shape[3]) * 2 != head_dim: raise ValueError( @@ -1470,8 +1472,8 @@ def minimax_m3_sparse_attn_decode( if measured_chunks is not None else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) ) - elif num_topk_chunks & (num_topk_chunks - 1): - raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + elif num_topk_chunks <= 0 or num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a positive power of two; got {num_topk_chunks}.") # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one # program per (request, block) covers all of them and dequantizes each block once. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py index 048d31b82d68..34794a95aadd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py @@ -38,6 +38,7 @@ def _gqa_sparse_decode_mq_kernel( sm_scale, stride_qn, stride_qh, + stride_qd, stride_k_blk, stride_k_pos, stride_v_blk, @@ -96,7 +97,7 @@ def _gqa_sparse_decode_mq_kernel( n_mask = off_n < n_valid q = tl.load( - q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :], + q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :] * stride_qd, mask=r_mask[:, None], other=0.0, ).to(tl.float16) @@ -208,6 +209,7 @@ def launch_mq_producer( sm_scale, q.stride(0), q.stride(1), + q.stride(2), k_paged.stride(0), k_paged.stride(2), v_paged.stride(0), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py index 76fc68fa5903..5bafb312b873 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py @@ -109,6 +109,7 @@ def sparse_attention_reference( valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) logits = (q64[tok] @ keys.T) * sm_scale logits = logits.masked_fill(~valid[None, :], float("-inf")) + values = values.masked_fill(~valid[:, None], 0) out[tok] = torch.softmax(logits, dim=-1) @ values return out @@ -327,6 +328,8 @@ def reference( rounded once. ``linear_softmax``, ``use_pdl``, ``early_trigger``, ``one_div`` and ``bt_prewait`` do not change what is computed; the hand-off buffers are checked separately with ``mxfp8_handoff_reference``. """ + if identity: + topk = _identity_table(seq_lens, decode_query_len, int(q.shape[0]), topk.device) attn = sparse_attention_reference( q.view(-1, NUM_HEADS, HEAD_DIM), k, @@ -416,7 +419,11 @@ def is_valid( req, offset = divmod(tok, decode_query_len) kv_len = max(seq[req] - decode_query_len + offset + 1, 0) nblocks = (kv_len + PAGE - 1) // PAGE - ids = topk[tok, : min(TOPK, nblocks)] + ids = ( + torch.arange(min(TOPK, nblocks), device=topk.device) + if identity + else topk[tok, : min(TOPK, nblocks)] + ) assert bool(((ids >= -1) & (ids < min(nblocks, num_cols))).all()), ( f"token {tok}: selected block ids must be -1 or below the request's block count; the block " "table is indexed with them unchecked" @@ -434,3 +441,14 @@ def is_valid( m3_sparse_decode_lc = M3SparseDecodeLc() + + +def _identity_table(seq_lens, decode_query_len, total_q, device): + """[T, 16] ids (i < min(16, blocks of the token's causal length) ? i : -1), the selection identity mode reads.""" + ids = torch.full((total_q, 16), -1, dtype=torch.int32, device=device) + for tok in range(total_q): + req, qoff = divmod(tok, decode_query_len) + kv_len = max(int(seq_lens[req]) - decode_query_len + qoff + 1, 0) + n = min(16, (kv_len + PAGE - 1) // PAGE) + ids[tok, :n] = torch.arange(n, dtype=torch.int32, device=device) + return ids diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py index 11ef3a496132..a5b626f3fe15 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py @@ -636,6 +636,7 @@ def kernel( split_rank, sD, bidx, + n_idx * self.cta_n, ar_ws, ar_flags, ar_rank, @@ -774,6 +775,7 @@ def epilog_warp( split_rank: cutlass.Int32, sD: cute.Tensor, bidx: cutlass.Int32, + row_offset: cutlass.Int32, ar_ws: cutlass.Int64, ar_flags: cutlass.Int64, ar_rank: cutlass.Int32, @@ -900,6 +902,7 @@ def epilog_warp( sD, epi_tid, bidx, + row_offset, ar_ws, ar_flags, ar_rank, @@ -926,6 +929,7 @@ def _push_tile( sD, epi_tid: cutlass.Int32, bidx: cutlass.Int32, + row_offset: cutlass.Int32, ar_ws: cutlass.Int64, ar_flags: cutlass.Int64, ar_rank: cutlass.Int32, @@ -946,7 +950,8 @@ def _push_tile( half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): word = epi_tid + w0 - t = word // pairs + local_t = word // pairs + t = row_offset + local_t p = word % pairs if t < ar_tokens: n = col0 + 2 * p @@ -961,7 +966,8 @@ def _push_tile( _st_global_b32( ar_ws + off * cutlass.Int64(4), _bf16x2_word( - sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + sD[2 * p, local_t].to(cutlass.Float32), + sD[2 * p + 1, local_t].to(cutlass.Float32), ), ) if cutlass.const_expr(self.use_pdl): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py index ab60e7d96370..281b56869436 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py @@ -428,7 +428,10 @@ def __init__( "M3MoeState allocates and compiles: build it before CUDA-graph capture" ) sms_all = torch.cuda.get_device_properties(device).multi_processor_count - if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + resolved_ctas = ( + (sms_all - SHARED_CTAS if shared else sms_all) if num_ctas is None else num_ctas + ) + if shared and resolved_ctas + SHARED_CTAS > sms_all: raise ValueError( f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " f"(all resident at once)" @@ -440,7 +443,7 @@ def __init__( if not ok: raise RuntimeError(f"m3_moe: {why}") sms = torch.cuda.get_device_properties(self.device).multi_processor_count - self.num_ctas = num_ctas or sms + self.num_ctas = resolved_ctas if not MIN_CTAS <= self.num_ctas <= sms: raise ValueError( f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py index 446776419315..24a73860a7f0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py @@ -151,9 +151,9 @@ def __init__(self, device: torch.device, num_tokens: int, num_ctas: Optional[int raise RuntimeError(f"m3_moe16: {why}") sms = torch.cuda.get_device_properties(self.device).multi_processor_count self.num_ctas = num_ctas or M16_NUM_CTAS - if not 128 <= self.num_ctas <= sms - SHARED_CTAS: + if self.num_ctas != 128 or self.num_ctas + SHARED_CTAS > sms: raise ValueError( - f"m3_moe16 runs 128 .. {sms - SHARED_CTAS} routed CTAs + {SHARED_CTAS} shared (all resident)" + f"m3_moe16 requires exactly 128 routed CTAs + {SHARED_CTAS} shared (all resident)" ) self.m = num_tokens plan, nb_cap = make_phase_plan16(self.num_ctas, num_tokens, M16_LAG) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py index d146b9d05d2e..240246d4548e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py @@ -1189,6 +1189,8 @@ def _check_nvfp4_inputs( ) if head_dim != triton.next_power_of_2(head_dim): raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + if v_paged.shape[:3] != k_paged.shape[:3]: + raise ValueError("NVFP4 K and V caches must have matching page, head and token extents.") for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): if int(data.shape[3]) * 2 != head_dim: raise ValueError( @@ -1460,8 +1462,8 @@ def minimax_m3_sparse_attn_decode( if measured_chunks is not None else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) ) - elif num_topk_chunks & (num_topk_chunks - 1): - raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + elif num_topk_chunks <= 0 or num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a positive power of two; got {num_topk_chunks}.") # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one # program per (request, block) covers all of them and dequantizes each block once. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py index 048d31b82d68..34794a95aadd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py @@ -38,6 +38,7 @@ def _gqa_sparse_decode_mq_kernel( sm_scale, stride_qn, stride_qh, + stride_qd, stride_k_blk, stride_k_pos, stride_v_blk, @@ -96,7 +97,7 @@ def _gqa_sparse_decode_mq_kernel( n_mask = off_n < n_valid q = tl.load( - q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :], + q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :] * stride_qd, mask=r_mask[:, None], other=0.0, ).to(tl.float16) @@ -208,6 +209,7 @@ def launch_mq_producer( sm_scale, q.stride(0), q.stride(1), + q.stride(2), k_paged.stride(0), k_paged.stride(2), v_paged.stride(0), diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py index a845ac68f257..76a99144322a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py @@ -113,6 +113,7 @@ def sparse_attention_reference( valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) logits = (q64[tok] @ keys.T) * sm_scale logits = logits.masked_fill(~valid[None, :], float("-inf")) + values = values.masked_fill(~valid[:, None], 0) out[tok] = torch.softmax(logits, dim=-1) @ values return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py index 2f1a599c45aa..1632da162493 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py @@ -749,6 +749,7 @@ def kernel( split_rank, sD, bidx, + n_idx * self.cta_n, ar_ws, ar_flags, ar_rank, @@ -942,6 +943,7 @@ def epilog_warp( split_rank: cutlass.Int32, sD: cute.Tensor, bidx: cutlass.Int32, + row_offset: cutlass.Int32, ar_ws: cutlass.Int64, ar_flags: cutlass.Int64, ar_rank: cutlass.Int32, @@ -1068,6 +1070,7 @@ def epilog_warp( sD, epi_tid, bidx, + row_offset, ar_ws, ar_flags, ar_rank, @@ -1094,6 +1097,7 @@ def _push_tile( sD, epi_tid: cutlass.Int32, bidx: cutlass.Int32, + row_offset: cutlass.Int32, ar_ws: cutlass.Int64, ar_flags: cutlass.Int64, ar_rank: cutlass.Int32, @@ -1114,7 +1118,8 @@ def _push_tile( half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): word = epi_tid + w0 - t = word // pairs + local_t = word // pairs + t = row_offset + local_t p = word % pairs if t < ar_tokens: n = col0 + 2 * p @@ -1129,7 +1134,8 @@ def _push_tile( _st_global_b32( ar_ws + off * cutlass.Int64(4), _bf16x2_word( - sD[2 * p, t].to(cutlass.Float32), sD[2 * p + 1, t].to(cutlass.Float32) + sD[2 * p, local_t].to(cutlass.Float32), + sD[2 * p + 1, local_t].to(cutlass.Float32), ), ) if cutlass.const_expr(self.use_pdl): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py index 70deb357f4a8..9f176116d722 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py @@ -381,7 +381,10 @@ def __init__( "M3MoeState allocates and compiles: build it before CUDA-graph capture" ) sms_all = torch.cuda.get_device_properties(device).multi_processor_count - if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + resolved_ctas = ( + (sms_all - SHARED_CTAS if shared else sms_all) if num_ctas is None else num_ctas + ) + if shared and resolved_ctas + SHARED_CTAS > sms_all: raise ValueError( f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " f"(all resident at once)" @@ -393,7 +396,7 @@ def __init__( if not ok: raise RuntimeError(f"m3_moe: {why}") sms = torch.cuda.get_device_properties(self.device).multi_processor_count - self.num_ctas = num_ctas or sms + self.num_ctas = resolved_ctas if not MIN_CTAS <= self.num_ctas <= sms: raise ValueError( f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py index e3694274c124..5b979f361997 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py @@ -1204,6 +1204,8 @@ def _check_nvfp4_inputs( ) if head_dim != triton.next_power_of_2(head_dim): raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + if v_paged.shape[:3] != k_paged.shape[:3]: + raise ValueError("NVFP4 K and V caches must have matching page, head and token extents.") for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): if int(data.shape[3]) * 2 != head_dim: raise ValueError( @@ -1479,8 +1481,8 @@ def minimax_m3_sparse_attn_decode( if measured_chunks is not None else resolve_num_topk_chunks(total_q, num_kv_heads, max_topk) ) - elif num_topk_chunks & (num_topk_chunks - 1): - raise ValueError(f"num_topk_chunks must be a power of two; got {num_topk_chunks}.") + elif num_topk_chunks <= 0 or num_topk_chunks & (num_topk_chunks - 1): + raise ValueError(f"num_topk_chunks must be a positive power of two; got {num_topk_chunks}.") # Every query token of a request reads the same blocks under identity selection (short-sequence mode): one # program per (request, block) covers all of them and dequantizes each block once. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py index 048d31b82d68..34794a95aadd 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py @@ -38,6 +38,7 @@ def _gqa_sparse_decode_mq_kernel( sm_scale, stride_qn, stride_qh, + stride_qd, stride_k_blk, stride_k_pos, stride_v_blk, @@ -96,7 +97,7 @@ def _gqa_sparse_decode_mq_kernel( n_mask = off_n < n_valid q = tl.load( - q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :], + q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :] * stride_qd, mask=r_mask[:, None], other=0.0, ).to(tl.float16) @@ -208,6 +209,7 @@ def launch_mq_producer( sm_scale, q.stride(0), q.stride(1), + q.stride(2), k_paged.stride(0), k_paged.stride(2), v_paged.stride(0), diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index 12164dcefaaf..5a9c0fbe4a10 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -586,8 +586,12 @@ def _extra_bytes_per_token(cls, model_config, mapping, num_layers, kwargs) -> in if mapping.pp_size > 1 else list(range(num_target_layers)) ) - sparse_layers = [layer for layer in local_layers if layer >= 3] - dense_layers = [layer for layer in local_layers if layer < 3] + sparse_layer_ids = kwargs.get("sparse_layer_ids") + if sparse_layer_ids is None: + sparse_layer_ids = default_sparse_layer_ids(config, num_target_layers) + sparse_layer_ids = set(int(layer) for layer in sparse_layer_ids) + sparse_layers = [layer for layer in local_layers if layer in sparse_layer_ids] + dense_layers = [layer for layer in local_layers if layer not in sparse_layer_ids] tp = 1 if mapping.enable_attention_dp else mapping.tp_size kv_heads = max(1, int(config.num_key_value_heads) // tp) head_dim = int( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py index 0acfdecb196e..bcac8e71786f 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/embedding_firstnorm/op.py @@ -36,7 +36,9 @@ def arg(tensor): def _check(ids, table, weight, raw, out) -> None: rows = ids.shape[0] if ids.dim() == 1 else 0 if not ( - 1 <= rows <= MAX_ROWS + all(t.is_cuda and t.device == ids.device for t in (ids, table, weight, raw, out)) + and all(t.data_ptr() % 16 == 0 for t in (table, weight, raw, out)) + and 1 <= rows <= MAX_ROWS and ids.dtype in (torch.int32, torch.int64) and ids.is_contiguous() and table.dtype == torch.bfloat16 @@ -45,6 +47,7 @@ def _check(ids, table, weight, raw, out) -> None: and table.is_contiguous() and weight.dtype == torch.bfloat16 and tuple(weight.shape) == (HIDDEN,) + and weight.is_contiguous() and all( t.dtype == torch.bfloat16 and tuple(t.shape) == (rows, HIDDEN) and t.is_contiguous() for t in (raw, out) @@ -52,7 +55,7 @@ def _check(ids, table, weight, raw, out) -> None: ): raise ValueError( "m3_embed_norm: expects ids int32/int64 [1..8], a contiguous bf16 table [V, 6144], bf16 weight [6144] " - "and contiguous bf16 raw / out [rows, 6144]" + "and contiguous bf16 raw / out [rows, 6144], all on one CUDA device with 16-byte aligned bf16 pointers" ) @@ -137,6 +140,7 @@ def __init__(self, table, norm, rows, id_dtype, use_pdl=True): self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) self.output = torch.empty_like(self.raw) ids = torch.empty(rows, device=table.device, dtype=id_dtype) + _check(ids, table, norm.weight, self.raw, self.output) _compile(ids, table, norm.weight, self.eps, self.raw, self.output, use_pdl) def __call__(self, ids): diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py index 12f6008d439f..0b29aba57c89 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op.py @@ -219,12 +219,40 @@ def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: key = (gu8.device, gu8.data_ptr(), gusf8.data_ptr(), dn8.data_ptr(), dnsf8.data_ptr()) + sources = (gu8, gusf8, dn8, dnsf8) + versions = ( + None if any(t.is_inference() for t in sources) else tuple(t._version for t in sources) + ) hit = _PRETILED.get(key) - if hit is None: - hit = _PRETILED[key] = ((gu8, gusf8, dn8, dnsf8), _pretile_shared(gu8, gusf8, dn8, dnsf8)) + if hit is None or versions is None or hit[2] != versions: + tiled = _pretile_shared(*sources) + if hit is not None: + # Reuse the layer's allocation across token-count states. Inference tensors + # have no mutation version, so rebuilding must refresh their contents. + with torch.inference_mode(): + hit[1].copy_(tiled) + tiled = hit[1] + hit = _PRETILED[key] = (sources, tiled, versions) return hit[1] +def _validate_output(out: torch.Tensor, m: int, device: torch.device) -> None: + device = torch.device(device) + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + if ( + out.dtype != torch.bfloat16 + or tuple(out.shape) != (m, HIDDEN) + or not out.is_contiguous() + or not out.is_cuda + or out.device != device + or out.data_ptr() % 16 != 0 + ): + raise ValueError( + f"m3_moe: out must be 16-byte aligned contiguous bf16 [{m}, {HIDDEN}] on {device}" + ) + + def make_phase_plan( num_ctas: int, m_max: int, lag: int = 2, tail_w: float = 1.0, early0: bool = False ): @@ -312,7 +340,10 @@ def __init__( "M3MoeState allocates and compiles: build it before CUDA-graph capture" ) sms_all = torch.cuda.get_device_properties(device).multi_processor_count - if shared and (num_ctas or sms_all) + SHARED_CTAS > sms_all: + resolved_ctas = ( + (sms_all - SHARED_CTAS if shared else sms_all) if num_ctas is None else num_ctas + ) + if shared and resolved_ctas + SHARED_CTAS > sms_all: raise ValueError( f"m3_moe (shared): num_ctas + {SHARED_CTAS} shared-expert CTAs must fit the {sms_all} SMs " f"(all resident at once)" @@ -324,7 +355,7 @@ def __init__( if not ok: raise RuntimeError(f"m3_moe: {why}") sms = torch.cuda.get_device_properties(self.device).multi_processor_count - self.num_ctas = num_ctas or sms + self.num_ctas = resolved_ctas if not MIN_CTAS <= self.num_ctas <= sms: raise ValueError( f"m3_moe runs {MIN_CTAS} .. {sms} CTAs (all resident at once), not {self.num_ctas}" @@ -430,6 +461,8 @@ def layer( def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: m = self.m + if out is not None: + _validate_output(out, m, self.device) if ( hidden is None or hidden.dtype != torch.bfloat16 @@ -575,7 +608,12 @@ def __init__( [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] ).float() self.escal = self.escal.reshape(6, e).contiguous() - self.bias = bias.float().contiguous() + device = torch.device(state.device) + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + if bias.numel() != e or not bias.is_cuda or bias.device != device: + raise ValueError(f"m3_moe: routing bias must have {e} elements on {device}") + self.bias = bias.float().reshape(e).contiguous() if ( router_weight is None or fc31_input_scale is None diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py index 446776419315..8617f1327d7a 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/moe/op16.py @@ -47,6 +47,7 @@ _compiled, _kernel_module, _shared_views, + _validate_output, _view, is_supported, ) @@ -229,6 +230,8 @@ def layer( def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: m = self.m + if out is not None: + _validate_output(out, m, self.device) if ( hidden is None or hidden.dtype != torch.bfloat16 diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py index 041dc12b1ffd..1807ae64d0f2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/minimax_m3/mx_gemv/op.py @@ -120,7 +120,16 @@ def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" split = cfg[0] - if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: + if ( + not 0 < m <= MAX_TOKENS + or n <= 0 + or k <= 0 + or n % 128 + or k % 128 + or split <= 0 + or 128 % split != 0 + or (k // 128) % split != 0 + ): return True tiles = k // 128 // split mma_n = _mma_n(m) @@ -316,6 +325,8 @@ def _run( last_uncompiled = None import cuda.bindings.driver as cuda_driver + if x.dim() != 2 or weight.dim() != 2 or not x.is_contiguous() or not weight.is_contiguous(): + return None m, k = x.shape if x_swiglu: k //= 2 @@ -363,6 +374,8 @@ def precompile( runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" import cuda.bindings.driver as cuda_driver + if weight.dim() != 2 or not weight.is_contiguous(): + return None n, k = weight.shape push = None if ar is None else (0, 0, 0, *ar) resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu) @@ -495,27 +508,14 @@ def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: b if ( weight.dim() != 2 or weight.dtype != torch.float8_e4m3fn + or not weight.is_contiguous() or get_config(*weight.shape) is None ): return n, k = weight.shape if (n, k, use_pdl) in _ensured: return + for width in (1, *MMA_WIDTHS[1:]): + precompile(weight, weight_scale, width, x_bf16=False, use_pdl=use_pdl) + precompile(weight, weight_scale, 1, x_bf16=True, use_pdl=use_pdl) _ensured.add((n, k, use_pdl)) - xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) - xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) - m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) - for width in MMA_WIDTHS[1:]: - m3_mx_gemv( - torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), - xsf, - weight, - weight_scale, - use_pdl=use_pdl, - ) - m3_mx_gemv_bf16( - torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), - weight, - weight_scale, - use_pdl=use_pdl, - ) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 43bffb485255..22a37483d135 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1049,9 +1049,10 @@ def replay( # anything below reads the static inputs on the device; a graph captured # with the pack's prologue issues them itself at the start of the replay. prologue_pack = self._step_input_prologue_keys.get(key) - step_inputs.flush(in_graph=prologue_pack is not None) - if prologue_pack is not None and not step_inputs.is_active( - prologue_pack): + captured_pack_active = (prologue_pack is not None + and step_inputs.is_active(prologue_pack)) + step_inputs.flush(in_graph=captured_pack_active) + if prologue_pack is not None and not captured_pack_active: # Replayed outside the pack's context: nothing filled its window. prologue_pack.quiesce_graph_window() if self._pre_replay_hooks: @@ -1072,7 +1073,11 @@ def replay( for hook in self._pre_replay_hooks: hook(key, seqlen, hook_tensors, current_inputs) - self.graphs[key].replay() + try: + self.graphs[key].replay() + finally: + if prologue_pack is not None: + prologue_pack.on_graph_replay() output_ref = self.graph_outputs[key] return output_ref diff --git a/tensorrt_llm/_torch/pyexecutor/step_inputs.py b/tensorrt_llm/_torch/pyexecutor/step_inputs.py index 448064f4968c..517778e4b6f5 100644 --- a/tensorrt_llm/_torch/pyexecutor/step_inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/step_inputs.py @@ -215,6 +215,7 @@ def __init__( self.d2d: list[tuple[int, int, int]] = [] # Every tensor a staged copy reads or writes, kept alive until the pack is issued. self.staged: list[torch.Tensor] = [] + self._graph_staged: list[torch.Tensor] = [] self.used = self.data_start if parallel_scatter: self._warm_up(device) @@ -362,6 +363,7 @@ def flush(self, in_graph: bool = False) -> None: if self._fill_graph_window(): if self.early_copy: self._upload_window(self.used) + self._graph_staged.extend(self.staged) self._clear() return self.quiesce_graph_window() @@ -391,9 +393,20 @@ def flush(self, in_graph: bool = False) -> None: event = self.events[self.parity] = torch.cuda.Event() event.record() self.parity = (self.parity + 1) % self.RING - # The scatter is queued on the current stream: the caching allocator keeps the staged tensors' memory for it. + self._release_after_submission(self.staged) self._clear() + def _release_after_submission(self, tensors: list[torch.Tensor]) -> None: + if tensors: + stream = torch.cuda.current_stream(self.dev.device) + for tensor in tensors: + tensor.record_stream(stream) + tensors.clear() + + def on_graph_replay(self) -> None: + """Release window operands after the graph reading their addresses has been queued.""" + self._release_after_submission(self._graph_staged) + def _fill_graph_window(self) -> bool: """Write the pack into the pinned window the replay's prologue issues; False when it does not fit or is not order-free (the caller issues it eagerly).""" diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 1b89fddf0061..20cd5290bf09 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -12,6 +12,7 @@ l0_cpu: backend: generic orchestrator: mpi tests: + - unittest/_torch/modeling_v2/models/minimax_m3/test_shared_kernel_contracts.py # MiniMax-M3 Rubin target unit tests. - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py # MiniMax-M3 target unit tests. diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml index 94bc49c05f3e..7b9e394088ff 100644 --- a/tests/integration/test_lists/test-db/l0_gr100.yml +++ b/tests/integration/test_lists/test-db/l0_gr100.yml @@ -13,6 +13,7 @@ l0_gr100: stage: post_merge backend: pytorch tests: + - unittest/_torch/modeling_v2/models/minimax_m3/test_shared_kernel_contracts.py # MiniMax-M3 Rubin target unit tests. - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py index fac388a85ec1..aec42648e25a 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_kv_size_estimate.py @@ -18,6 +18,7 @@ def _model_config(nvfp4=True, index_dtype="fp8"): return types.SimpleNamespace( pretrained_config=types.SimpleNamespace( num_hidden_layers=60, + sparse_attention_config={"sparse_attention_freq": [0] * 3 + [1] * 57}, num_key_value_heads=4, head_dim=128, hidden_size=6144, @@ -82,3 +83,46 @@ def base(cls, model_config, mapping, num_layers=None, **kwargs): monkeypatch.setattr(cache_manager, "get_sm_version", lambda: 107) monkeypatch.setattr(cache_manager, "_SM107_EXACT_KV_SIZE_ESTIMATE", False) assert KVM.get_cache_size_per_token(None, None) == (17280, 0) + + +@pytest.mark.parametrize("dense_layers", [4, 10, 60]) +@pytest.mark.parametrize("explicit", [False, True]) +def test_estimate_matches_nondefault_cache_layout(monkeypatch, dense_layers, explicit): + config = _model_config() + sparse_ids = list(range(dense_layers, 60)) + kwargs = {} + if explicit: + kwargs["sparse_layer_ids"] = sparse_ids + else: + config.pretrained_config.sparse_attention_config = { + "sparse_attention_freq": [0] * dense_layers + [1] * (60 - dense_layers) + } + manager = object.__new__(KVM) + manager.sparse_layer_ids = ( + sparse_ids + if explicit + else cache_manager.default_sparse_layer_ids(config.pretrained_config, 60) + ) + assert manager.sparse_layer_ids == sparse_ids + manager.dtype = cache_manager.DataType.NVFP4 + manager.pp_layers = list(range(60)) + manager.layer_offsets = {layer: layer for layer in range(60)} + manager.num_kv_heads_per_layer = [2] * 60 + manager.head_dim_per_layer = [128] * 60 + manager.kv_factor = 2 + manager.sparse_index_dim = 128 + manager.indexer_kv_dtype = "fp8" + # The generic NVFP4 K/V allocation: 512 values / 2 + 512 scale bytes / 16. + generic_layer_bytes = 288 + monkeypatch.setattr( + cache_manager.KVCacheManagerV2, + "get_layer_bytes_per_token", + lambda self, layer, role: generic_layer_bytes, + ) + kv_bytes = sum(manager.get_layer_bytes_per_token(i, cache_manager.Role.ALL) for i in range(60)) + index_buffers = manager._extra_buffers_per_layer(tokens_per_block=128) + index_bytes = ( + sum(buffer.size for buffers in index_buffers.values() for buffer in buffers) // 128 + ) + extra = KVM._extra_bytes_per_token(config, _mapping(2, False), None, kwargs) + assert 60 * generic_layer_bytes + extra == kv_bytes + index_bytes diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py index 4052191ded74..c6c91681361c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_silu_and_mul requires a CUDA device", allow_module_level=True) OAI = (7.0, 1.702, 1.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py index 0ebd8285aa9e..3d56b92c342e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -9,7 +9,8 @@ import tensorrt_llm # noqa: F401 (registers the trtllm ops) -assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the dense FP8 q/k/v op requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py index e3de00587e67..29aa10b54689 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention.py @@ -23,7 +23,8 @@ import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_attention requires a CUDA device", allow_module_level=True) # Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. # The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py index d401a9d532e4..6fdd007fb536 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_draft_attention_paged_context.py @@ -26,7 +26,8 @@ import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above from _validating import ULP, assert_within_ulp # noqa: E402 -assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_attention requires a CUDA device", allow_module_level=True) pytestmark = pytest.mark.skipif( get_sm_version() != 107, reason="the SM107 thop paged-context path (no TRTLLM-Gen FMHA)" ) @@ -88,6 +89,20 @@ def _prepare(metadata, request_ids, qlens, cached, num_contexts: int) -> None: metadata.prepare() +def _rope(x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """NeoX RoPE over the whole head at the drafter's theta, float64.""" + half = rt.HEAD_DIM // 2 + inv_freq = rt.DRAFT_ROPE_THETA ** ( + -torch.arange(0, rt.HEAD_DIM, 2, dtype=torch.float64, device=x.device) / rt.HEAD_DIM + ) + angles = positions.to(torch.float64)[:, None] * inv_freq[None, :] + x = x.to(torch.float64) + x1, x2 = x[..., :half], x[..., half:] + return torch.cat( + [x1 * angles.cos() - x2 * angles.sin(), x2 * angles.cos() + x1 * angles.sin()], -1 + ) + + @pytest.mark.parametrize("case", _CASES, ids=[c[0] for c in _CASES]) def test_paged_context_at_128_token_pages(case) -> None: name, cached, qlens, contexts = case @@ -147,6 +162,13 @@ def test_paged_context_at_128_token_pages(case) -> None: slots = rt.flat_slots(manager, rt.DRAFT, rid, positions).long() rows = qkv[first : first + n] for h in range(rt.KV_HEADS): + k_cols = q_size + h * rt.HEAD_DIM + want_k = _rope(rows[:, k_cols : k_cols + rt.HEAD_DIM].float(), positions) + got_k = pool[slots // rt.PAGE, 0, h, slots % rt.PAGE].to(torch.float64) + step = want_k.abs().clamp_min(2.0**-6) * 2.0**-3 + assert bool(((got_k - want_k).abs() <= step).all()), ( + f"K head {h} is appended rotated, in E4M3" + ) v_cols = q_size + (rt.KV_HEADS + h) * rt.HEAD_DIM want_v = rows[:, v_cols : v_cols + rt.HEAD_DIM].to(torch.float8_e4m3fn) got_v = pool[slots // rt.PAGE, 1, h, slots % rt.PAGE] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py index 6a42e13000de..cda2c3cf64bb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -27,7 +27,8 @@ m3_fullctx_decode as fullctx, ) -assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_drafter_attn requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py index 3d77b932b3e7..c380f655de47 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -28,7 +28,8 @@ ) from tensorrt_llm.bindings.internal import thop -assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the drafter preprocess op requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp2_drafter_rope_fp8_kv"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py index 918174c1295f..e44a6333a68a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -16,7 +16,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fp8_indexer_qk_norm_rope requires a CUDA device", allow_module_level=True) DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py index 6a21693bcf9a..22b9d3ee7a8e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fullctx_decode.py @@ -29,7 +29,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fullctx_decode requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py index 85c5e6b746c1..fabae6422ba1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py @@ -23,7 +23,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_index_score requires a CUDA device", allow_module_level=True) HEADS = 2 # index heads per rank at TP2 HEAD_DIM = 128 @@ -92,7 +93,7 @@ def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> Non for seq_lens, split in (((length,), 1536), ((length, max(qlen, length // 3)), 768)): idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) - op(idx_q, index_k, block_table, seq, score, 4, split, True) + op(idx_q, index_k, block_table, seq, score, qlen, split, True) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, qlen ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py index 4f7e66dd8910..84f8ba6da693 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_msa_attention.py @@ -28,7 +28,8 @@ import _m3_attn_runtime_tp2 as rt # noqa: E402 — needs the path declared above from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_msa_attention requires a CUDA device", allow_module_level=True) # Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, # per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py index dcdfc973c8ba..4dca333f6718 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_producer.py @@ -30,7 +30,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_producer requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_producer requires a CUDA device", allow_module_level=True) PAGES = 1024 MAX_POSITIONS = 262144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py index bcc0ad67c843..3ffe4321fac1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_qk_norm_rope_fp8 requires a CUDA device", allow_module_level=True) HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py index 4be27dc1a2ff..4a3fe4690227 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py @@ -18,7 +18,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_qkv_insert requires a CUDA device", allow_module_level=True) PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 32, 2, 2, 128, 64 MAX_POSITIONS = 262144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py index 19123d839dea..b3b2d97b1497 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py @@ -15,7 +15,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_blocks requires a CUDA device", allow_module_level=True) TOPK, INIT, LOCAL = 16, 0, 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py index 9fc6e361a03e..233fdb92573a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long.py @@ -20,7 +20,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_long requires a CUDA device", allow_module_level=True) # Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key # widths, past the row (clamped), and zero (a CUDA-graph padding row). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py index 60b04b9dde33..011cf1b9a690 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_long_dsl.py @@ -22,7 +22,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_long_dsl requires a CUDA device", allow_module_level=True) # Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key # widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py index 3a19b3bf5a7d..616eb6e1a7b3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py @@ -21,7 +21,8 @@ ) from tensorrt_llm.bindings.internal import thop -assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_tokmap_guard requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp2 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py index d35556356264..67e30c0143b4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py @@ -20,7 +20,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above -assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("msa_kv_update requires a CUDA device", allow_module_level=True) #: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. _SPARE = 3 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py index 822a3b2a0ffc..e1c013b90b16 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py @@ -31,10 +31,12 @@ from pathlib import Path from typing import Any, Callable, List, Optional, Tuple +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("the fused all-reduce + RMSNorm requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 900 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py index 8e0e595bdc1c..551c954c6b10 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -32,10 +32,12 @@ from pathlib import Path from typing import Any, Dict, List, Optional +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_mc_allreduce_norm requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 900 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py index c6baaa215341..d80024bb76c8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py @@ -32,10 +32,12 @@ from pathlib import Path from typing import Any, Optional +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "m3_moe_engine requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_moe_engine requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 1200 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py index 75f31cdb994b..7944e1514c0d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -12,7 +12,8 @@ import pytest import torch -assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_mc_allreduce_norm requires CUDA devices", allow_module_level=True) # Spawns 2 rank processes over the visible devices; under xdist workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py index b77aabe64e20..2f667fb353f3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py @@ -15,7 +15,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_embedding requires a CUDA device", allow_module_level=True) VOCAB, HIDDEN = 4096, 6144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py index b455f6f2a049..ae88229a79b6 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_bf16_linear requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py index f1d036a31909..16f37f33b86a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py @@ -22,7 +22,8 @@ op as gemv, ) -assert torch.cuda.is_available(), "the dense seam test requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the dense seam test requires a CUDA device", allow_module_level=True) ALPHA, LIMIT, BETA = 1.702, 7.0, 1.0 ROWS = list(range(1, modeling.TINY_MAX_TOKENS + 1)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py index 2b8b3a4f2e0c..13f745c07dfe 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -20,7 +20,8 @@ m3_dense_swiglu, ) -assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the dense swigluoai op requires a CUDA device", allow_module_level=True) if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp2_swiglu_oai"): pytest.skip( "needs the TP2 op library with minimax_m3_sm107_tp2_swiglu_oai", allow_module_level=True diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py index b44554a8d717..dd9f8bbaa9f2 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_draft_lowm.py @@ -30,7 +30,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_lowm requires a CUDA device", allow_module_level=True) # Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). # diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py index ebfee3b50ee0..9f8a268b1495 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fp32_linear requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py index 9057207107af..f7d0ff7eab42 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv.py @@ -20,7 +20,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mx_gemv requires a CUDA device", allow_module_level=True) # Element band, in ULP units of the row scale: see compare. diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py index 4055f2d8db4e..c9da3b928106 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -23,7 +23,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mx_gemv_bf16 requires a CUDA device", allow_module_level=True) # The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. ELEMENT_ULP = 1.5 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py index f59ca4d385b4..cbeb86cbb1e9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -16,7 +16,8 @@ from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mxfp8_gemm requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py index 5d272004698e..9a67a31a5200 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -16,7 +16,8 @@ from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mxfp8_linear requires a CUDA device", allow_module_level=True) _FLAGS = { "native": (False, False, False, False), diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py index e685bc5df277..964ecf3f28d5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_pinned_gemm.py @@ -13,7 +13,8 @@ ) from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner -assert torch.cuda.is_available(), "MX-TAC requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("MX-TAC requires a CUDA device", allow_module_level=True) # The TP2 projections at decode rows the GEMVs do not serve (12, 16) and a few beyond. CASES = [ diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py index 46548f91c1d7..164a36b93f8c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_ugpu_head.py @@ -13,7 +13,8 @@ ) from tensorrt_llm._torch.modules.linear import Linear -assert torch.cuda.is_available(), "the uGPU head requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the uGPU head requires a CUDA device", allow_module_level=True) def _locality(): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py index 1d6de982a67a..245b39f825b9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_engine.py @@ -34,7 +34,8 @@ ) from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_moe_engine requires a CUDA device", allow_module_level=True) _STATES = {} diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py index 59559dc01d89..18a2aa0e2599 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_epoch_wrap.py @@ -21,7 +21,8 @@ __extra_import_path__ = ["../../../..", "."] from _m3_moe_operands import hidden_states, layer_for # noqa: E402 -assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_moe_engine requires a CUDA device", allow_module_level=True) EPOCH0 = 0x7FFFFFF0 CALLS = 32 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py index 0530ad9cc24e..eafe82183c3f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_moe_sh_take.py @@ -25,7 +25,8 @@ __extra_import_path__ = ["../../../..", "."] from _m3_moe_operands import hidden_states, layer_for # noqa: E402 -assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_moe_engine requires a CUDA device", allow_module_level=True) SHIPPED = engine_op.SH_TAKE_BY_M GATED = (0, 0, 0, 0, 0, 1, 1, 1) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py index 59201301c623..d27bd5c90e29 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -22,7 +22,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_embedding_firstnorm requires a CUDA device", allow_module_level=True) H = 6144 EPS = 1e-6 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py index 15758e2645e2..d22dd99fecbb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py @@ -20,7 +20,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_capture_norm requires a CUDA device", allow_module_level=True) LAYERS = 3 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py index b255f786fc5e..2bda1f79c148 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py @@ -19,7 +19,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_chunk_norm requires a CUDA device", allow_module_level=True) def _case(spec, seed): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py index deeb0b10c5d3..3efc753d50bd 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py @@ -24,9 +24,10 @@ tactic_for, ) -assert torch.cuda.is_available(), "the push variants require a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the push variants require a CUDA device", allow_module_level=True) -H, MAXT, EMPTY = 6144, 8, -2147483648 +H, MAXT, EMPTY = 6144, 16, -2147483648 # The boundary linears (output width H): o_proj and down, each certified row count on split-K or an even direct width, # up to the 8-token bucket the push arms on (ops/m3_draft_lowm.arm_push); the batch-4/8 rows past it run unpushed. CELLS = [ @@ -34,9 +35,11 @@ for nk, rows in sorted(ROWS.items()) if nk[0] == H and nk[1] != 18432 # o_proj / down, not the fc for m in rows - if m <= MAXT + if m <= 8 and (tactic_for(nk[0], nk[1], m)[0] == "s" or tactic_for(nk[0], nk[1], m)[1][1] % 2 == 0) ] +# Two output-row tiles must populate distinct exchange rows. +CELLS.append((6144, 16, "s", (64, 8, 2, 10))) def _expected_words(y: torch.Tensor, flags: torch.Tensor, world: int, rank: int) -> torch.Tensor: diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py index 964e62f86558..76a4f7dc449e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py @@ -27,7 +27,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_input requires a CUDA device", allow_module_level=True) def _case(spec, seed): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py index a817e8ce00dc..8b40a25a7ef1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_silu.py @@ -14,7 +14,8 @@ from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_silu requires a CUDA device", allow_module_level=True) ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py index 0528b9f1abb3..6742c23ff618 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py @@ -57,7 +57,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_step requires a CUDA device", allow_module_level=True) T = DRAFT_LEN + 1 _LOW32 = 0xFFFFFFFF diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py index 7c9a826c1d3d..7d8180a999fb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_fold_tail.py @@ -20,7 +20,8 @@ m3_spec_tail, ) -assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fold_tail requires a CUDA device", allow_module_level=True) VOCAB = 200064 NUM_SLOTS = 9 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py index 3bad6ccc5d2a..aa664ff99747 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_force_vec requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py index 53a58df60bef..09a40616268f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_force_vec_req requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py index 6d585d8afa32..0fae391d6b99 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py @@ -18,7 +18,8 @@ run_splitk_dense, ) -assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_lowm_splitk requires a CUDA device", allow_module_level=True) # Every certified (N, K, rows) of the TP2 drafter table that a split-K tactic serves: (N, K, rows, tactic). CELLS = [ diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py index 7a59742e1302..c62f29d1f0a8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_sampled_accept.py @@ -17,7 +17,8 @@ from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_sampled_accept requires a CUDA device", allow_module_level=True) DRAFT_LEN = 3 #: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py index 5c6a99bb495a..5ca4be168a4c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_shard_accept.py @@ -27,7 +27,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_shard_accept requires a CUDA device", allow_module_level=True) def _case(spec, seed, edit_logits=None): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py index f0417a37a805..c0aae0581d0f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py @@ -13,7 +13,8 @@ from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_slot_scatter requires a CUDA device", allow_module_level=True) STALE = -4 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py index b023cbd47add..2e39c56397c9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_spec_offsets.py @@ -10,7 +10,8 @@ m3_spec_offsets as op, ) -assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_spec_offsets requires a CUDA device", allow_module_level=True) I32 = torch.iinfo(torch.int32) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py index e6e732174a8b..271b407f83f4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py @@ -13,7 +13,8 @@ from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_step_gather requires a CUDA device", allow_module_level=True) #: Rows of the engine's flat buffers (the max batch), and their stale fill. MAX_ROWS = 8 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py index 13fc115c0afa..8332f85183dd 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_prologue.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_step_prologue requires a CUDA device", allow_module_level=True) MAX_DESC = StepInputPack.MAX_DESC WINDOW = StepInputPack.GRAPH_WINDOW diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py index 37b04013094a..5b1d1603346e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_rank_agreement.py @@ -53,7 +53,7 @@ class _FakeComm: def __init__(self, world): self.world = world - self.barrier = threading.Barrier(world) + self.barrier = threading.Barrier(world, timeout=10) self.lock = threading.Lock() self.vals = [] self.calls = 0 @@ -78,7 +78,7 @@ def _mapping(tp_size=2, adp=False): return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) -def _run_ranks(fns): +def _run_ranks(fns, comm): out, errs = [None] * len(fns), [] def run(i): @@ -86,12 +86,18 @@ def run(i): out[i] = fns[i]() except BaseException as e: # noqa: BLE001 errs.append(e) + comm.barrier.abort() threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] for t in threads: t.start() for t in threads: t.join(timeout=30) + if any(t.is_alive() for t in threads): + comm.barrier.abort() + for t in threads: + t.join(timeout=10) + assert not any(t.is_alive() for t in threads), "rank worker did not terminate" assert not errs, errs return out @@ -123,7 +129,9 @@ def test_tp_agree_refuses_capture(monkeypatch): @pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) def test_tp_agree_built_two_ranks(two_rank_comm, oks): built = [object() if ok else None for ok in oks] - out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + out = _run_ranks( + [lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)], two_rank_comm + ) assert out == (built if all(oks) else [None, None]) @@ -147,7 +155,7 @@ def build(i, num_tokens, shared=True): @pytest.mark.parametrize("fail", [(False, False), (False, True)]) def test_engine_build_agreed_across_ranks(two_rank_comm, fail): cores = [_core(f) for f in fail] - _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores], two_rank_comm) keys = sorted(cores[0].moe_engines[3]) assert keys and two_rank_comm.calls == 2 * len(keys) a, b = cores[0].moe_engines[3], cores[1].moe_engines[3] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py index b58bcfacd648..747a8c3931e8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_silu_and_mul requires a CUDA device", allow_module_level=True) OAI = (7.0, 1.702, 1.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py index 824cb6488e58..0e05e12f6707 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py @@ -31,7 +31,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_dense_decode requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_dense_decode requires a CUDA device", allow_module_level=True) PAGE = 32 SUBPAGES = 2 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py index 98c6a27fc039..b4ccbea622d4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -10,7 +10,8 @@ import tensorrt_llm # noqa: F401 (registers the trtllm ops) -assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the dense FP8 q/k/v op requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py index a594c099c2ed..a0be7f0dd6f4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_draft_attention.py @@ -23,7 +23,8 @@ import _m3_attn_runtime_tp4 as rt # noqa: E402 — needs the path declared above from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_attention requires a CUDA device", allow_module_level=True) # Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. # The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py index e5cec6259f9a..b468359b42cc 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -30,7 +30,8 @@ m3_tokmap_guard as tokmap_guard, ) -assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_drafter_attn requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py index aefbea361b74..02713095ef13 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -28,7 +28,8 @@ ) from tensorrt_llm.bindings.internal import thop -assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the drafter preprocess op requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp4_drafter_rope_fp8_kv"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py index 529a556ea4c3..75d9aa31ccd3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -16,7 +16,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fp8_indexer_qk_norm_rope requires a CUDA device", allow_module_level=True) DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py index 0014a17a7b01..9e7ea8c951d5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fullctx_decode.py @@ -29,7 +29,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fullctx_decode requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py index 1d4353c466cf..829e76748ddf 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py @@ -29,7 +29,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_index_score requires a CUDA device", allow_module_level=True) HEADS = 1 # index heads per rank at TP4 HEAD_DIM = 128 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py index be30baa9bd5b..942bf95e80ce 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score_t2sc.py @@ -24,7 +24,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_index_score requires a CUDA device", allow_module_level=True) HEADS = 1 # index heads per rank at TP4 (4 index heads over 4 KV groups) HEAD_DIM = 128 @@ -97,7 +98,7 @@ def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> Non ): idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) - op(idx_q, index_k, block_table, seq, score, 4, split, True) + op(idx_q, index_k, block_table, seq, score, qlen, split, True) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, qlen ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py index 70cbd7b25768..5df724b0e37f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_msa_attention.py @@ -28,7 +28,8 @@ import _m3_attn_runtime_tp4 as rt # noqa: E402 — needs the path declared above from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_msa_attention requires a CUDA device", allow_module_level=True) # Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, # per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py index f759d0d20d3a..46464d79f6a4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_producer.py @@ -32,7 +32,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_producer requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_producer requires a CUDA device", allow_module_level=True) PAGES = 1024 MAX_POSITIONS = 262144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py index 071795f8c3fa..33e0780fde08 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_qk_norm_rope_fp8 requires a CUDA device", allow_module_level=True) HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py index a7ee9f877707..87023e25c703 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py @@ -18,7 +18,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_qkv_insert requires a CUDA device", allow_module_level=True) PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 16, 1, 1, 128, 64 MAX_POSITIONS = 262144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py index 6f23d21ce10f..1ea6abdf1dc3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py @@ -15,7 +15,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_blocks requires a CUDA device", allow_module_level=True) TOPK, INIT, LOCAL = 16, 0, 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py index 6f4d512ba75e..88def7b245d0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long.py @@ -20,7 +20,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_long requires a CUDA device", allow_module_level=True) # Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key # widths, past the row (clamped), and zero (a CUDA-graph padding row). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py index 816351089a99..3df3c3049bd6 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_long_dsl.py @@ -22,7 +22,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_long_dsl requires a CUDA device", allow_module_level=True) # Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key # widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py index 1c72256dac47..523a09addea8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py @@ -31,7 +31,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_sparse_decode requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_sparse_decode requires a CUDA device", allow_module_level=True) # Band, in BF16 ulps (2^-8) of each [16, 128] head row's largest magnitude, and of the relative RMS. Every route # computes from exact operands (E4M3 q and E2M1 x E4M3 K/V are exact in FP16), so the gap to the float64 reference diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py index 0c0becea2433..f88a42670313 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode_lc.py @@ -25,7 +25,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_sparse_decode_lc requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_sparse_decode_lc requires a CUDA device", allow_module_level=True) #: Hand-off buffers are sized as MERGE_MXFP8 sizes them. _HANDOFF_ROWS = 128 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py index 013e12cd3703..780a53fda069 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py @@ -21,7 +21,8 @@ ) from tensorrt_llm.bindings.internal import thop -assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_tokmap_guard requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp4 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py index a1682f829082..10003353bcbf 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py @@ -25,7 +25,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("msa_kv_update requires a CUDA device", allow_module_level=True) #: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. _SPARE = 3 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py index 127859c5cb2a..3fb3c2bbe14e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py @@ -31,10 +31,12 @@ from pathlib import Path from typing import Any, Callable, List, Optional, Tuple +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("the fused all-reduce + RMSNorm requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 900 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py index db9dba58f108..be06233e492e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py @@ -31,10 +31,12 @@ from pathlib import Path from typing import Any, Dict, Optional +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_mc_allreduce_norm requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 900 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py index 13f7835c7d66..09b0038dceb3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py @@ -47,10 +47,12 @@ from pathlib import Path from typing import Any, Callable, Iterator, List, Optional, Tuple +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("the fused all-reduce + RMSNorm requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 900 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py index 852149f12cc7..1ee5855270b9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -32,10 +32,12 @@ from pathlib import Path from typing import Any, Dict, List, Optional +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_mc_allreduce_norm requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 900 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py index ad7539d27829..ba72d9394bf6 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_capture_in_boundary.py @@ -12,7 +12,8 @@ import pytest import torch -assert torch.cuda.is_available(), "the multicast all-reduce requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("the multicast all-reduce requires CUDA devices", allow_module_level=True) # Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py index 3cc1592bc3c5..fcbdb199320b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_draft_ar_norm.py @@ -16,7 +16,8 @@ import pytest import torch -assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("the fused all-reduce + RMSNorm requires CUDA devices", allow_module_level=True) _RESULT_TAG = "[ulp]" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py index 31a054642e89..b1ae305834b4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -11,7 +11,8 @@ import pytest import torch -assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_mc_allreduce_norm requires CUDA devices", allow_module_level=True) # Spawns 4 rank processes over the visible devices; under xdist workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py index 4258edbc4638..1d2d922612ba 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py @@ -15,7 +15,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_embedding requires a CUDA device", allow_module_level=True) VOCAB, HIDDEN = 4096, 6144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py index f17d7ec63189..cf2493c69fc3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_bf16_linear requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py index c53155b7c02b..729612c53161 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -20,7 +20,8 @@ m3_dense_swiglu, ) -assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the dense swigluoai op requires a CUDA device", allow_module_level=True) if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp4_swiglu_oai"): pytest.skip( "needs the TP4 op library with minimax_m3_sm107_tp4_swiglu_oai", allow_module_level=True diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py index ebc36807e49b..3f84423fd2c5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_draft_lowm.py @@ -33,7 +33,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_lowm requires a CUDA device", allow_module_level=True) # Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). # diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py index fc71fc9f522a..7e6aff6798d8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fp32_linear requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py index 0cd67d512573..ebb388279924 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv.py @@ -25,7 +25,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mx_gemv requires a CUDA device", allow_module_level=True) # Element band, in ULP units of the row scale: see compare. diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py index 2c3947b5cc81..17e5ff6b7109 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -28,7 +28,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mx_gemv_bf16 requires a CUDA device", allow_module_level=True) # The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. ELEMENT_ULP = 1.5 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py index 54379ac1f3e2..d28ea1ad8dd0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -16,7 +16,8 @@ from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mxfp8_gemm requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py index b869ca3aef13..f0199cbc31d0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -16,7 +16,8 @@ from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mxfp8_linear requires a CUDA device", allow_module_level=True) _FLAGS = { "native": (False, False, False, False), diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py index a2ae95a27105..1616fa9bd7d3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py @@ -22,7 +22,8 @@ from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mxfp8_pinned_gemm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mxfp8_pinned_gemm requires a CUDA device", allow_module_level=True) TACTIC = modeling.MXFP8_DECODE_GEMM_TACTIC diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py index 22e987e17979..6f31b14c1c38 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_quantize.py @@ -20,7 +20,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mxfp8_pinned_quantize requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mxfp8_pinned_quantize requires a CUDA device", allow_module_level=True) TACTIC = modeling.MXFP8_DECODE_QUANT_TACTIC diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py index 467cfd888e5e..8bb9fc8fd1ec 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py @@ -38,7 +38,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_target_head requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_target_head requires a CUDA device", allow_module_level=True) def _seed(why: str) -> int: diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py index 4ad3e870f33a..c4818ec67f8c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe16_engine.py @@ -27,7 +27,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_moe16_engine requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_moe16_engine requires a CUDA device", allow_module_level=True) # As test_modeling_v2_m3_moe_engine.compare (same arithmetic, same quantization boundaries). BAND_ELEMENT = 4.0 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py index c77b345c4bc7..480be0f4f794 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_moe_engine.py @@ -27,7 +27,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_moe_engine requires a CUDA device", allow_module_level=True) _STATES = {} diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py index c207756c1b49..c0fde3f9e81d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -27,7 +27,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_embedding_firstnorm requires a CUDA device", allow_module_level=True) H = 6144 EPS = 1e-6 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py index d31aebfe2610..d2a171790f64 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py @@ -25,7 +25,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_capture_norm requires a CUDA device", allow_module_level=True) LAYERS = 3 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py index 2caf1da9304d..8848ed3a2b6f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py @@ -24,7 +24,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_chunk_norm requires a CUDA device", allow_module_level=True) def _case(spec, seed): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py index 2f82a2a9d63b..a4c8090bf4d2 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py @@ -32,7 +32,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_input requires a CUDA device", allow_module_level=True) def _case(spec, seed): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py index 6d1183397871..d43cbaad5d7a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_silu.py @@ -24,7 +24,8 @@ def gen(seed: int) -> torch.Generator: return torch.Generator(device="cuda").manual_seed(seed) -assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_silu requires a CUDA device", allow_module_level=True) ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py index 8bb087bf3cd4..1a17fa8ee539 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -62,7 +62,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_step requires a CUDA device", allow_module_level=True) T = DRAFT_LEN + 1 _LOW32 = 0xFFFFFFFF diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py index ef33aecf3a6a..edcfd8a451f4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -55,7 +55,8 @@ seed_of, ) -assert torch.cuda.is_available(), "the exact-numerics checks require a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the exact-numerics checks require a CUDA device", allow_module_level=True) DRAFT_LEN = 3 CAPTURED_LAYERS = 3 @@ -255,6 +256,8 @@ def test_draft_tokens_through_the_full_vocab_head(rows) -> None: top2 = stock_logits.topk(2, dim=-1).values decidable = (top2[:, 0] - top2[:, 1]) > 2 * perturb stock_tok, ours_tok = stock_logits.argmax(-1), ours_logits.argmax(-1) + assert torch.isfinite(ours_logits).all() + assert bool(decidable.any()), "no draft-token rows satisfy the reference margin" assert torch.equal(ours_tok[decidable], stock_tok[decidable]) print( f"[exact-numerics] rows={rows} draft head: {int((ours_tok == stock_tok).sum())}/{rows} tokens equal, " diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py index d3e117c326e7..fced3c8d7785 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_fold_tail.py @@ -20,7 +20,8 @@ m3_spec_tail, ) -assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fold_tail requires a CUDA device", allow_module_level=True) VOCAB = 200064 NUM_SLOTS = 9 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py index 717deed4c1aa..122201ef3e65 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_force_vec requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py index 8a4d03c2fd03..ba7e9b58accb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec_req.py @@ -19,7 +19,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_force_vec_req requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py index 01aa7eb2af53..7b457c6bdc3e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py @@ -18,7 +18,8 @@ run_splitk_dense, ) -assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_lowm_splitk requires a CUDA device", allow_module_level=True) # Every certified (N, K, rows) of the TP4 drafter table that a split-K tactic serves: (N, K, rows, tactic). CELLS = [ diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py index 4ed33c8a7251..68db7d98209e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_sampled_accept.py @@ -21,7 +21,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_sampled_accept requires a CUDA device", allow_module_level=True) DRAFT_LEN = 3 #: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py index 0c90ceac2e70..3b81085ff29d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_shard_accept.py @@ -32,7 +32,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_shard_accept requires a CUDA device", allow_module_level=True) def _case(spec, seed, edit_logits=None): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py index 4b27e42af50b..7f48d822e327 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py @@ -18,7 +18,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_slot_scatter requires a CUDA device", allow_module_level=True) STALE = -4 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py index 4a2a242816d3..74e63d9e8e5e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_spec_offsets.py @@ -10,7 +10,8 @@ m3_spec_offsets as op, ) -assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_spec_offsets requires a CUDA device", allow_module_level=True) I32 = torch.iinfo(torch.int32) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py index c57435ce1f30..ce5562f0d404 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py @@ -18,7 +18,8 @@ reason="the sm_107 TP4 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_step_gather requires a CUDA device", allow_module_level=True) #: Rows of the engine's flat buffers (the max batch), and their stale fill. MAX_ROWS = 8 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py index a5d0cec0ed6c..496fe3ff12ae 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_prologue.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_step_prologue requires a CUDA device", allow_module_level=True) MAX_DESC = StepInputPack.MAX_DESC WINDOW = StepInputPack.GRAPH_WINDOW diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py index 3046e9bc61bd..0e915fbcb18e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py @@ -39,7 +39,8 @@ dispatch = modeling -assert torch.cuda.is_available(), "the prebuild checks require a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the prebuild checks require a CUDA device", allow_module_level=True) INDEX_HEADS = 4 INDEX_DIM = 128 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py index 0b5d60b19523..433090091875 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_rank_agreement.py @@ -53,7 +53,7 @@ class _FakeComm: def __init__(self, world): self.world = world - self.barrier = threading.Barrier(world) + self.barrier = threading.Barrier(world, timeout=10) self.lock = threading.Lock() self.vals = [] self.calls = 0 @@ -78,7 +78,7 @@ def _mapping(tp_size=2, adp=False): return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) -def _run_ranks(fns): +def _run_ranks(fns, comm): out, errs = [None] * len(fns), [] def run(i): @@ -86,12 +86,18 @@ def run(i): out[i] = fns[i]() except BaseException as e: # noqa: BLE001 errs.append(e) + comm.barrier.abort() threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] for t in threads: t.start() for t in threads: t.join(timeout=30) + if any(t.is_alive() for t in threads): + comm.barrier.abort() + for t in threads: + t.join(timeout=10) + assert not any(t.is_alive() for t in threads), "rank worker did not terminate" assert not errs, errs return out @@ -123,7 +129,9 @@ def test_tp_agree_refuses_capture(monkeypatch): @pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) def test_tp_agree_built_two_ranks(two_rank_comm, oks): built = [object() if ok else None for ok in oks] - out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + out = _run_ranks( + [lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)], two_rank_comm + ) assert out == (built if all(oks) else [None, None]) @@ -148,7 +156,7 @@ def build(i, num_tokens, shared=True): @pytest.mark.parametrize("fail", [(False, False), (False, True)]) def test_engine_build_agreed_across_ranks(two_rank_comm, fail): cores = [_core(f) for f in fail] - _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores], two_rank_comm) keys = sorted(cores[0].moe_engines[3]) + sorted(cores[0].moe16_engines[3]) assert keys and two_rank_comm.calls == 2 * len(keys) for engines in ("moe_engines", "moe16_engines"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py index 1153eaee47bd..891ffce1d511 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_silu_and_mul requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_silu_and_mul requires a CUDA device", allow_module_level=True) OAI = (7.0, 1.702, 1.0) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py index bef0db5ce41b..07d91aef8ad1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_dense_qk_norm_rope_fp8.py @@ -10,7 +10,8 @@ import tensorrt_llm # noqa: F401 (registers the trtllm ops) -assert torch.cuda.is_available(), "the dense FP8 q/k/v op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the dense FP8 q/k/v op requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py index c89183995648..7ee756142c77 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_draft_attention.py @@ -23,7 +23,8 @@ import _m3_attn_runtime_tp8 as rt # noqa: E402 — needs the path declared above from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_attention requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_attention requires a CUDA device", allow_module_level=True) # Band, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS. # The reference reads the same FP8 cache rows the call attends, so what is left is the FP8 FMHA's own arithmetic: diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py index 174be4bef5a1..41a0f02ddfb1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -27,7 +27,8 @@ m3_fullctx_decode as fullctx, ) -assert torch.cuda.is_available(), "m3_drafter_attn requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_drafter_attn requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py index b1759f370b7a..44188d52a4bb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -28,7 +28,8 @@ ) from tensorrt_llm.bindings.internal import thop -assert torch.cuda.is_available(), "the drafter preprocess op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the drafter preprocess op requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) if not hasattr(torch.ops.trtllm, "minimax_m3_sm107tp8_drafter_rope_fp8_kv"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py index 2179f1720e0d..baba0c067512 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -16,7 +16,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_fp8_indexer_qk_norm_rope requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fp8_indexer_qk_norm_rope requires a CUDA device", allow_module_level=True) DIM, ROTARY, EPS, THETA, PAGE = 128, 64, 1e-6, 5e6, 128 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py index 48d1bba922ac..5c3d90809f0f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fullctx_decode.py @@ -29,7 +29,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_fullctx_decode requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fullctx_decode requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py index e768b7442a15..9693051474f2 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py @@ -23,7 +23,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_index_score requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_index_score requires a CUDA device", allow_module_level=True) HEADS = 1 # index heads per rank at TP8 (4 over 4 KV groups, each group on a rank pair) HEAD_DIM = 128 @@ -96,7 +97,7 @@ def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> Non ): idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) - op(idx_q, index_k, block_table, seq, score, 4, split, True) + op(idx_q, index_k, block_table, seq, score, qlen, split, True) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, qlen ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py index 0539d9639125..d03456c05cb6 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_msa_attention.py @@ -28,7 +28,8 @@ import _m3_attn_runtime_tp8 as rt # noqa: E402 — needs the path declared above from _validating import ULP, assert_within_ulp, certifying # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_msa_attention requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_msa_attention requires a CUDA device", allow_module_level=True) # Bands, in BF16 ulps (2^-8) of each token's [heads x 128] row's largest magnitude (element) and of the relative RMS, # per kernel the step reaches. The reference reads the same quantized cache rows (and casts Q to E4M3 where the FP8 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py index 124942104b82..b3a28b962c7c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_producer.py @@ -32,7 +32,8 @@ reason="the sm_107 TP8 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_producer requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_producer requires a CUDA device", allow_module_level=True) PAGES = 1024 MAX_POSITIONS = 262144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py index 6e58bacee51c..e48bf7aabc0d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_qk_norm_rope_fp8 requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_qk_norm_rope_fp8 requires a CUDA device", allow_module_level=True) HEAD_DIM, ROTARY_DIM, EPS, THETA = 128, 64, 1e-6, 5e6 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py index 514567365f8c..85c1a03c59be 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py @@ -18,7 +18,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_qkv_insert requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_qkv_insert requires a CUDA device", allow_module_level=True) PAGES, PAGE, HEADS, KV, INDEX, D, ROTARY = 256, 128, 8, 1, 1, 128, 64 MAX_POSITIONS = 262144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py index bf3c3c271148..af3b3c3d4f0e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py @@ -15,7 +15,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_blocks requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_blocks requires a CUDA device", allow_module_level=True) TOPK, INIT, LOCAL = 16, 0, 1 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py index 945025a650ca..d202e89141b9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long.py @@ -20,7 +20,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_long requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_long requires a CUDA device", allow_module_level=True) # Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 2-, 4- and 9-key # widths, past the row (clamped), and zero (a CUDA-graph padding row). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py index 4022b2cfb133..265765b92f3c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_long_dsl.py @@ -22,7 +22,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above -assert torch.cuda.is_available(), "m3_select_long_dsl requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_select_long_dsl requires a CUDA device", allow_module_level=True) # Valid-block counts that land a row on each of the kernel's paths: <= 16 (all picked), the 4-, 11- and 24-key # widths and their edges, past the row (clamped), and zero (a CUDA-graph padding row). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py index ac61687d6968..296b43e23c30 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode_lc8.py @@ -32,7 +32,8 @@ reason="the sm_107 TP8 target runs on SM107 only", ) -assert torch.cuda.is_available(), "m3_sparse_decode_lc8 requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_sparse_decode_lc8 requires a CUDA device", allow_module_level=True) #: Hand-off buffers are sized as MERGE_MXFP8 sizes them. _HANDOFF_ROWS = 128 @@ -328,7 +329,7 @@ def test_one_div_merge_is_deterministic_and_in_band( def run(out, q8, sf, use_pdl, one_div): args = list(_args(case, out, q8, sf, linear, prewait)) - args[18] = use_pdl + args[17] = use_pdl op(*args, merge_once=True, one_div=one_div) runs = [] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py index e21326054b2e..461e9695b6b2 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py @@ -21,7 +21,8 @@ ) from tensorrt_llm.bindings.internal import thop -assert torch.cuda.is_available(), "m3_tokmap_guard requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_tokmap_guard requires a CUDA device", allow_module_level=True) if torch.cuda.get_device_capability() != (10, 7): pytest.skip("the tp8 target is certified on SM 10.7 only", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py index c29fd0085966..a9aea90d3eaa 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py @@ -20,7 +20,8 @@ __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 — needs the path declared above -assert torch.cuda.is_available(), "msa_kv_update requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("msa_kv_update requires a CUDA device", allow_module_level=True) #: Spare capacity past the live rows, as the metadata's buffers have: the op must leave it alone. _SPARE = 3 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py index 45b043d24e95..006771133817 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py @@ -34,10 +34,12 @@ from pathlib import Path from typing import Any, Callable, List, Optional, Tuple +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "the fused all-reduce + RMSNorm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("the fused all-reduce + RMSNorm requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 900 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py index 6d7def71212a..b493b3a4afd9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -42,10 +42,12 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Tuple +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_mc_allreduce_norm requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 900 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py index 95374cf56848..4f0c30322f4a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py @@ -19,7 +19,7 @@ in CUDA_VISIBLE_DEVICES and kills them all on the first failure or on the deadline (a rank whose peer never pushes polls forever). The collected entry point is ``test_modeling_v2_m3_moe_push_op_matrix.py``. - CUDA_VISIBLE_DEVICES=0,1 python _m3_moe_push_op_matrix.py + CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python _m3_moe_push_op_matrix.py """ import os @@ -31,10 +31,12 @@ from pathlib import Path from typing import Any, Optional +import pytest import torch import torch.distributed as dist -assert torch.cuda.is_available(), "m3_moe_engine requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_moe_engine requires CUDA devices", allow_module_level=True) _WORKER_FLAG = "--rank-worker" DEADLINE_S = 1200 @@ -324,10 +326,10 @@ def _spawn_ranks() -> None: """One worker process per claimed device, each in its own session; on the first nonzero exit or on the deadline every session is killed.""" visible = os.environ.get("CUDA_VISIBLE_DEVICES") - assert visible, "set CUDA_VISIBLE_DEVICES to the 2 devices this run owns" + assert visible, "set CUDA_VISIBLE_DEVICES to the 8 devices this run owns" world_size = len([d for d in visible.split(",") if d.strip()]) - assert world_size == 2, ( - f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 2" + assert world_size == 8, ( + f"CUDA_VISIBLE_DEVICES names {world_size} device(s); the matrix runs at 8" ) with socket.socket() as s: s.bind(("127.0.0.1", 0)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py index 0a9a7d39c177..0f12b78cb73a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/test_modeling_v2_m3_mc_allreduce_norm_op_matrix.py @@ -22,7 +22,8 @@ finalize_partial, ) -assert torch.cuda.is_available(), "m3_mc_allreduce_norm requires CUDA devices" +if not torch.cuda.is_available(): + pytest.skip("m3_mc_allreduce_norm requires CUDA devices", allow_module_level=True) # Spawns one rank process per device of the world size; under xdist workers would fight for them. diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py index 546ad3e11ef7..2ae718a482a1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py @@ -15,7 +15,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_embedding requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_embedding requires a CUDA device", allow_module_level=True) VOCAB, HIDDEN = 4096, 6144 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py index 509c1c09cebd..4386d7434b8a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_bf16_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_bf16_linear requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py index 67796dc0302f..ca565059786a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -20,7 +20,8 @@ m3_dense_swiglu, ) -assert torch.cuda.is_available(), "the dense swigluoai op requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the dense swigluoai op requires a CUDA device", allow_module_level=True) if not hasattr(torch.ops.trtllm, "minimax_m3_sm107_tp8_swiglu_oai"): pytest.skip( "needs the TP8 op library with minimax_m3_sm107_tp8_swiglu_oai", allow_module_level=True diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py index e63d7f3f11ed..919d76f7b539 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_draft_lowm.py @@ -30,7 +30,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_lowm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_lowm requires a CUDA device", allow_module_level=True) # Element band, in BF16 ulps of the row scale, and the relative RMS (``assert_within_ulp``). # diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py index 149d230a929c..94fd72e397fb 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py @@ -14,7 +14,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_fp32_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fp32_linear requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py index 918aa9e5ffab..100114112b01 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv.py @@ -20,7 +20,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mx_gemv requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mx_gemv requires a CUDA device", allow_module_level=True) # Element band, in ULP units of the row scale: see compare. diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py index c8c9461b44da..ebd80213194f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mx_gemv_bf16.py @@ -23,7 +23,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mx_gemv_bf16 requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mx_gemv_bf16 requires a CUDA device", allow_module_level=True) # The element band of test_modeling_v2_m3_mx_gemv.compare (derived there): the same fp32 sums against fp64. ELEMENT_ULP = 1.5 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py index 4374a72aa8a0..f3b9d7c10cdf 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -16,7 +16,8 @@ from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mxfp8_gemm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mxfp8_gemm requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py index dcff36cd9b93..82938bf53a14 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -16,7 +16,8 @@ from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_mxfp8_linear requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_mxfp8_linear requires a CUDA device", allow_module_level=True) _FLAGS = { "native": (False, False, False, False), diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py index 457d56dd8698..5ffb8f58d48c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_pinned_gemm.py @@ -13,7 +13,8 @@ ) from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner -assert torch.cuda.is_available(), "MX-TAC requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("MX-TAC requires a CUDA device", allow_module_level=True) # The TP8 per-rank projections (ops/gemm/m3_mx_gemv) at decode rows the GEMVs do not serve (9..16) and a few beyond. CASES = [ diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py index 1826394596eb..70ff6e13cdc4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py @@ -16,7 +16,8 @@ _swiglu_oai_bf16x2, ) -assert torch.cuda.is_available(), "the swigluoai prologue oracle requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("the swigluoai prologue oracle requires a CUDA device", allow_module_level=True) THREADS = 256 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py index 8836a2e7726d..4f90dda67ae3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_moe_engine.py @@ -33,7 +33,8 @@ ) from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_moe_engine requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_moe_engine requires a CUDA device", allow_module_level=True) pytestmark = pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7), diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py index 05f9e80edbd9..09cee232deb9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/norm/test_modeling_v2_m3_embedding_firstnorm.py @@ -22,7 +22,8 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_embedding_firstnorm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_embedding_firstnorm requires a CUDA device", allow_module_level=True) H = 6144 EPS = 1e-6 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py index 605c96c773cc..121001424351 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py @@ -20,7 +20,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_capture_norm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_capture_norm requires a CUDA device", allow_module_level=True) LAYERS = 3 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py index 28d5cbefc927..b28c1a5057ed 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_normed.py @@ -30,7 +30,13 @@ def _worker(normed): def _scaling_norm(monkeypatch): """The RMSNorm the worker applies, as ``x * w`` (CPU): enough to tell which blocks were normalized here.""" monkeypatch.setattr(modeling, "flashinfer_rmsnorm", lambda x, w, eps: x * w) + require = modeling.require_mode() + monkeypatch.setattr(modeling, "_logged", set()) modeling.set_require(False) + try: + yield + finally: + modeling.set_require(require) def _drafter(): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py index 8e8982524c8b..bced642e6b8b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py @@ -19,7 +19,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_chunk_norm requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_chunk_norm requires a CUDA device", allow_module_level=True) def _case(spec, seed): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py index f1f5f6160e32..972204f0d653 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py @@ -27,7 +27,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_input requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_input requires a CUDA device", allow_module_level=True) def _case(spec, seed): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py index 5ae74445fe33..cf8dfc4b76d0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_silu.py @@ -14,7 +14,8 @@ from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_silu requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_silu requires a CUDA device", allow_module_level=True) ALL = torch.arange(-(1 << 15), 1 << 15, dtype=torch.int32).to(torch.int16).view(torch.bfloat16) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py index 877e52628100..ba136429a30e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -58,7 +58,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_draft_step requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_draft_step requires a CUDA device", allow_module_level=True) T = DRAFT_LEN + 1 _LOW32 = 0xFFFFFFFF diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py index 5e6ae369f8fa..6ff42bc84120 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_fold_tail.py @@ -20,7 +20,8 @@ m3_spec_tail, ) -assert torch.cuda.is_available(), "m3_fold_tail requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_fold_tail requires a CUDA device", allow_module_level=True) VOCAB = 200064 NUM_SLOTS = 9 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py index 6b63ba265112..16fe304dadf8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_force_vec requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_force_vec requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py index 871dfa22d617..84ef40ee731b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py @@ -17,7 +17,8 @@ __extra_import_path__ = ["../../../.."] from _validating import validating # noqa: E402 — needs the path declared above -assert torch.cuda.is_available(), "m3_force_vec_req requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_force_vec_req requires a CUDA device", allow_module_level=True) @pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py index 821c6eeb08ae..02a447e0e857 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py @@ -18,7 +18,8 @@ run_splitk_dense_swiglu, ) -assert torch.cuda.is_available(), "m3_lowm_splitk requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_lowm_splitk requires a CUDA device", allow_module_level=True) # Every split-K cell of the TP8 drafter table: (N, K, max rows, tactic) for kind "s". CELLS = [ diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py index 6066e23c92a3..1610c34192ca 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_sampled_accept.py @@ -17,7 +17,8 @@ from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_sampled_accept requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_sampled_accept requires a CUDA device", allow_module_level=True) DRAFT_LEN = 3 #: Hot tokens of the distribution cells; every other logit sits 22 below the top (mass ~ 1e-5 of the total). diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py index ca795ef4f3ca..213de2535bf7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_shard_accept.py @@ -27,7 +27,8 @@ ) from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_shard_accept requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_shard_accept requires a CUDA device", allow_module_level=True) def _case(spec, seed, edit_logits=None): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py index fab3b7fd7853..f29ff1ddb175 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py @@ -13,7 +13,8 @@ from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_slot_scatter requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_slot_scatter requires a CUDA device", allow_module_level=True) STALE = -4 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py index c4ed9c08c2c3..2e8b405a4e0b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_spec_offsets.py @@ -10,7 +10,8 @@ m3_spec_offsets as op, ) -assert torch.cuda.is_available(), "m3_spec_offsets requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_spec_offsets requires a CUDA device", allow_module_level=True) I32 = torch.iinfo(torch.int32) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py index 1c23a9357876..fa284cfc029f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py @@ -13,7 +13,8 @@ from _m3_spec_inputs import VOCAB, gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_step_gather requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_step_gather requires a CUDA device", allow_module_level=True) #: Rows of the engine's flat buffers (the max batch), and their stale fill. MAX_ROWS = 8 diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py index a4076e46f9d1..c76f24abeb0e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_prologue.py @@ -17,7 +17,8 @@ from _m3_spec_inputs import gen, seed_of # noqa: E402 — needs the path declared above from _validating import certifying, validating # noqa: E402, F401 -assert torch.cuda.is_available(), "m3_step_prologue requires a CUDA device" +if not torch.cuda.is_available(): + pytest.skip("m3_step_prologue requires a CUDA device", allow_module_level=True) # The reference and the pack tests run this branch's step_inputs (the step epoch, prologue_op): a run that mounts only # the target's files gets the image's older module, whose graph scatter has no epoch argument. assert "epoch" in step_inputs._graph_scatter_kernel.arg_names, ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py index 8f28a67e722f..c4e4af933407 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/test_modeling_v2_m3_rank_agreement.py @@ -53,7 +53,7 @@ class _FakeComm: def __init__(self, world): self.world = world - self.barrier = threading.Barrier(world) + self.barrier = threading.Barrier(world, timeout=10) self.lock = threading.Lock() self.vals = [] self.calls = 0 @@ -78,7 +78,7 @@ def _mapping(tp_size=2, adp=False): return SimpleNamespace(tp_size=tp_size, enable_attention_dp=adp, tp_group=list(range(tp_size))) -def _run_ranks(fns): +def _run_ranks(fns, comm): out, errs = [None] * len(fns), [] def run(i): @@ -86,12 +86,18 @@ def run(i): out[i] = fns[i]() except BaseException as e: # noqa: BLE001 errs.append(e) + comm.barrier.abort() threads = [threading.Thread(target=run, args=(i,)) for i in range(len(fns))] for t in threads: t.start() for t in threads: t.join(timeout=30) + if any(t.is_alive() for t in threads): + comm.barrier.abort() + for t in threads: + t.join(timeout=10) + assert not any(t.is_alive() for t in threads), "rank worker did not terminate" assert not errs, errs return out @@ -127,7 +133,9 @@ def test_tp_agree_refuses_capture(monkeypatch): @pytest.mark.parametrize("oks", [(True, True), (True, False), (False, True), (False, False)]) def test_tp_agree_built_two_ranks(two_rank_comm, oks): built = [object() if ok else None for ok in oks] - out = _run_ranks([lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)]) + out = _run_ranks( + [lambda i=i: mc.tp_agree_built(_mapping(2), built[i], "t") for i in range(2)], two_rank_comm + ) assert out == (built if all(oks) else [None, None]) @@ -156,7 +164,7 @@ def build(i, num_tokens, shared=True): def test_engine_build_agreed_across_ranks(monkeypatch, world, failing): comm = _patch_comm(monkeypatch, world) cores = [_core(rank in failing, world) for rank in range(world)] - _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores]) + _run_ranks([lambda c=c: c._prebuild_engines(3) for c in cores], comm) keys = sorted(cores[0].moe_engines[3]) assert keys and comm.calls == world * len(keys) assert not any(hasattr(c, "moe16_engines") for c in cores) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_shared_kernel_contracts.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_shared_kernel_contracts.py new file mode 100644 index 000000000000..44d90b4cf39c --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_shared_kernel_contracts.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.embedding_firstnorm import op as embedding +from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.moe import op as moe +from tensorrt_llm._torch.cute_dsl_kernels.minimax_m3.mx_gemv import op as gemv + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("n,k", [(1537, 6144), (1536, 6145), (0, 6144), (1536, 0)]) +def test_explicit_gemv_config_declines_incomplete_tiles(n, k): + assert gemv._declines(n, k, 1, False, (2, 12)) + + +@pytest.mark.cpu_only +def test_frozen_warmup_retries_failed_compilation(monkeypatch): + monkeypatch.setattr(gemv, "_ensured", set()) + monkeypatch.setattr(gemv, "_frozen", [True]) + weight = torch.empty(1536, 6144, dtype=torch.float8_e4m3fn) + scales = torch.empty(0) + calls = [] + + def fail(*args, **kwargs): + raise RuntimeError("compile failed") + + monkeypatch.setattr(gemv, "precompile", fail) + with pytest.raises(RuntimeError, match="compile failed"): + gemv.ensure_compiled(weight, scales, False) + monkeypatch.setattr( + gemv, "precompile", lambda *args, **kw: calls.append((args[2], kw["x_bf16"])) + ) + gemv.ensure_compiled(weight, scales, False) + assert calls == [(1, False), (16, False), (32, False), (48, False), (1, True)] + gemv.ensure_compiled(weight, scales, False) + assert len(calls) == 5 + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("inference", [False, True]) +def test_reloaded_shared_weights_refresh_the_cached_tiles(monkeypatch, inference): + monkeypatch.setattr(moe, "_PRETILED", {}) + monkeypatch.setattr(moe, "_pretile_shared", lambda *sources: torch.cat(sources)) + with torch.inference_mode(inference): + sources = [torch.arange(4) for _ in range(4)] + before = moe._pretiled(*sources) + sources[0].add_(10) + after = moe._pretiled(*sources) + assert torch.equal(after, torch.cat(sources)) + assert after.data_ptr() == before.data_ptr() + # A later state can be constructed outside the context that created its weights. + assert torch.equal(moe._pretiled(*sources), torch.cat(sources)) + + +@pytest.mark.cpu_only +def test_shared_state_default_reserves_shared_ctas(monkeypatch): + class ReachedKernelSelection(Exception): + pass + + def kernel(config, path): + assert config["num_ctas"] == 148 - moe.SHARED_CTAS + raise ReachedKernelSelection() + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda device: Mock(spec_set=["multi_processor_count"], multi_processor_count=148), + ) + monkeypatch.setattr(moe, "is_supported", lambda device: (True, "")) + monkeypatch.setattr(moe, "make_phase_plan", lambda *args: (torch.zeros(1), 1)) + monkeypatch.setattr(moe, "_kernel_module", kernel) + with pytest.raises(ReachedKernelSelection): + moe.M3MoeState(torch.device("cuda", 0), num_tokens=6, shared=True) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA tensor storage required") +@pytest.mark.parametrize("bad", ["offset", "stride", "cpu"]) +def test_embedding_rejects_unsupported_buffers(bad): + h = embedding.HIDDEN + ids = torch.zeros(1, dtype=torch.int32, device="cuda") + table = torch.empty(2, h, dtype=torch.bfloat16, device="cuda") + weight = torch.empty(h, dtype=torch.bfloat16, device="cuda") + raw = torch.empty(1, h, dtype=torch.bfloat16, device="cuda") + out = torch.empty_like(raw) + embedding._check(ids, table, weight, raw, out) + if bad == "offset": + out = torch.empty(h + 1, dtype=torch.bfloat16, device="cuda")[1:].view(1, h) + elif bad == "stride": + weight = torch.empty(2 * h, dtype=torch.bfloat16, device="cuda")[::2] + else: + ids = ids.cpu() + with pytest.raises(ValueError): + embedding._check(ids, table, weight, raw, out) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA tensor storage required") +@pytest.mark.parametrize("bad", ["shape", "dtype", "stride", "device", "offset"]) +def test_moe_rejects_invalid_output_storage(bad): + device = torch.device("cuda", torch.cuda.current_device()) + out = torch.empty(2, moe.HIDDEN, device=device, dtype=torch.bfloat16) + moe._validate_output(out, 2, device) + if bad == "offset": + out = torch.empty(2 * moe.HIDDEN + 1, device=device, dtype=torch.bfloat16)[1:].view( + 2, moe.HIDDEN + ) + elif bad == "shape": + out = out[:1] + elif bad == "dtype": + out = out.float() + elif bad == "stride": + out = torch.empty(2, 2 * moe.HIDDEN, device=device, dtype=torch.bfloat16)[:, ::2] + else: + out = out.cpu() + with pytest.raises(ValueError): + moe._validate_output(out, 2, device) + + +@pytest.mark.cpu_only +def test_gemv_declines_transposed_weights_and_strided_activations(): + x = torch.empty(1, 6144, dtype=torch.float8_e4m3fn) + weight = torch.empty(6144, 1536, dtype=torch.float8_e4m3fn).T + scale = torch.empty(0) + assert gemv.m3_mx_gemv(x, scale, weight, scale) is None + x = torch.empty(1, 12288, dtype=torch.float8_e4m3fn)[:, ::2] + assert gemv.m3_mx_gemv(x, scale, weight.contiguous(), scale) is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA binding validation required") +@pytest.mark.parametrize("geometry", [(0, 64, 1), (1, 0, 1), (1, 64, 0), (1 << 32, 64, 1)]) +@pytest.mark.parametrize("tp", [2, 4, 8]) +def test_full_context_reduce_rejects_invalid_geometry(geometry, tp): + partial = torch.zeros(1, 1, 1, 1, 128, device="cuda", dtype=torch.bfloat16) + stats = torch.zeros(1, 1, 1, 1, device="cuda") + lens = torch.ones(1, device="cuda", dtype=torch.int32) + out = torch.empty(1, 128, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="producer geometry"): + prefix = "sm107_tp2" if tp == 2 else f"sm107tp{tp}" + getattr(torch.ops.trtllm, f"minimax_m3_{prefix}_fullctx_reduce")( + partial, stats, lens, out, 1, *geometry, True, False + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA binding validation required") +@pytest.mark.parametrize("one_warp", [False, True]) +@pytest.mark.parametrize( + "undersized", ["src_new", "slots", "ids", "draft_ids", "pos_off", "kv_off"] +) +@pytest.mark.parametrize("tp", [4, 8]) +def test_fold_tail_rejects_undersized_buffers(one_warp, undersized, tp): + def ints(*shape): + return torch.zeros(*shape, device="cuda", dtype=torch.int32) + + args = dict( + ctrl=ints(1), + slots=torch.arange(2, device="cuda", dtype=torch.int64), + rows=2, + src_new=ints(2, 1), + dst_new=ints(2, 1), + src_next=ints(2, 2), + dst_next=ints(2, 2), + src_draft=ints(2, 1), + dst_draft=ints(2, 1), + src_lens=ints(2), + dst_lens=ints(2), + ids=ints(4), + draft_ids=ints(2), + pos_off=ints(4), + kv_off=ints(2), + tokens_per_step=2, + ) + args[undersized] = args[undersized][:1] + suffix = "_one_warp" if one_warp else "" + op = getattr(torch.ops.trtllm, f"minimax_m3_sm107_tp{tp}_fold_tail{suffix}") + with pytest.raises(RuntimeError, match="must hold rows|staging outputs"): + op(**args) diff --git a/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py b/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py index 374a712d3325..b593ab930904 100644 --- a/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py +++ b/tests/unittest/_torch/pyexecutor/test_cuda_graph_runner.py @@ -643,3 +643,51 @@ def short(threshold: int, host_tokens: int) -> bool: short(span, t) and worst_attended(t) > span for t in range(span - 4 * tokens_per_step, span + 1) ) + + +@pytest.mark.parametrize("same_pack", [False, True]) +@pytest.mark.parametrize("replay_fails", [False, True]) +def test_replay_defers_copies_only_to_its_captured_pack(monkeypatch, same_pack, replay_fails): + key = KeyType(batch_size=1, draft_len=0) + captured = Mock() + active = captured if same_pack else Mock() + queued = [] + metadata = object() + inputs = { + "attn_metadata": metadata, + "input_ids": torch.zeros(1), + "position_ids": torch.zeros(1, 1), + } + runner = object.__new__(CUDAGraphRunner) + runner.graph_metadata = {key: {"attn_metadata": metadata, "spec_metadata": None}} + runner.shared_static_tensors = inputs + runner.config = SimpleNamespace(use_mrope=False) + runner._check_strict_buffer_stability = lambda meta: None + runner._get_num_tokens_for_key = lambda key: 1 + runner._step_input_prologue_keys = {key: captured} + runner._pre_replay_hooks = [] + + def replay(): + queued.append("replay") + if replay_fails: + raise RuntimeError("replay failed") + + runner.graphs = {key: SimpleNamespace(replay=replay)} + runner.graph_outputs = {key: None} + captured.on_graph_replay.side_effect = lambda: queued.append("release") + monkeypatch.setattr( + cuda_graph_runner_module.step_inputs, "is_active", lambda pack: pack is active + ) + monkeypatch.setattr( + cuda_graph_runner_module.step_inputs, "flush", lambda *, in_graph: queued.append(in_graph) + ) + if replay_fails: + with pytest.raises(RuntimeError, match="replay failed"): + runner.replay(key, inputs) + else: + runner.replay(key, inputs) + assert queued == [same_pack, "replay", "release"] + if same_pack: + captured.quiesce_graph_window.assert_not_called() + else: + captured.quiesce_graph_window.assert_called_once_with() diff --git a/tests/unittest/_torch/pyexecutor/test_step_inputs.py b/tests/unittest/_torch/pyexecutor/test_step_inputs.py index 9ca87ba292fd..a00652c45547 100644 --- a/tests/unittest/_torch/pyexecutor/test_step_inputs.py +++ b/tests/unittest/_torch/pyexecutor/test_step_inputs.py @@ -3,11 +3,14 @@ """Packing a decode step's small input copies into one upload and one scatter launch.""" import threading +import weakref +from types import SimpleNamespace import pytest import torch from tensorrt_llm._torch.pyexecutor import h2d_cache, step_inputs +from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDAGraphRunner, KeyType from tensorrt_llm._torch.pyexecutor.h2d_cache import H2DSkipCache, copy_h2d_if_changed from tensorrt_llm._torch.pyexecutor.step_inputs import StepInputPack @@ -723,6 +726,7 @@ def _step(pack, graph, dsts, values, overlap=False, flush=True): step_inputs.discard() step_inputs.flush(in_graph=True) graph.replay() + pack.on_graph_replay() def test_in_graph_pack_is_issued_by_the_replay(): @@ -942,3 +946,53 @@ def test_eager_flush_does_not_advance_the_step_epoch(): torch.cuda.synchronize() assert dst.tolist() == [1, 2, 3, 4] assert int(pack.epoch.item()) == start + + +@pytest.mark.parametrize("early", [False, True]) +def test_graph_window_keeps_temporary_device_sources_until_replay(early): + pack, graph, (a, b, c), out = _graph_case(early) + source = torch.arange(16, device="cuda", dtype=a.dtype) + source_ref = weakref.ref(source) + with step_inputs.active(pack): + step_inputs.begin() + step_inputs.copy_(a, source) + del source + step_inputs.flush(in_graph=True) + assert source_ref() is not None + overwrite = torch.full((16,), -7, device="cuda", dtype=a.dtype) + graph.replay() + pack.on_graph_replay() + torch.cuda.synchronize() + assert out[:16].tolist() == list(range(16)) + assert overwrite.eq(-7).all() + assert source_ref() is None + + +@pytest.mark.parametrize("early", [False, True]) +def test_replay_with_a_different_pack_issues_its_copies_eagerly(early): + captured, graph, (a, b, c), out = _graph_case(early) + _step(captured, graph, (a, b, c), (list(range(16)), list(range(16)))) + active = StepInputPack(_device(), parallel_scatter=True, in_graph=True, early_copy=early) + key = KeyType(batch_size=16, draft_len=0) + metadata = object() + runner = object.__new__(CUDAGraphRunner) + runner.graph_metadata = {key: {"attn_metadata": metadata, "spec_metadata": None}} + runner.shared_static_tensors = {"input_ids": a, "position_ids": b.view(1, 16)} + runner.config = SimpleNamespace(use_mrope=False) + runner._check_strict_buffer_stability = lambda meta: None + runner._get_num_tokens_for_key = lambda key: 16 + runner._step_input_prologue_keys = {key: captured} + runner._pre_replay_hooks = [] + runner.graphs = {key: graph} + runner.graph_outputs = {key: out} + inputs = { + "attn_metadata": metadata, + "input_ids": torch.full_like(a, 77), + "position_ids": torch.full_like(b.view(1, 16), 88), + } + with step_inputs.active(active): + step_inputs.begin() + runner.replay(key, inputs) + torch.cuda.synchronize() + assert out[:16].tolist() == [77] * 16 + assert out[16:32].tolist() == [88] * 16 diff --git a/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py index 5a56b77e1659..13de7bdd0377 100644 --- a/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py +++ b/tests/unittest/_torch/speculative/test_spec_sampler_step_fold.py @@ -276,7 +276,14 @@ def test_store_copies_go_on_the_side_stream_when_the_sampler_stored_the_rows(): @pytest.mark.parametrize("draft_width", [MAX_DRAFT_LEN, MAX_DRAFT_LEN - 1]) def test_side_stream_store_copies_match_the_main_stream_copies(draft_width): scheduled, outputs = _step(draft_width) - on_main = _host_outputs(_sampler().sample_async(scheduled, outputs, [])) + main_sampler = _sampler() + expected = _reference_store(_store_snapshot(main_sampler), outputs) + for name, tensor in expected.items(): + getattr(main_sampler.store, name).copy_(tensor) + main_sampler.bind_spec_step_fold(_fold(rows_stored=True)) + main_state = main_sampler.sample_async(scheduled, outputs, []) + assert main_state.sampler_event.side_stream_event is None + on_main = _host_outputs(main_state) sampler = _sampler() sampler.d2h_side_stream_copies = True From 950592f0addddf034472fb56dc18197cbac34a5a Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Sat, 31 Oct 2026 13:42:11 -0700 Subject: [PATCH 5/8] [None][test] Cover strided Rubin multi-query attention inputs Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../test_modeling_v2_m3_sparse_decode.py | 42 +++++++++++++++++++ .../test_modeling_v2_m3_sparse_decode.py | 42 +++++++++++++++++++ .../test_modeling_v2_m3_sparse_decode.py | 42 +++++++++++++++++++ 3 files changed, 126 insertions(+) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py index feb450a5d350..1667f6034162 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py @@ -27,6 +27,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode import ( # noqa: E501 m3_sparse_decode as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.triton_sparse_decode_mq import ( # noqa: E501 + launch_mq_producer, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( minimax_m3_sparse_attn_decode as stock_decode, ) @@ -449,3 +452,42 @@ def test_identity_route_reads_its_own_lengths_not_the_table(requests, fault) -> fault, case.seq_lens.tolist(), ) + + +def test_mq_producer_preserves_strided_query_elements() -> None: + case = build_step((129, 257), 4, identity=True, seed=71) + # MQ processes one KV head; use the first head's query group for every TP target. + heads = NUM_HEADS // case.k_paged.shape[1] + q = case.q[:, :heads].contiguous() + storage = torch.zeros((*q.shape[:-1], 2 * HEAD_DIM), dtype=q.dtype, device=q.device) + strided = storage[..., ::2] + strided.copy_(q) + assert strided.stride(2) == 2 + results = [] + for query in (q, strided): + partial = torch.empty( + case.topk.shape[-1], q.shape[0], heads, HEAD_DIM, dtype=torch.float32, device=q.device + ) + lse = torch.empty(partial.shape[:-1], dtype=torch.float32, device=q.device) + launch_mq_producer( + query, + case.k_paged[:, :1], + case.v_paged[:, :1], + case.k_block_scale[:, :1], + case.v_block_scale[:, :1], + case.k_global_scale, + case.v_global_scale, + case.topk[:1], + case.block_table, + case.seq_lens, + partial, + lse, + sm_scale=case.sm_scale, + decode_query_len=case.decode_query_len, + use_pdl=False, + early_trigger=False, + ) + results.append((partial, lse)) + torch.cuda.synchronize() + assert torch.equal(results[0][0], results[1][0]) + assert torch.equal(results[0][1], results[1][1]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py index 523a09addea8..1c88c39e824f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py @@ -14,6 +14,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode import ( # noqa: E501 m3_sparse_decode as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.triton_sparse_decode_mq import ( # noqa: E501 + launch_mq_producer, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( minimax_m3_sparse_attn_decode as stock_decode, ) @@ -441,3 +444,42 @@ def test_lc_fails_closed_on_ids_past_the_row(requests) -> None: torch.cuda.synchronize() assert torch.isfinite(outs[0].float()).all() assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)) + + +def test_mq_producer_preserves_strided_query_elements() -> None: + case = build_step((129, 257), 4, identity=True, seed=71) + # MQ processes one KV head; use the first head's query group for every TP target. + heads = NUM_HEADS // case.k_paged.shape[1] + q = case.q[:, :heads].contiguous() + storage = torch.zeros((*q.shape[:-1], 2 * HEAD_DIM), dtype=q.dtype, device=q.device) + strided = storage[..., ::2] + strided.copy_(q) + assert strided.stride(2) == 2 + results = [] + for query in (q, strided): + partial = torch.empty( + case.topk.shape[-1], q.shape[0], heads, HEAD_DIM, dtype=torch.float32, device=q.device + ) + lse = torch.empty(partial.shape[:-1], dtype=torch.float32, device=q.device) + launch_mq_producer( + query, + case.k_paged[:, :1], + case.v_paged[:, :1], + case.k_block_scale[:, :1], + case.v_block_scale[:, :1], + case.k_global_scale, + case.v_global_scale, + case.topk[:1], + case.block_table, + case.seq_lens, + partial, + lse, + sm_scale=case.sm_scale, + decode_query_len=case.decode_query_len, + use_pdl=False, + early_trigger=False, + ) + results.append((partial, lse)) + torch.cuda.synchronize() + assert torch.equal(results[0][0], results[1][0]) + assert torch.equal(results[0][1], results[1][1]) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py index 36dbdd14d592..195c5cb5c567 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py @@ -26,6 +26,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_sparse_decode import ( # noqa: E501 m3_sparse_decode as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.triton_sparse_decode_mq import ( # noqa: E501 + launch_mq_producer, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( minimax_m3_sparse_attn_decode as stock_decode, ) @@ -624,3 +627,42 @@ def test_lc_fails_closed_on_ids_past_the_row(requests) -> None: torch.cuda.synchronize() assert torch.isfinite(outs[0].float()).all() assert torch.equal(outs[0].view(torch.int16), outs[1].view(torch.int16)) + + +def test_mq_producer_preserves_strided_query_elements() -> None: + case = build_step((129, 257), 4, identity=True, seed=71) + # MQ processes one KV head; use the first head's query group for every TP target. + heads = NUM_HEADS // case.k_paged.shape[1] + q = case.q[:, :heads].contiguous() + storage = torch.zeros((*q.shape[:-1], 2 * HEAD_DIM), dtype=q.dtype, device=q.device) + strided = storage[..., ::2] + strided.copy_(q) + assert strided.stride(2) == 2 + results = [] + for query in (q, strided): + partial = torch.empty( + case.topk.shape[-1], q.shape[0], heads, HEAD_DIM, dtype=torch.float32, device=q.device + ) + lse = torch.empty(partial.shape[:-1], dtype=torch.float32, device=q.device) + launch_mq_producer( + query, + case.k_paged[:, :1], + case.v_paged[:, :1], + case.k_block_scale[:, :1], + case.v_block_scale[:, :1], + case.k_global_scale, + case.v_global_scale, + case.topk[:1], + case.block_table, + case.seq_lens, + partial, + lse, + sm_scale=case.sm_scale, + decode_query_len=case.decode_query_len, + use_pdl=False, + early_trigger=False, + ) + results.append((partial, lse)) + torch.cuda.synchronize() + assert torch.equal(results[0][0], results[1][0]) + assert torch.equal(results[0][1], results[1][1]) From 04f4a3f200f55e4cbf4064ebdbdc7e0150462b55 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 12:44:57 -0800 Subject: [PATCH 6/8] [None][refactor] Share MiniMax-M3 SM107 implementations across TP targets Move common Python operators, kernel implementations, native code and test helpers to the MiniMax-M3 family while preserving TP-specific assembly, shape/tactic tables, topology and runtime state. Keep existing numerical and dispatch policies, and add shared distributed sampling coverage. Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- cpp/tensorrt_llm/kernels/CMakeLists.txt | 6 +- .../drafterRopeFp8Kv.cu | 14 +- .../drafterRopeFp8Kv.h | 6 +- .../foldTail.cu | 6 +- .../foldTail.h | 7 +- .../kernels/minimaxM3Sm107/fullCtxReduce.cu | 47 + .../fullCtxReduce.cuh} | 116 +- .../fullCtxReduce.h | 8 +- .../mcAllreduceNorm.cu} | 184 +- .../mcAllreduceNorm.h} | 8 +- .../minimaxM3Sm107/mcAllreduceUtils.cuh | 192 + .../producerPw.cu} | 50 +- .../producerPw.h} | 26 +- .../kernels/minimaxM3Sm107/sampledAccept.cu | 965 ++++ .../kernels/minimaxM3Sm107/sampledAccept.h | 87 + .../siluAndMul.cu | 6 +- .../siluAndMul.h | 4 +- .../specOffsets.cu | 6 +- .../specOffsets.h | 4 +- .../stepPrologue.cu | 6 +- .../stepPrologue.h | 4 +- .../swigluOai.cu | 6 +- .../swigluOai.h | 4 +- .../minimaxM3Sm107Tp2/drafterRopeFp8Kv.h | 79 - .../kernels/minimaxM3Sm107Tp2/foldTail.cu | 35 +- .../minimaxM3Sm107Tp2/fullCtxReduce.cu | 182 +- .../minimaxM3Sm107Tp2/mcAllreduceNormTp2.h | 83 - .../kernels/minimaxM3Sm107Tp2/siluAndMul.cu | 89 - .../kernels/minimaxM3Sm107Tp2/siluAndMul.h | 40 - .../kernels/minimaxM3Sm107Tp2/specOffsets.cu | 58 - .../kernels/minimaxM3Sm107Tp2/specOffsets.h | 41 - .../kernels/minimaxM3Sm107Tp2/stepPrologue.cu | 111 - .../kernels/minimaxM3Sm107Tp2/swigluOai.cu | 100 - .../kernels/minimaxM3Sm107Tp2/swigluOai.h | 43 - .../kernels/minimaxM3Sm107Tp4/draftStep.cu | 2 +- .../kernels/minimaxM3Sm107Tp4/draftStep.h | 4 +- .../minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu | 189 - .../minimaxM3Sm107Tp4/drafterRopeFp8Kv.h | 79 - .../kernels/minimaxM3Sm107Tp4/foldTail.h | 130 - .../minimaxM3Sm107Tp4/fullCtxReduce.cu | 270 -- .../kernels/minimaxM3Sm107Tp4/fullCtxReduce.h | 77 - .../minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu | 174 +- .../kernels/minimaxM3Sm107Tp4/producerPw4.cu | 340 -- .../kernels/minimaxM3Sm107Tp4/producerPw4.h | 76 - .../minimaxM3Sm107Tp4/sampledAccept.cu | 922 +--- .../kernels/minimaxM3Sm107Tp4/sampledAccept.h | 35 +- .../kernels/minimaxM3Sm107Tp4/specOffsets.cu | 58 - .../kernels/minimaxM3Sm107Tp4/stepPrologue.cu | 111 - .../kernels/minimaxM3Sm107Tp4/stepPrologue.h | 47 - .../kernels/minimaxM3Sm107Tp4/swigluOai.cu | 100 - .../kernels/minimaxM3Sm107Tp8/draftStep.cu | 2 +- .../kernels/minimaxM3Sm107Tp8/draftStep.h | 4 +- .../minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu | 189 - .../kernels/minimaxM3Sm107Tp8/foldTail.cu | 64 - .../minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu | 1096 ----- .../kernels/minimaxM3Sm107Tp8/producerPw8.cu | 340 -- .../kernels/minimaxM3Sm107Tp8/producerPw8.h | 76 - .../minimaxM3Sm107Tp8/sampledAccept.cu | 924 +--- .../kernels/minimaxM3Sm107Tp8/sampledAccept.h | 35 +- .../kernels/minimaxM3Sm107Tp8/siluAndMul.cu | 89 - .../kernels/minimaxM3Sm107Tp8/siluAndMul.h | 40 - .../kernels/minimaxM3Sm107Tp8/specOffsets.h | 41 - .../kernels/minimaxM3Sm107Tp8/stepPrologue.h | 47 - .../kernels/minimaxM3Sm107Tp8/swigluOai.h | 43 - ...lOp.cpp => minimaxM3Sm107SiluAndMulOp.cpp} | 48 +- ...Op.cpp => minimaxM3Sm107SpecOffsetsOp.cpp} | 54 +- ...p.cpp => minimaxM3Sm107StepPrologueOp.cpp} | 46 +- ...aiOp.cpp => minimaxM3Sm107SwigluOaiOp.cpp} | 46 +- .../minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp | 6 +- .../minimaxM3Sm107Tp2McAllreduceNormOp.cpp | 8 +- .../minimaxM3Sm107Tp2ProducerPwOp.cpp | 12 +- .../minimaxM3Sm107Tp2SiluAndMulOp.cpp | 66 - .../minimaxM3Sm107Tp2SpecOffsetsOp.cpp | 74 - .../minimaxM3Sm107Tp2StepPrologueOp.cpp | 68 - .../minimaxM3Sm107Tp2SwigluOaiOp.cpp | 65 - .../minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp | 6 +- .../minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp | 6 +- .../minimaxM3Sm107Tp4FoldTailParams.h | 6 +- .../minimaxM3Sm107Tp4FullCtxReduceOp.cpp | 4 +- .../minimaxM3Sm107Tp4ProducerPwOp.cpp | 12 +- .../minimaxM3Sm107Tp4StepPrologueOp.cpp | 68 - .../minimaxM3Sm107Tp4SwigluOaiOp.cpp | 65 - .../minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp | 6 +- .../minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp | 6 +- .../minimaxM3Sm107Tp8FoldTailParams.h | 6 +- .../minimaxM3Sm107Tp8FullCtxReduceOp.cpp | 4 +- .../minimaxM3Sm107Tp8McAllreduceNormOp.cpp | 8 +- .../minimaxM3Sm107Tp8ProducerPwOp.cpp | 12 +- .../minimaxM3Sm107Tp8SiluAndMulOp.cpp | 66 - .../minimaxM3Sm107Tp8SpecOffsetsOp.cpp | 74 - .../modeling_v2/models/minimax_m3/config.py | 93 + .../minimax_m3/kernels/attention/__init__.py | 2 + .../kernels/attention/index_score_runtime.py | 127 + .../kernels/attention/m3_index_score.py | 549 +++ .../kernels/attention/m3_select_long_dsl.py | 285 ++ .../attention/producer_tp8.py} | 0 .../kernels/attention/triton_sparse_decode.py | 892 ++++ .../attention}/triton_sparse_decode_mq.py | 0 .../kernels/embedding_firstnorm/__init__.py | 2 + .../m3_embedding_firstnorm_kernel.py | 0 .../minimax_m3/kernels/gemm/__init__.py | 2 + .../ops => kernels/gemm}/m3_lowm_splitk.py | 0 .../gemm}/m3_mx_gemv_kernel.py | 87 +- .../models/minimax_m3/kernels/moe/__init__.py | 2 + .../m3_moe => kernels/moe}/m3_moe_kernel.py | 82 +- .../minimax_m3/kernels/spec/__init__.py | 2 + .../minimax_m3/kernels/spec/m3_fold_tail.py | 64 + .../ops => kernels/spec}/m3_force_vec.py | 0 .../minimax_m3_nvfp4__sm_107__tp2/modeling.py | 518 +-- .../ops/attention/m3_draft_attention.py | 108 +- .../ops/attention/m3_fullctx_decode.py | 5 +- .../ops/attention/m3_index_score.py | 95 +- .../ops/attention/m3_msa_attention.py | 212 +- .../ops/attention/m3_producer.py | 96 +- .../ops/attention/m3_qk_norm_rope_fp8.py | 137 - .../ops/attention/m3_select_long.py | 49 +- .../ops/attention/m3_sparse_decode.py | 33 +- .../ops/attention/m3_sparse_decode_lc.py | 127 +- .../ops/comm/m3_allreduce.py | 61 - .../ops/comm/m3_mc_allreduce_norm.py | 105 +- .../ops/comm/m3_mc_finalize_allreduce_norm.py | 58 +- .../ops/embedding/m3_embedding.py | 48 - .../ops/embedding_firstnorm/op.py | 91 +- .../ops/gemm/m3_dense_swiglu.py | 72 - .../ops/gemm/m3_draft_lowm.py | 8 +- .../ops/gemm/m3_mxfp8_gemm.py | 58 - .../ops/gemm/m3_mxfp8_linear.py | 137 - .../ops/m3_draft_step.py | 5 +- .../ops/m3_drafter_attn.py | 402 +- .../ops/m3_fold_tail.py | 156 +- .../ops/m3_fullctx_decode.py | 158 +- .../ops/m3_index_score.py | 649 +-- .../ops/m3_locality.py | 429 +- .../ops/m3_lowm_splitk.py | 1255 ------ .../ops/m3_mc_allreduce.py | 75 +- .../ops/m3_moe/op.py | 507 +-- .../ops/m3_prequant.py | 105 +- .../ops/m3_router_gemv.py | 12 +- .../ops/m3_select_long_dsl.py | 370 +- .../ops/m3_shard_accept.py | 48 +- .../ops/m3_target_head.py | 84 +- .../ops/m3_tokmap_guard.py | 138 +- .../ops/moe/m3_moe_engine.py | 241 +- .../ops/moe/m3_router_gemv_fp4.py | 63 - .../ops/msa_kv_update.py | 89 - .../ops/mx_gemv/m3_mx_gemv_kernel.py | 1237 ------ .../ops/mx_gemv/op.py | 359 +- .../ops/norm/m3_embedding_firstnorm.py | 11 +- .../ops/spec/m3_draft_silu.py | 2 +- .../ops/spec/m3_draft_step.py | 42 +- .../ops/spec/m3_force_vec.py | 136 - .../ops/spec/m3_sampled_accept.py | 120 +- .../ops/spec/m3_shard_accept.py | 97 +- .../ops/spec/m3_slot_scatter.py | 72 - .../ops/spec/m3_spec_offsets.py | 2 +- .../ops/spec/m3_step_prologue.py | 2 +- .../ops/triton_sparse_decode.py | 1186 +---- .../minimax_m3_nvfp4__sm_107__tp4/modeling.py | 480 +- .../ops/activation/m3_silu_and_mul.py | 83 - .../ops/attention/m3_dense_decode.py | 190 - .../ops/attention/m3_draft_attention.py | 30 +- .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 - .../ops/attention/m3_fullctx_decode.py | 5 +- .../ops/attention/m3_index_score.py | 95 +- .../ops/attention/m3_msa_attention.py | 86 +- .../ops/attention/m3_producer.py | 87 +- .../ops/attention/m3_qkv_insert.py | 134 - .../ops/attention/m3_select_blocks.py | 82 - .../ops/attention/m3_select_long.py | 49 +- .../ops/attention/m3_sparse_decode.py | 33 +- .../ops/attention/m3_sparse_decode_lc.py | 122 +- .../ops/attention/msa_kv_update.py | 200 - .../ops/comm/m3_mc_allreduce_norm.py | 105 +- .../ops/comm/m3_mc_finalize_allreduce_norm.py | 58 +- .../m3_embedding_firstnorm_kernel.py | 128 - .../ops/embedding_firstnorm/op.py | 91 +- .../ops/gemm/m3_bf16_linear.py | 41 - .../ops/gemm/m3_dense_swiglu.py | 72 - .../ops/gemm/m3_draft_lowm.py | 8 +- .../ops/gemm/m3_fp32_linear.py | 46 - .../ops/gemm/m3_mxfp8_gemm.py | 58 - .../ops/gemm/m3_mxfp8_pinned_gemm.py | 5 +- .../ops/gemm/m3_target_head.py | 346 -- .../ops/m3_draft_lowm.py | 4 +- .../ops/m3_draft_step.py | 6 +- .../ops/m3_drafter_attn.py | 397 +- .../ops/m3_fold_tail.py | 158 +- .../ops/m3_force_vec.py | 64 - .../ops/m3_fullctx_decode.py | 158 +- .../ops/m3_index_score.py | 624 +-- .../ops/m3_lowm_splitk.py | 1255 ------ .../ops/m3_mc_allreduce.py | 75 +- .../ops/m3_moe/m3_moe_kernel.py | 3914 ----------------- .../ops/m3_moe/op.py | 510 +-- .../ops/m3_moe/op16.py | 2 + .../ops/m3_prequant.py | 105 +- .../ops/m3_select_long_dsl.py | 371 +- .../ops/m3_shard_accept.py | 48 +- .../ops/m3_spec_tail.py | 20 +- .../ops/m3_target_head.py | 84 +- .../ops/m3_tokmap_guard.py | 138 +- .../ops/moe/m3_moe_engine.py | 241 +- .../ops/moe/m3_router_tc.py | 119 - .../ops/msa_kv_update.py | 89 - .../ops/mx_gemv/m3_mx_gemv_kernel.py | 1299 ------ .../ops/mx_gemv/op.py | 359 +- .../ops/norm/m3_embedding_firstnorm.py | 11 +- .../ops/spec/m3_capture_norm.py | 80 - .../ops/spec/m3_chunk_norm.py | 70 - .../ops/spec/m3_draft_input.py | 140 - .../ops/spec/m3_draft_silu.py | 2 +- .../ops/spec/m3_draft_step.py | 42 +- .../ops/spec/m3_force_vec.py | 86 - .../ops/spec/m3_force_vec_req.py | 5 +- .../ops/spec/m3_sampled_accept.py | 120 +- .../ops/spec/m3_shard_accept.py | 97 +- .../ops/spec/m3_slot_scatter.py | 72 - .../ops/spec/m3_spec_offsets.py | 2 +- .../ops/spec/m3_step_gather.py | 152 - .../ops/spec/m3_step_prologue.py | 2 +- .../ops/triton_sparse_decode.py | 1175 +---- .../ops/triton_sparse_decode_mq.py | 239 - .../minimax_m3_nvfp4__sm_107__tp8/modeling.py | 511 +-- .../ops/activation/m3_silu_and_mul.py | 83 - .../ops/attention/m3_dense_decode.py | 190 - .../ops/attention/m3_draft_attention.py | 108 +- .../attention/m3_fp8_indexer_qk_norm_rope.py | 114 - .../ops/attention/m3_fullctx_decode.py | 5 +- .../ops/attention/m3_index_score.py | 95 +- .../ops/attention/m3_msa_attention.py | 212 +- .../ops/attention/m3_producer.py | 87 +- .../ops/attention/m3_qk_norm_rope_fp8.py | 137 - .../ops/attention/m3_qkv_insert.py | 134 - .../ops/attention/m3_select_blocks.py | 82 - .../ops/attention/m3_select_long.py | 49 +- .../ops/attention/m3_sparse_decode.py | 33 +- .../ops/attention/m3_sparse_decode_lc8.py | 122 +- .../ops/attention/msa_kv_update.py | 200 - .../ops/comm/m3_allreduce.py | 60 - .../ops/comm/m3_mc_allreduce_norm.py | 105 +- .../ops/comm/m3_mc_finalize_allreduce_norm.py | 58 +- .../ops/embedding/m3_embedding.py | 48 - .../m3_embedding_firstnorm_kernel.py | 128 - .../ops/embedding_firstnorm/op.py | 91 +- .../ops/gemm/m3_bf16_linear.py | 41 - .../ops/gemm/m3_draft_lowm.py | 6 +- .../ops/gemm/m3_fp32_linear.py | 46 - .../ops/gemm/m3_mxfp8_linear.py | 137 - .../ops/gemm/m3_target_head.py | 346 -- .../ops/m3_draft_step.py | 5 +- .../ops/m3_drafter_attn.py | 396 +- .../ops/m3_fold_tail.py | 170 +- .../ops/m3_force_vec.py | 63 - .../ops/m3_fullctx_decode.py | 158 +- .../ops/m3_index_score.py | 663 +-- .../ops/m3_locality.py | 429 +- .../ops/m3_mc_allreduce.py | 74 +- .../ops/m3_moe/m3_moe_kernel.py | 3841 ---------------- .../ops/m3_moe/op.py | 487 +- .../ops/m3_prequant.py | 105 +- .../ops/m3_router_gemv.py | 12 +- .../ops/m3_select_long_dsl.py | 371 +- .../ops/m3_shard_accept.py | 48 +- .../ops/m3_target_head.py | 84 +- .../ops/m3_tokmap_guard.py | 138 +- .../ops/moe/m3_moe_engine.py | 241 +- .../ops/moe/m3_router_gemv.py | 101 - .../ops/moe/m3_router_tc.py | 106 - .../ops/mx_gemv/op.py | 396 +- .../ops/norm/m3_embedding_firstnorm.py | 11 +- .../ops/spec/m3_capture_norm.py | 80 - .../ops/spec/m3_chunk_norm.py | 74 - .../ops/spec/m3_draft_input.py | 140 - .../ops/spec/m3_draft_silu.py | 2 +- .../ops/spec/m3_draft_step.py | 42 +- .../ops/spec/m3_sampled_accept.py | 120 +- .../ops/spec/m3_shard_accept.py | 97 +- .../ops/spec/m3_spec_offsets.py | 2 +- .../ops/spec/m3_step_gather.py | 160 - .../ops/spec/m3_step_prologue.py | 2 +- .../ops/triton_sparse_decode.py | 1177 +---- .../ops/triton_sparse_decode_mq.py | 239 - .../models/minimax_m3/modeling_support.py | 287 ++ .../models/minimax_m3/ops/__init__.py | 2 + .../minimax_m3/ops/activation/__init__.py | 2 + .../ops/activation/m3_silu_and_mul.py | 61 +- .../minimax_m3/ops/attention/__init__.py | 2 + .../ops/attention/draft_attention.py | 105 + .../minimax_m3/ops/attention/drafter.py | 407 ++ .../minimax_m3/ops/attention/fullctx.py | 182 + .../ops/attention/long_select_host.py | 84 + .../ops/attention/m3_dense_decode.py | 0 .../attention/m3_fp8_indexer_qk_norm_rope.py | 25 +- .../ops/attention/m3_index_score.py | 104 + .../ops/attention/m3_qk_norm_rope_fp8.py | 25 +- .../ops/attention/m3_qkv_insert.py | 0 .../ops/attention/m3_select_blocks.py | 0 .../models/minimax_m3/ops/attention/msa.py | 198 + .../ops/attention/msa_kv_update.py | 0 .../attention/msa_kv_update_hook.py} | 4 +- .../ops/attention/producer_reference.py | 179 + .../ops/attention/select_reference.py | 62 + .../ops/attention/sparse_decode_policy.py | 267 ++ .../ops/attention/sparse_reference.py | 184 + .../minimax_m3/ops/attention/token_map.py | 181 + .../models/minimax_m3/ops/comm/__init__.py | 2 + .../ops/comm/m3_allreduce.py | 34 +- .../minimax_m3/ops/comm/mc_workspace.py | 80 + .../models/minimax_m3/ops/comm/reference.py | 159 + .../minimax_m3/ops/embedding/__init__.py | 2 + .../ops/embedding/decode_embedding.py | 157 + .../ops/embedding/m3_embedding.py | 53 +- .../models/minimax_m3/ops/gemm/__init__.py | 2 + .../models/minimax_m3/ops/gemm/locality.py | 431 ++ .../ops/gemm/m3_bf16_linear.py | 33 +- .../ops/gemm/m3_dense_swiglu.py | 45 +- .../ops/gemm/m3_fp32_linear.py | 0 .../ops/gemm/m3_mxfp8_gemm.py | 32 +- .../ops/gemm/m3_mxfp8_linear.py | 80 +- .../ops/gemm/m3_target_head.py | 0 .../models/minimax_m3/ops/gemm/mxfp8_core.py | 384 ++ .../models/minimax_m3/ops/gemm/mxfp8_host.py | 212 + .../models/minimax_m3/ops/gemm/mxfp8_state.py | 61 + .../models/minimax_m3/ops/gemm/prequant.py | 85 + .../minimax_m3/ops/gemm/target_head_state.py | 109 + .../models/minimax_m3/ops/moe/__init__.py | 2 + .../ops/moe/m3_router_gemv.py | 0 .../ops/moe/m3_router_gemv_fp4.py | 6 +- .../ops/moe/m3_router_tc.py | 85 +- .../models/minimax_m3/ops/moe/phase_plan.py | 206 + .../models/minimax_m3/ops/moe/reference.py | 233 + .../models/minimax_m3/ops/moe/runtime.py | 232 + .../models/minimax_m3/ops/moe/state.py | 192 + .../models/minimax_m3/ops/spec/__init__.py | 2 + .../minimax_m3/ops/spec/accept_workspace.py | 64 + .../ops/spec/acceptance_reference.py | 97 + .../ops/spec/draft_head_reference.py | 73 + .../models/minimax_m3/ops/spec/fold_tail.py | 117 + .../ops/spec/m3_capture_norm.py | 58 +- .../ops/spec/m3_chunk_norm.py | 50 +- .../ops/spec/m3_draft_input.py | 0 .../ops/spec/m3_force_vec.py | 6 +- .../ops/spec/m3_slot_scatter.py | 16 +- .../ops/spec/m3_step_gather.py | 81 +- .../minimax_m3/ops/spec/sampling_reference.py | 132 + .../models/minimax_m3/reporting.py | 92 + .../_torch/cute_dsl_kernels/__init__.py | 2 + tensorrt_llm/_torch/distributed/ops.py | 6 +- .../engine/runners/decoder/inputs.py | 11 +- .../accuracy/modeling_v2/test_minimax_m3.py | 52 +- .../integration/test_lists/test-db/l0_cpu.yml | 3 + .../test_allreduce_tune_nccl_symmetric.py | 93 +- .../gemm => }/_m3_mxfp8_operands.py | 0 .../models/minimax_m3/_m3_test_inputs.py | 49 + .../test_distributed_sampled_reference.py | 255 ++ .../comm/test_sampled_accept_tp8_exchange.py | 225 + .../test_sampled_accept_tp8_graph_exchange.py | 138 + .../test_modeling_v2_m3_silu_and_mul.py | 6 +- ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 4 +- .../test_modeling_v2_m3_index_score.py | 33 +- .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 4 +- .../test_modeling_v2_m3_qkv_insert.py | 2 +- .../test_modeling_v2_m3_select_blocks.py | 2 +- .../test_modeling_v2_m3_sparse_decode.py | 6 +- .../test_modeling_v2_m3_tokmap_guard.py | 13 +- .../test_modeling_v2_msa_kv_update.py | 6 +- .../comm/_m3_allreduce_op_matrix.py | 61 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 55 +- .../comm/_m3_moe_push_op_matrix.py | 55 +- .../test_modeling_v2_m3_embedding.py | 4 +- .../gemm/test_modeling_v2_m3_bf16_linear.py | 4 +- .../gemm/test_modeling_v2_m3_dense_seam.py | 7 +- .../gemm/test_modeling_v2_m3_dense_swiglu.py | 16 +- .../gemm/test_modeling_v2_m3_fp32_linear.py | 2 +- .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 11 +- .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 6 +- .../moe/_m3_router.py | 8 +- .../test_modeling_v2_m3_fp32_router_rows.py | 2 +- .../moe/test_modeling_v2_m3_router_tc.py | 2 +- .../spec/_m3_spec_inputs.py | 52 +- .../spec/test_modeling_v2_m3_capture_norm.py | 11 +- .../spec/test_modeling_v2_m3_chunk_norm.py | 6 +- ...test_modeling_v2_m3_draft_ar_push_words.py | 6 +- .../spec/test_modeling_v2_m3_draft_input.py | 2 +- .../spec/test_modeling_v2_m3_draft_step.py | 8 +- .../spec/test_modeling_v2_m3_force_al.py | 2 +- .../spec/test_modeling_v2_m3_force_vec.py | 2 +- .../spec/test_modeling_v2_m3_force_vec_req.py | 2 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 8 +- .../spec/test_modeling_v2_m3_slot_scatter.py | 2 +- .../spec/test_modeling_v2_m3_step_gather.py | 6 +- .../test_modeling_v2_m3_silu_and_mul.py | 6 +- .../test_modeling_v2_m3_dense_decode.py | 9 +- .../test_modeling_v2_m3_drafter_fullctx.py | 14 +- ...test_modeling_v2_m3_drafter_rope_fp8_kv.py | 18 + ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 4 +- .../test_modeling_v2_m3_index_score.py | 25 +- .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 4 +- .../test_modeling_v2_m3_qkv_insert.py | 2 +- .../test_modeling_v2_m3_select_blocks.py | 2 +- .../test_modeling_v2_m3_sparse_decode.py | 6 +- .../test_modeling_v2_m3_tokmap_guard.py | 13 +- .../test_modeling_v2_msa_kv_update.py | 6 +- .../comm/_m3_allreduce_op_matrix.py | 61 +- .../comm/_m3_capture_in_boundary_op_matrix.py | 53 +- .../comm/_m3_draft_ar_norm_op_matrix.py | 53 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 55 +- .../test_modeling_v2_m3_embedding.py | 4 +- .../gemm/_m3_mxfp8_operands.py | 30 - .../gemm/test_modeling_v2_m3_bf16_linear.py | 4 +- .../gemm/test_modeling_v2_m3_dense_swiglu.py | 16 +- .../gemm/test_modeling_v2_m3_fp32_linear.py | 2 +- .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 11 +- .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 6 +- .../test_modeling_v2_m3_mxfp8_pinned_gemm.py | 2 +- .../gemm/test_modeling_v2_m3_target_head.py | 4 +- .../moe/_m3_router.py | 2 +- .../moe/test_modeling_v2_m3_router_tc.py | 2 +- .../spec/_m3_spec_inputs.py | 52 +- .../spec/test_modeling_v2_m3_capture_norm.py | 11 +- .../spec/test_modeling_v2_m3_chunk_norm.py | 6 +- .../spec/test_modeling_v2_m3_draft_input.py | 2 +- .../spec/test_modeling_v2_m3_draft_step.py | 8 +- .../test_modeling_v2_m3_exact_numerics.py | 10 +- .../spec/test_modeling_v2_m3_force_vec.py | 2 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 8 +- .../spec/test_modeling_v2_m3_slot_scatter.py | 2 +- .../spec/test_modeling_v2_m3_step_gather.py | 6 +- .../test_modeling_v2_m3_prebuild.py | 4 + .../test_modeling_v2_m3_silu_and_mul.py | 6 +- ...modeling_v2_m3_fp8_indexer_qk_norm_rope.py | 4 +- .../test_modeling_v2_m3_index_score.py | 33 +- .../test_modeling_v2_m3_qk_norm_rope_fp8.py | 4 +- .../test_modeling_v2_m3_qkv_insert.py | 2 +- .../test_modeling_v2_m3_select_blocks.py | 2 +- .../test_modeling_v2_m3_sparse_decode.py | 6 +- .../test_modeling_v2_m3_tokmap_guard.py | 13 +- .../test_modeling_v2_msa_kv_update.py | 6 +- .../comm/_m3_allreduce_op_matrix.py | 61 +- .../comm/_m3_mc_allreduce_norm_op_matrix.py | 55 +- .../comm/_m3_moe_push_op_matrix.py | 55 +- .../test_modeling_v2_m3_embedding.py | 4 +- .../gemm/_m3_mxfp8_operands.py | 30 - .../gemm/test_modeling_v2_m3_bf16_linear.py | 4 +- .../gemm/test_modeling_v2_m3_dense_swiglu.py | 16 +- .../gemm/test_modeling_v2_m3_fp32_linear.py | 2 +- .../gemm/test_modeling_v2_m3_mxfp8_gemm.py | 11 +- .../gemm/test_modeling_v2_m3_mxfp8_linear.py | 6 +- .../test_modeling_v2_m3_swiglu_oai_bf16x2.py | 2 +- .../moe/_m3_router.py | 8 +- .../moe/test_modeling_v2_m3_router_tc.py | 2 +- .../spec/_m3_spec_inputs.py | 52 +- .../spec/test_modeling_v2_m3_capture_norm.py | 11 +- .../spec/test_modeling_v2_m3_chunk_norm.py | 6 +- .../spec/test_modeling_v2_m3_di_fc_flag.py | 1 + .../spec/test_modeling_v2_m3_draft_input.py | 2 +- .../spec/test_modeling_v2_m3_draft_step.py | 8 +- .../spec/test_modeling_v2_m3_force_al.py | 2 +- .../spec/test_modeling_v2_m3_force_vec.py | 2 +- .../spec/test_modeling_v2_m3_force_vec_req.py | 2 +- .../spec/test_modeling_v2_m3_lowm_splitk.py | 8 +- .../spec/test_modeling_v2_m3_slot_scatter.py | 2 +- .../spec/test_modeling_v2_m3_step_gather.py | 6 +- .../test_modeling_v2_eagle3_worker_api.py | 136 + .../test_modeling_v2_spec_tail_contracts.py | 96 + .../test_collective_launcher_cleanup.py | 89 + .../pyexecutor/test_first_content_emission.py | 19 + .../_torch/pyexecutor/test_runtime_hooks.py | 37 + 468 files changed, 14546 insertions(+), 43779 deletions(-) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp2 => minimaxM3Sm107}/drafterRopeFp8Kv.cu (91%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp8 => minimaxM3Sm107}/drafterRopeFp8Kv.h (93%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp4 => minimaxM3Sm107}/foldTail.cu (93%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp8 => minimaxM3Sm107}/foldTail.h (96%) create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cu rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp8/fullCtxReduce.cu => minimaxM3Sm107/fullCtxReduce.cuh} (72%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp8 => minimaxM3Sm107}/fullCtxReduce.h (92%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu => minimaxM3Sm107/mcAllreduceNorm.cu} (87%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp8/mcAllreduceNormTp8.h => minimaxM3Sm107/mcAllreduceNorm.h} (92%) create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceUtils.cuh rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp2/producerPw2.cu => minimaxM3Sm107/producerPw.cu} (87%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp2/producerPw2.h => minimaxM3Sm107/producerPw.h} (77%) create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.cu create mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.h rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp4 => minimaxM3Sm107}/siluAndMul.cu (96%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp4 => minimaxM3Sm107}/siluAndMul.h (95%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp8 => minimaxM3Sm107}/specOffsets.cu (93%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp4 => minimaxM3Sm107}/specOffsets.h (95%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp8 => minimaxM3Sm107}/stepPrologue.cu (96%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp2 => minimaxM3Sm107}/stepPrologue.h (96%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp8 => minimaxM3Sm107}/swigluOai.cu (96%) rename cpp/tensorrt_llm/kernels/{minimaxM3Sm107Tp4 => minimaxM3Sm107}/swigluOai.h (96%) delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h delete mode 100644 cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h rename cpp/tensorrt_llm/thop/minimaxM3/{minimaxM3Sm107Tp4SiluAndMulOp.cpp => minimaxM3Sm107SiluAndMulOp.cpp} (51%) rename cpp/tensorrt_llm/thop/minimaxM3/{minimaxM3Sm107Tp4SpecOffsetsOp.cpp => minimaxM3Sm107SpecOffsetsOp.cpp} (52%) rename cpp/tensorrt_llm/thop/minimaxM3/{minimaxM3Sm107Tp8StepPrologueOp.cpp => minimaxM3Sm107StepPrologueOp.cpp} (57%) rename cpp/tensorrt_llm/thop/minimaxM3/{minimaxM3Sm107Tp8SwigluOaiOp.cpp => minimaxM3Sm107SwigluOaiOp.cpp} (56%) delete mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp delete mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp delete mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp delete mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp delete mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp delete mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp delete mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp delete mode 100644 cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/config.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/index_score_runtime.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/m3_index_score.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/m3_select_long_dsl.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py => kernels/attention/producer_tp8.py} (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/triton_sparse_decode.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2/ops => kernels/attention}/triton_sparse_decode_mq.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/embedding_firstnorm/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2/ops => kernels}/embedding_firstnorm/m3_embedding_firstnorm_kernel.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8/ops => kernels/gemm}/m3_lowm_splitk.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv => kernels/gemm}/m3_mx_gemv_kernel.py (95%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/moe/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe => kernels/moe}/m3_moe_kernel.py (98%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/m3_fold_tail.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2/ops => kernels/spec}/m3_force_vec.py (100%) delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py delete mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/modeling_support.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/activation/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/activation/m3_silu_and_mul.py (67%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/draft_attention.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/drafter.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/fullctx.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/long_select_host.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/attention/m3_dense_decode.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/attention/m3_fp8_indexer_qk_norm_rope.py (85%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_index_score.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp4 => }/ops/attention/m3_qk_norm_rope_fp8.py (84%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/attention/m3_qkv_insert.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/attention/m3_select_blocks.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/attention/msa_kv_update.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py => ops/attention/msa_kv_update_hook.py} (95%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/producer_reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/select_reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/sparse_decode_policy.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/sparse_reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/token_map.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp4 => }/ops/comm/m3_allreduce.py (77%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/mc_workspace.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/decode_embedding.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp4 => }/ops/embedding/m3_embedding.py (65%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/locality.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/gemm/m3_bf16_linear.py (61%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8 => }/ops/gemm/m3_dense_swiglu.py (66%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/gemm/m3_fp32_linear.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8 => }/ops/gemm/m3_mxfp8_gemm.py (74%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp4 => }/ops/gemm/m3_mxfp8_linear.py (77%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/gemm/m3_target_head.py (100%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_core.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_host.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_state.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/prequant.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/target_head_state.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/__init__.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/moe/m3_router_gemv.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8 => }/ops/moe/m3_router_gemv_fp4.py (95%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/moe/m3_router_tc.py (67%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/phase_plan.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/runtime.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/state.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/accept_workspace.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/acceptance_reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/draft_head_reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/fold_tail.py rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/spec/m3_capture_norm.py (71%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/spec/m3_chunk_norm.py (76%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/spec/m3_draft_input.py (100%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8 => }/ops/spec/m3_force_vec.py (96%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp8 => }/ops/spec/m3_slot_scatter.py (87%) rename tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2 => }/ops/spec/m3_step_gather.py (71%) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/sampling_reference.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/reporting.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/__init__.py rename tests/unittest/_torch/modeling_v2/models/minimax_m3/{minimax_m3_nvfp4__sm_107__tp2/gemm => }/_m3_mxfp8_operands.py (100%) create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/_m3_test_inputs.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_distributed_sampled_reference.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_exchange.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py delete mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py delete mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_spec_tail_contracts.py create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/test_collective_launcher_cleanup.py diff --git a/cpp/tensorrt_llm/kernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/CMakeLists.txt index b79d6abee1bc..937917791cb1 100644 --- a/cpp/tensorrt_llm/kernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/CMakeLists.txt @@ -93,7 +93,7 @@ list(FILTER SRC_CU EXCLUDE REGEX "minimaxM3/(fusedDecodeAttn|attnMega|attnC1)\\.cu$") # The MiniMax-M3 sm_107 decode kernels are built as their own sm_107-only OBJECT # library below. -list(FILTER SRC_CU EXCLUDE REGEX "minimaxM3Sm107Tp[248]/.*") +list(FILTER SRC_CU EXCLUDE REGEX "minimaxM3Sm107(Tp[248])?/.*") if(NOT ENABLE_MULTI_DEVICE) list(FILTER SRC_CU EXCLUDE REGEX "customAllReduceKernels*.*cu$") @@ -151,8 +151,8 @@ set_cuda_architectures(minimax_m3_attn_src 100 103) # the older ones. A build without an sm_10x architecture compiles them for # sm_107a: their torch ops link against the host launchers, and only the sm_107 # targets launch them. -file(GLOB MINIMAX_M3_SM107_SRC minimaxM3Sm107Tp2/*.cu minimaxM3Sm107Tp4/*.cu - minimaxM3Sm107Tp8/*.cu) +file(GLOB MINIMAX_M3_SM107_SRC minimaxM3Sm107/*.cu minimaxM3Sm107Tp2/*.cu + minimaxM3Sm107Tp4/*.cu minimaxM3Sm107Tp8/*.cu) add_library(minimax_m3_sm107_src OBJECT ${MINIMAX_M3_SM107_SRC}) set_property(TARGET minimax_m3_sm107_src PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET minimax_m3_sm107_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.cu similarity index 91% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.cu index 179770d14a0e..32610bcbd887 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.cu @@ -15,19 +15,19 @@ * limitations under the License. */ -// The TP2 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's GPT-NeoX -// path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin negated, -// dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale E4M3 -// store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query +// The sm_107 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's +// GPT-NeoX path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin +// negated, dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale +// E4M3 store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query // saturates the same way (448, not NaN). -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.h" TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp2 +namespace minimaxM3Sm107 { namespace drafterRopeFp8Kv { @@ -183,7 +183,7 @@ cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaS } } // namespace drafterRopeFp8Kv -} // namespace minimaxM3Sm107Tp2 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.h similarity index 93% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.h index f5316504784a..6131221a5668 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.h @@ -27,7 +27,7 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp8 +namespace minimaxM3Sm107 { namespace drafterRopeFp8Kv { @@ -63,7 +63,7 @@ struct Params int ropeMode; }; -//! The MiniMax-M3 TP8 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, +//! The MiniMax-M3 sm_107 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, //! rounded to BF16 (the in-tree preprocess's q output and cache input), then Q, K and V to saturating E4M3: K / V //! as the in-tree preprocess stores them at unit KV scale, Q as torch's CUDA cast of the BF16 query. Rows map to //! (request token / qLen, row token % qLen) at cache position seqLen - qLen + row. @@ -73,7 +73,7 @@ struct Params cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream); } // namespace drafterRopeFp8Kv -} // namespace minimaxM3Sm107Tp8 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/foldTail.cu similarity index 93% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/foldTail.cu index 309db8a47f3d..d4c119f0272f 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/foldTail.cu @@ -15,13 +15,13 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h" TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp4 +namespace minimaxM3Sm107 { namespace @@ -58,7 +58,7 @@ cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream) return cudaGetLastError(); } -} // namespace minimaxM3Sm107Tp4 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h similarity index 96% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h index 2c565e6c0c7f..42746073630b 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h @@ -26,7 +26,7 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp8 +namespace minimaxM3Sm107 { //! Operands of the speculative step's forward-tail store and next-step staging (kernels/m3_fold_tail.py). Every tensor @@ -76,7 +76,8 @@ cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream); //! own kernel (e.g. the drafter's last-pass draft_step in its last CTA, after its final tokens are written and fenced). //! The caller runs rows i < *p.ctrl (one warp per row, or one warp looping over the rows); same code as the kernel, so //! the writes are bit for bit the standalone tail's. -__device__ __forceinline__ void foldTailRow(FoldTailParams const& p, int i, int c) +template +__device__ __forceinline__ void foldTailRow(ParamsT const& p, int i, int c) { int64_t const s = p.slots[i]; int32_t const nw = c < p.snW ? p.srcNew[i * p.sn0 + c * p.sn1] : 0; @@ -124,7 +125,7 @@ __device__ __forceinline__ void foldTailWarp(FoldTailParams const& p, int lane) } #endif -} // namespace minimaxM3Sm107Tp8 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cu new file mode 100644 index 000000000000..3e04960f150f --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cu @@ -0,0 +1,47 @@ +/* + * 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. + */ + +// The split-KV merge for the MiniMax-M3 NVFP4 TP4/TP8 (Rubin) full-context decode. The producer is the PrimTS split-KV +// decode with its separate-reduction scratch: per (request, KV head, split, row) one normalized 16-bit partial O and +// one fp32 log2-LSE. Active splits per row follow the producer's split ranges (the causal union of the row's grouped +// Q CTA, instruction-group aligned): splits past them hold nothing and are not read. +// +// One CTA per (request, KV head, kRowsPerCta rows); per row kGroups split groups x kThreadsPerRow threads x 8 dims. +// A thread issues the loads of every split it owns (strided over the groups) before using any, folds them under its +// group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to +// BF16 once. + +#include "tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cuh" +#include "tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107 +{ +namespace fullCtxReduce +{ +cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) +{ + return minimaxM3Sm107::fullCtxReduce::detail::invoke<1, true>(params, pdl, stream); +} +} // namespace fullCtxReduce +} // namespace minimaxM3Sm107 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cuh similarity index 72% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cuh index ea33729aed7a..a3a050f5c848 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cuh @@ -25,9 +25,9 @@ // group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to // BF16 once. -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h" +#pragma once -#include +#include "tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.h" #include #include #include @@ -36,20 +36,19 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp8 +namespace minimaxM3Sm107 { namespace fullCtxReduce { -namespace +namespace detail { constexpr int kGroups = 8; constexpr int kThreadsPerRow = kHeadDim / 8; -constexpr int kRowsPerCta = 1; constexpr int kSplitsPerGroup = kMaxSplits / kGroups; -constexpr int kThreads = kGroups * kThreadsPerRow * kRowsPerCta; -__device__ int activeSplits(Params const& p, int b, int row) +template +__device__ int activeSplits(ParamsT const& p, int b, int row) { int const seqKv = p.seqLensKv[b]; int visible = seqKv; @@ -80,14 +79,17 @@ __device__ float2 unpack(uint32_t word) } } -template -__global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) +template +__global__ void __launch_bounds__(kGroups* kThreadsPerRow* kRowsPerCta) fullCtxReduceKernel(ParamsT const p) { #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 asm volatile("griddepcontrol.wait;" ::: "memory"); // The dependent (the o_proj activation quantization / GEMV) reads `out` only after its own griddepcontrol.wait, // which waits for this whole grid, so it is released here and runs its weight-only prologue under the merge. - asm volatile("griddepcontrol.launch_dependents;"); + if constexpr (kQuant) + { + asm volatile("griddepcontrol.launch_dependents;"); + } #endif int const b = blockIdx.z; int const h = blockIdx.y; @@ -197,53 +199,55 @@ __global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) size_t const outRow = (static_cast(b) * p.seqLenQ + token) * (p.numKvHeads * p.headsPerKv) + h * p.headsPerKv + head; reinterpret_cast(p.out + outRow * kHeadDim)[t] = *reinterpret_cast(packed); - if (p.q8 == nullptr) + if constexpr (kQuant) { - return; - } - // MXFP8 copy of the BF16 row: a 32-column block is 4 neighbouring threads x 8 columns (lanes 0..15 of warp 0 - // hold the row). Block scale = the UE8M0 exponent of ceil(amax / 448), values = satfinite E4M3 of x * 2^-scale. - float x[8]; - float amax = 0.f; + if (p.q8 == nullptr) + { + return; + } + // MXFP8 copy of the BF16 row: a 32-column block is 4 neighbouring threads x 8 columns (lanes 0..15 of warp 0 + // hold the row). Block scale = the UE8M0 exponent of ceil(amax / 448), values = satfinite E4M3 of x * 2^-scale. + float x[8]; + float amax = 0.f; #pragma unroll - for (int j = 0; j < 4; ++j) - { - float2 const f = __bfloat1622float2(packed[j]); - x[2 * j] = f.x; - x[2 * j + 1] = f.y; - amax = fmaxf(amax, fmaxf(fabsf(f.x), fabsf(f.y))); - } - amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 1)); - amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 2)); - float rcp; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); - float const sfv = amax * rcp; - int const bits = __float_as_int(sfv); - int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); - e = min(e, 254); - float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; - uint16_t q2[4]; + for (int j = 0; j < 4; ++j) + { + float2 const f = __bfloat1622float2(packed[j]); + x[2 * j] = f.x; + x[2 * j + 1] = f.y; + amax = fmaxf(amax, fmaxf(fabsf(f.x), fabsf(f.y))); + } + amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 1)); + amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 2)); + float rcp; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); + float const sfv = amax * rcp; + int const bits = __float_as_int(sfv); + int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); + e = min(e, 254); + float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; + uint16_t q2[4]; #pragma unroll - for (int j = 0; j < 4; ++j) - { - asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2[j]) : "f"(x[2 * j + 1] * qs), "f"(x[2 * j] * qs)); - } - int const tok = b * p.seqLenQ + token; - int const col = (h * p.headsPerKv + head) * kHeadDim + t * 8; - *reinterpret_cast(p.q8 + static_cast(tok) * (p.numSfKTiles * 128) + col) - = *reinterpret_cast(q2); - if (t % 4 == 0) - { - int const kSf = col / 32; - int64_t const off = static_cast(tok / 128) * p.numSfKTiles * 512 + (kSf / 4) * 512 + (tok % 32) * 16 - + ((tok % 128) / 32) * 4 + kSf % 4; - p.sf[off] = static_cast(e); + for (int j = 0; j < 4; ++j) + { + asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2[j]) : "f"(x[2 * j + 1] * qs), "f"(x[2 * j] * qs)); + } + int const tok = b * p.seqLenQ + token; + int const col = (h * p.headsPerKv + head) * kHeadDim + t * 8; + *reinterpret_cast(p.q8 + static_cast(tok) * (p.numSfKTiles * 128) + col) + = *reinterpret_cast(q2); + if (t % 4 == 0) + { + int const kSf = col / 32; + int64_t const off = static_cast(tok / 128) * p.numSfKTiles * 512 + (kSf / 4) * 512 + + (tok % 32) * 16 + ((tok % 128) / 32) * 4 + kSf % 4; + p.sf[off] = static_cast(e); + } } } -} // namespace - -cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) +template +cudaError_t invoke(ParamsT const& params, bool pdl, cudaStream_t stream) { if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1 || params.qTokensPerCta <= 0 || params.tileKv <= 0 || params.instsKv <= 0 || params.instsKv > std::numeric_limits::max() / kMaxSplits) @@ -252,19 +256,21 @@ cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t str } cudaLaunchConfig_t config = {}; config.gridDim = dim3((params.rows + kRowsPerCta - 1) / kRowsPerCta, params.numKvHeads, params.batch); - config.blockDim = dim3(kThreads); + config.blockDim = dim3(kGroups * kThreadsPerRow * kRowsPerCta); config.stream = stream; cudaLaunchAttribute attribute[1]; attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; config.attrs = attribute; config.numAttrs = 1; - return params.bf16Partial ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) - : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); + return params.bf16Partial + ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) + : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); } +} // namespace detail } // namespace fullCtxReduce -} // namespace minimaxM3Sm107Tp8 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.h similarity index 92% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.h index 1095ec4e020d..cded8b184e68 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.h @@ -27,10 +27,10 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp8 +namespace minimaxM3Sm107 { -//! The split-KV merge of the MiniMax-M3 TP2 full-context decode (PrimTS split-KV decode, separate-reduction scratch) -//! straight to BF16. See fullCtxReduce.cu. +//! The split-KV merge of the MiniMax-M3 TP4/TP8 full-context decode (PrimTS split-KV decode, separate-reduction +//! scratch) straight to BF16. See fullCtxReduce.cu. namespace fullCtxReduce { @@ -71,7 +71,7 @@ struct Params cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream); } // namespace fullCtxReduce -} // namespace minimaxM3Sm107Tp8 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.cu similarity index 87% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.cu index cad2b93558ab..947b29538d27 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.cu @@ -15,7 +15,8 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceUtils.cuh" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/cudaUtils.h" @@ -52,160 +53,23 @@ namespace cg = cooperative_groups; namespace { -constexpr int kEmpty = static_cast(0x80000000u); -constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k - -// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], -// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c -// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp -// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its -// arithmetic is the same source. -constexpr int kCapThreads = 256; -constexpr int kCapIters = 3; -constexpr int kCapWarps = kCapThreads / 32; -constexpr int kCapHidden = kCapThreads * kCapIters * 8; - -// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). -__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) -{ - __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); -#pragma unroll - for (int q = 0; q < 8; ++q) - { - float const v = __bfloat162float(e[q]); - ss += v * v; - } - return ss; -} - -__device__ __forceinline__ float cap_warp_sum(float v) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v += __shfl_xor_sync(0xffffffffu, v, o); - return v; -} - -// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. -__device__ __forceinline__ void cap_store( - float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) -{ - uint4 const wr = *reinterpret_cast(w); - __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); - uint4 o; - __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); -#pragma unroll - for (int q = 0; q < 8; ++q) - oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); - *reinterpret_cast(cap) = o; -} - -__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) -{ - if ((w & 0xFFFFu) == 0x8000u) - w &= 0xFFFF0000u; - if ((w >> 16) == 0x8000u) - w &= 0x0000FFFFu; - return w; -} - -__device__ __forceinline__ float bf16_lo(uint32_t w) -{ - return __uint_as_float(w << 16); -} - -// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's -// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals -// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the -// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. -__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) -{ - float value; - asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); - uint32_t result; - asm("{\n" - " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" - " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" - " setp.le.f32 p_zero, %1, 0f00000000;\n" - " mov.b32 bits, %1;\n" - " shr.b32 exp_biased, bits, 23;\n" - " and.b32 exp_biased, exp_biased, 255;\n" - " and.b32 mantissa, bits, 0x7FFFFF;\n" - " setp.ne.u32 p_has_mant, mantissa, 0;\n" - " selp.u32 bump, 1, 0, p_has_mant;\n" - " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" - " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" - " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" - " @p_tiny_sub mov.u32 bump, 0;\n" - " add.u32 r, exp_biased, bump;\n" - " setp.gt.u32 p_ovf, r, 254;\n" - " selp.u32 r, 254, r, p_ovf;\n" - " selp.u32 %0, 0, r, p_zero;\n" - "}\n" - : "=r"(result) - : "f"(value)); - return result; -} - -__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) -{ - float inv; - asm("{\n" - " .reg .s32 new_exp;\n" - " .reg .b32 float_bits;\n" - " .reg .pred p_zero;\n" - " setp.eq.u32 p_zero, %1, 0;\n" - " sub.s32 new_exp, 254, %1;\n" - " max.s32 new_exp, new_exp, 0;\n" - " shl.b32 float_bits, new_exp, 23;\n" - " mov.b32 %0, float_bits;\n" - " @p_zero mov.b32 %0, 0;\n" - "}\n" - : "=f"(inv) - : "r"(ue8m0)); - return inv; -} - -// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. -__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) -{ - uint16_t pair; - asm("{\n" - " .reg .f32 f0, f1;\n" - " mul.f32 f0, %1, %3;\n" - " mul.f32 f1, %2, %3;\n" - " min.f32 f0, f0, 0f43E00000;\n" - " max.f32 f0, f0, 0fC3E00000;\n" - " min.f32 f1, f1, 0f43E00000;\n" - " max.f32 f1, f1, 0fC3E00000;\n" - " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" - "}\n" - : "=h"(pair) - : "f"(lo), "f"(hi), "f"(inv)); - return pair; -} - -__device__ __forceinline__ float rcp_approx_ftz(float a) -{ - float b; - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); - return b; -} - -__device__ __forceinline__ float bf16_hi(uint32_t w) -{ - return __uint_as_float(w & 0xFFFF0000u); -} - -__device__ __forceinline__ int4 ld_volatile_v4(int const* p) -{ - int4 v; - asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" - : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) - : "l"(p) - : "memory"); - return v; -} +using minimaxM3Sm107::mcDetail::kEmpty; +using minimaxM3Sm107::mcDetail::kFinMaxTopk; +using minimaxM3Sm107::mcDetail::kCapThreads; +using minimaxM3Sm107::mcDetail::kCapIters; +using minimaxM3Sm107::mcDetail::kCapWarps; +using minimaxM3Sm107::mcDetail::kCapHidden; +using minimaxM3Sm107::mcDetail::cap_chunk_sq; +using minimaxM3Sm107::mcDetail::cap_warp_sum; +using minimaxM3Sm107::mcDetail::cap_store; +using minimaxM3Sm107::mcDetail::canon_bf16x2; +using minimaxM3Sm107::mcDetail::bf16_lo; +using minimaxM3Sm107::mcDetail::mxfp8_ue8m0; +using minimaxM3Sm107::mcDetail::mxfp8_inv_scale; +using minimaxM3Sm107::mcDetail::mxfp8_e4m3x2; +using minimaxM3Sm107::mcDetail::rcp_approx_ftz; +using minimaxM3Sm107::mcDetail::bf16_hi; +using minimaxM3Sm107::mcDetail::ld_volatile_v4; struct FinalizeArgs { @@ -673,7 +537,7 @@ bool checkBinary(Kernel kernel, char const* name) } template -void launch(MinimaxM3Sm107Tp2McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) +void launch(MinimaxM3Sm107McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) { int const M = p.numTokens, H = p.hidden; TLLM_CHECK_WITH_INFO(H % (CLUSTER * 8) == 0 && H / CLUSTER / 8 <= THREADS, "unsupported hidden size %d", H); @@ -997,7 +861,7 @@ __global__ void __launch_bounds__(THREADS, 1) mc_allreduce_norm_row_kernel(__nv_ } void launchRow( - MinimaxM3Sm107Tp2McAllreduceNormParams const& p, int32_t maxTokens, bool drafter, bool pdl, cudaStream_t stream) + MinimaxM3Sm107McAllreduceNormParams const& p, int32_t maxTokens, bool drafter, bool pdl, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(p.hidden == kRowThreads * 8 && p.world == 2, "the row kernel takes H 6144 at 2 ranks"); static bool const binaryOk @@ -1043,14 +907,14 @@ constexpr int kThreads = 128; } // namespace -bool isMinimaxM3Sm107Tp2McAllreduceNormSupported(int32_t maxTokens, int32_t opts) +bool isMinimaxM3Sm107McAllreduceNormSupported(int32_t maxTokens, int32_t opts) { return (maxTokens == 8 && (opts == 7 || opts == 15 || opts == 23 || opts == 31)) || (maxTokens == 16 && opts == 15) || (maxTokens == 48 && opts == 15); } -void invokeMinimaxM3Sm107Tp2McAllreduceNorm(MinimaxM3Sm107Tp2McAllreduceNormParams const& params, int32_t maxTokens, - int32_t opts, bool pdl, cudaStream_t stream) +void invokeMinimaxM3Sm107McAllreduceNorm( + MinimaxM3Sm107McAllreduceNormParams const& params, int32_t maxTokens, int32_t opts, bool pdl, cudaStream_t stream) { if (params.fc2 == nullptr && params.world == 2 && params.hidden == kRowThreads * 8) { diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.h similarity index 92% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.h index 4b36daf66813..d0e2dce84013 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.h @@ -35,7 +35,7 @@ namespace kernels //! pushes it through the multicast mapping `wsMc` into every rank's exchange, polls its own unicast view `wsUc`, //! sums the ranks in FP32 in rank order, adds the residual (BF16 round) and applies RMSNorm. With `outQuant` it //! also writes the MXFP8 copy of the normed rows (E4M3 [numTokens, hidden] and swizzled UE8M0 scales). -struct MinimaxM3Sm107Tp8McAllreduceNormParams +struct MinimaxM3Sm107McAllreduceNormParams { __nv_bfloat16 const* x{nullptr}; //!< [numTokens, hidden] partial; in finalize mode the FC2 rows [rows, hidden]. __nv_bfloat16 const* x2{nullptr}; //!< Optional second addend [numTokens, hidden]. @@ -71,12 +71,12 @@ struct MinimaxM3Sm107Tp8McAllreduceNormParams }; //! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). -bool isMinimaxM3Sm107Tp8McAllreduceNormSupported(int32_t maxTokens, int32_t opts); +bool isMinimaxM3Sm107McAllreduceNormSupported(int32_t maxTokens, int32_t opts); //! Launch with 8-CTA clusters of 128 threads, one cluster per token. `maxTokens` selects the exchange bucket; //! `opts` the epilogue variant (every variant is bit-identical). -void invokeMinimaxM3Sm107Tp8McAllreduceNorm(MinimaxM3Sm107Tp8McAllreduceNormParams const& params, int32_t maxTokens, - int32_t opts, bool pdl, cudaStream_t stream); +void invokeMinimaxM3Sm107McAllreduceNorm( + MinimaxM3Sm107McAllreduceNormParams const& params, int32_t maxTokens, int32_t opts, bool pdl, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceUtils.cuh b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceUtils.cuh new file mode 100644 index 000000000000..0121f6ecc066 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceUtils.cuh @@ -0,0 +1,192 @@ +/* + * 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 +{ +namespace minimaxM3Sm107 +{ +namespace mcDetail +{ +constexpr int kEmpty = static_cast(0x80000000u); +constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k + +// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], +// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c +// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp +// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its +// arithmetic is the same source. +constexpr int kCapThreads = 256; +constexpr int kCapIters = 3; +constexpr int kCapWarps = kCapThreads / 32; +constexpr int kCapHidden = kCapThreads * kCapIters * 8; + +// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). +__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) +{ + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); +#pragma unroll + for (int q = 0; q < 8; ++q) + { + float const v = __bfloat162float(e[q]); + ss += v * v; + } + return ss; +} + +__device__ __forceinline__ float cap_warp_sum(float v) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +} + +// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. +__device__ __forceinline__ void cap_store( + float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) +{ + uint4 const wr = *reinterpret_cast(w); + __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); + uint4 o; + __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); +#pragma unroll + for (int q = 0; q < 8; ++q) + oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); + *reinterpret_cast(cap) = o; +} + +__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) +{ + if ((w & 0xFFFFu) == 0x8000u) + w &= 0xFFFF0000u; + if ((w >> 16) == 0x8000u) + w &= 0x0000FFFFu; + return w; +} + +__device__ __forceinline__ float bf16_lo(uint32_t w) +{ + return __uint_as_float(w << 16); +} + +// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's +// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals +// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the +// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. +__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) +{ + float value; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); + uint32_t result; + asm("{\n" + " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" + " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" + " setp.le.f32 p_zero, %1, 0f00000000;\n" + " mov.b32 bits, %1;\n" + " shr.b32 exp_biased, bits, 23;\n" + " and.b32 exp_biased, exp_biased, 255;\n" + " and.b32 mantissa, bits, 0x7FFFFF;\n" + " setp.ne.u32 p_has_mant, mantissa, 0;\n" + " selp.u32 bump, 1, 0, p_has_mant;\n" + " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" + " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" + " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" + " @p_tiny_sub mov.u32 bump, 0;\n" + " add.u32 r, exp_biased, bump;\n" + " setp.gt.u32 p_ovf, r, 254;\n" + " selp.u32 r, 254, r, p_ovf;\n" + " selp.u32 %0, 0, r, p_zero;\n" + "}\n" + : "=r"(result) + : "f"(value)); + return result; +} + +__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) +{ + float inv; + asm("{\n" + " .reg .s32 new_exp;\n" + " .reg .b32 float_bits;\n" + " .reg .pred p_zero;\n" + " setp.eq.u32 p_zero, %1, 0;\n" + " sub.s32 new_exp, 254, %1;\n" + " max.s32 new_exp, new_exp, 0;\n" + " shl.b32 float_bits, new_exp, 23;\n" + " mov.b32 %0, float_bits;\n" + " @p_zero mov.b32 %0, 0;\n" + "}\n" + : "=f"(inv) + : "r"(ue8m0)); + return inv; +} + +// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. +__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) +{ + uint16_t pair; + asm("{\n" + " .reg .f32 f0, f1;\n" + " mul.f32 f0, %1, %3;\n" + " mul.f32 f1, %2, %3;\n" + " min.f32 f0, f0, 0f43E00000;\n" + " max.f32 f0, f0, 0fC3E00000;\n" + " min.f32 f1, f1, 0f43E00000;\n" + " max.f32 f1, f1, 0fC3E00000;\n" + " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" + "}\n" + : "=h"(pair) + : "f"(lo), "f"(hi), "f"(inv)); + return pair; +} + +__device__ __forceinline__ float rcp_approx_ftz(float a) +{ + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +__device__ __forceinline__ float bf16_hi(uint32_t w) +{ + return __uint_as_float(w & 0xFFFF0000u); +} + +__device__ __forceinline__ int4 ld_volatile_v4(int const* p) +{ + int4 v; + asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p) + : "memory"); + return v; +} + +} // namespace mcDetail +} // namespace minimaxM3Sm107 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/producerPw.cu similarity index 87% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/producerPw.cu index da30d840cc2e..4657dcb2220d 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/producerPw.cu @@ -15,13 +15,14 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/producerPw.h" #include #include #include -// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP2 heads. +// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP2/TP4/TP8 +// heads. // // The SM107 TP2 target's copy of minimaxM3/producerPw.cu (the sm_103 TP4 / TP8 targets keep theirs). Same math and // outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with --use_fast_math like @@ -42,9 +43,9 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp2 +namespace minimaxM3Sm107 { -namespace producerPw2 +namespace producerPw { namespace { @@ -127,11 +128,11 @@ __device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float se return static_cast(fp32x8_to_e2m1(q)); } -template +template __global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) { #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) if constexpr (kEarly) { @@ -146,10 +147,10 @@ __global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) if (tokenIdx >= p.numTokens) return; - constexpr int kBegin = kHeadsQ; - constexpr int vBegin = kBegin + kHeadsKV; - constexpr int indexQBegin = vBegin + kHeadsKV; - constexpr int indexKHead = indexQBegin + kHeadsIndex; + constexpr int kBegin = kHeadsQ; + constexpr int vBegin = kBegin + kHeadsKV; + constexpr int indexQBegin = vBegin + kHeadsKV; + constexpr int indexKHead = indexQBegin + kHeadsIndex; bool const isQ = localHead < kBegin; bool const isK = localHead >= kBegin && localHead < vBegin; bool const isV = localHead >= vBegin && localHead < indexQBegin; @@ -258,13 +259,15 @@ __global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) if (isQ) { - store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); + store_fp8x4( + p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); return; } if (isIndexQ) { store_fp8x4(p.indexQOut, - (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); + (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, + e); return; } @@ -301,10 +304,10 @@ __global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) #endif } -template +template cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) { - constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; + constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; int const warps = numTokens * kTotalHeads; cudaLaunchConfig_t cfg = {}; cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); @@ -315,25 +318,32 @@ cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; cfg.attrs = attr; cfg.numAttrs = 1; - return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); + return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); } } // namespace -cudaError_t invokeProducerPw2( +template +cudaError_t invokeProducerPw( Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream) { if (numTokens <= 0) return cudaSuccess; if (prewaitStepInputs) { - return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); + return early ? launch(p, numTokens, pdl, stream) + : launch(p, numTokens, pdl, stream); } - return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); + return early ? launch(p, numTokens, pdl, stream) + : launch(p, numTokens, pdl, stream); } -} // namespace producerPw2 -} // namespace minimaxM3Sm107Tp2 +template cudaError_t invokeProducerPw<2>(Params const&, int, bool, bool, bool, cudaStream_t); +template cudaError_t invokeProducerPw<4>(Params const&, int, bool, bool, bool, cudaStream_t); +template cudaError_t invokeProducerPw<8>(Params const&, int, bool, bool, bool, cudaStream_t); + +} // namespace producerPw +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/producerPw.h similarity index 77% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/producerPw.h index 9e94868af843..74aef9ee80d2 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/producerPw.h @@ -28,12 +28,12 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp2 +namespace minimaxM3Sm107 { -//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP2 heads (32 query, 2 KV and -//! 2 index heads per rank) with its constant loads before the PDL grid-dependency wait and an optional early dependents -//! trigger. Same outputs, bit for bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw2.cu. -namespace producerPw2 +//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP2/TP4/TP8 heads with its +//! constant loads before the PDL grid-dependency wait and an optional early dependents trigger. Same outputs, bit for +//! bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw.cu. +namespace producerPw { struct Params @@ -60,17 +60,21 @@ struct Params }; //! The heads per rank this producer is built for. -constexpr int kHeadsQ = 32; -constexpr int kHeadsKV = 2; -constexpr int kHeadsIndex = 2; +template +inline constexpr int kHeadsQ = 64 / kTp; +template +inline constexpr int kHeadsKV = kTp == 2 ? 2 : 1; +template +inline constexpr int kHeadsIndex = kTp == 2 ? 2 : 1; //! Launch over `numTokens` tokens, one warp per (token, packed head). `prewaitStepInputs` reads the position ids, the //! cache slots and the RoPE rows before the grid-dependency wait (they must not be written by a PDL predecessor). -cudaError_t invokeProducerPw2( +template +cudaError_t invokeProducerPw( Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream); -} // namespace producerPw2 -} // namespace minimaxM3Sm107Tp2 +} // namespace producerPw +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.cu new file mode 100644 index 000000000000..f1a4abe338ff --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.cu @@ -0,0 +1,965 @@ +/* + * 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/minimaxM3Sm107/sampledAccept.h" + +#include "tensorrt_llm/common/assert.h" + +#include +#include + +#include +#include + +// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). +// +// Shared sm_107 (Rubin) TP4/TP8 sampled verification with the SA-1R decision variants. The algorithm and +// arithmetic are unchanged. Each target keeps its existing kMaxVs specialization (TP4: 51200; TP8: 100352), +// which sizes the per-CTA logit staging and the launch's dynamic shared memory. +// +// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the +// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p +// active must stay on the generic path; the caller decides. +// +// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL +// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z +// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it +// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens +// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the +// nucleus with probability >= top_p, so most rows decide on their second exchange. +// +// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, +// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local +// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the +// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based +// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one +// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and +// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the +// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing +// and two call halves keep a rank one call ahead from overwriting a slot still being read. +// +// SA-1R (Args.one_round_k 9 / 82; MoE prototype moe/sa1r/6c753632ec, exact on 39,968 rows, eager + graph): each CTA +// also multicasts its largest bf16 logits with the round-0 record (3 per extra word). After round 0 every CTA knows M, +// the candidate logit cv and wc = __expf(cv - M); round 1 would accept iff S <= target, S = the fp32 sum (its tree) of +// w = __expf(__fmul_rn(x, tInv) - M) over every token with w > wc. x -> w is monotone (__fmul_rn by tInv > 0, fsub, +// __expf: checked exhaustively over all non-positive fp32 inputs), so if no CTA's unlisted value can beat wc (82: the +// bound tau; 9: the 9th listed), the listed w > wc are exactly round 1's summands, bitwise. Two fp32 sums of the same +// m non-negative terms differ by < m * 2^-23 * S, so accepting iff S_1r + m * 2^-21 * S_1r (round up) <= target +// implies round 1 accepts; m <= 1 is exact. Otherwise round 1 runs as before. An accepted row outputs the round-0 +// candidate, as the two-round kernel does, and consumes no Philox draw it would not. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ +namespace minimaxM3Sm107 +{ +namespace sampledAccept +{ +namespace +{ +constexpr int kThreads = 256; +constexpr int kWarps = kThreads / 32; +constexpr int kMaxRec = kMaxWorld * kMaxSplits; +constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) +constexpr int kTagBits = 64 - 32 - kIdBits; +constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard +// Record polls trap after this many spins (the multicast all-reduce's bound): a missing writer becomes a CUDA error, +// never a hang. +constexpr long long kMaxSpins = 1LL << 31; + +// SA-1R variant kTopK: 0 none; K exact top-K list; 10 * K + P pooled (each thread keeps its top P plus its (P + 1)-th +// as a bound; the CTA lists the pool's top K plus tau = max(K-th listed, max thread bound)). +__host__ __device__ constexpr int mode_k(int m) +{ + return m >= 10 ? m / 10 : m; +} + +__host__ __device__ constexpr int mode_pool(int m) +{ + return m >= 10 ? m % 10 : 0; +} + +__host__ __device__ constexpr int mode_vals(int m) +{ + return mode_k(m) + (mode_pool(m) ? 1 : 0); +} + +// Words per record: key, logit of key, f2, f3, then the round-0 list (3 bf16 + a 16-bit tag per word). +template +struct Rec +{ + static constexpr int kXw = (mode_vals(kTopK) + 2) / 3; + static constexpr int kEx = 4 + kXw; +}; + +// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. +template +__host__ __device__ inline long long half_words(int world) +{ + return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * Rec::kEx; +} + +__device__ __forceinline__ uint32_t ordered(float v) +{ + uint32_t const b = __float_as_uint(v); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +// bf16 bits <-> an order-preserving 16-bit key. +__device__ __forceinline__ uint32_t ord16(uint32_t b) +{ + return (b & 0x8000u) ? (~b & 0xFFFFu) : (b | 0x8000u); +} + +__device__ __forceinline__ uint32_t unord16(uint32_t k) +{ + return (k & 0x8000u) ? (k & 0x7FFFu) : (~k & 0xFFFFu); +} + +// Philox4x32-10. +__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) +{ +#pragma unroll + for (int i = 0; i < 10; ++i) + { + uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; + uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; + c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); + k.x += 0x9E3779B9u; + k.y += 0xBB67AE85u; + } + return c; +} + +// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. +__device__ __forceinline__ float uniform(uint32_t bits) +{ + return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); +} + +// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. +__device__ __forceinline__ unsigned long long gkey(float lg, int id) +{ + return (static_cast(ordered(lg)) << kIdBits) + | static_cast((1u << kIdBits) - 1u - static_cast(id)); +} + +__device__ __forceinline__ int key_token(unsigned long long k) +{ + return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); +} + +__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) +{ + asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); +} + +__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) +{ + unsigned long long v; + asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); + return v; +} + +template +struct Shared +{ + __nv_bfloat16 row[kMaxVs / kSplits]; + float redf[kWarps]; + unsigned long long redk[kWarps]; + unsigned long long rk[kMaxRec]; + float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; + float wsum[kWarps]; + int owner; + int pick; + float pickx; + bool last; + float dS[kWarps]; // SA-1R: per-warp partial mass above the candidate + int dM[kWarps]; // SA-1R: per-warp count | incomplete << 16 + uint32_t wtop[kWarps][kTopK > 0 ? mode_k(kTopK) : 1]; // SA-1R: per-warp top-K keys + uint32_t wbound[kWarps]; // SA-1R pooled: per-warp max thread bound + uint32_t ctop[kTopK > 0 ? mode_vals(kTopK) : 1]; // SA-1R: the CTA's top-K keys (descending), by rank + unsigned long long rx[kMaxRec][Rec::kXw > 0 ? Rec::kXw : 1]; +}; + +template +__device__ __forceinline__ float block_max(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = s.redf[0]; + for (int i = 1; i < kWarps; ++i) + r = fmaxf(r, s.redf[i]); + __syncthreads(); + return r; +} + +// Fixed-order block sum: identical result for identical inputs. +template +__device__ __forceinline__ float block_sum(float v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + if ((threadIdx.x & 31) == 0) + s.redf[threadIdx.x >> 5] = v; + __syncthreads(); + float r = 0.f; + for (int i = 0; i < kWarps; ++i) + r += s.redf[i]; + __syncthreads(); + return r; +} + +template +__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) +{ +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + if ((threadIdx.x & 31) == 0) + s.redk[threadIdx.x >> 5] = v; + __syncthreads(); + unsigned long long r = s.redk[0]; + for (int i = 1; i < kWarps; ++i) + r = s.redk[i] > r ? s.redk[i] : r; + __syncthreads(); + return r; +} + +// Exclusive prefix of v over the threads in thread order, and the total. +template +__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + float inc = v; +#pragma unroll + for (int o = 1; o < 32; o <<= 1) + { + float const t = __shfl_up_sync(0xffffffffu, inc, o); + if (lane >= o) + inc += t; + } + if (lane == 31) + s.wsum[warp] = inc; + __syncthreads(); + float base = 0.f; + total = 0.f; + for (int i = 0; i < kWarps; ++i) + { + base += i < warp ? s.wsum[i] : 0.f; + total += s.wsum[i]; + } + __syncthreads(); + return base + inc - v; +} + +// Warp-wide integer max / sum / or (the redux instructions from sm_80; a butterfly otherwise, same results). +__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_max_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + { + uint32_t const w = __shfl_xor_sync(0xffffffffu, v, o); + v = w > v ? w : v; + } + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_add_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_add_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v += __shfl_xor_sync(0xffffffffu, v, o); + return v; +#endif +} + +__device__ __forceinline__ uint32_t warp_or_u32(uint32_t v) +{ +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + return __reduce_or_sync(0xffffffffu, v); +#else +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + v |= __shfl_xor_sync(0xffffffffu, v, o); + return v; +#endif +} + +// Insert key into a descending register list of K (fully unrolled: no local memory). +template +__device__ __forceinline__ void topk_insert(uint32_t (&a)[K], uint32_t key) +{ +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const hi = key > a[j] ? key : a[j]; + key = key > a[j] ? a[j] : key; + a[j] = hi; + } +} + +template +__device__ __forceinline__ void topk_pop(uint32_t (&a)[K]) +{ +#pragma unroll + for (int j = 0; j + 1 < K; ++j) + a[j] = a[j + 1]; + a[K - 1] = 0u; +} + +// Per-thread insertion with an early reject: after the first few elements almost every key fails the K-th test. +template +__device__ __forceinline__ void thread_insert(uint32_t (&a)[K], uint32_t key) +{ + if (key > a[K - 1]) + topk_insert(a, key); +} + +// The warp's K largest keys (descending) into s.wtop[warp]; published by the caller's next block barrier. +template +__device__ __forceinline__ void warp_topk(uint32_t (&a)[K], S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const m = warp_max_u32(a[0]); + if (a[0] == m && m != 0u) + topk_pop(a); + if (lane == 0) + s.wtop[warp][j] = m; + } +} + +// Pooled: K rounds over each lane's P pool keys (descending) and the warp's max thread bound. +template +__device__ __forceinline__ void warp_pool(uint32_t (&a)[P + 1], S& s) +{ + int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + uint32_t const bound = warp_max_u32(a[P]); + uint32_t pl[P]; +#pragma unroll + for (int j = 0; j < P; ++j) + pl[j] = a[j]; +#pragma unroll + for (int j = 0; j < K; ++j) + { + uint32_t const m = warp_max_u32(pl[0]); + if (pl[0] == m && m != 0u) + topk_pop

(pl); + if (lane == 0) + s.wtop[warp][j] = m; + } + if (lane == 0) + s.wbound[warp] = bound; +} + +// A row's kSplits CTAs wait for each other's records, so they launch as one thread-block cluster: the hardware +// schedules a cluster's CTAs together, whatever other streams hold. Rows never wait on other rows. +template +__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) +{ + using R = Rec; + extern __shared__ __align__(16) unsigned char smem[]; + Shared& s = *reinterpret_cast*>(smem); + float constexpr kNegInf = -__builtin_huge_valf(); + int const tid = threadIdx.x; + int const row = blockIdx.x / kSplits; + int const split = blockIdx.x % kSplits; + int const world = a.world; + int const nrec = world * kSplits; + // The row's sampling parameters are step inputs (the spec metadata, written before the target forward), not the + // predecessor head GEMM's output, so with prewait_params they load before the dependency wait and overlap the + // head's tail; the same values either way. + float temperature = 0.f, topP = 0.f; + long long seed_in = 0, offset_in = 0; + auto load_params = [&]() + { + temperature = a.temperatures[row]; + topP = a.top_ps[row]; + seed_in = a.seeds[row]; + offset_in = a.offsets[row]; + }; + if (a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("griddepcontrol.wait;" ::: "memory"); +#endif + if (!a.prewait_params) + load_params(); +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // early_trigger: the dependent grid (the first draft pass's input) may launch now and park at its own + // griddepcontrol.wait, which still returns only once this grid has completed and flushed; it skips the launch + // latency after this kernel ends. Otherwise dependents are released at the acceptance epilogue (the CTAs of + // splits 1.. exit without releasing). + if (a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + float const tempInv = 1.0f / (temperature + kTempEpsilon); + unsigned long long const seed = static_cast(seed_in); + unsigned long long const offset = static_cast(offset_in); + unsigned int const call = *reinterpret_cast(a.counters + 1); + // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of + // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round + // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten + // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. + auto ktag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; + }; + auto ftag_of = [&](int r) -> unsigned long long + { + unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; + return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; + }; + // SA-1R round-0 list words: a 16-bit tag, never 0. + unsigned long long const xtag = (((call >> 1) << 1) & 0x7FFFu) | 0x8000u; + long long const rec_stride = static_cast(world) * kMaxSplits * R::kEx; + unsigned long long* const xu = a.ws_uc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride; + unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) + + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * R::kEx; + long long const round_stride = static_cast(kMaxRows) * rec_stride; + + // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- + __nv_bfloat16 const* src = a.logits + row * a.ld; + int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; + int const s0 = min(split * sl_len, a.Vs); + int const n = min(s0 + sl_len, a.Vs) - s0; + int const id0 = a.vocab_offset + s0; + float tmax = kNegInf; + // SA-1R: each thread's largest (ord16(bf16) << 16 | ~index) keys while staging (only rows with a nucleus cut). + bool const want_x = kTopK > 0 && topP < 1.0f; + constexpr int kK = mode_k(kTopK), kPool = mode_pool(kTopK), kVals = mode_vals(kTopK); + constexpr int kTk = kPool ? kPool + 1 : (kK > 0 ? kK : 1); + uint32_t tk[kTk]; +#pragma unroll + for (int j = 0; j < kTk; ++j) + tk[j] = 0u; + bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; + int const nvec = vec ? n / 8 : 0; + for (int v = tid; v < nvec; v += kThreads) + { + uint4 const w = *reinterpret_cast(src + s0 + 8 * v); + *reinterpret_cast(&s.row[8 * v]) = w; + __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); +#pragma unroll + for (int j = 0; j < 8; ++j) + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); + if constexpr (kTopK > 0) + { + if (want_x) + { +#pragma unroll + for (int j = 0; j < 8; ++j) + { + uint32_t const key + = (ord16(__bfloat16_as_ushort(e[j])) << 16) | (0xFFFFu - static_cast(8 * v + j)); + if constexpr (kPool > 0) + topk_insert(tk, key); // branchless: P + 1 compare-swaps + else + thread_insert(tk, key); + } + } + } + } + for (int i = nvec * 8 + tid; i < n; i += kThreads) + { + s.row[i] = src[s0 + i]; + tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); + if constexpr (kTopK > 0) + { + if (want_x) + topk_insert( + tk, (ord16(__bfloat16_as_ushort(src[s0 + i])) << 16) | (0xFFFFu - static_cast(i))); + } + } + if constexpr (kTopK > 0) + { + if (tid < kVals) + s.ctop[tid] = 0u; // published with s.wtop by block_max's barrier + if (want_x) + { + if constexpr (kPool > 0) + warp_pool(tk, s); // s.wtop / s.wbound published by block_max's barrier + else + warp_topk(tk, s); // s.wtop published by block_max's barrier + } + } + float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier + if constexpr (kTopK > 0) + { + // SA-1R: the CTA's top K of the warp lists by rank (keys are unique): a lane quad per key, no serial merge; + // s.ctop is read after the round-0 draw's first barrier + constexpr int kN = kWarps * kK; + constexpr int kSp = kN * 4 <= kThreads ? 4 : 1; // threads per key + if (want_x && tid < kN * kSp) + { + uint32_t const* const wl = &s.wtop[0][0]; + int const part = tid % kSp; + uint32_t const key = wl[tid / kSp]; + int rank = 0; +#pragma unroll + for (int j = part; j < kN; j += kSp) + rank += wl[j] > key ? 1 : 0; + if constexpr (kSp == 4) + { + rank += __shfl_xor_sync(0xffffffffu, rank, 1); + rank += __shfl_xor_sync(0xffffffffu, rank, 2); + } + if (part == 0 && key != 0u && rank < kK) + s.ctop[rank] = key; + } + } + float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z + unsigned long long ckey = 0ULL; + int const gid = a.rank * kMaxSplits + split; + for (int r = 0;; ++r) + { + // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) + float const ref = r == 0 ? smax : M; + float const wmin = r == 0 ? 0.f : wc; + float ts = 0.f; + for (int i = tid; i < n; i += kThreads) + { + float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); + ts += w > wmin ? w : 0.f; + } + float total; + float const excl = block_excl_scan(ts, total, s); + if constexpr (kTopK > 0) + { + // SA-1R round 0: the list words depend only on the staged slice, so they go out now (s.ctop / s.wbound are + // published by the barrier above), one lane per word, ahead of the main record: off the readers' path + if (r == 0 && want_x && tid < R::kXw) + { + uint32_t tau = 0u; + if constexpr (kPool > 0) + { + tau = s.ctop[kK - 1]; // the largest unlisted pool key is <= the K-th listed +#pragma unroll + for (int wq = 0; wq < kWarps; ++wq) + tau = s.wbound[wq] > tau ? s.wbound[wq] : tau; + } + unsigned long long w = xtag; +#pragma unroll + for (int t = 0; t < 3; ++t) + { + int const e = 3 * tid + t; + uint32_t const kv = e < kK ? s.ctop[e] : (kPool > 0 && e == kK ? tau : 0u); + // missing entries: -inf (bf16 0xFF80), w = 0 + uint32_t const b = kv != 0u ? unord16(kv >> 16) : 0xFF80u; + w |= static_cast(b) << (16 * (3 - t)); + } + st_sys(xm + 4 + tid, w); // round 0 = slot 0 + } + } + uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), + static_cast(offset), static_cast(offset >> 32)), + make_uint2(static_cast(seed), static_cast(seed >> 32))); + float const tgt = uniform(rnd.x) * total; + if (tid == 0) + s.owner = -1; + __syncthreads(); + if (ts > 0.f && excl <= tgt) + atomicMax(&s.owner, tid); + __syncthreads(); + if (tid == s.owner) + { + float acc = excl; + int pick = -1; + float px = 0.f; + if constexpr (kIlpWalk) + { + // SA-PW: the owner's strided tokens 8 at a time. The 8 x and w (the same expressions) are + // independent, so their MUFU work overlaps; the serial sum then adds them in the same order with + // the same early exit: the same acc sequence, pick and px as the one-token walk, bit for bit. + constexpr int kW = 8; + bool done = false; + for (int i0 = tid; i0 < n && !done; i0 += kW * kThreads) + { + float xs[kW], wv[kW]; +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + xs[k] = i < n ? __fmul_rn(__bfloat162float(s.row[i]), tempInv) : 0.f; + wv[k] = __expf(xs[k] - ref); + } +#pragma unroll + for (int k = 0; k < kW; ++k) + { + int const i = i0 + k * kThreads; + if (!done && i < n && wv[k] > wmin) + { + acc += wv[k]; + pick = i; + px = xs[k]; + if (acc > tgt) + done = true; + } + } + } + } + else + { + for (int i = tid; i < n; i += kThreads) + { + float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); + float const w = __expf(x - ref); + if (w > wmin) + { + acc += w; + pick = i; + px = x; + if (acc > tgt) + break; + } + } + } + s.pick = pick; + s.pickx = px; + } + __syncthreads(); + unsigned long long sb = 0ULL; + if (total > 0.f && s.owner >= 0) + sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); + float const smab = total; + unsigned long long const ktag = ktag_of(r); + unsigned long long const ftag = ftag_of(r); + unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; + bool const with_x = want_x && r == 0; + if (tid == 0) + { + unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; + st_sys(xw + 0, (sb << kTagBits) | ktag); + st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); + st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); + st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); + } + // ---- every record of the row, every rank --------------------------------------------------------------- + int const nw = with_x ? R::kEx : 4; + for (int t = tid; t < nrec * nw; t += kThreads) + { + int const q = t / nw, f = t % nw; + unsigned long long const* const p = xr + (q / kSplits * kMaxSplits + q % kSplits) * R::kEx + f; + unsigned long long v; + long long spins = 0; + if (f == 0) + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); + s.rk[q] = v >> kTagBits; + } + else if (f < 4) + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFFFFF00000000ULL) != ftag); + float const fv = __uint_as_float(static_cast(v)); + (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; + } + else + { + do + { + v = ld_sys(p); + if (++spins > kMaxSpins) + __trap(); + } while ((v & 0xFFFFULL) != xtag); + s.rx[q][f - 4] = v; + } + } + __syncthreads(); + // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- + unsigned long long gk = 0ULL; + float gv = 0.f; + for (int q = 0; q < nrec; ++q) + { + if (s.rk[q] > gk) + { + gk = s.rk[q]; + gv = s.rv[q]; + } + } + bool accept; + if (r == 0) + { + M = kNegInf; + for (int q = 0; q < nrec; ++q) + M = fmaxf(M, s.r2[q]); + float Z = 0.f; + for (int q = 0; q < nrec; ++q) + Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); + target = topP * Z; + accept = !(topP < 1.0f); // no nucleus cut: the first draw stands + ckey = gk; + cv = gv; + if constexpr (kTopK > 0) + { + // SA-1R: every thread takes (record, entry) items; the per-warp partials are published by the round's + // barrier below and every thread combines them in warp order after it (no new barrier) + if (with_x) + { + float const wcand = __expf(cv - M); // = round 1's wmin + float part = 0.f; + int cnt = 0, inc = 0; + for (int t = tid; t < nrec * kVals; t += kThreads) + { + int const q = t / kVals, e = t % kVals; + uint32_t const b = static_cast(s.rx[q][e / 3] >> (16 * (3 - e % 3))) & 0xFFFFu; + float const w = __expf( + __fmul_rn(__bfloat162float(__ushort_as_bfloat16(static_cast(b))), tempInv) + - M); // w bitwise as round 1 computes it + if (w > wcand) + { + if (kPool > 0 && e == kK) + inc = 1; // tau (the largest unlisted value) is heavier than the candidate + else + { + part += w; + ++cnt; + if (kPool == 0 && e == kK - 1) + inc = 1; // the K-th is heavier: unlisted heavier tokens may exist + } + } + } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) + part += __shfl_xor_sync(0xffffffffu, part, o); // any order: the margin covers it + cnt = static_cast(warp_add_u32(static_cast(cnt))); + inc = static_cast(warp_or_u32(static_cast(inc))); + if ((tid & 31) == 0) + { + s.dS[tid >> 5] = part; + s.dM[tid >> 5] = cnt | (inc << 16); + } + } + } + } + else + { + float S = 0.f; // mass strictly above the current candidate + for (int q = 0; q < nrec; ++q) + S += s.r2[q]; + accept = S <= target || gk == 0ULL; + if (!accept) + { + ckey = gk; + cv = gv; + } + } + wc = __expf(cv - M); + __syncthreads(); // records consumed before the next round's writes to s + bool decided = false; + if constexpr (kTopK > 0) + { + if (r == 0 && with_x) + { + float S1 = 0.f; + int m = 0, inc = 0; + for (int wq = 0; wq < kWarps; ++wq) + { + S1 += s.dS[wq]; + m += s.dM[wq] & 0xFFFF; + inc |= s.dM[wq] >> 16; + } + if (!inc) + { + float const margin = m <= 1 ? 0.f : S1 * static_cast(m) * (1.0f / 2097152.0f); + decided = __fadd_ru(S1, margin) <= target; + accept = accept || decided; + } + } + } + if (accept) + { + if (split != 0) + return; + if (tid == 0) + { + a.tokens[row] = key_token(ckey); + if (a.extra_attempts != nullptr && r > 1) + atomicAdd(a.extra_attempts, r - 1); + if (a.decided != nullptr && decided) + atomicAdd(a.decided, 1); + } + break; + } + } + + // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + if (!a.early_trigger) + asm volatile("griddepcontrol.launch_dependents;"); +#endif + __threadfence(); + __syncthreads(); + if (tid == 0) + s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; + __syncthreads(); + if (!s.last) + return; + __threadfence(); + if (tid == 0) + { + a.counters[0] = 0u; + a.counters[1] = call + 1u; + } + if (tid < a.nreq) + { + int const* tok = a.tokens; + int const b = tid; + int const k1 = a.draft_len + 1; + int nacc = 1; + bool run = true; + for (int k = 0; k < a.draft_len; ++k) + { + int const tt = *reinterpret_cast(tok + b * k1 + k); + a.acc_tokens[b * k1 + k] = tt; + run = run && a.draft[b * a.draft_len + k] == tt; + nacc += run ? 1 : 0; + } + a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); + if (a.force_vec != nullptr) + nacc = a.force_vec[b]; + else if (a.force_total > 0) + nacc = a.force_total; + a.num_acc[b] = nacc; + a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; + if (a.next_new != nullptr) + a.next_new[b * a.next_stride] + = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); + } +} + +template +cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) +{ + int device; + cudaError_t const deviceStatus = cudaGetDevice(&device); + if (deviceStatus != cudaSuccess) + return deviceStatus; + static std::mutex configMutex; + static std::unordered_map deviceClusters; + int max_clusters = 0; + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(a.rows * kSplits); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = sizeof(Shared); + cfg.stream = stream; + cudaLaunchAttribute attrs[2]; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = kSplits; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; + cfg.attrs = attrs; + cfg.numAttrs = 2; + { + std::lock_guard lock(configMutex); + max_clusters = deviceClusters[device]; + if (max_clusters == 0) + { + cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); + if (err == cudaSuccess) + { + err = cudaOccupancyMaxActiveClusters( + &max_clusters, sampled_accept_kernel, &cfg); + } + if (err != cudaSuccess) + { + return err; + } + } + deviceClusters[device] = max_clusters; + } + TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); + return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); +} + +template +cudaError_t launchRows(Args const& a, bool use_pdl, cudaStream_t stream) +{ + // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); + // more rows: 4 CTAs per row. + if (a.ilp_walk) + { + return a.rows <= 32 ? launch(a, use_pdl, stream) + : launch(a, use_pdl, stream); + } + return a.rows <= 32 ? launch(a, use_pdl, stream) + : launch(a, use_pdl, stream); +} + +} // namespace + +long long sampledAcceptWsWords(int world, int oneRoundK) +{ + TLLM_CHECK_WITH_INFO( + oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "SA-1R variant %d is not built", oneRoundK); + return 2 * (oneRoundK == 9 ? half_words<9>(world) : oneRoundK == 82 ? half_words<82>(world) : half_words<0>(world)); +} + +template +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) +{ + switch (a.one_round_k) + { + case 0: return launchRows(a, usePdl, stream); + case 9: return launchRows(a, usePdl, stream); + case 82: return launchRows(a, usePdl, stream); + default: TLLM_CHECK_WITH_INFO(false, "SA-1R variant %d is not built", a.one_round_k); + } + return cudaErrorInvalidValue; +} + +template cudaError_t invokeSampledAccept<51200>(Args const&, bool, cudaStream_t); +template cudaError_t invokeSampledAccept<100352>(Args const&, bool, cudaStream_t); + +} // namespace sampledAccept +} // namespace minimaxM3Sm107 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.h new file mode 100644 index 000000000000..ea19bb4f374a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.h @@ -0,0 +1,87 @@ +/* + * 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 +{ +namespace minimaxM3Sm107 +{ +//! Fused sampled (temperature + top-p) verification of the MiniMax-M3 Eagle3 target head over the TP-sharded +//! vocabulary, decode rows <= kMaxRows; see sampledAccept.cu. +namespace sampledAccept +{ + +constexpr int kMaxSplits = 8; // CTAs per verify row (at most) +constexpr int kMaxRows = 48; // verify rows per call +constexpr int kMaxWorld = 8; // TP ranks +constexpr int kIdBits = 18; // token ids < 2^18 + +struct Args +{ + __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard + long long ld; + int rows, Vs, vocab_offset; + float const* temperatures; + float const* top_ps; + long long const* seeds; + long long const* offsets; + unsigned long long* ws_uc; + unsigned long long* ws_mc; + unsigned int* counters; // [0] finished rows, [1] call + int* tokens; // [rows] sampled token per row + int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows + int rank, world; + int const* draft; + int const* batch_idx; + int* acc_tokens; + int* num_acc; + long long* gather_ids; + int draft_len, nreq, force_total; + int* next_new; + int next_stride; + int const* force_vec; + int prewait_params; // nonzero: the row's sampling parameters load before the PDL dependency wait + int early_trigger; // nonzero: dependents are released right after the dependency wait, not at the epilogue + // SA-1R (one-round decision): 0 the two-round kernel; 9 each CTA's top-9 logits ride its round-0 record; 82 pooled + // (each thread's top 2, the CTA lists the pool's top 8 plus the bound of everything unlisted). A row whose first + // candidate is provably in the nucleus then decides in round 0, with the same token as the two-round kernel. + int one_round_k; + int* decided; // optional [1]: rows decided in round 0 (SA-1R), summed + int ilp_walk; // nonzero: the owner walk takes its tokens 8 at a time (same weights, same summation order) +}; + +//! Words (uint64) of the multicast exchange workspace for `world` ranks and SA-1R variant `oneRoundK`: two call halves. +long long sampledAcceptWsWords(int world, int oneRoundK = 0); + +//! Launch the sampled verification over a.rows * (8 or 4) co-resident CTAs (8 per row up to 32 rows). +template +cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream); + +} // namespace sampledAccept +} // namespace minimaxM3Sm107 +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/siluAndMul.cu similarity index 96% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/siluAndMul.cu index 970b2aa21cf4..bebc12f822b9 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/siluAndMul.cu @@ -15,13 +15,13 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/siluAndMul.h" TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp4 +namespace minimaxM3Sm107 { namespace @@ -83,7 +83,7 @@ cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int return cudaGetLastError(); } -} // namespace minimaxM3Sm107Tp4 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/siluAndMul.h similarity index 95% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/siluAndMul.h index dccbb44572ef..e3c409df8b72 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/siluAndMul.h @@ -26,7 +26,7 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp4 +namespace minimaxM3Sm107 { //! out[m, j] = bf16(silu(g) * (u + 0)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both @@ -34,7 +34,7 @@ namespace minimaxM3Sm107Tp4 //! (no limit, alpha 1, beta 0), the plain SwiGLU of the Eagle3 drafter MLP. cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream); -} // namespace minimaxM3Sm107Tp4 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/specOffsets.cu similarity index 93% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/specOffsets.cu index fd101e7f521c..8a315c4ae41d 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/specOffsets.cu @@ -15,13 +15,13 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/specOffsets.h" TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp8 +namespace minimaxM3Sm107 { namespace @@ -52,7 +52,7 @@ cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, return cudaGetLastError(); } -} // namespace minimaxM3Sm107Tp8 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/specOffsets.h similarity index 95% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/specOffsets.h index 878ca0197e9e..33b2ad4046a5 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/specOffsets.h @@ -26,7 +26,7 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp4 +namespace minimaxM3Sm107 { //! The overlap scheduler's speculative step correction in one launch: pos[i] += posOff[i] for i < nPos and @@ -35,7 +35,7 @@ namespace minimaxM3Sm107Tp4 cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, int64_t nKv, cudaStream_t stream); -} // namespace minimaxM3Sm107Tp4 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/stepPrologue.cu similarity index 96% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/stepPrologue.cu index b1ab46a36088..e6a009f50ca9 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/stepPrologue.cu @@ -15,13 +15,13 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/stepPrologue.h" TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp8 +namespace minimaxM3Sm107 { namespace @@ -105,7 +105,7 @@ cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, return cudaGetLastError(); } -} // namespace minimaxM3Sm107Tp8 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/stepPrologue.h similarity index 96% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/stepPrologue.h index 44a7b7bcb44d..acd1ae191589 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/stepPrologue.h @@ -26,7 +26,7 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp2 +namespace minimaxM3Sm107 { //! The decode graph's step-input prologue in one launch (step_inputs.StepInputPack with early_copy): the work of the @@ -41,7 +41,7 @@ namespace minimaxM3Sm107Tp2 //! incremented once. cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream); -} // namespace minimaxM3Sm107Tp2 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/swigluOai.cu similarity index 96% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/swigluOai.cu index 4fad864f2d64..2fe60fdc1f5c 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/swigluOai.cu @@ -15,13 +15,13 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/swigluOai.h" TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp8 +namespace minimaxM3Sm107 { namespace @@ -94,7 +94,7 @@ cudaError_t invokeSwigluOai( : cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD); } -} // namespace minimaxM3Sm107Tp8 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/swigluOai.h similarity index 96% rename from cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h rename to cpp/tensorrt_llm/kernels/minimaxM3Sm107/swigluOai.h index f2898e100006..846e62a5af56 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107/swigluOai.h @@ -26,7 +26,7 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels { -namespace minimaxM3Sm107Tp4 +namespace minimaxM3Sm107 { //! The dense MLP's swigluoai activation: out[m, j] = bf16(min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + @@ -37,7 +37,7 @@ namespace minimaxM3Sm107Tp4 cudaError_t invokeSwigluOai( uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream); -} // namespace minimaxM3Sm107Tp4 +} // namespace minimaxM3Sm107 } // namespace kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h deleted file mode 100644 index 3302c94390f3..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h +++ /dev/null @@ -1,79 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp2 -{ -namespace drafterRopeFp8Kv -{ - -constexpr int kHeadDim = 128; - -struct Params -{ - //! [tokens, (numHeadsQ + 2 * numHeadsKv) * 128] BF16: the drafter layer's packed Q, K, V rows. - __nv_bfloat16 const* qkv; - //! [positions, rotaryDim / 2] (cos, sin) pairs: the layer's rotary_cos_sin table. - float2 const* cosSin; - //! [requests] cache lengths including this step's rows. - int const* seqLens; - //! [requests, qLen] spec-decoding position offsets of a multi-token pass, or null. - int const* specPositionOffsets; - //! [tokens, numHeadsQ, 128] E4M3 query out. - __nv_fp8_e4m3* qOut; - //! Base of the layer's K / V sub-pages, [*, numHeadsKv, tokensPerBlock, 128] E4M3. - __nv_fp8_e4m3* pool; - //! [requests, 2, maxBlocks] sub-page indices: K in row 0, V in row 1. - int const* blockTable; - long long blockTableRequestStride; - long long blockTableKvStride; - int numTokens; - int qLen; - int numHeadsQ; - int numHeadsKv; - int rotaryDim; - int tokensPerBlock; - //! How the rotation's two products are summed: 0 two products then the add, 1 fma(cos, x, sin * pair), - //! 2 fma(sin, pair, cos * x). - int ropeMode; -}; - -//! The MiniMax-M3 TP2 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, -//! rounded to BF16 (the in-tree preprocess's q output and cache input), then Q, K and V to saturating E4M3: K / V -//! as the in-tree preprocess stores them at unit KV scale, Q as torch's CUDA cast of the BF16 query. Rows map to -//! (request token / qLen, row token % qLen) at cache position seqLen - qLen + row. -//! -//! PDL (pdl): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction and the kernel -//! releases its dependent right after it. The dependent may read only step-invariant data before its own wait. -cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream); - -} // namespace drafterRopeFp8Kv -} // namespace minimaxM3Sm107Tp2 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu index dd0ad05437a6..d3d5ace1c50c 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.cu @@ -15,6 +15,7 @@ * limitations under the License. */ +#include "tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h" #include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/foldTail.h" TRTLLM_NAMESPACE_BEGIN @@ -37,39 +38,7 @@ __global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) { return; } - int64_t const s = p.slots[i]; - int32_t const nw = c < p.snW ? p.srcNew[i * p.sn0 + c * p.sn1] : 0; - if (c < p.dnW) - { - p.dstNew[s * p.dn0 + c * p.dn1] = nw; - } - int32_t const nxt = c < p.sxW ? p.srcNext[i * p.sx0 + c * p.sx1] : 0; - if (c < p.dxW) - { - p.dstNext[s * p.dx0 + c * p.dx1] = nxt; - } - int32_t const drf = c < p.sdW ? p.srcDraft[i * p.sd0 + c * p.sd1] : 0; - if (c < p.ddW) - { - p.dstDraft[s * p.dd0 + c * p.dd1] = drf; - } - int32_t const ln = p.srcLens[i]; - int32_t const t = p.tokensPerStep; - int32_t const d = t - 1; - if (c == 0) - { - p.dstLens[s] = ln; - p.kvOff[i] = ln - t; - } - if (c < t) - { - p.ids[i * t + c] = c < p.dxW ? nxt : 0; - p.posOff[i * t + c] = ln; - } - if (c < d) - { - p.draftIds[i * d + c] = c < p.ddW ? drf : 0; - } + minimaxM3Sm107::foldTailRow(p, i, c); } } // namespace diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu index 270b6ca53865..924533b55b83 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.cu @@ -25,13 +25,9 @@ // group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to // BF16 once. +#include "tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.cuh" #include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/fullCtxReduce.h" -#include -#include -#include -#include - TRTLLM_NAMESPACE_BEGIN namespace kernels @@ -40,184 +36,10 @@ namespace minimaxM3Sm107Tp2 { namespace fullCtxReduce { -namespace -{ - -constexpr int kGroups = 8; -constexpr int kThreadsPerRow = kHeadDim / 8; -constexpr int kRowsPerCta = 2; -constexpr int kSplitsPerGroup = kMaxSplits / kGroups; -constexpr int kThreads = kGroups * kThreadsPerRow * kRowsPerCta; - -__device__ int activeSplits(Params const& p, int b, int row) -{ - int const seqKv = p.seqLensKv[b]; - int visible = seqKv; - if (p.causal) - { - int const token = row / p.headsPerKv; - int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; - int const end = base + min(p.qTokensPerCta, p.seqLenQ - base); - visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); - } - int64_t const tiles = (static_cast(visible) + p.tileKv - 1) / p.tileKv; - int64_t const tilesPerGroup = static_cast(p.splits) * p.instsKv; - int64_t const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; - int64_t const local = (groups > 1 ? groups : 1) * p.instsKv; - return static_cast((tiles + local - 1) / local); -} - -template -__device__ float2 unpack(uint32_t word) -{ - if constexpr (kBf16Partial) - { - return __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&word)); - } - else - { - return __half22float2(*reinterpret_cast<__half2 const*>(&word)); - } -} - -template -__global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("griddepcontrol.wait;" ::: "memory"); -#endif - int const b = blockIdx.z; - int const h = blockIdx.y; - int const localRow = threadIdx.x / (kGroups * kThreadsPerRow); - int const row = blockIdx.x * kRowsPerCta + localRow; - int const g = (threadIdx.x / kThreadsPerRow) % kGroups; - int const t = threadIdx.x % kThreadsPerRow; - bool const validRow = row < p.rows; - int const active = validRow ? activeSplits(p, b, row) : 0; - - __shared__ float sMax[kRowsPerCta][kGroups]; - __shared__ float sSum[kRowsPerCta][kGroups]; - __shared__ float sAcc[kRowsPerCta][kGroups][kHeadDim]; - - uint4 const* partialO = static_cast(p.partialO); - size_t const kvIdx = static_cast(b) * p.numKvHeads + h; - float lse[kSplitsPerGroup]; - uint4 vec[kSplitsPerGroup]; -#pragma unroll - for (int i = 0; i < kSplitsPerGroup; ++i) - { - int const s = g + i * kGroups; - if (s < active) - { - size_t const r = (kvIdx * p.splits + s) * p.rows + row; - lse[i] = __ldcg(p.partialStats + r); - vec[i] = __ldcg(partialO + r * kThreadsPerRow + t); - } - } - float groupMax = -CUDART_INF_F; -#pragma unroll - for (int i = 0; i < kSplitsPerGroup; ++i) - { - if (g + i * kGroups < active) - { - groupMax = fmaxf(groupMax, lse[i]); - } - } - // A split of the CTA's causal union can hold no key of this row (log2-LSE -inf). Such a split, and a group of - // only such splits, is skipped rather than weighted 0: its partial O carries no defined value. - float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; - float sum = 0.f; -#pragma unroll - for (int i = 0; i < kSplitsPerGroup; ++i) - { - if (g + i * kGroups < active && lse[i] != -CUDART_INF_F) - { - float const w = exp2f(lse[i] - groupMax); - sum += w; - uint32_t const* words = reinterpret_cast(&vec[i]); -#pragma unroll - for (int j = 0; j < 4; ++j) - { - float2 const f = unpack(words[j]); - acc[2 * j] += w * f.x; - acc[2 * j + 1] += w * f.y; - } - } - } - if (t == 0) - { - sMax[localRow][g] = groupMax; - sSum[localRow][g] = sum; - } -#pragma unroll - for (int j = 0; j < 8; ++j) - { - sAcc[localRow][g][t * 8 + j] = acc[j]; - } - __syncthreads(); - if (g != 0 || !validRow) - { - return; - } - float rowMax = -CUDART_INF_F; -#pragma unroll - for (int k = 0; k < kGroups; ++k) - { - rowMax = fmaxf(rowMax, sMax[localRow][k]); - } - float total = 0.f; - float o[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; -#pragma unroll - for (int k = 0; k < kGroups; ++k) - { - if (sMax[localRow][k] == -CUDART_INF_F) - { - continue; - } - float const scale = exp2f(sMax[localRow][k] - rowMax); - total += scale * sSum[localRow][k]; -#pragma unroll - for (int j = 0; j < 8; ++j) - { - o[j] += scale * sAcc[localRow][k][t * 8 + j]; - } - } - float const inv = total > 0.f ? 1.f / total : 0.f; - int const token = row / p.headsPerKv; - int const head = row % p.headsPerKv; - __nv_bfloat162 packed[4]; -#pragma unroll - for (int j = 0; j < 4; ++j) - { - packed[j] = __floats2bfloat162_rn(o[2 * j] * inv, o[2 * j + 1] * inv); - } - size_t const outRow - = (static_cast(b) * p.seqLenQ + token) * (p.numKvHeads * p.headsPerKv) + h * p.headsPerKv + head; - reinterpret_cast(p.out + outRow * kHeadDim)[t] = *reinterpret_cast(packed); -} - -} // namespace - cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) { - if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1 || params.qTokensPerCta <= 0 - || params.tileKv <= 0 || params.instsKv <= 0 || params.instsKv > std::numeric_limits::max() / kMaxSplits) - { - return cudaErrorInvalidValue; - } - cudaLaunchConfig_t config = {}; - config.gridDim = dim3((params.rows + kRowsPerCta - 1) / kRowsPerCta, params.numKvHeads, params.batch); - config.blockDim = dim3(kThreads); - config.stream = stream; - cudaLaunchAttribute attribute[1]; - attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - config.attrs = attribute; - config.numAttrs = 1; - return params.bf16Partial ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) - : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); + return minimaxM3Sm107::fullCtxReduce::detail::invoke<2, false>(params, pdl, stream); } - } // namespace fullCtxReduce } // namespace minimaxM3Sm107Tp2 } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h deleted file mode 100644 index e8ade9c62c3e..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h +++ /dev/null @@ -1,83 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -//! Operands of the MiniMax-M3 multicast all-reduce + residual + RMSNorm kernel (decode, at most 48 tokens). -//! -//! Per rank the kernel forms the partial `bf16(x + x2)` (or, in finalize mode, -//! `bf16(bf16(sum_k scale * fc2[perm]) + x2)`, the routed sum rounded before the shared expert is added, as stock), -//! pushes it through the multicast mapping `wsMc` into every rank's exchange, polls its own unicast view `wsUc`, -//! sums the ranks in FP32 in rank order, adds the residual (BF16 round) and applies RMSNorm. With `outQuant` it -//! also writes the MXFP8 copy of the normed rows (E4M3 [numTokens, hidden] and swizzled UE8M0 scales). -struct MinimaxM3Sm107Tp2McAllreduceNormParams -{ - __nv_bfloat16 const* x{nullptr}; //!< [numTokens, hidden] partial; in finalize mode the FC2 rows [rows, hidden]. - __nv_bfloat16 const* x2{nullptr}; //!< Optional second addend [numTokens, hidden]. - __nv_bfloat16 const* fc2{nullptr}; //!< Finalize mode: the permuted FC2 output (same buffer as x); null otherwise. - int32_t const* perm{nullptr}; //!< Finalize mode: [numTokens, topk] expanded index -> permuted row, -1 dropped. - void const* scales{nullptr}; //!< Finalize mode: [numTokens, topk] routing weights (fp32 or bf16). - bool scaleFp32{false}; - int32_t topk{0}; - int32_t numTokens{0}; - int32_t hidden{0}; - __nv_bfloat16 const* residual{nullptr}; //!< [numTokens, hidden] - __nv_bfloat16 const* normWeight{nullptr}; //!< [hidden] - float eps{0.F}; - int32_t* wsUc{nullptr}; //!< This rank's unicast view of the exchange (int32 words, re-armed after reading). - int32_t* wsMc{nullptr}; //!< The multicast mapping of the exchange. - int32_t* flags{nullptr}; //!< Per-CTA call-parity counters, [maxTokens * 8]. - int32_t rank{0}; - int32_t world{0}; - __nv_bfloat16* outNormed{nullptr}; - __nv_bfloat16* outResidual{nullptr}; - uint8_t* outQuant{nullptr}; //!< Optional E4M3 [numTokens, hidden]. - uint8_t* outSf{nullptr}; //!< Optional UE8M0 scales, 128x4-swizzled. - //! 0: push x (+ x2) here; 2: a producer already pushed x and launched this grid after its pushes; 3: as 2, with x - //! in exchange set A and x2 in set B. - int32_t prepushed{0}; - //! Optional Eagle3 capture of a capture layer's boundary: capOut[t * capStride + c] = the m3_spec_tail - //! capture_norm of the updated residual row with weight capWeight [hidden] and capEps, bit for bit (hidden 6144, - //! not the drafter boundary). - __nv_bfloat16* capOut{nullptr}; - int64_t capStride{0}; - __nv_bfloat16 const* capWeight{nullptr}; - float capEps{0.F}; -}; - -//! The (bucket, epilogue options) pairs this launcher is built for: (8, 7), (8, 15), (16, 15) and (48, 15). -bool isMinimaxM3Sm107Tp2McAllreduceNormSupported(int32_t maxTokens, int32_t opts); - -//! Launch with 8-CTA clusters of 128 threads, one cluster per token. `maxTokens` selects the exchange bucket; -//! `opts` the epilogue variant (every variant is bit-identical). -void invokeMinimaxM3Sm107Tp2McAllreduceNorm(MinimaxM3Sm107Tp2McAllreduceNormParams const& params, int32_t maxTokens, - int32_t opts, bool pdl, cudaStream_t stream); - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu deleted file mode 100644 index cb7d83b89bb4..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.cu +++ /dev/null @@ -1,89 +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/minimaxM3Sm107Tp2/siluAndMul.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp2 -{ - -namespace -{ - -constexpr int kThreads = 128; - -// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves). The -// instruction sequence of the Triton kernel's PTX at its defaults (host/silu/silu_and_mul_plain.ptx): g * -log2(e) as -// one product, ex2.approx, + 1, div.full 1 / that, g * s, (u + 0.0) as an IEEE add, the product, one rounding to bf16. -__device__ __forceinline__ uint32_t siluMulBf16x2(uint32_t gate, uint32_t up) -{ - uint32_t r; - asm("{\n\t.reg .b32 gl, gh, ul, uh, t0, t1, e0, e1, s0, s1, p0, p1, q0, q1, r0, r1, one, zero;\n\t" - ".reg .b16 hl, hh;\n\t" - "mov.b32 {hl, hh}, %1;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" - "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" - "mul.rn.f32 t0, gl, 0fBFB8AA3B;\n\tmul.rn.f32 t1, gh, 0fBFB8AA3B;\n\t" - "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" - "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" - "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" - "mul.rn.f32 p0, gl, s0;\n\tmul.rn.f32 p1, gh, s1;\n\t" - "mov.b32 zero, 0f00000000;\n\tadd.rn.f32 q0, ul, zero;\n\tadd.rn.f32 q1, uh, zero;\n\t" - "mul.rn.f32 r0, q0, p0;\n\tmul.rn.f32 r1, q1, p1;\n\t" - "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" - : "=r"(r) - : "r"(gate), "r"(up)); - return r; -} - -// Eight outputs per thread: 16-byte loads of the gate and the up half, one 16-byte store. -__global__ void __launch_bounds__(kThreads) - siluAndMulKernel(uint4 const* __restrict__ x, uint4* __restrict__ out, int64_t rows, int64_t vecs) -{ - int64_t const row = blockIdx.y; - int64_t const v = static_cast(blockIdx.x) * kThreads + threadIdx.x; - if (row >= rows || v >= vecs) - { - return; - } - uint4 const g = x[row * 2 * vecs + v]; - uint4 const u = x[row * 2 * vecs + vecs + v]; - uint4 o; - o.x = siluMulBf16x2(g.x, u.x); - o.y = siluMulBf16x2(g.y, u.y); - o.z = siluMulBf16x2(g.z, u.z); - o.w = siluMulBf16x2(g.w, u.w); - out[row * vecs + v] = o; -} - -} // namespace - -cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream) -{ - int64_t const vecs = d / 8; - dim3 const grid(static_cast((vecs + kThreads - 1) / kThreads), static_cast(rows)); - siluAndMulKernel<<>>( - reinterpret_cast(x), reinterpret_cast(out), rows, vecs); - return cudaGetLastError(); -} - -} // namespace minimaxM3Sm107Tp2 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h deleted file mode 100644 index 01b32137ac4a..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/siluAndMul.h +++ /dev/null @@ -1,40 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp2 -{ - -//! out[m, j] = bf16(silu(g) * (u + 0)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both -//! contiguous, d a multiple of 8, 16-byte aligned rows): bit for bit the Triton silu_and_mul kernel at its defaults -//! (no limit, alpha 1, beta 0), the plain SwiGLU of the Eagle3 drafter MLP. -cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream); - -} // namespace minimaxM3Sm107Tp2 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu deleted file mode 100644 index 22c073f74b3c..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.cu +++ /dev/null @@ -1,58 +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/minimaxM3Sm107Tp2/specOffsets.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp2 -{ - -namespace -{ - -constexpr int kThreads = 256; - -__global__ void __launch_bounds__(kThreads) - specOffsetsKernel(int32_t* __restrict__ pos, int32_t const* __restrict__ posOff, int64_t nPos, - int32_t* __restrict__ kv, int32_t const* __restrict__ kvOff, int64_t nKv) -{ - for (int64_t i = threadIdx.x; i < nPos; i += kThreads) - { - pos[i] = static_cast(static_cast(pos[i]) + static_cast(posOff[i])); - } - for (int64_t j = threadIdx.x; j < nKv; j += kThreads) - { - kv[j] = static_cast(static_cast(kv[j]) + static_cast(kvOff[j])); - } -} - -} // namespace - -cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, - int64_t nKv, cudaStream_t stream) -{ - specOffsetsKernel<<<1, kThreads, 0, stream>>>(pos, posOff, nPos, kv, kvOff, nKv); - return cudaGetLastError(); -} - -} // namespace minimaxM3Sm107Tp2 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h deleted file mode 100644 index 1de212a37179..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/specOffsets.h +++ /dev/null @@ -1,41 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp2 -{ - -//! The overlap scheduler's speculative step correction in one launch: pos[i] += posOff[i] for i < nPos and -//! kv[j] += kvOff[j] for j < nKv (int32, wrapping as the ATen adds do). A plain launch: it is a step-input writer that -//! pre-wait readers rely on (rubin-20261015/PREWAIT_AUDIT.md item b). -cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, - int64_t nKv, cudaStream_t stream); - -} // namespace minimaxM3Sm107Tp2 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu deleted file mode 100644 index 9b6dbb42faed..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/stepPrologue.cu +++ /dev/null @@ -1,111 +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/minimaxM3Sm107Tp2/stepPrologue.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp2 -{ - -namespace -{ - -constexpr int kThreads = 512; -constexpr int kWarps = kThreads / 32; - -// One warp per copy. Plain (coherent) loads: a device copy may read bytes another warp of this CTA just stored, which -// the CTA barrier between the two kinds makes visible; a non-coherent load could return the stale line. -__device__ __forceinline__ void copyWarp(uint8_t* dst, uint8_t const* src, int64_t bytes, int lane) -{ - auto const d = reinterpret_cast(dst); - auto const s = reinterpret_cast(src); - if (((d | s | static_cast(bytes)) & 15) == 0) - { - auto* d16 = reinterpret_cast(dst); - auto const* s16 = reinterpret_cast(src); - for (int64_t i = lane; i < bytes / 16; i += 32) - { - d16[i] = s16[i]; - } - } - else if (((d | s | static_cast(bytes)) & 3) == 0) - { - auto* d4 = reinterpret_cast(dst); - auto const* s4 = reinterpret_cast(src); - for (int64_t i = lane; i < bytes / 4; i += 32) - { - d4[i] = s4[i]; - } - } - else - { - for (int64_t i = lane; i < bytes; i += 32) - { - dst[i] = src[i]; - } - } -} - -// One CTA: a step's window is a few dozen copies of a few hundred bytes, so the launch is the cost; a CTA barrier -// orders the host copies before the device copies where the parallel scatter needed a second launch. -__global__ void __launch_bounds__(kThreads) stepPrologueKernel(uint8_t* window, int64_t* epoch, int64_t maxDesc) -{ - auto* header = reinterpret_cast(window); - int64_t const numHost = min(header[0], maxDesc); - int64_t const numDevice = min(header[1], maxDesc); - int64_t const* desc = header + 2; - int const warp = threadIdx.x / 32; - int const lane = threadIdx.x % 32; - if (epoch != nullptr && threadIdx.x == 0) - { - *epoch += 1; - } - for (int64_t i = warp; i < numHost; i += kWarps) - { - int64_t const* c = desc + 3 * i; - copyWarp(reinterpret_cast(c[0]), window + c[1], c[2], lane); - } - __syncthreads(); - for (int64_t i = warp; i < numDevice; i += kWarps) - { - // The device copies' slots follow all header[0] host copies, as the parallel scatter indexes them. - int64_t const* c = desc + 3 * (header[0] + i); - copyWarp(reinterpret_cast(c[0]), reinterpret_cast(c[1]), c[2], lane); - } - // Every thread has read the header and the descriptors: consume the window. - __syncthreads(); - if (threadIdx.x < 2) - { - header[threadIdx.x] = 0; - } -} - -} // namespace - -cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream) -{ - stepPrologueKernel<<<1, kThreads, 0, stream>>>(window, epoch, maxDesc); - return cudaGetLastError(); -} - -} // namespace minimaxM3Sm107Tp2 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu deleted file mode 100644 index 5da59fcec01a..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.cu +++ /dev/null @@ -1,100 +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/minimaxM3Sm107Tp2/swigluOai.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp2 -{ - -namespace -{ - -constexpr int kThreads = 256; - -// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves): the -// instruction sequence of the Triton kernel at limit 7, alpha 1.702, beta 1 (Indexer kdev17, bitwise on SM107): the -// gate clamp as min.bf16x2 before the widening, the up clamp as min.xorsign.abs, (u + 1), (g * -1.702) * log2(e), -// ex2.approx, + 1, div.full 1 / that, s * g, the product, one rounding to bf16. -__device__ __forceinline__ uint32_t swigluOaiBf16x2(uint32_t gate, uint32_t up) -{ - uint32_t r; - asm("{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" - "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, %1, glim;\n\t" - "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" - "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" - "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" - "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" - "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" - "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" - "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" - "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" - "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" - "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" - "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" - "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" - : "=r"(r) - : "r"(gate), "r"(up)); - return r; -} - -template -__global__ void __launch_bounds__(kThreads) - swigluOaiKernel(uint32_t const* __restrict__ x, uint32_t* __restrict__ out, int64_t rows, int64_t halfD) -{ - if constexpr (kPdl) - { - asm volatile("griddepcontrol.wait;" ::: "memory"); - asm volatile("griddepcontrol.launch_dependents;"); - } - int64_t const i = blockIdx.x * static_cast(kThreads) + threadIdx.x; - if (i >= rows * halfD) - { - return; - } - int64_t const row = i / halfD; - int64_t const col = i - row * halfD; - uint32_t const* xr = x + row * 2 * halfD; - out[i] = swigluOaiBf16x2(xr[col], xr[halfD + col]); -} - -} // namespace - -cudaError_t invokeSwigluOai( - uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream) -{ - int64_t const n = rows * halfD; - cudaLaunchConfig_t cfg{}; - cfg.gridDim = dim3(static_cast((n + kThreads - 1) / kThreads)); - cfg.blockDim = dim3(kThreads); - cfg.stream = stream; - cudaLaunchAttribute attr[1]; - attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - cfg.attrs = attr; - cfg.numAttrs = 1; - return pdl ? cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD) - : cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD); -} - -} // namespace minimaxM3Sm107Tp2 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h deleted file mode 100644 index c8db72a5efd3..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp2/swigluOai.h +++ /dev/null @@ -1,43 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp2 -{ - -//! The dense MLP's swigluoai activation: out[m, j] = bf16(min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + -//! 1)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both contiguous, d even): bit for bit -//! the Triton silu_and_mul kernel at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With `pdl` the launch allows -//! programmatic serialization and the grid waits on its dependency, then releases its dependents before it computes -//! (the down GEMV stages its weights during the activation and reads it only after its own grid wait). -cudaError_t invokeSwigluOai( - uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream); - -} // namespace minimaxM3Sm107Tp2 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu index 2e67b3ae7ac1..7f2a04422c1c 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.cu @@ -585,7 +585,7 @@ __global__ void __launch_bounds__(kWarps * 32, 1) __syncthreads(); if (threadIdx.x < 32) { - foldTailWarp(p.fold, static_cast(threadIdx.x)); + minimaxM3Sm107::foldTailWarp(p.fold, static_cast(threadIdx.x)); } } #endif diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h index 85296a3fd6d7..5261a7c4e1f2 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/draftStep.h @@ -18,7 +18,7 @@ #pragma once #include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h" #include #include @@ -90,7 +90,7 @@ struct Params CodedArgs cd; // The step's forward-tail store + next-step staging (foldTail.h), run by the last CTA after this pass's // tokens and bookkeeping: set on the drafter's last pass only (dispatch.DRAFT_STEP_FOLD_TAIL). - FoldTailParams fold{}; + minimaxM3Sm107::FoldTailParams fold{}; bool hasFold = false; // D-FIX: sum the pieces of a tile that several warps stream in warp order (deterministic) instead of by // shared-memory atomics in arrival order (dispatch.DRAFT_STEP_FIXED_ORDER). diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu deleted file mode 100644 index a506548c1aec..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.cu +++ /dev/null @@ -1,189 +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. - */ - -// The TP4 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's GPT-NeoX -// path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin negated, -// dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale E4M3 -// store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query -// saturates the same way (448, not NaN). - -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ -namespace drafterRopeFp8Kv -{ -namespace -{ - -constexpr int kElemsPerThread = kHeadDim / 32; -constexpr int kBlockSize = 256; - -template -__device__ __forceinline__ float rotate(float c, float x, float s, float pair) -{ - if constexpr (kRopeMode == 0) - { - return __fadd_rn(__fmul_rn(c, x), __fmul_rn(s, pair)); - } - else if constexpr (kRopeMode == 1) - { - return __fmaf_rn(c, x, __fmul_rn(s, pair)); - } - else - { - return __fmaf_rn(s, pair, __fmul_rn(c, x)); - } -} - -__device__ __forceinline__ void storeE4m3Pairs(__nv_fp8_e4m3* dst, __nv_bfloat16 const (&v)[kElemsPerThread]) -{ -#pragma unroll - for (int i = 0; i < kElemsPerThread; i += 2) - { - __nv_bfloat162 pair; - pair.x = v[i]; - pair.y = v[i + 1]; - __nv_fp8x2_e4m3 const packed(pair); - reinterpret_cast<__nv_fp8x2_storage_t*>(dst)[i / 2] = packed.__x; - } -} - -template -__global__ void __launch_bounds__(kBlockSize) drafterRopeFp8KvKernel(Params const p) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("griddepcontrol.wait;" ::: "memory"); - asm volatile("griddepcontrol.launch_dependents;"); -#endif - int const lane = threadIdx.x % 32; - int const warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; - int const slots = p.numHeadsQ + p.numHeadsKv; - int const token = warp / slots; - if (token >= p.numTokens) - { - return; - } - int const slot = warp % slots; - bool const isQ = slot < p.numHeadsQ; - int const head = isQ ? slot : slot - p.numHeadsQ; - int const request = token / p.qLen; - int const row = token % p.qLen; - int const cacheLen = p.seqLens[request]; - int const kvPos = cacheLen - p.qLen + row; - int const ropePos - = p.specPositionOffsets != nullptr ? p.specPositionOffsets[request * p.qLen + row] + cacheLen - p.qLen : kvPos; - int const rowWidth = (p.numHeadsQ + 2 * p.numHeadsKv) * kHeadDim; - __nv_bfloat16 const* src - = p.qkv + static_cast(token) * rowWidth + (isQ ? head * kHeadDim : (p.numHeadsQ + head) * kHeadDim); - int const half = p.rotaryDim / 2; - float2 const* coef = p.cosSin + static_cast(ropePos) * half; - - __nv_bfloat16 rotated[kElemsPerThread]; -#pragma unroll - for (int i = 0; i < kElemsPerThread; ++i) - { - int const d = lane * kElemsPerThread + i; - bool const firstHalf = d < half; - float const x = __bfloat162float(src[d]); - float const pair = __bfloat162float(src[firstHalf ? d + half : d - half]); - float2 const cs = d < p.rotaryDim ? coef[firstHalf ? d : d - half] : make_float2(1.0f, 0.0f); - float const s = firstHalf ? -cs.y : cs.y; - rotated[i] = __float2bfloat16(rotate(cs.x, x, s, pair)); - } - - if (isQ) - { - storeE4m3Pairs( - p.qOut + (static_cast(token) * p.numHeadsQ + head) * kHeadDim + lane * kElemsPerThread, rotated); - return; - } - - int const block = kvPos / p.tokensPerBlock; - long long const inPage = (static_cast(head) * p.tokensPerBlock + kvPos % p.tokensPerBlock) * kHeadDim - + lane * kElemsPerThread; - long long const pageElems = static_cast(p.numHeadsKv) * p.tokensPerBlock * kHeadDim; - int const* table = p.blockTable + request * p.blockTableRequestStride; - long long const kPage = table[block]; - long long const vPage = table[p.blockTableKvStride + block]; - storeE4m3Pairs(p.pool + kPage * pageElems + inPage, rotated); - - __nv_bfloat16 const* vSrc = p.qkv + static_cast(token) * rowWidth - + (p.numHeadsQ + p.numHeadsKv + head) * kHeadDim + lane * kElemsPerThread; - __nv_bfloat16 v[kElemsPerThread]; -#pragma unroll - for (int i = 0; i < kElemsPerThread; ++i) - { - v[i] = vSrc[i]; - } - storeE4m3Pairs(p.pool + vPage * pageElems + inPage, v); -} - -//! With PDL and maxSharedCarveout, the maximum shared-memory carveout lets the full-context decode's large-smem CTAs -//! (the dependent this kernel releases right after its wait) co-reside on the SMs this small kernel occupies. Set once -//! per instantiation. -template -cudaError_t launchWithMaxSharedCarveout( - cudaLaunchConfig_t const& config, Params const& params, bool pdl, bool maxSharedCarveout) -{ - auto const kernel = drafterRopeFp8KvKernel; - if (pdl && maxSharedCarveout) - { - static cudaError_t const carveout = cudaFuncSetAttribute( - kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); - if (carveout != cudaSuccess) - { - return carveout; - } - } - return cudaLaunchKernelEx(&config, kernel, params); -} - -} // namespace - -cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream) -{ - int const warps = params.numTokens * (params.numHeadsQ + params.numHeadsKv); - int const warpsPerBlock = kBlockSize / 32; - cudaLaunchConfig_t config = {}; - config.gridDim = dim3((warps + warpsPerBlock - 1) / warpsPerBlock); - config.blockDim = dim3(kBlockSize); - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attribute[1]; - attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - config.attrs = attribute; - config.numAttrs = 1; - switch (params.ropeMode) - { - case 0: return launchWithMaxSharedCarveout<0>(config, params, pdl, maxSharedCarveout); - case 1: return launchWithMaxSharedCarveout<1>(config, params, pdl, maxSharedCarveout); - case 2: return launchWithMaxSharedCarveout<2>(config, params, pdl, maxSharedCarveout); - default: return cudaErrorInvalidValue; - } -} - -} // namespace drafterRopeFp8Kv -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h deleted file mode 100644 index 2a39d03d9d76..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h +++ /dev/null @@ -1,79 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ -namespace drafterRopeFp8Kv -{ - -constexpr int kHeadDim = 128; - -struct Params -{ - //! [tokens, (numHeadsQ + 2 * numHeadsKv) * 128] BF16: the drafter layer's packed Q, K, V rows. - __nv_bfloat16 const* qkv; - //! [positions, rotaryDim / 2] (cos, sin) pairs: the layer's rotary_cos_sin table. - float2 const* cosSin; - //! [requests] cache lengths including this step's rows. - int const* seqLens; - //! [requests, qLen] spec-decoding position offsets of a multi-token pass, or null. - int const* specPositionOffsets; - //! [tokens, numHeadsQ, 128] E4M3 query out. - __nv_fp8_e4m3* qOut; - //! Base of the layer's K / V sub-pages, [*, numHeadsKv, tokensPerBlock, 128] E4M3. - __nv_fp8_e4m3* pool; - //! [requests, 2, maxBlocks] sub-page indices: K in row 0, V in row 1. - int const* blockTable; - long long blockTableRequestStride; - long long blockTableKvStride; - int numTokens; - int qLen; - int numHeadsQ; - int numHeadsKv; - int rotaryDim; - int tokensPerBlock; - //! How the rotation's two products are summed: 0 two products then the add, 1 fma(cos, x, sin * pair), - //! 2 fma(sin, pair, cos * x). - int ropeMode; -}; - -//! The MiniMax-M3 TP4 Eagle3 drafter layer's generation preprocess: GPT-NeoX RoPE from the cos/sin table on Q and K, -//! rounded to BF16 (the in-tree preprocess's q output and cache input), then Q, K and V to saturating E4M3: K / V -//! as the in-tree preprocess stores them at unit KV scale, Q as torch's CUDA cast of the BF16 query. Rows map to -//! (request token / qLen, row token % qLen) at cache position seqLen - qLen + row. -//! -//! PDL (pdl): the launch is programmatic; griddepcontrol.wait is the kernel's first instruction and the kernel -//! releases its dependent right after it. The dependent may read only step-invariant data before its own wait. -cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream); - -} // namespace drafterRopeFp8Kv -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h deleted file mode 100644 index 92b658cdc193..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h +++ /dev/null @@ -1,130 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ - -//! Operands of the speculative step's forward-tail store and next-step staging (kernels/m3_fold_tail.py). Every tensor -//! is int32 except `slots` (int64); strides are in elements. The store views are [slots, width], the sources -//! [rows, width]. -struct FoldTailParams -{ - int32_t const* ctrl; // ctrl[0]: the rows the step armed (0: write nothing) - int64_t const* slots; - int32_t const* srcNew; - int64_t sn0, sn1; - int32_t snW; - int32_t* dstNew; - int64_t dn0, dn1; - int32_t dnW; - int32_t const* srcNext; - int64_t sx0, sx1; - int32_t sxW; - int32_t* dstNext; - int64_t dx0, dx1; - int32_t dxW; - int32_t const* srcDraft; - int64_t sd0, sd1; - int32_t sdW; - int32_t* dstDraft; - int64_t dd0, dd1; - int32_t ddW; - int32_t const* srcLens; - int32_t* dstLens; - int32_t* ids; // [rows * T] - int32_t* draftIds; // [rows * (T - 1)] - int32_t* posOff; // [rows * T] - int32_t* kvOff; // [rows] - int32_t tokensPerStep; - int32_t maxRows; // validated capacity of the source and staging buffers -}; - -//! One CTA of 32 lanes per row, `rows` CTAs; the first ctrl[0] rows write. Widths and tokensPerStep must be <= 32. -cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream); - -//! The same writes from ONE warp (one CTA of 32 lanes looping over the rows, foldTailWarp): the in-kernel form a caller -//! such as the drafter's last-pass draft_step runs in its last CTA; launched alone it is that form's oracle. -cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream); - -#ifdef __CUDACC__ -//! Row i of the tail at lane c (0..31), the work of foldTailKernel's CTA i, for a caller that runs the tail inside its -//! own kernel (e.g. the drafter's last-pass draft_step in its last CTA, after its final tokens are written and fenced). -//! The caller runs rows i < *p.ctrl (one warp per row, or one warp looping over the rows); same code as the kernel, so -//! the writes are bit for bit the standalone tail's. -__device__ __forceinline__ void foldTailRow(FoldTailParams const& p, int i, int c) -{ - int64_t const s = p.slots[i]; - int32_t const nw = c < p.snW ? p.srcNew[i * p.sn0 + c * p.sn1] : 0; - if (c < p.dnW) - { - p.dstNew[s * p.dn0 + c * p.dn1] = nw; - } - int32_t const nxt = c < p.sxW ? p.srcNext[i * p.sx0 + c * p.sx1] : 0; - if (c < p.dxW) - { - p.dstNext[s * p.dx0 + c * p.dx1] = nxt; - } - int32_t const drf = c < p.sdW ? p.srcDraft[i * p.sd0 + c * p.sd1] : 0; - if (c < p.ddW) - { - p.dstDraft[s * p.dd0 + c * p.dd1] = drf; - } - int32_t const ln = p.srcLens[i]; - int32_t const t = p.tokensPerStep; - int32_t const d = t - 1; - if (c == 0) - { - p.dstLens[s] = ln; - p.kvOff[i] = ln - t; - } - if (c < t) - { - p.ids[i * t + c] = c < p.dxW ? nxt : 0; - p.posOff[i * t + c] = ln; - } - if (c < d) - { - p.draftIds[i * d + c] = c < p.ddW ? drf : 0; - } -} - -//! The whole tail from one warp of a running kernel (lane = threadIdx.x % 32): rows i < *p.ctrl in order. -__device__ __forceinline__ void foldTailWarp(FoldTailParams const& p, int lane) -{ - int const rows = min(*p.ctrl, p.maxRows); - for (int i = 0; i < rows; ++i) - { - foldTailRow(p, i, lane); - } -} -#endif - -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu deleted file mode 100644 index ffab1d36db50..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.cu +++ /dev/null @@ -1,270 +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. - */ - -// The split-KV merge for the MiniMax-M3 NVFP4 TP2 (Rubin) full-context decode. The producer is the PrimTS split-KV -// decode with its separate-reduction scratch: per (request, KV head, split, row) one normalized 16-bit partial O and -// one fp32 log2-LSE. Active splits per row follow the producer's split ranges (the causal union of the row's grouped -// Q CTA, instruction-group aligned): splits past them hold nothing and are not read. -// -// One CTA per (request, KV head, kRowsPerCta rows); per row kGroups split groups x kThreadsPerRow threads x 8 dims. -// A thread issues the loads of every split it owns (strided over the groups) before using any, folds them under its -// group's max, and the groups merge through SMEM with one max rescale; the row is normalized once and rounded to -// BF16 once. - -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ -namespace fullCtxReduce -{ -namespace -{ - -constexpr int kGroups = 8; -constexpr int kThreadsPerRow = kHeadDim / 8; -constexpr int kRowsPerCta = 1; -constexpr int kSplitsPerGroup = kMaxSplits / kGroups; -constexpr int kThreads = kGroups * kThreadsPerRow * kRowsPerCta; - -__device__ int activeSplits(Params const& p, int b, int row) -{ - int const seqKv = p.seqLensKv[b]; - int visible = seqKv; - if (p.causal) - { - int const token = row / p.headsPerKv; - int const base = (token / p.qTokensPerCta) * p.qTokensPerCta; - int const end = base + min(p.qTokensPerCta, p.seqLenQ - base); - visible = min(max(seqKv - p.seqLenQ + end, 0), seqKv); - } - int64_t const tiles = (static_cast(visible) + p.tileKv - 1) / p.tileKv; - int64_t const tilesPerGroup = static_cast(p.splits) * p.instsKv; - int64_t const groups = (tiles + tilesPerGroup - 1) / tilesPerGroup; - int64_t const local = (groups > 1 ? groups : 1) * p.instsKv; - return static_cast((tiles + local - 1) / local); -} - -template -__device__ float2 unpack(uint32_t word) -{ - if constexpr (kBf16Partial) - { - return __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&word)); - } - else - { - return __half22float2(*reinterpret_cast<__half2 const*>(&word)); - } -} - -template -__global__ void __launch_bounds__(kThreads) fullCtxReduceKernel(Params const p) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("griddepcontrol.wait;" ::: "memory"); - // The dependent (the o_proj activation quantization / GEMV) reads `out` only after its own griddepcontrol.wait, - // which waits for this whole grid, so it is released here and runs its weight-only prologue under the merge. - asm volatile("griddepcontrol.launch_dependents;"); -#endif - int const b = blockIdx.z; - int const h = blockIdx.y; - int const localRow = threadIdx.x / (kGroups * kThreadsPerRow); - int const row = blockIdx.x * kRowsPerCta + localRow; - int const g = (threadIdx.x / kThreadsPerRow) % kGroups; - int const t = threadIdx.x % kThreadsPerRow; - bool const validRow = row < p.rows; - int const active = validRow ? activeSplits(p, b, row) : 0; - - __shared__ float sMax[kRowsPerCta][kGroups]; - __shared__ float sSum[kRowsPerCta][kGroups]; - __shared__ float sAcc[kRowsPerCta][kGroups][kHeadDim]; - - uint4 const* partialO = static_cast(p.partialO); - size_t const kvIdx = static_cast(b) * p.numKvHeads + h; - float lse[kSplitsPerGroup]; - uint4 vec[kSplitsPerGroup]; -#pragma unroll - for (int i = 0; i < kSplitsPerGroup; ++i) - { - int const s = g + i * kGroups; - if (s < active) - { - size_t const r = (kvIdx * p.splits + s) * p.rows + row; - lse[i] = __ldcg(p.partialStats + r); - vec[i] = __ldcg(partialO + r * kThreadsPerRow + t); - } - } - float groupMax = -CUDART_INF_F; -#pragma unroll - for (int i = 0; i < kSplitsPerGroup; ++i) - { - if (g + i * kGroups < active) - { - groupMax = fmaxf(groupMax, lse[i]); - } - } - // A split of the CTA's causal union can hold no key of this row (log2-LSE -inf). Such a split, and a group of - // only such splits, is skipped rather than weighted 0: its partial O carries no defined value. - float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; - float sum = 0.f; -#pragma unroll - for (int i = 0; i < kSplitsPerGroup; ++i) - { - if (g + i * kGroups < active && lse[i] != -CUDART_INF_F) - { - float const w = exp2f(lse[i] - groupMax); - sum += w; - uint32_t const* words = reinterpret_cast(&vec[i]); -#pragma unroll - for (int j = 0; j < 4; ++j) - { - float2 const f = unpack(words[j]); - acc[2 * j] += w * f.x; - acc[2 * j + 1] += w * f.y; - } - } - } - if (t == 0) - { - sMax[localRow][g] = groupMax; - sSum[localRow][g] = sum; - } -#pragma unroll - for (int j = 0; j < 8; ++j) - { - sAcc[localRow][g][t * 8 + j] = acc[j]; - } - __syncthreads(); - if (g != 0 || !validRow) - { - return; - } - float rowMax = -CUDART_INF_F; -#pragma unroll - for (int k = 0; k < kGroups; ++k) - { - rowMax = fmaxf(rowMax, sMax[localRow][k]); - } - float total = 0.f; - float o[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; -#pragma unroll - for (int k = 0; k < kGroups; ++k) - { - if (sMax[localRow][k] == -CUDART_INF_F) - { - continue; - } - float const scale = exp2f(sMax[localRow][k] - rowMax); - total += scale * sSum[localRow][k]; -#pragma unroll - for (int j = 0; j < 8; ++j) - { - o[j] += scale * sAcc[localRow][k][t * 8 + j]; - } - } - float const inv = total > 0.f ? 1.f / total : 0.f; - int const token = row / p.headsPerKv; - int const head = row % p.headsPerKv; - __nv_bfloat162 packed[4]; -#pragma unroll - for (int j = 0; j < 4; ++j) - { - packed[j] = __floats2bfloat162_rn(o[2 * j] * inv, o[2 * j + 1] * inv); - } - size_t const outRow - = (static_cast(b) * p.seqLenQ + token) * (p.numKvHeads * p.headsPerKv) + h * p.headsPerKv + head; - reinterpret_cast(p.out + outRow * kHeadDim)[t] = *reinterpret_cast(packed); - if (p.q8 == nullptr) - { - return; - } - // MXFP8 copy of the BF16 row: a 32-column block is 4 neighbouring threads x 8 columns (lanes 0..15 of warp 0 - // hold the row). Block scale = the UE8M0 exponent of ceil(amax / 448), values = satfinite E4M3 of x * 2^-scale. - float x[8]; - float amax = 0.f; -#pragma unroll - for (int j = 0; j < 4; ++j) - { - float2 const f = __bfloat1622float2(packed[j]); - x[2 * j] = f.x; - x[2 * j + 1] = f.y; - amax = fmaxf(amax, fmaxf(fabsf(f.x), fabsf(f.y))); - } - amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 1)); - amax = fmaxf(amax, __shfl_xor_sync(0x0000ffffu, amax, 2)); - float rcp; - asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(rcp) : "f"(448.0F)); - float const sfv = amax * rcp; - int const bits = __float_as_int(sfv); - int e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0); - e = min(e, 254); - float const qs = amax != 0.F ? __int_as_float((254 - e) << 23) : 0.F; - uint16_t q2[4]; -#pragma unroll - for (int j = 0; j < 4; ++j) - { - asm("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(q2[j]) : "f"(x[2 * j + 1] * qs), "f"(x[2 * j] * qs)); - } - int const tok = b * p.seqLenQ + token; - int const col = (h * p.headsPerKv + head) * kHeadDim + t * 8; - *reinterpret_cast(p.q8 + static_cast(tok) * (p.numSfKTiles * 128) + col) - = *reinterpret_cast(q2); - if (t % 4 == 0) - { - int const kSf = col / 32; - int64_t const off = static_cast(tok / 128) * p.numSfKTiles * 512 + (kSf / 4) * 512 + (tok % 32) * 16 - + ((tok % 128) / 32) * 4 + kSf % 4; - p.sf[off] = static_cast(e); - } -} - -} // namespace - -cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream) -{ - if (params.splits < 1 || params.splits > kMaxSplits || params.rows < 1 || params.qTokensPerCta <= 0 - || params.tileKv <= 0 || params.instsKv <= 0 || params.instsKv > std::numeric_limits::max() / kMaxSplits) - { - return cudaErrorInvalidValue; - } - cudaLaunchConfig_t config = {}; - config.gridDim = dim3((params.rows + kRowsPerCta - 1) / kRowsPerCta, params.numKvHeads, params.batch); - config.blockDim = dim3(kThreads); - config.stream = stream; - cudaLaunchAttribute attribute[1]; - attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - config.attrs = attribute; - config.numAttrs = 1; - return params.bf16Partial ? cudaLaunchKernelEx(&config, fullCtxReduceKernel, params) - : cudaLaunchKernelEx(&config, fullCtxReduceKernel, params); -} - -} // namespace fullCtxReduce -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h deleted file mode 100644 index 95af79e2b2ce..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h +++ /dev/null @@ -1,77 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ -//! The split-KV merge of the MiniMax-M3 TP2 full-context decode (PrimTS split-KV decode, separate-reduction scratch) -//! straight to BF16. See fullCtxReduce.cu. -namespace fullCtxReduce -{ - -constexpr int kHeadDim = 128; -constexpr int kMaxSplits = 128; - -struct Params -{ - void const* partialO; // [B, Hkv, S, R, 128] normalized 16-bit partial O per split - float const* partialStats; // [B, Hkv, S, R] log2-LSE per split - int const* seqLensKv; // [B] - __nv_bfloat16* out; // [B, SQ, Hkv * hR, 128] - int batch; - int numKvHeads; - int splits; // S, the planned split capacity (<= kMaxSplits) - int rows; // R = hR * SQ, rows [token, local head] - int seqLenQ; // SQ - int headsPerKv; // hR - int qTokensPerCta; // query tokens one producer CTA owns (grouped Q) - int tileKv; // producer K/V tile - int instsKv; // producer K/V instructions per loop step - bool causal; - bool bf16Partial; - uint8_t* q8 = nullptr; // optional E4M3 copy of out [B * SQ, Hkv * hR * 128] (the o_proj MXFP8 hand-off) - uint8_t* sf = nullptr; // its UE8M0 block scales, 128x4-swizzled as trtllm::mxfp8_quantize(x, True) lays them out - int numSfKTiles = 0; // Hkv * hR * 128 / 128 -}; - -//! Merges the active splits of every row in fp32 with one max rescale and writes BF16; with q8 set, also the MXFP8 -//! copy of the BF16 rows (bit for bit trtllm::mxfp8_quantize of out). With pdl the launch is -//! programmatic: the grid waits on the producer (griddepcontrol.wait) before reading its scratch. -//! -//! PDL contract: griddepcontrol.wait is the kernel's first instruction, so nothing (scratch, lengths, out) is read or -//! written before every prerequisite grid has completed and flushed; a producer's early trigger -//! (griddepcontrol.launch_dependents) only lets this grid launch and wait sooner. The kernel triggers its own -//! dependents right after its wait: the next kernel's griddepcontrol.wait still waits for this whole grid, so its -//! post-wait code sees all of out, and its pre-wait prologue must not read out. -cudaError_t invokeFullCtxReduce(Params const& params, bool pdl, cudaStream_t stream); - -} // namespace fullCtxReduce -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu index 1942082c6df9..471137ccc0d6 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.cu @@ -15,6 +15,7 @@ * limitations under the License. */ +#include "tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceUtils.cuh" #include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/mcAllreduceNormTp4.h" #include "tensorrt_llm/common/assert.h" @@ -26,7 +27,7 @@ #include #include -// SM107 (Rubin) TP4 copy of minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu for an sm_107 TP4 target: the row kernel at 4 +// SM107 (Rubin) TP4 copy of minimaxM3Sm107/mcAllreduceNorm.cu for an sm_107 TP4 target: the row kernel at 4 // ranks. Two changes in the arithmetic: the rank sum rounds to BF16 before the residual add, the rounding points of the // stock all-reduce + fused residual RMSNorm on this target (bf16 all-reduce output, then bf16(residual + sum)); and // every call but the MoE finalize runs the row kernel below, whose RMS statistic is the stock fused norm's reduction @@ -53,160 +54,23 @@ namespace cg = cooperative_groups; namespace { -constexpr int kEmpty = static_cast(0x80000000u); -constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k - -// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], -// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c -// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp -// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its -// arithmetic is the same source. (The TP2 target's mcAllreduceNormTp2.cu has the same epilogue.) -constexpr int kCapThreads = 256; -constexpr int kCapIters = 3; -constexpr int kCapWarps = kCapThreads / 32; -constexpr int kCapHidden = kCapThreads * kCapIters * 8; - -// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). -__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) -{ - __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); -#pragma unroll - for (int q = 0; q < 8; ++q) - { - float const v = __bfloat162float(e[q]); - ss += v * v; - } - return ss; -} - -__device__ __forceinline__ float cap_warp_sum(float v) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v += __shfl_xor_sync(0xffffffffu, v, o); - return v; -} - -// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. -__device__ __forceinline__ void cap_store( - float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) -{ - uint4 const wr = *reinterpret_cast(w); - __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); - uint4 o; - __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); -#pragma unroll - for (int q = 0; q < 8; ++q) - oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); - *reinterpret_cast(cap) = o; -} - -__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) -{ - if ((w & 0xFFFFu) == 0x8000u) - w &= 0xFFFF0000u; - if ((w >> 16) == 0x8000u) - w &= 0x0000FFFFu; - return w; -} - -__device__ __forceinline__ float bf16_lo(uint32_t w) -{ - return __uint_as_float(w << 16); -} - -// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's -// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals -// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the -// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. -__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) -{ - float value; - asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); - uint32_t result; - asm("{\n" - " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" - " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" - " setp.le.f32 p_zero, %1, 0f00000000;\n" - " mov.b32 bits, %1;\n" - " shr.b32 exp_biased, bits, 23;\n" - " and.b32 exp_biased, exp_biased, 255;\n" - " and.b32 mantissa, bits, 0x7FFFFF;\n" - " setp.ne.u32 p_has_mant, mantissa, 0;\n" - " selp.u32 bump, 1, 0, p_has_mant;\n" - " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" - " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" - " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" - " @p_tiny_sub mov.u32 bump, 0;\n" - " add.u32 r, exp_biased, bump;\n" - " setp.gt.u32 p_ovf, r, 254;\n" - " selp.u32 r, 254, r, p_ovf;\n" - " selp.u32 %0, 0, r, p_zero;\n" - "}\n" - : "=r"(result) - : "f"(value)); - return result; -} - -__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) -{ - float inv; - asm("{\n" - " .reg .s32 new_exp;\n" - " .reg .b32 float_bits;\n" - " .reg .pred p_zero;\n" - " setp.eq.u32 p_zero, %1, 0;\n" - " sub.s32 new_exp, 254, %1;\n" - " max.s32 new_exp, new_exp, 0;\n" - " shl.b32 float_bits, new_exp, 23;\n" - " mov.b32 %0, float_bits;\n" - " @p_zero mov.b32 %0, 0;\n" - "}\n" - : "=f"(inv) - : "r"(ue8m0)); - return inv; -} - -// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. -__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) -{ - uint16_t pair; - asm("{\n" - " .reg .f32 f0, f1;\n" - " mul.f32 f0, %1, %3;\n" - " mul.f32 f1, %2, %3;\n" - " min.f32 f0, f0, 0f43E00000;\n" - " max.f32 f0, f0, 0fC3E00000;\n" - " min.f32 f1, f1, 0f43E00000;\n" - " max.f32 f1, f1, 0fC3E00000;\n" - " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" - "}\n" - : "=h"(pair) - : "f"(lo), "f"(hi), "f"(inv)); - return pair; -} - -__device__ __forceinline__ float rcp_approx_ftz(float a) -{ - float b; - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); - return b; -} - -__device__ __forceinline__ float bf16_hi(uint32_t w) -{ - return __uint_as_float(w & 0xFFFF0000u); -} - -__device__ __forceinline__ int4 ld_volatile_v4(int const* p) -{ - int4 v; - asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" - : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) - : "l"(p) - : "memory"); - return v; -} +using minimaxM3Sm107::mcDetail::kEmpty; +using minimaxM3Sm107::mcDetail::kFinMaxTopk; +using minimaxM3Sm107::mcDetail::kCapThreads; +using minimaxM3Sm107::mcDetail::kCapIters; +using minimaxM3Sm107::mcDetail::kCapWarps; +using minimaxM3Sm107::mcDetail::kCapHidden; +using minimaxM3Sm107::mcDetail::cap_chunk_sq; +using minimaxM3Sm107::mcDetail::cap_warp_sum; +using minimaxM3Sm107::mcDetail::cap_store; +using minimaxM3Sm107::mcDetail::canon_bf16x2; +using minimaxM3Sm107::mcDetail::bf16_lo; +using minimaxM3Sm107::mcDetail::mxfp8_ue8m0; +using minimaxM3Sm107::mcDetail::mxfp8_inv_scale; +using minimaxM3Sm107::mcDetail::mxfp8_e4m3x2; +using minimaxM3Sm107::mcDetail::rcp_approx_ftz; +using minimaxM3Sm107::mcDetail::bf16_hi; +using minimaxM3Sm107::mcDetail::ld_volatile_v4; struct FinalizeArgs { diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu deleted file mode 100644 index a0d5ac50025a..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.cu +++ /dev/null @@ -1,340 +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/minimaxM3Sm107Tp4/producerPw4.h" - -#include -#include -#include - -// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP4 heads. -// -// The SM107 TP4 target's copy of the SM107 TP2 producerPw2.cu (itself a copy of minimaxM3/producerPw.cu, which the -// sm_103 targets keep), for its pre-wait step inputs. Same math and -// outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with --use_fast_math like -// the rest of the library), specialized to 16 query, 1 KV and 1 index head per rank (one index-K head), head_dim 128, -// rotary_dim 64, page 128, NVFP4 main KV and FP8 index-K. One warp per (token, packed head); K and V rows go to their -// head's block of the HND pools (the native kernel's head strides). -// -// The decode producer is latency bound: after griddepcontrol.wait the native kernel runs a chain of dependent global -// loads (packed row -> norm weight -> position id -> RoPE row -> cache slot -> quant scale). This kernel reads the -// constant norm weights and KV quant scale before griddepcontrol.wait and issues the packed row, position id and cache -// slot loads together after it, followed by the RoPE row; with `prewaitStepInputs` the position id, the slot and the -// RoPE row are read before the wait too, leaving only the packed row after it. With `early` it also calls -// griddepcontrol.launch_dependents on entry: consumers read its outputs only after their own griddepcontrol.wait, which -// waits for this grid to complete, so the early trigger only lets the next kernel launch and run its pre-wait prologue -// while this one runs. - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ -namespace producerPw4 -{ -namespace -{ - -constexpr int kHeadDim = 128; -constexpr int kRotaryDim = 64; -constexpr int kPageSize = 128; -constexpr int kElems = kHeadDim / 32; -constexpr int kBlock = 256; - -__device__ __forceinline__ float rcp_approx_ftz(float a) -{ - float b; - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); - return b; -} - -__device__ __forceinline__ uint32_t sf_group_mask4() -{ - uint32_t laneId; - asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); - return 0xFu << (laneId & ~3u); -} - -__device__ __forceinline__ uint32_t fp32x8_to_e2m1(float (&a)[8]) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - uint32_t val; - asm volatile( - "{\n" - ".reg .b8 byte0;\n" - ".reg .b8 byte1;\n" - ".reg .b8 byte2;\n" - ".reg .b8 byte3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" - "}" - : "=r"(val) - : "f"(a[0]), "f"(a[1]), "f"(a[2]), "f"(a[3]), "f"(a[4]), "f"(a[5]), "f"(a[6]), "f"(a[7])); - return val; -#else - return 0; -#endif -} - -__device__ __forceinline__ void store_fp8x4(__nv_fp8_e4m3* out, int64_t off, float const (&e)[kElems]) -{ - __nv_fp8x2_e4m3 const low(make_float2(e[0], e[1])); - __nv_fp8x2_e4m3 const high(make_float2(e[2], e[3])); - *reinterpret_cast(out + off) = static_cast(low.__x) | (static_cast(high.__x) << 16); -} - -__device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float secondLevelScale, uint8_t* scaleOut) -{ - __nv_bfloat16 rounded[4]; - float localMax = 0.0F; -#pragma unroll - for (int i = 0; i < 4; ++i) - { - rounded[i] = __float2bfloat16_rn(e[i]); - localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); - } - uint32_t const groupMask = sf_group_mask4(); - localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); - localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); - float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); - __nv_fp8_e4m3 const encodedScale(scaleValue); - if (scaleOut != nullptr) - *scaleOut = encodedScale.__x; - float const narrowedScale = static_cast(encodedScale); - float const outputScale - = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; - float q[8] = {}; -#pragma unroll - for (int i = 0; i < 4; ++i) - q[i] = __bfloat162float(rounded[i]) * outputScale; - return static_cast(fp32x8_to_e2m1(q)); -} - -template -__global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - if constexpr (kEarly) - { - asm volatile("griddepcontrol.launch_dependents;"); - } -#endif - int const warpId = threadIdx.x / 32; - int const laneId = threadIdx.x % 32; - int const globalWarp = blockIdx.x * (kBlock / 32) + warpId; - int const tokenIdx = globalWarp / kTotalHeads; - int const localHead = globalWarp % kTotalHeads; - if (tokenIdx >= p.numTokens) - return; - - constexpr int kBegin = kHeadsQ; - constexpr int vBegin = kBegin + kHeadsKV; - constexpr int indexQBegin = vBegin + kHeadsKV; - constexpr int indexKHead = indexQBegin + kHeadsIndex; - bool const isQ = localHead < kBegin; - bool const isK = localHead >= kBegin && localHead < vBegin; - bool const isV = localHead >= vBegin && localHead < indexQBegin; - bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; - bool const isIndexK = localHead == indexKHead; - - // Constant data read before the wait: norm weights and the KV quant scale. - float w[kElems] = {}; - float quantScale = 0.0F; - __nv_bfloat16 const* normWeight - = isQ ? p.qWeight : (isK ? p.kWeight : (isIndexQ ? p.indexQWeight : p.indexKWeight)); - if (!isV) - { - uint2 const wv = *reinterpret_cast(normWeight + laneId * kElems); - float2 const w01 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.x)); - float2 const w23 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.y)); - w[0] = w01.x; - w[1] = w01.y; - w[2] = w23.x; - w[3] = w23.y; - } - if (isK || isV) - { - quantScale = p.kvQuantScale[isV ? 2 : 1]; - } - - int64_t const inputOffset = (static_cast(tokenIdx) * kTotalHeads + localHead) * kHeadDim + laneId * kElems; - bool const rotates = !isV && laneId * kElems < kRotaryDim; - bool const slotted = isK || isV || isIndexK; - int const coefficient = (laneId * kElems) % (kRotaryDim / 2); - // kPrewait: the position id, the cache slot and the RoPE row are read before the wait. They are per-step metadata - // written before the replay (host-to-device) or by plain-launched kernels earlier in the step (position_ids by the - // step's input preprocessing, the slots by msa_kv_update or its torch fallback), so they are final before this - // grid can start whatever its predecessors trigger; the RoPE table is constant. - int positionId = 0; - int slot = 0; - float4 c = make_float4(0.0F, 0.0F, 0.0F, 0.0F); - float4 s = c; - auto const loadStepInputs = [&]() - { - positionId = rotates ? p.positionIds[tokenIdx] : 0; - slot = slotted ? p.outCacheLoc[tokenIdx] : 0; - if (rotates) - { - int64_t const ropeRow = static_cast(positionId) * kRotaryDim; - c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); - s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); - } - }; - if constexpr (kPrewait) - { - loadStepInputs(); - } -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.wait;" ::: "memory"); -#endif - // The packed row comes from the predecessor grid; without kPrewait the step inputs are issued with it, so the - // native kernel's serial chain of dependent loads becomes two round trips. - uint2 const packed = *reinterpret_cast(p.packed + inputOffset); - if constexpr (!kPrewait) - { - loadStepInputs(); - } - float cosine[kElems] = {c.x, c.y, c.z, c.w}; - float sine[kElems] = {s.x, s.y, s.z, s.w}; - - float e[kElems]; - float sumSquares = 0.0F; -#pragma unroll - for (int pair = 0; pair < 2; ++pair) - { - float2 const v = __bfloat1622float2( - *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); - e[2 * pair] = v.x; - e[2 * pair + 1] = v.y; - if (!isV) - sumSquares += v.x * v.x + v.y * v.y; - } - - if (!isV) - { -#pragma unroll - for (int mask = 16; mask > 0; mask >>= 1) - sumSquares += __shfl_xor_sync(0xffffffff, sumSquares, mask, 32); - float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kHeadDim) + p.eps); -#pragma unroll - for (int i = 0; i < kElems; ++i) - { - e[i] *= rmsReciprocal * (1.0F + w[i]); - } - constexpr int kPairOffset = (kRotaryDim / 2) / kElems; -#pragma unroll - for (int i = 0; i < kElems; ++i) - { - int const dim = laneId * kElems + i; - float paired = __shfl_xor_sync(0xffffffff, e[i], kPairOffset); - if (dim < kRotaryDim) - { - if (dim < kRotaryDim / 2) - paired = -paired; - // The native kernel's contraction of e * cos + paired * sin; spelled out so this build matches it. - e[i] = __fmaf_rn(paired, sine[i], __fmul_rn(e[i], cosine[i])); - } - } - } - - if (isQ) - { - store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); - return; - } - if (isIndexQ) - { - store_fp8x4(p.indexQOut, - (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); - return; - } - - if (slot < 0) - return; - int const page = slot / kPageSize; - if (page >= p.numPages) - return; - int const withinPage = slot & (kPageSize - 1); - if (isIndexK) - { - store_fp8x4(p.indexK, - static_cast(page) * p.indexPageStride + static_cast(withinPage) * p.indexTokenStride - + laneId * kElems, - e); - return; - } - - int const head = isK ? localHead - kBegin : localHead - vBegin; - int const plane = isV ? 1 : 0; - int const scaleColumn = laneId / 4; - int64_t const scaleWithinPage = isV - ? static_cast(withinPage / 4) * 32 + static_cast(scaleColumn) * 4 + withinPage % 4 - : static_cast(withinPage) * 8 + scaleColumn; - int64_t const scaleOffset = static_cast(page) * p.scalePageStride - + static_cast(plane) * p.scalePlaneStride + static_cast(head) * p.scaleHeadStride - + scaleWithinPage; - uint8_t* scaleOut = laneId % 4 == 0 ? p.kvScale + scaleOffset : nullptr; - uint16_t const packedFp4 = quantize_fp4x4(e, quantScale, scaleOut); - int64_t const dataOffset = static_cast(page) * p.dataPageStride - + static_cast(plane) * p.dataPlaneStride + static_cast(head) * p.dataHeadStride - + static_cast(withinPage) * p.dataTokenStride + laneId * 2; - *reinterpret_cast(p.kvData + dataOffset) = packedFp4; -#endif -} - -template -cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) -{ - constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; - int const warps = numTokens * kTotalHeads; - cudaLaunchConfig_t cfg = {}; - cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); - cfg.blockDim = dim3(kBlock); - cfg.stream = stream; - cudaLaunchAttribute attr[1]; - attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - cfg.attrs = attr; - cfg.numAttrs = 1; - return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); -} - -} // namespace - -cudaError_t invokeProducerPw4( - Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream) -{ - if (numTokens <= 0) - return cudaSuccess; - if (prewaitStepInputs) - { - return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); - } - return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); -} - -} // namespace producerPw4 -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h deleted file mode 100644 index 6f8aa54a812b..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h +++ /dev/null @@ -1,76 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ -//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP4 heads (16 query, 1 KV and -//! 1 index head per rank) with its constant loads before the PDL grid-dependency wait and an optional early dependents -//! trigger. Same outputs, bit for bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw4.cu. -namespace producerPw4 -{ - -struct Params -{ - __nv_bfloat16 const* packed; - __nv_fp8_e4m3* qOut; - __nv_fp8_e4m3* indexQOut; - uint8_t* kvData; - uint8_t* kvScale; - __nv_fp8_e4m3* indexK; - int const* outCacheLoc; - float const* kvQuantScale; - int64_t dataPageStride, dataPlaneStride, dataHeadStride, dataTokenStride; - int64_t scalePageStride, scalePlaneStride, scaleHeadStride; - int64_t indexPageStride, indexTokenStride, numPages; - int numTokens; - float eps; - __nv_bfloat16 const* qWeight; - __nv_bfloat16 const* kWeight; - __nv_bfloat16 const* indexQWeight; - __nv_bfloat16 const* indexKWeight; - float const* rotaryCosSin; - int const* positionIds; -}; - -//! The heads per rank this producer is built for. -constexpr int kHeadsQ = 16; -constexpr int kHeadsKV = 1; -constexpr int kHeadsIndex = 1; - -//! Launch over `numTokens` tokens, one warp per (token, packed head). `prewaitStepInputs` reads the position ids, the -//! cache slots and the RoPE rows before the grid-dependency wait (they must not be written by a PDL predecessor). -cudaError_t invokeProducerPw4( - Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream); - -} // namespace producerPw4 -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu index ce95ebb5b71e..08f7149c87b7 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.cu @@ -17,51 +17,6 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h" -#include "tensorrt_llm/common/assert.h" - -#include -#include - -#include -#include - -// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). -// -// sm_107 (Rubin) TP4 copy of kernels/minimaxM3Sm107Tp2/sampledAccept.cu: identical algorithm and arithmetic as the -// shared kernels/minimaxM3/sampledAccept.cu, with the TP2 copy's launch (a row's CTAs as one thread-block cluster) and -// bounded record polls. The vocab shard limit stays the shared kernel's (kMaxVs 51200 >= the TP4 shard 50016). -// -// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the -// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p -// active must stay on the generic path; the caller decides. -// -// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL -// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z -// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it -// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens -// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the -// nucleus with probability >= top_p, so most rows decide on their second exchange. -// -// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, -// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local -// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the -// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based -// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one -// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and -// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the -// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing -// and two call halves keep a rank one call ahead from overwriting a slot still being read. -// -// SA-1R (Args.one_round_k 9 / 82; MoE prototype moe/sa1r/6c753632ec, exact on 39,968 rows, eager + graph): each CTA -// also multicasts its largest bf16 logits with the round-0 record (3 per extra word). After round 0 every CTA knows M, -// the candidate logit cv and wc = __expf(cv - M); round 1 would accept iff S <= target, S = the fp32 sum (its tree) of -// w = __expf(__fmul_rn(x, tInv) - M) over every token with w > wc. x -> w is monotone (__fmul_rn by tInv > 0, fsub, -// __expf: checked exhaustively over all non-positive fp32 inputs), so if no CTA's unlisted value can beat wc (82: the -// bound tau; 9: the 9th listed), the listed w > wc are exactly round 1's summands, bitwise. Two fp32 sums of the same -// m non-negative terms differ by < m * 2^-23 * S, so accepting iff S_1r + m * 2^-21 * S_1r (round up) <= target -// implies round 1 accepts; m <= 1 is exact. Otherwise round 1 runs as before. An accepted row outputs the round-0 -// candidate, as the two-round kernel does, and consumes no Philox draw it would not. - TRTLLM_NAMESPACE_BEGIN namespace kernels @@ -70,888 +25,15 @@ namespace minimaxM3Sm107Tp4 { namespace sampledAccept { -namespace -{ -constexpr int kThreads = 256; -constexpr int kWarps = kThreads / 32; -constexpr int kMaxRec = kMaxWorld * kMaxSplits; -constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) -constexpr int kTagBits = 64 - 32 - kIdBits; -constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard -// Record polls trap after this many spins (the multicast all-reduce's bound): a missing writer becomes a CUDA error, -// never a hang. -constexpr long long kMaxSpins = 1LL << 31; - -// SA-1R variant kTopK: 0 none; K exact top-K list; 10 * K + P pooled (each thread keeps its top P plus its (P + 1)-th -// as a bound; the CTA lists the pool's top K plus tau = max(K-th listed, max thread bound)). -__host__ __device__ constexpr int mode_k(int m) -{ - return m >= 10 ? m / 10 : m; -} - -__host__ __device__ constexpr int mode_pool(int m) -{ - return m >= 10 ? m % 10 : 0; -} - -__host__ __device__ constexpr int mode_vals(int m) -{ - return mode_k(m) + (mode_pool(m) ? 1 : 0); -} - -// Words per record: key, logit of key, f2, f3, then the round-0 list (3 bf16 + a 16-bit tag per word). -template -struct Rec -{ - static constexpr int kXw = (mode_vals(kTopK) + 2) / 3; - static constexpr int kEx = 4 + kXw; -}; - -// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. -template -__host__ __device__ inline long long half_words(int world) -{ - return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * Rec::kEx; -} - -__device__ __forceinline__ uint32_t ordered(float v) -{ - uint32_t const b = __float_as_uint(v); - return (b & 0x80000000u) ? ~b : (b | 0x80000000u); -} - -// bf16 bits <-> an order-preserving 16-bit key. -__device__ __forceinline__ uint32_t ord16(uint32_t b) -{ - return (b & 0x8000u) ? (~b & 0xFFFFu) : (b | 0x8000u); -} - -__device__ __forceinline__ uint32_t unord16(uint32_t k) -{ - return (k & 0x8000u) ? (k & 0x7FFFu) : (~k & 0xFFFFu); -} - -// Philox4x32-10. -__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) -{ -#pragma unroll - for (int i = 0; i < 10; ++i) - { - uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; - uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; - c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); - k.x += 0x9E3779B9u; - k.y += 0xBB67AE85u; - } - return c; -} - -// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. -__device__ __forceinline__ float uniform(uint32_t bits) -{ - return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); -} - -// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. -__device__ __forceinline__ unsigned long long gkey(float lg, int id) -{ - return (static_cast(ordered(lg)) << kIdBits) - | static_cast((1u << kIdBits) - 1u - static_cast(id)); -} - -__device__ __forceinline__ int key_token(unsigned long long k) -{ - return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); -} - -__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) -{ - asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); -} - -__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) -{ - unsigned long long v; - asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); - return v; -} - -template -struct Shared -{ - __nv_bfloat16 row[kMaxVs / kSplits]; - float redf[kWarps]; - unsigned long long redk[kWarps]; - unsigned long long rk[kMaxRec]; - float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; - float wsum[kWarps]; - int owner; - int pick; - float pickx; - bool last; - float dS[kWarps]; // SA-1R: per-warp partial mass above the candidate - int dM[kWarps]; // SA-1R: per-warp count | incomplete << 16 - uint32_t wtop[kWarps][kTopK > 0 ? mode_k(kTopK) : 1]; // SA-1R: per-warp top-K keys - uint32_t wbound[kWarps]; // SA-1R pooled: per-warp max thread bound - uint32_t ctop[kTopK > 0 ? mode_vals(kTopK) : 1]; // SA-1R: the CTA's top-K keys (descending), by rank - unsigned long long rx[kMaxRec][Rec::kXw > 0 ? Rec::kXw : 1]; -}; - -template -__device__ __forceinline__ float block_max(float v, S& s) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); - if ((threadIdx.x & 31) == 0) - s.redf[threadIdx.x >> 5] = v; - __syncthreads(); - float r = s.redf[0]; - for (int i = 1; i < kWarps; ++i) - r = fmaxf(r, s.redf[i]); - __syncthreads(); - return r; -} - -// Fixed-order block sum: identical result for identical inputs. -template -__device__ __forceinline__ float block_sum(float v, S& s) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v += __shfl_xor_sync(0xffffffffu, v, o); - if ((threadIdx.x & 31) == 0) - s.redf[threadIdx.x >> 5] = v; - __syncthreads(); - float r = 0.f; - for (int i = 0; i < kWarps; ++i) - r += s.redf[i]; - __syncthreads(); - return r; -} - -template -__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - { - unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); - v = w > v ? w : v; - } - if ((threadIdx.x & 31) == 0) - s.redk[threadIdx.x >> 5] = v; - __syncthreads(); - unsigned long long r = s.redk[0]; - for (int i = 1; i < kWarps; ++i) - r = s.redk[i] > r ? s.redk[i] : r; - __syncthreads(); - return r; -} - -// Exclusive prefix of v over the threads in thread order, and the total. -template -__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) -{ - int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; - float inc = v; -#pragma unroll - for (int o = 1; o < 32; o <<= 1) - { - float const t = __shfl_up_sync(0xffffffffu, inc, o); - if (lane >= o) - inc += t; - } - if (lane == 31) - s.wsum[warp] = inc; - __syncthreads(); - float base = 0.f; - total = 0.f; - for (int i = 0; i < kWarps; ++i) - { - base += i < warp ? s.wsum[i] : 0.f; - total += s.wsum[i]; - } - __syncthreads(); - return base + inc - v; -} - -// Warp-wide integer max / sum / or (the redux instructions from sm_80; a butterfly otherwise, same results). -__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - return __reduce_max_sync(0xffffffffu, v); -#else -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - { - uint32_t const w = __shfl_xor_sync(0xffffffffu, v, o); - v = w > v ? w : v; - } - return v; -#endif -} - -__device__ __forceinline__ uint32_t warp_add_u32(uint32_t v) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - return __reduce_add_sync(0xffffffffu, v); -#else -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v += __shfl_xor_sync(0xffffffffu, v, o); - return v; -#endif -} - -__device__ __forceinline__ uint32_t warp_or_u32(uint32_t v) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - return __reduce_or_sync(0xffffffffu, v); -#else -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v |= __shfl_xor_sync(0xffffffffu, v, o); - return v; -#endif -} - -// Insert key into a descending register list of K (fully unrolled: no local memory). -template -__device__ __forceinline__ void topk_insert(uint32_t (&a)[K], uint32_t key) -{ -#pragma unroll - for (int j = 0; j < K; ++j) - { - uint32_t const hi = key > a[j] ? key : a[j]; - key = key > a[j] ? a[j] : key; - a[j] = hi; - } -} - -template -__device__ __forceinline__ void topk_pop(uint32_t (&a)[K]) -{ -#pragma unroll - for (int j = 0; j + 1 < K; ++j) - a[j] = a[j + 1]; - a[K - 1] = 0u; -} - -// Per-thread insertion with an early reject: after the first few elements almost every key fails the K-th test. -template -__device__ __forceinline__ void thread_insert(uint32_t (&a)[K], uint32_t key) -{ - if (key > a[K - 1]) - topk_insert(a, key); -} - -// The warp's K largest keys (descending) into s.wtop[warp]; published by the caller's next block barrier. -template -__device__ __forceinline__ void warp_topk(uint32_t (&a)[K], S& s) -{ - int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; -#pragma unroll - for (int j = 0; j < K; ++j) - { - uint32_t const m = warp_max_u32(a[0]); - if (a[0] == m && m != 0u) - topk_pop(a); - if (lane == 0) - s.wtop[warp][j] = m; - } -} - -// Pooled: K rounds over each lane's P pool keys (descending) and the warp's max thread bound. -template -__device__ __forceinline__ void warp_pool(uint32_t (&a)[P + 1], S& s) -{ - int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; - uint32_t const bound = warp_max_u32(a[P]); - uint32_t pl[P]; -#pragma unroll - for (int j = 0; j < P; ++j) - pl[j] = a[j]; -#pragma unroll - for (int j = 0; j < K; ++j) - { - uint32_t const m = warp_max_u32(pl[0]); - if (pl[0] == m && m != 0u) - topk_pop

(pl); - if (lane == 0) - s.wtop[warp][j] = m; - } - if (lane == 0) - s.wbound[warp] = bound; -} - -// A row's kSplits CTAs wait for each other's records, so they launch as one thread-block cluster: the hardware -// schedules a cluster's CTAs together, whatever other streams hold. Rows never wait on other rows. -template -__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) -{ - using R = Rec; - extern __shared__ __align__(16) unsigned char smem[]; - Shared& s = *reinterpret_cast*>(smem); - float constexpr kNegInf = -__builtin_huge_valf(); - int const tid = threadIdx.x; - int const row = blockIdx.x / kSplits; - int const split = blockIdx.x % kSplits; - int const world = a.world; - int const nrec = world * kSplits; - // The row's sampling parameters are step inputs (the spec metadata, written before the target forward), not the - // predecessor head GEMM's output, so with prewait_params they load before the dependency wait and overlap the - // head's tail; the same values either way. - float temperature = 0.f, topP = 0.f; - long long seed_in = 0, offset_in = 0; - auto load_params = [&]() - { - temperature = a.temperatures[row]; - topP = a.top_ps[row]; - seed_in = a.seeds[row]; - offset_in = a.offsets[row]; - }; - if (a.prewait_params) - load_params(); -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.wait;" ::: "memory"); -#endif - if (!a.prewait_params) - load_params(); -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - // early_trigger: the dependent grid (the first draft pass's input) may launch now and park at its own - // griddepcontrol.wait, which still returns only once this grid has completed and flushed; it skips the launch - // latency after this kernel ends. Otherwise dependents are released at the acceptance epilogue (the CTAs of - // splits 1.. exit without releasing). - if (a.early_trigger) - asm volatile("griddepcontrol.launch_dependents;"); -#endif - float const tempInv = 1.0f / (temperature + kTempEpsilon); - unsigned long long const seed = static_cast(seed_in); - unsigned long long const offset = static_cast(offset_in); - unsigned int const call = *reinterpret_cast(a.counters + 1); - // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of - // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round - // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten - // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. - auto ktag_of = [&](int r) -> unsigned long long - { - unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; - return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; - }; - auto ftag_of = [&](int r) -> unsigned long long - { - unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; - return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; - }; - // SA-1R round-0 list words: a 16-bit tag, never 0. - unsigned long long const xtag = (((call >> 1) << 1) & 0x7FFFu) | 0x8000u; - long long const rec_stride = static_cast(world) * kMaxSplits * R::kEx; - unsigned long long* const xu = a.ws_uc + static_cast(call & 1u) * half_words(world) - + static_cast(row) * rec_stride; - unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) - + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * R::kEx; - long long const round_stride = static_cast(kMaxRows) * rec_stride; - - // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- - __nv_bfloat16 const* src = a.logits + row * a.ld; - int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; - int const s0 = min(split * sl_len, a.Vs); - int const n = min(s0 + sl_len, a.Vs) - s0; - int const id0 = a.vocab_offset + s0; - float tmax = kNegInf; - // SA-1R: each thread's largest (ord16(bf16) << 16 | ~index) keys while staging (only rows with a nucleus cut). - bool const want_x = kTopK > 0 && topP < 1.0f; - constexpr int kK = mode_k(kTopK), kPool = mode_pool(kTopK), kVals = mode_vals(kTopK); - constexpr int kTk = kPool ? kPool + 1 : (kK > 0 ? kK : 1); - uint32_t tk[kTk]; -#pragma unroll - for (int j = 0; j < kTk; ++j) - tk[j] = 0u; - bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; - int const nvec = vec ? n / 8 : 0; - for (int v = tid; v < nvec; v += kThreads) - { - uint4 const w = *reinterpret_cast(src + s0 + 8 * v); - *reinterpret_cast(&s.row[8 * v]) = w; - __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); -#pragma unroll - for (int j = 0; j < 8; ++j) - tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); - if constexpr (kTopK > 0) - { - if (want_x) - { -#pragma unroll - for (int j = 0; j < 8; ++j) - { - uint32_t const key - = (ord16(__bfloat16_as_ushort(e[j])) << 16) | (0xFFFFu - static_cast(8 * v + j)); - if constexpr (kPool > 0) - topk_insert(tk, key); // branchless: P + 1 compare-swaps - else - thread_insert(tk, key); - } - } - } - } - for (int i = nvec * 8 + tid; i < n; i += kThreads) - { - s.row[i] = src[s0 + i]; - tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); - if constexpr (kTopK > 0) - { - if (want_x) - topk_insert( - tk, (ord16(__bfloat16_as_ushort(src[s0 + i])) << 16) | (0xFFFFu - static_cast(i))); - } - } - if constexpr (kTopK > 0) - { - if (tid < kVals) - s.ctop[tid] = 0u; // published with s.wtop by block_max's barrier - if (want_x) - { - if constexpr (kPool > 0) - warp_pool(tk, s); // s.wtop / s.wbound published by block_max's barrier - else - warp_topk(tk, s); // s.wtop published by block_max's barrier - } - } - float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier - if constexpr (kTopK > 0) - { - // SA-1R: the CTA's top K of the warp lists by rank (keys are unique): a lane quad per key, no serial merge; - // s.ctop is read after the round-0 draw's first barrier - constexpr int kN = kWarps * kK; - constexpr int kSp = kN * 4 <= kThreads ? 4 : 1; // threads per key - if (want_x && tid < kN * kSp) - { - uint32_t const* const wl = &s.wtop[0][0]; - int const part = tid % kSp; - uint32_t const key = wl[tid / kSp]; - int rank = 0; -#pragma unroll - for (int j = part; j < kN; j += kSp) - rank += wl[j] > key ? 1 : 0; - if constexpr (kSp == 4) - { - rank += __shfl_xor_sync(0xffffffffu, rank, 1); - rank += __shfl_xor_sync(0xffffffffu, rank, 2); - } - if (part == 0 && key != 0u && rank < kK) - s.ctop[rank] = key; - } - } - float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z - unsigned long long ckey = 0ULL; - int const gid = a.rank * kMaxSplits + split; - for (int r = 0;; ++r) - { - // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) - float const ref = r == 0 ? smax : M; - float const wmin = r == 0 ? 0.f : wc; - float ts = 0.f; - for (int i = tid; i < n; i += kThreads) - { - float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); - ts += w > wmin ? w : 0.f; - } - float total; - float const excl = block_excl_scan(ts, total, s); - if constexpr (kTopK > 0) - { - // SA-1R round 0: the list words depend only on the staged slice, so they go out now (s.ctop / s.wbound are - // published by the barrier above), one lane per word, ahead of the main record: off the readers' path - if (r == 0 && want_x && tid < R::kXw) - { - uint32_t tau = 0u; - if constexpr (kPool > 0) - { - tau = s.ctop[kK - 1]; // the largest unlisted pool key is <= the K-th listed -#pragma unroll - for (int wq = 0; wq < kWarps; ++wq) - tau = s.wbound[wq] > tau ? s.wbound[wq] : tau; - } - unsigned long long w = xtag; -#pragma unroll - for (int t = 0; t < 3; ++t) - { - int const e = 3 * tid + t; - uint32_t const kv = e < kK ? s.ctop[e] : (kPool > 0 && e == kK ? tau : 0u); - // missing entries: -inf (bf16 0xFF80), w = 0 - uint32_t const b = kv != 0u ? unord16(kv >> 16) : 0xFF80u; - w |= static_cast(b) << (16 * (3 - t)); - } - st_sys(xm + 4 + tid, w); // round 0 = slot 0 - } - } - uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), - static_cast(offset), static_cast(offset >> 32)), - make_uint2(static_cast(seed), static_cast(seed >> 32))); - float const tgt = uniform(rnd.x) * total; - if (tid == 0) - s.owner = -1; - __syncthreads(); - if (ts > 0.f && excl <= tgt) - atomicMax(&s.owner, tid); - __syncthreads(); - if (tid == s.owner) - { - float acc = excl; - int pick = -1; - float px = 0.f; - if constexpr (kIlpWalk) - { - // SA-PW: the owner's strided tokens 8 at a time. The 8 x and w (the same expressions) are - // independent, so their MUFU work overlaps; the serial sum then adds them in the same order with - // the same early exit: the same acc sequence, pick and px as the one-token walk, bit for bit. - constexpr int kW = 8; - bool done = false; - for (int i0 = tid; i0 < n && !done; i0 += kW * kThreads) - { - float xs[kW], wv[kW]; -#pragma unroll - for (int k = 0; k < kW; ++k) - { - int const i = i0 + k * kThreads; - xs[k] = i < n ? __fmul_rn(__bfloat162float(s.row[i]), tempInv) : 0.f; - wv[k] = __expf(xs[k] - ref); - } -#pragma unroll - for (int k = 0; k < kW; ++k) - { - int const i = i0 + k * kThreads; - if (!done && i < n && wv[k] > wmin) - { - acc += wv[k]; - pick = i; - px = xs[k]; - if (acc > tgt) - done = true; - } - } - } - } - else - { - for (int i = tid; i < n; i += kThreads) - { - float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); - float const w = __expf(x - ref); - if (w > wmin) - { - acc += w; - pick = i; - px = x; - if (acc > tgt) - break; - } - } - } - s.pick = pick; - s.pickx = px; - } - __syncthreads(); - unsigned long long sb = 0ULL; - if (total > 0.f && s.owner >= 0) - sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); - float const smab = total; - unsigned long long const ktag = ktag_of(r); - unsigned long long const ftag = ftag_of(r); - unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; - bool const with_x = want_x && r == 0; - if (tid == 0) - { - unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; - st_sys(xw + 0, (sb << kTagBits) | ktag); - st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); - st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); - st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); - } - // ---- every record of the row, every rank --------------------------------------------------------------- - int const nw = with_x ? R::kEx : 4; - for (int t = tid; t < nrec * nw; t += kThreads) - { - int const q = t / nw, f = t % nw; - unsigned long long const* const p = xr + (q / kSplits * kMaxSplits + q % kSplits) * R::kEx + f; - unsigned long long v; - long long spins = 0; - if (f == 0) - { - do - { - v = ld_sys(p); - if (++spins > kMaxSpins) - __trap(); - } while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); - s.rk[q] = v >> kTagBits; - } - else if (f < 4) - { - do - { - v = ld_sys(p); - if (++spins > kMaxSpins) - __trap(); - } while ((v & 0xFFFFFFFF00000000ULL) != ftag); - float const fv = __uint_as_float(static_cast(v)); - (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; - } - else - { - do - { - v = ld_sys(p); - if (++spins > kMaxSpins) - __trap(); - } while ((v & 0xFFFFULL) != xtag); - s.rx[q][f - 4] = v; - } - } - __syncthreads(); - // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- - unsigned long long gk = 0ULL; - float gv = 0.f; - for (int q = 0; q < nrec; ++q) - { - if (s.rk[q] > gk) - { - gk = s.rk[q]; - gv = s.rv[q]; - } - } - bool accept; - if (r == 0) - { - M = kNegInf; - for (int q = 0; q < nrec; ++q) - M = fmaxf(M, s.r2[q]); - float Z = 0.f; - for (int q = 0; q < nrec; ++q) - Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); - target = topP * Z; - accept = !(topP < 1.0f); // no nucleus cut: the first draw stands - ckey = gk; - cv = gv; - if constexpr (kTopK > 0) - { - // SA-1R: every thread takes (record, entry) items; the per-warp partials are published by the round's - // barrier below and every thread combines them in warp order after it (no new barrier) - if (with_x) - { - float const wcand = __expf(cv - M); // = round 1's wmin - float part = 0.f; - int cnt = 0, inc = 0; - for (int t = tid; t < nrec * kVals; t += kThreads) - { - int const q = t / kVals, e = t % kVals; - uint32_t const b = static_cast(s.rx[q][e / 3] >> (16 * (3 - e % 3))) & 0xFFFFu; - float const w = __expf( - __fmul_rn(__bfloat162float(__ushort_as_bfloat16(static_cast(b))), tempInv) - - M); // w bitwise as round 1 computes it - if (w > wcand) - { - if (kPool > 0 && e == kK) - inc = 1; // tau (the largest unlisted value) is heavier than the candidate - else - { - part += w; - ++cnt; - if (kPool == 0 && e == kK - 1) - inc = 1; // the K-th is heavier: unlisted heavier tokens may exist - } - } - } -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - part += __shfl_xor_sync(0xffffffffu, part, o); // any order: the margin covers it - cnt = static_cast(warp_add_u32(static_cast(cnt))); - inc = static_cast(warp_or_u32(static_cast(inc))); - if ((tid & 31) == 0) - { - s.dS[tid >> 5] = part; - s.dM[tid >> 5] = cnt | (inc << 16); - } - } - } - } - else - { - float S = 0.f; // mass strictly above the current candidate - for (int q = 0; q < nrec; ++q) - S += s.r2[q]; - accept = S <= target || gk == 0ULL; - if (!accept) - { - ckey = gk; - cv = gv; - } - } - wc = __expf(cv - M); - __syncthreads(); // records consumed before the next round's writes to s - bool decided = false; - if constexpr (kTopK > 0) - { - if (r == 0 && with_x) - { - float S1 = 0.f; - int m = 0, inc = 0; - for (int wq = 0; wq < kWarps; ++wq) - { - S1 += s.dS[wq]; - m += s.dM[wq] & 0xFFFF; - inc |= s.dM[wq] >> 16; - } - if (!inc) - { - float const margin = m <= 1 ? 0.f : S1 * static_cast(m) * (1.0f / 2097152.0f); - decided = __fadd_ru(S1, margin) <= target; - accept = accept || decided; - } - } - } - if (accept) - { - if (split != 0) - return; - if (tid == 0) - { - a.tokens[row] = key_token(ckey); - if (a.extra_attempts != nullptr && r > 1) - atomicAdd(a.extra_attempts, r - 1); - if (a.decided != nullptr && decided) - atomicAdd(a.decided, 1); - } - break; - } - } - - // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - if (!a.early_trigger) - asm volatile("griddepcontrol.launch_dependents;"); -#endif - __threadfence(); - __syncthreads(); - if (tid == 0) - s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; - __syncthreads(); - if (!s.last) - return; - __threadfence(); - if (tid == 0) - { - a.counters[0] = 0u; - a.counters[1] = call + 1u; - } - if (tid < a.nreq) - { - int const* tok = a.tokens; - int const b = tid; - int const k1 = a.draft_len + 1; - int nacc = 1; - bool run = true; - for (int k = 0; k < a.draft_len; ++k) - { - int const tt = *reinterpret_cast(tok + b * k1 + k); - a.acc_tokens[b * k1 + k] = tt; - run = run && a.draft[b * a.draft_len + k] == tt; - nacc += run ? 1 : 0; - } - a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); - if (a.force_vec != nullptr) - nacc = a.force_vec[b]; - else if (a.force_total > 0) - nacc = a.force_total; - a.num_acc[b] = nacc; - a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; - if (a.next_new != nullptr) - a.next_new[b * a.next_stride] - = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); - } -} - -template -cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) -{ - int device; - cudaError_t const deviceStatus = cudaGetDevice(&device); - if (deviceStatus != cudaSuccess) - return deviceStatus; - static std::mutex configMutex; - static std::unordered_map deviceClusters; - int max_clusters = 0; - cudaLaunchConfig_t cfg = {}; - cfg.gridDim = dim3(a.rows * kSplits); - cfg.blockDim = dim3(kThreads); - cfg.dynamicSmemBytes = sizeof(Shared); - cfg.stream = stream; - cudaLaunchAttribute attrs[2]; - attrs[0].id = cudaLaunchAttributeClusterDimension; - attrs[0].val.clusterDim.x = kSplits; - attrs[0].val.clusterDim.y = 1; - attrs[0].val.clusterDim.z = 1; - attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; - cfg.attrs = attrs; - cfg.numAttrs = 2; - { - std::lock_guard lock(configMutex); - max_clusters = deviceClusters[device]; - if (max_clusters == 0) - { - cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); - if (err == cudaSuccess) - { - err = cudaOccupancyMaxActiveClusters( - &max_clusters, sampled_accept_kernel, &cfg); - } - if (err != cudaSuccess) - { - return err; - } - } - deviceClusters[device] = max_clusters; - } - TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); - return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); -} - -template -cudaError_t launchRows(Args const& a, bool use_pdl, cudaStream_t stream) -{ - // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); - // more rows: 4 CTAs per row. - if (a.ilp_walk) - { - return a.rows <= 32 ? launch<8, kTopK, true>(a, use_pdl, stream) : launch<4, kTopK, true>(a, use_pdl, stream); - } - return a.rows <= 32 ? launch<8, kTopK, false>(a, use_pdl, stream) : launch<4, kTopK, false>(a, use_pdl, stream); -} - -} // namespace - long long sampledAcceptWsWords(int world, int oneRoundK) { - TLLM_CHECK_WITH_INFO( - oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "SA-1R variant %d is not built", oneRoundK); - return 2 * (oneRoundK == 9 ? half_words<9>(world) : oneRoundK == 82 ? half_words<82>(world) : half_words<0>(world)); + return minimaxM3Sm107::sampledAccept::sampledAcceptWsWords(world, oneRoundK); } cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) { - switch (a.one_round_k) - { - case 0: return launchRows<0>(a, usePdl, stream); - case 9: return launchRows<9>(a, usePdl, stream); - case 82: return launchRows<82>(a, usePdl, stream); - default: TLLM_CHECK_WITH_INFO(false, "SA-1R variant %d is not built", a.one_round_k); - } - return cudaErrorInvalidValue; + return minimaxM3Sm107::sampledAccept::invokeSampledAccept(a, usePdl, stream); } - } // namespace sampledAccept } // namespace minimaxM3Sm107Tp4 } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h index 0ebd153992d0..09f377cececb 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/sampledAccept.h @@ -18,6 +18,7 @@ #pragma once #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.h" #include #include @@ -40,39 +41,7 @@ constexpr int kMaxWorld = 8; // TP ranks constexpr int kMaxVs = 51200; // vocab shard width (TP4: 200064 / 4 = 50016) constexpr int kIdBits = 18; // token ids < 2^18 -struct Args -{ - __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard - long long ld; - int rows, Vs, vocab_offset; - float const* temperatures; - float const* top_ps; - long long const* seeds; - long long const* offsets; - unsigned long long* ws_uc; - unsigned long long* ws_mc; - unsigned int* counters; // [0] finished rows, [1] call - int* tokens; // [rows] sampled token per row - int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows - int rank, world; - int const* draft; - int const* batch_idx; - int* acc_tokens; - int* num_acc; - long long* gather_ids; - int draft_len, nreq, force_total; - int* next_new; - int next_stride; - int const* force_vec; - int prewait_params; // nonzero: the row's sampling parameters load before the PDL dependency wait - int early_trigger; // nonzero: dependents are released right after the dependency wait, not at the epilogue - // SA-1R (one-round decision): 0 the two-round kernel; 9 each CTA's top-9 logits ride its round-0 record; 82 pooled - // (each thread's top 2, the CTA lists the pool's top 8 plus the bound of everything unlisted). A row whose first - // candidate is provably in the nucleus then decides in round 0, with the same token as the two-round kernel. - int one_round_k; - int* decided; // optional [1]: rows decided in round 0 (SA-1R), summed - int ilp_walk; // nonzero: the owner walk takes its tokens 8 at a time (same weights, same summation order) -}; +using Args = minimaxM3Sm107::sampledAccept::Args; //! Words (uint64) of the multicast exchange workspace for `world` ranks and SA-1R variant `oneRoundK`: two call halves. long long sampledAcceptWsWords(int world, int oneRoundK = 0); diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu deleted file mode 100644 index 679bcb002d1d..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.cu +++ /dev/null @@ -1,58 +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/minimaxM3Sm107Tp4/specOffsets.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ - -namespace -{ - -constexpr int kThreads = 256; - -__global__ void __launch_bounds__(kThreads) - specOffsetsKernel(int32_t* __restrict__ pos, int32_t const* __restrict__ posOff, int64_t nPos, - int32_t* __restrict__ kv, int32_t const* __restrict__ kvOff, int64_t nKv) -{ - for (int64_t i = threadIdx.x; i < nPos; i += kThreads) - { - pos[i] = static_cast(static_cast(pos[i]) + static_cast(posOff[i])); - } - for (int64_t j = threadIdx.x; j < nKv; j += kThreads) - { - kv[j] = static_cast(static_cast(kv[j]) + static_cast(kvOff[j])); - } -} - -} // namespace - -cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, - int64_t nKv, cudaStream_t stream) -{ - specOffsetsKernel<<<1, kThreads, 0, stream>>>(pos, posOff, nPos, kv, kvOff, nKv); - return cudaGetLastError(); -} - -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu deleted file mode 100644 index 42361b5ef47f..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.cu +++ /dev/null @@ -1,111 +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/minimaxM3Sm107Tp4/stepPrologue.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ - -namespace -{ - -constexpr int kThreads = 512; -constexpr int kWarps = kThreads / 32; - -// One warp per copy. Plain (coherent) loads: a device copy may read bytes another warp of this CTA just stored, which -// the CTA barrier between the two kinds makes visible; a non-coherent load could return the stale line. -__device__ __forceinline__ void copyWarp(uint8_t* dst, uint8_t const* src, int64_t bytes, int lane) -{ - auto const d = reinterpret_cast(dst); - auto const s = reinterpret_cast(src); - if (((d | s | static_cast(bytes)) & 15) == 0) - { - auto* d16 = reinterpret_cast(dst); - auto const* s16 = reinterpret_cast(src); - for (int64_t i = lane; i < bytes / 16; i += 32) - { - d16[i] = s16[i]; - } - } - else if (((d | s | static_cast(bytes)) & 3) == 0) - { - auto* d4 = reinterpret_cast(dst); - auto const* s4 = reinterpret_cast(src); - for (int64_t i = lane; i < bytes / 4; i += 32) - { - d4[i] = s4[i]; - } - } - else - { - for (int64_t i = lane; i < bytes; i += 32) - { - dst[i] = src[i]; - } - } -} - -// One CTA: a step's window is a few dozen copies of a few hundred bytes, so the launch is the cost; a CTA barrier -// orders the host copies before the device copies where the parallel scatter needed a second launch. -__global__ void __launch_bounds__(kThreads) stepPrologueKernel(uint8_t* window, int64_t* epoch, int64_t maxDesc) -{ - auto* header = reinterpret_cast(window); - int64_t const numHost = min(header[0], maxDesc); - int64_t const numDevice = min(header[1], maxDesc); - int64_t const* desc = header + 2; - int const warp = threadIdx.x / 32; - int const lane = threadIdx.x % 32; - if (epoch != nullptr && threadIdx.x == 0) - { - *epoch += 1; - } - for (int64_t i = warp; i < numHost; i += kWarps) - { - int64_t const* c = desc + 3 * i; - copyWarp(reinterpret_cast(c[0]), window + c[1], c[2], lane); - } - __syncthreads(); - for (int64_t i = warp; i < numDevice; i += kWarps) - { - // The device copies' slots follow all header[0] host copies, as the parallel scatter indexes them. - int64_t const* c = desc + 3 * (header[0] + i); - copyWarp(reinterpret_cast(c[0]), reinterpret_cast(c[1]), c[2], lane); - } - // Every thread has read the header and the descriptors: consume the window. - __syncthreads(); - if (threadIdx.x < 2) - { - header[threadIdx.x] = 0; - } -} - -} // namespace - -cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream) -{ - stepPrologueKernel<<<1, kThreads, 0, stream>>>(window, epoch, maxDesc); - return cudaGetLastError(); -} - -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h deleted file mode 100644 index c934decff161..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/stepPrologue.h +++ /dev/null @@ -1,47 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ - -//! The decode graph's step-input prologue in one launch (step_inputs.StepInputPack with early_copy): the work of the -//! two parallel scatter launches and the header clear that follow them, in that order. -//! -//! The window is the pack's staging layout (kernels/stepInputsScatter.h): an int64 header `[numHostCopies, -//! numDeviceCopies]`, then one `(dst, src, bytes)` int64 triple per copy, host copies first; a host copy's `src` is a -//! byte offset into the window, a device copy's a device address. Up to `maxDesc` copies of each kind are written (the -//! slots the parallel scatter's grid covers), every host copy before any device copy, so a device copy may read a -//! destination a host copy filled; the copies of one kind must not overlap (the pack issues the window only then). -//! Afterwards the header is zeroed, so a replay with no new upload writes nothing. With `epoch` non-null, it is -//! incremented once. -cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream); - -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu deleted file mode 100644 index 507520ec0de6..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp4/swigluOai.cu +++ /dev/null @@ -1,100 +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/minimaxM3Sm107Tp4/swigluOai.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp4 -{ - -namespace -{ - -constexpr int kThreads = 256; - -// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves): the -// instruction sequence of the Triton kernel at limit 7, alpha 1.702, beta 1 (Indexer kdev17, bitwise on SM107): the -// gate clamp as min.bf16x2 before the widening, the up clamp as min.xorsign.abs, (u + 1), (g * -1.702) * log2(e), -// ex2.approx, + 1, div.full 1 / that, s * g, the product, one rounding to bf16. -__device__ __forceinline__ uint32_t swigluOaiBf16x2(uint32_t gate, uint32_t up) -{ - uint32_t r; - asm("{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" - "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, %1, glim;\n\t" - "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" - "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" - "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" - "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" - "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" - "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" - "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" - "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" - "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" - "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" - "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" - "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" - : "=r"(r) - : "r"(gate), "r"(up)); - return r; -} - -template -__global__ void __launch_bounds__(kThreads) - swigluOaiKernel(uint32_t const* __restrict__ x, uint32_t* __restrict__ out, int64_t rows, int64_t halfD) -{ - if constexpr (kPdl) - { - asm volatile("griddepcontrol.wait;" ::: "memory"); - asm volatile("griddepcontrol.launch_dependents;"); - } - int64_t const i = blockIdx.x * static_cast(kThreads) + threadIdx.x; - if (i >= rows * halfD) - { - return; - } - int64_t const row = i / halfD; - int64_t const col = i - row * halfD; - uint32_t const* xr = x + row * 2 * halfD; - out[i] = swigluOaiBf16x2(xr[col], xr[halfD + col]); -} - -} // namespace - -cudaError_t invokeSwigluOai( - uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream) -{ - int64_t const n = rows * halfD; - cudaLaunchConfig_t cfg{}; - cfg.gridDim = dim3(static_cast((n + kThreads - 1) / kThreads)); - cfg.blockDim = dim3(kThreads); - cfg.stream = stream; - cudaLaunchAttribute attr[1]; - attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - cfg.attrs = attr; - cfg.numAttrs = 1; - return pdl ? cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD) - : cudaLaunchKernelEx(&cfg, swigluOaiKernel, x, out, rows, halfD); -} - -} // namespace minimaxM3Sm107Tp4 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu index fb28ce95250b..eeedda902a09 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.cu @@ -629,7 +629,7 @@ __global__ void __launch_bounds__(kWarps * 32, 1) __syncthreads(); if (threadIdx.x < 32) { - foldTailWarp(p.fold, static_cast(threadIdx.x)); + minimaxM3Sm107::foldTailWarp(p.fold, static_cast(threadIdx.x)); } } #endif diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h index 485f383732fe..332df6d7369e 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/draftStep.h @@ -18,7 +18,7 @@ #pragma once #include "tensorrt_llm/common/config.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h" #include #include @@ -90,7 +90,7 @@ struct Params CodedArgs cd; // The step's forward-tail store + next-step staging (foldTail.h), run by the last CTA after this pass's // tokens and bookkeeping: set on the drafter's last pass only (dispatch.DRAFT_STEP_FOLD_TAIL). - FoldTailParams fold{}; + minimaxM3Sm107::FoldTailParams fold{}; bool hasFold = false; // D-FIX: sum the pieces of a tile that several warps stream in warp order (deterministic) instead of by // shared-memory atomics in arrival order (dispatch.DRAFT_STEP_FIXED_ORDER). diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu deleted file mode 100644 index 68279dc3bfaa..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.cu +++ /dev/null @@ -1,189 +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. - */ - -// The TP8 drafter layer's generation preprocess in one launch. The rotation is applyBiasRopeUpdateKVCacheV2's GPT-NeoX -// path at BF16 (apply_rotary_embedding_gptneox: cos * x + sin * pair in fp32, the first rotary half's sin negated, -// dims past the rotary width with (cos, sin) = (1, 0), rounded to BF16); the K / V store is its unit-scale E4M3 -// store (__nv_fp8x2_e4m3 from BF16 pairs), which is also the query's E4M3: torch's CUDA cast of the BF16 query -// saturates the same way (448, not NaN). - -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ -namespace drafterRopeFp8Kv -{ -namespace -{ - -constexpr int kElemsPerThread = kHeadDim / 32; -constexpr int kBlockSize = 256; - -template -__device__ __forceinline__ float rotate(float c, float x, float s, float pair) -{ - if constexpr (kRopeMode == 0) - { - return __fadd_rn(__fmul_rn(c, x), __fmul_rn(s, pair)); - } - else if constexpr (kRopeMode == 1) - { - return __fmaf_rn(c, x, __fmul_rn(s, pair)); - } - else - { - return __fmaf_rn(s, pair, __fmul_rn(c, x)); - } -} - -__device__ __forceinline__ void storeE4m3Pairs(__nv_fp8_e4m3* dst, __nv_bfloat16 const (&v)[kElemsPerThread]) -{ -#pragma unroll - for (int i = 0; i < kElemsPerThread; i += 2) - { - __nv_bfloat162 pair; - pair.x = v[i]; - pair.y = v[i + 1]; - __nv_fp8x2_e4m3 const packed(pair); - reinterpret_cast<__nv_fp8x2_storage_t*>(dst)[i / 2] = packed.__x; - } -} - -template -__global__ void __launch_bounds__(kBlockSize) drafterRopeFp8KvKernel(Params const p) -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("griddepcontrol.wait;" ::: "memory"); - asm volatile("griddepcontrol.launch_dependents;"); -#endif - int const lane = threadIdx.x % 32; - int const warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; - int const slots = p.numHeadsQ + p.numHeadsKv; - int const token = warp / slots; - if (token >= p.numTokens) - { - return; - } - int const slot = warp % slots; - bool const isQ = slot < p.numHeadsQ; - int const head = isQ ? slot : slot - p.numHeadsQ; - int const request = token / p.qLen; - int const row = token % p.qLen; - int const cacheLen = p.seqLens[request]; - int const kvPos = cacheLen - p.qLen + row; - int const ropePos - = p.specPositionOffsets != nullptr ? p.specPositionOffsets[request * p.qLen + row] + cacheLen - p.qLen : kvPos; - int const rowWidth = (p.numHeadsQ + 2 * p.numHeadsKv) * kHeadDim; - __nv_bfloat16 const* src - = p.qkv + static_cast(token) * rowWidth + (isQ ? head * kHeadDim : (p.numHeadsQ + head) * kHeadDim); - int const half = p.rotaryDim / 2; - float2 const* coef = p.cosSin + static_cast(ropePos) * half; - - __nv_bfloat16 rotated[kElemsPerThread]; -#pragma unroll - for (int i = 0; i < kElemsPerThread; ++i) - { - int const d = lane * kElemsPerThread + i; - bool const firstHalf = d < half; - float const x = __bfloat162float(src[d]); - float const pair = __bfloat162float(src[firstHalf ? d + half : d - half]); - float2 const cs = d < p.rotaryDim ? coef[firstHalf ? d : d - half] : make_float2(1.0f, 0.0f); - float const s = firstHalf ? -cs.y : cs.y; - rotated[i] = __float2bfloat16(rotate(cs.x, x, s, pair)); - } - - if (isQ) - { - storeE4m3Pairs( - p.qOut + (static_cast(token) * p.numHeadsQ + head) * kHeadDim + lane * kElemsPerThread, rotated); - return; - } - - int const block = kvPos / p.tokensPerBlock; - long long const inPage = (static_cast(head) * p.tokensPerBlock + kvPos % p.tokensPerBlock) * kHeadDim - + lane * kElemsPerThread; - long long const pageElems = static_cast(p.numHeadsKv) * p.tokensPerBlock * kHeadDim; - int const* table = p.blockTable + request * p.blockTableRequestStride; - long long const kPage = table[block]; - long long const vPage = table[p.blockTableKvStride + block]; - storeE4m3Pairs(p.pool + kPage * pageElems + inPage, rotated); - - __nv_bfloat16 const* vSrc = p.qkv + static_cast(token) * rowWidth - + (p.numHeadsQ + p.numHeadsKv + head) * kHeadDim + lane * kElemsPerThread; - __nv_bfloat16 v[kElemsPerThread]; -#pragma unroll - for (int i = 0; i < kElemsPerThread; ++i) - { - v[i] = vSrc[i]; - } - storeE4m3Pairs(p.pool + vPage * pageElems + inPage, v); -} - -//! With PDL and maxSharedCarveout, the maximum shared-memory carveout lets the full-context decode's large-smem CTAs -//! (the dependent this kernel releases right after its wait) co-reside on the SMs this small kernel occupies. Set once -//! per instantiation. -template -cudaError_t launchWithMaxSharedCarveout( - cudaLaunchConfig_t const& config, Params const& params, bool pdl, bool maxSharedCarveout) -{ - auto const kernel = drafterRopeFp8KvKernel; - if (pdl && maxSharedCarveout) - { - static cudaError_t const carveout = cudaFuncSetAttribute( - kernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); - if (carveout != cudaSuccess) - { - return carveout; - } - } - return cudaLaunchKernelEx(&config, kernel, params); -} - -} // namespace - -cudaError_t launch(Params const& params, bool pdl, bool maxSharedCarveout, cudaStream_t stream) -{ - int const warps = params.numTokens * (params.numHeadsQ + params.numHeadsKv); - int const warpsPerBlock = kBlockSize / 32; - cudaLaunchConfig_t config = {}; - config.gridDim = dim3((warps + warpsPerBlock - 1) / warpsPerBlock); - config.blockDim = dim3(kBlockSize); - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attribute[1]; - attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - config.attrs = attribute; - config.numAttrs = 1; - switch (params.ropeMode) - { - case 0: return launchWithMaxSharedCarveout<0>(config, params, pdl, maxSharedCarveout); - case 1: return launchWithMaxSharedCarveout<1>(config, params, pdl, maxSharedCarveout); - case 2: return launchWithMaxSharedCarveout<2>(config, params, pdl, maxSharedCarveout); - default: return cudaErrorInvalidValue; - } -} - -} // namespace drafterRopeFp8Kv -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu deleted file mode 100644 index e80a7810b9ba..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.cu +++ /dev/null @@ -1,64 +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/minimaxM3Sm107Tp8/foldTail.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ - -namespace -{ - -// Row i of the step (lane c = column), as the Triton fold tail's program i: the store's padded rows, then the next -// step's inputs read from them (zero past each source width). -__global__ void __launch_bounds__(32) foldTailKernel(FoldTailParams const p) -{ - int const i = blockIdx.x; - if (i >= *p.ctrl || i >= p.maxRows) - { - return; - } - foldTailRow(p, i, static_cast(threadIdx.x)); -} - -__global__ void __launch_bounds__(32) foldTailOneWarpKernel(FoldTailParams const p) -{ - foldTailWarp(p, static_cast(threadIdx.x)); -} - -} // namespace - -cudaError_t invokeFoldTail(FoldTailParams const& p, int32_t rows, cudaStream_t stream) -{ - foldTailKernel<<>>(p); - return cudaGetLastError(); -} - -cudaError_t invokeFoldTailOneWarp(FoldTailParams const& p, cudaStream_t stream) -{ - foldTailOneWarpKernel<<<1, 32, 0, stream>>>(p); - return cudaGetLastError(); -} - -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu deleted file mode 100644 index a5a157e33fe4..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.cu +++ /dev/null @@ -1,1096 +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/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/tllmException.h" - -#include -#include -#include -#include - -// The sm_107 TP8 target's own copy of minimaxM3Sm107Tp2/mcAllreduceNormTp2.cu (the version with the Eagle3 capture -// epilogue, cap_out), verbatim but for the names: the TP8 ext carries it, so the TP8 target no longer needs an image -// whose TP2 op has cap_out (SpecTail S8). The TP2 file's notes follow unchanged. -// -// SM107 (Rubin) copy of minimaxM3/mcAllreduceNorm.cu for the sm_107 TP8 target. Two changes in the arithmetic: the -// rank sum rounds to BF16 before the residual add, the rounding points of the stock all-reduce + fused residual -// RMSNorm on this target (bf16 all-reduce output, then bf16(residual + sum)); and every call but the MoE finalize runs -// the row kernel below, whose RMS statistic is the stock fused norm's reduction tree (bitwise the stock outputs). -// -// The kernels use sm_100-family features (thread-block clusters, cp.async.bulk, st.async, mbarrier transaction -// counts, block-scaled conversions): they are compiled for every sm_10x target of the build, the sm_100f family target -// included (which runs on SM107). A device pass for any other architecture compiles them to a trap, never to a silent -// empty body, and the host entry point checks that the kernel binary it launches holds sm_10x code (see -// checkBinary). -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 -#define MINIMAX_M3_KERNELS_ENABLED 1 -#else -#define MINIMAX_M3_KERNELS_ENABLED 0 -#endif - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ - -namespace cg = cooperative_groups; - -namespace -{ - -constexpr int kEmpty = static_cast(0x80000000u); -constexpr int kFinMaxTopk = 8; // MoE finalize: top-k rows loaded together up to this k - -// The Eagle3 capture written by a capture layer's boundary (cap_out, optional): rmsnorm(updated residual) * fc_norm[i], -// bit for bit the m3_spec_tail capture_norm kernel (kernels/minimaxM3/specTail.cu at H 6144: 256 threads, thread c -// sums the squares of the 8-column chunks c, c + 256 and c + 512 in that order, the warp butterfly 16..1, the 8 warp -// sums added in warp order, rsqrtf(sum / H + eps), bf16(v * rs * w)). Only its layout is reproduced here; its -// arithmetic is the same source. -constexpr int kCapThreads = 256; -constexpr int kCapIters = 3; -constexpr int kCapWarps = kCapThreads / 32; -constexpr int kCapHidden = kCapThreads * kCapIters * 8; - -// capture_norm's per-thread chain over one 8-column chunk (continues `ss`). -__device__ __forceinline__ float cap_chunk_sq(uint4 const raw, float ss) -{ - __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&raw); -#pragma unroll - for (int q = 0; q < 8; ++q) - { - float const v = __bfloat162float(e[q]); - ss += v * v; - } - return ss; -} - -__device__ __forceinline__ float cap_warp_sum(float v) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v += __shfl_xor_sync(0xffffffffu, v, o); - return v; -} - -// bf16(v * rs * w) for one 8-column chunk of the updated residual (v as floats) into cap. -__device__ __forceinline__ void cap_store( - float const (&v)[8], float rs, __nv_bfloat16 const* __restrict__ w, __nv_bfloat16* __restrict__ cap) -{ - uint4 const wr = *reinterpret_cast(w); - __nv_bfloat16 const* we = reinterpret_cast<__nv_bfloat16 const*>(&wr); - uint4 o; - __nv_bfloat16* oe = reinterpret_cast<__nv_bfloat16*>(&o); -#pragma unroll - for (int q = 0; q < 8; ++q) - oe[q] = __float2bfloat16(v[q] * rs * __bfloat162float(we[q])); - *reinterpret_cast(cap) = o; -} - -__device__ __forceinline__ uint32_t canon_bf16x2(uint32_t w) -{ - if ((w & 0xFFFFu) == 0x8000u) - w &= 0xFFFF0000u; - if ((w >> 16) == 0x8000u) - w &= 0x0000FFFFu; - return w; -} - -__device__ __forceinline__ float bf16_lo(uint32_t w) -{ - return __uint_as_float(w << 16); -} - -// The MXFP8 publish quantizes exactly as FlashInfer's CuTe DSL mxfp8 quantize (MXFP8QuantizeRunner.CUTE_DSL, MX-TAC's -// pinned decode quantize; quantization_cute_dsl_utils.py), the three helpers below being its PTX verbatim. It equals -// trtllm::mxfp8_quantize bit for bit except on blocks with 0 < amax < 448 * 2^-127 (UE8M0 byte 0), where the -// TRT-LLM kernel's rcp.approx.ftz of the denormal scale gives inf (NaN / garbage E4M3) and this flushes the block to 0. -__device__ __forceinline__ uint32_t mxfp8_ue8m0(float amax) -{ - float value; - asm("mul.rn.f32 %0, %1, %2;" : "=f"(value) : "f"(amax), "f"(1.0f / 448.0f)); - uint32_t result; - asm("{\n" - " .reg .pred p_zero, p_has_mant, p_exp_zero, p_tiny_sub, p_ovf;\n" - " .reg .u32 bits, exp_biased, mantissa, bump, r;\n" - " setp.le.f32 p_zero, %1, 0f00000000;\n" - " mov.b32 bits, %1;\n" - " shr.b32 exp_biased, bits, 23;\n" - " and.b32 exp_biased, exp_biased, 255;\n" - " and.b32 mantissa, bits, 0x7FFFFF;\n" - " setp.ne.u32 p_has_mant, mantissa, 0;\n" - " selp.u32 bump, 1, 0, p_has_mant;\n" - " setp.eq.u32 p_exp_zero, exp_biased, 0;\n" - " setp.le.u32 p_tiny_sub, mantissa, 0x400000;\n" - " and.pred p_tiny_sub, p_exp_zero, p_tiny_sub;\n" - " @p_tiny_sub mov.u32 bump, 0;\n" - " add.u32 r, exp_biased, bump;\n" - " setp.gt.u32 p_ovf, r, 254;\n" - " selp.u32 r, 254, r, p_ovf;\n" - " selp.u32 %0, 0, r, p_zero;\n" - "}\n" - : "=r"(result) - : "f"(value)); - return result; -} - -__device__ __forceinline__ float mxfp8_inv_scale(uint32_t ue8m0) -{ - float inv; - asm("{\n" - " .reg .s32 new_exp;\n" - " .reg .b32 float_bits;\n" - " .reg .pred p_zero;\n" - " setp.eq.u32 p_zero, %1, 0;\n" - " sub.s32 new_exp, 254, %1;\n" - " max.s32 new_exp, new_exp, 0;\n" - " shl.b32 float_bits, new_exp, 23;\n" - " mov.b32 %0, float_bits;\n" - " @p_zero mov.b32 %0, 0;\n" - "}\n" - : "=f"(inv) - : "r"(ue8m0)); - return inv; -} - -// Two BF16-valued floats (lo, hi) times inv_scale -> two E4M3 bytes (lo in the low byte), saturating to +/-448. -__device__ __forceinline__ uint32_t mxfp8_e4m3x2(float lo, float hi, float inv) -{ - uint16_t pair; - asm("{\n" - " .reg .f32 f0, f1;\n" - " mul.f32 f0, %1, %3;\n" - " mul.f32 f1, %2, %3;\n" - " min.f32 f0, f0, 0f43E00000;\n" - " max.f32 f0, f0, 0fC3E00000;\n" - " min.f32 f1, f1, 0f43E00000;\n" - " max.f32 f1, f1, 0fC3E00000;\n" - " cvt.rn.satfinite.e4m3x2.f32 %0, f1, f0;\n" - "}\n" - : "=h"(pair) - : "f"(lo), "f"(hi), "f"(inv)); - return pair; -} - -__device__ __forceinline__ float rcp_approx_ftz(float a) -{ - float b; - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); - return b; -} - -__device__ __forceinline__ float bf16_hi(uint32_t w) -{ - return __uint_as_float(w & 0xFFFF0000u); -} - -__device__ __forceinline__ int4 ld_volatile_v4(int const* p) -{ - int4 v; - asm volatile("ld.volatile.global.v4.s32 {%0,%1,%2,%3}, [%4];\n" - : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) - : "l"(p) - : "memory"); - return v; -} - -struct FinalizeArgs -{ - __nv_bfloat16 const* fc2; // [permuted rows, H] routed-expert FC2 output (null: plain input x) - int const* perm; // [M, topk] expanded index -> permuted row, -1 = dropped - void const* scales; // [M, topk] routing weights (fp32 or bf16) - int scale_fp32; - int topk; -}; - -// grid = min(M, MAXT) * CLUSTER CTAs, cluster (CLUSTER, 1, 1), THREADS threads. -// OPTS (each bit-identical): 1 residual + norm weight loaded at entry (before the poll); 2 cluster RMS exchange by -// st.async + mbarrier complete_tx instead of cluster.sync; 4 call parity rotated right after the poll; 8 max-smem -// carveout so the next kernel's CTAs co-reside early. -// OPTS 16 (NOT bit-identical with the others; a different boundary): the Eagle3 drafter's boundary, whose stock is a -// plain BF16 all-reduce, then FlashInfer fused_add_rmsnorm: the RMS statistic and the normalization take the -// unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded to BF16). -// CAP: the Eagle3 capture epilogue (cap_out) is compiled in only where a call passes cap_out, so every other call -// runs the kernel without the capture branches, shared arrays or parameters in its body. -template -__global__ void __launch_bounds__(THREADS) mc_allreduce_norm_kernel(__nv_bfloat16 const* __restrict__ x, - __nv_bfloat16 const* __restrict__ x2, FinalizeArgs fin, int M, int H, __nv_bfloat16 const* __restrict__ residual, - __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, - int world, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, - uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, - __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) -{ -#if MINIMAX_M3_KERNELS_ENABLED - __shared__ float warp_sq[THREADS / 32]; - __shared__ alignas(16) float cluster_sq[CLUSTER]; - extern __shared__ __align__(16) uint4 cap_slice[]; // [cols / 8] this CTA's updated chunks (cap_out only) - __shared__ alignas(8) unsigned long long xbar; - cg::cluster_group cluster = cg::this_cluster(); - if constexpr ((OPTS & 2) != 0) - { - if (threadIdx.x == 0) - { - uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); - asm volatile("mbarrier.init.shared::cta.b64 [%0], 1;" ::"r"(b) : "memory"); - asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); - } - asm volatile("barrier.cluster.arrive.relaxed.aligned;" ::: "memory"); - } - int const crank = static_cast(cluster.block_rank()); - int const t = blockIdx.x / CLUSTER; - int const cols = H / CLUSTER; - int const c0 = crank * cols + threadIdx.x * 8; - bool const active = t < M && threadIdx.x * 8 < cols; - int64_t const half_words = static_cast(MAXT) * world * (H / 2); - - // prepushed 1: the producer (m3_mx_gemv) already pushed this rank's rows into the exchange slots; 2: and it - // launched this grid only after its pushes (its grid wait passed), so the slots are polled without a grid wait; - // 3: as 2, with x in set A and x2 in set B. -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - if (prepushed < 2) - asm volatile("griddepcontrol.wait;" ::: "memory"); -#endif - int const parity = *reinterpret_cast(flags + blockIdx.x) & 1; - uint4 rraw = make_uint4(0u, 0u, 0u, 0u); - uint4 wraw = make_uint4(0u, 0u, 0u, 0u); - if constexpr ((OPTS & 1) != 0) - { - // the residual row was written grids earlier (prepushed >= 2: the producer launched this grid only after its - // own grid wait; otherwise the grid wait above) - if (active) - { - rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); - wraw = *reinterpret_cast(norm_w + c0); - } - } - // Push this rank's slice into every rank's buffer. - if (active && !prepushed) - { - uint4 raw; - if (fin.fc2 != nullptr) - { - // MoE finalize: fp32 sum of the routing-weighted top-k expert rows plus the shared expert, rounded once. - float f[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; - if (fin.topk <= kFinMaxTopk) - { - // All top-k indices and scales, then every expert-row load in flight at once, then the same in-order - // FP32 accumulation as the loop below (one memory round trip instead of topk; bit-identical). - int prow[kFinMaxTopk]; - float sc[kFinMaxTopk]; -#pragma unroll - for (int k = 0; k < kFinMaxTopk; ++k) - { - prow[k] = -1; - sc[k] = 0.f; - if (k < fin.topk) - { - int const e = t * fin.topk + k; - prow[k] = fin.perm[e]; - sc[k] = fin.scale_fp32 ? static_cast(fin.scales)[e] - : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); - } - } - uint4 rv[kFinMaxTopk]; -#pragma unroll - for (int k = 0; k < kFinMaxTopk; ++k) - { - if (prow[k] >= 0) - rv[k] = *reinterpret_cast(fin.fc2 + static_cast(prow[k]) * H + c0); - } -#pragma unroll - for (int k = 0; k < kFinMaxTopk; ++k) - { - if (prow[k] < 0) - continue; - __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv[k]); -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const v = __bfloat1622float2(h[q]); - f[2 * q] += v.x * sc[k]; - f[2 * q + 1] += v.y * sc[k]; - } - } - } - else - for (int k = 0; k < fin.topk; ++k) - { - int const e = t * fin.topk + k; - int const prow = fin.perm[e]; - if (prow < 0) - continue; - float const sc = fin.scale_fp32 - ? static_cast(fin.scales)[e] - : __bfloat162float(static_cast<__nv_bfloat16 const*>(fin.scales)[e]); - uint4 const rv = *reinterpret_cast(fin.fc2 + static_cast(prow) * H + c0); - __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const v = __bfloat1622float2(h[q]); - f[2 * q] += v.x * sc; - f[2 * q + 1] += v.y * sc; - } - } - if (x2 != nullptr) - { - // As stock: the routed output is rounded to BF16 first, then the shared expert is added (one more - // BF16 rounding of the FP32 sum). - uint4 const rv = *reinterpret_cast(x2 + static_cast(t) * H + c0); - __nv_bfloat162 const* h = reinterpret_cast<__nv_bfloat162 const*>(&rv); -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const routed = __bfloat1622float2(__floats2bfloat162_rn(f[2 * q], f[2 * q + 1])); - float2 const v = __bfloat1622float2(h[q]); - f[2 * q] = routed.x + v.x; - f[2 * q + 1] = routed.y + v.y; - } - } - __nv_bfloat162* o = reinterpret_cast<__nv_bfloat162*>(&raw); -#pragma unroll - for (int q = 0; q < 4; ++q) - o[q] = __floats2bfloat162_rn(f[2 * q], f[2 * q + 1]); - } - else - { - raw = *reinterpret_cast(x + static_cast(t) * H + c0); - } - if (fin.fc2 == nullptr && x2 != nullptr) - { - uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); - __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); - __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const fa = __bfloat1622float2(a[q]); - float2 const fb = __bfloat1622float2(b[q]); - a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); - } - } - int64_t const slot = parity * half_words + (static_cast(t) * world + rank) * (H / 2) + c0 / 2; - *reinterpret_cast(ws_mc + slot) - = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), - static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); - } -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.launch_dependents;"); -#endif - float s[8] = {}; - float sq = 0.f; - if (active) - { - // The residual row (written grids earlier) and every rank's slot are loaded together; slots still empty are - // re-polled. Ranks are then summed in rank order, as before. - if constexpr ((OPTS & 1) == 0) - rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); - constexpr int kMaxWorld = 8; - int* const p0 = ws_uc + parity * half_words + static_cast(t) * world * (H / 2) + c0 / 2; - int4 v[kMaxWorld]; - int4 v2[kMaxWorld]; - bool const two = prepushed == 3; - int* const p1 = p0 + 2 * half_words; // set B - uint32_t pending = (1u << world) - 1u; - uint32_t pending2 = two ? pending : 0u; - long long spins = 0; - while ((pending | pending2) != 0u) - { -#pragma unroll - for (int r = 0; r < kMaxWorld; ++r) - { - if ((pending >> r) & 1u) - v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); - if ((pending2 >> r) & 1u) - v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); - } -#pragma unroll - for (int r = 0; r < kMaxWorld; ++r) - { - if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty - && v[r].w != kEmpty) - pending &= ~(1u << r); - if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty - && v2[r].w != kEmpty) - pending2 &= ~(1u << r); - } - if (++spins > (1LL << 31)) - __trap(); - } -#pragma unroll - for (int r = 0; r < kMaxWorld; ++r) - { - if (r < world) - { - *reinterpret_cast(p0 + static_cast(r) * (H / 2)) - = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); - uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), - static_cast(v[r].z), static_cast(v[r].w)}; - if (two) - { - *reinterpret_cast(p1 + static_cast(r) * (H / 2)) - = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); - uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), - static_cast(v2[r].z), static_cast(v2[r].w)}; -#pragma unroll - for (int q = 0; q < 4; ++q) - { - __nv_bfloat162 const u - = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); - ww[q] = *reinterpret_cast(&u); - } - } -#pragma unroll - for (int q = 0; q < 4; ++q) - { - s[2 * q] += bf16_lo(ww[q]); - s[2 * q + 1] += bf16_hi(ww[q]); - } - } - } - // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. -#pragma unroll - for (int i = 0; i < 8; ++i) - s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); - __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); - uint32_t upd[4]; -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const rf = __bfloat1622float2(rr[q]); - __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); - upd[q] = *reinterpret_cast(&u); - if constexpr ((OPTS & 16) != 0) - { - float const x0 = s[2 * q] + rf.x, x1 = s[2 * q + 1] + rf.y; - s[2 * q] = x0; - s[2 * q + 1] = x1; - sq += x0 * x0 + x1 * x1; - } - else - { - float2 const uf = __bfloat1622float2(u); - s[2 * q] = uf.x; - s[2 * q + 1] = uf.y; - sq += uf.x * uf.x + uf.y * uf.y; - } - } - *reinterpret_cast(out_residual + static_cast(t) * H + c0) - = make_uint4(upd[0], upd[1], upd[2], upd[3]); - if constexpr ((OPTS & 16) == 0) - { - if constexpr (CAP) - cap_slice[threadIdx.x] = make_uint4(upd[0], upd[1], upd[2], upd[3]); - } - } - if constexpr ((OPTS & 4) != 0) - { - // every CTA (active or not) rotates its call parity once its own poll is done (read only by later grids) - if (threadIdx.x == 0) - flags[blockIdx.x] = parity ^ 1; - } -#pragma unroll - for (int off = 16; off > 0; off >>= 1) - sq += __shfl_xor_sync(0xffffffffu, sq, off); - if ((threadIdx.x & 31) == 0) - warp_sq[threadIdx.x / 32] = sq; - __syncthreads(); - if constexpr ((OPTS & 2) != 0) - asm volatile("barrier.cluster.wait.aligned;" ::: "memory"); // peers' xbar initialized (long since) - if (threadIdx.x == 0) - { - float total = 0.f; -#pragma unroll - for (int w = 0; w < THREADS / 32; ++w) - total += warp_sq[w]; - if constexpr ((OPTS & 2) == 0) - for (int c = 0; c < CLUSTER; ++c) - *cluster.map_shared_rank(&cluster_sq[crank], c) = total; - else - { - uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); - uint32_t const d = static_cast(__cvta_generic_to_shared(&cluster_sq[crank])); - asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(b), "r"(CLUSTER * 4) - : "memory"); - for (int c = 0; c < CLUSTER; ++c) - { - uint32_t rd, rb; - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rd) : "r"(d), "r"(c)); - asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(rb) : "r"(b), "r"(c)); - asm volatile("st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" ::"r"(rd), - "r"(__float_as_uint(total)), "r"(rb) - : "memory"); - } - } - } - if constexpr ((OPTS & 2) == 0) - cluster.sync(); - else - { - // every peer's sum landed (st.async complete_tx on this CTA's barrier); no cluster-wide release-arrive - uint32_t const b = static_cast(__cvta_generic_to_shared(&xbar)); - uint32_t done = 0; - while (!done) - asm volatile( - "{\n .reg .pred p;\n mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 p, [%1], 0;\n" - " selp.u32 %0, 1, 0, p;\n}" - : "=r"(done) - : "r"(b) - : "memory"); - } - if constexpr ((OPTS & 16) == 0) - { - if constexpr (CAP) - { - // The row's chunks are spread over the cluster (CTA c holds chunks [c * per, (c + 1) * per)); every CTA - // forms capture_norm's statistic itself, its 128 threads standing in for capture threads ct and ct + 128. - __shared__ float cap_red[kCapWarps]; - int const per = cols / 8; - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - float ca = 0.f, cb = 0.f; -#pragma unroll - for (int it = 0; it < kCapIters; ++it) - { - int const ka = it * kCapThreads + static_cast(threadIdx.x); - int const kb = ka + kCapThreads / 2; - ca = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[ka % per], ka / per), ca); - cb = cap_chunk_sq(*cluster.map_shared_rank(&cap_slice[kb % per], kb / per), cb); - } - // Every peer read is done before any CTA moves on (and later exits). - asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); - asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); - ca = cap_warp_sum(ca); - cb = cap_warp_sum(cb); - if ((threadIdx.x & 31) == 0) - { - cap_red[threadIdx.x / 32] = ca; - cap_red[threadIdx.x / 32 + kCapWarps / 2] = cb; - } - __syncthreads(); - float tot = 0.f; -#pragma unroll - for (int w = 0; w < kCapWarps; ++w) - tot += cap_red[w]; - float const rs = rsqrtf(tot / H + cap_eps); - if (active) - cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); - } - } - if (active) - { - float total = 0.f; -#pragma unroll - for (int c = 0; c < CLUSTER; ++c) - total += cluster_sq[c]; - float const inv = rsqrtf(total / static_cast(H) + eps); - if constexpr ((OPTS & 1) == 0) - wraw = *reinterpret_cast(norm_w + c0); - __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); - uint32_t o[4]; -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const wf = __bfloat1622float2(wv[q]); - __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); - o[q] = *reinterpret_cast(&n); - float2 const nf = __bfloat1622float2(n); - s[2 * q] = nf.x; - s[2 * q + 1] = nf.y; - } - *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); - } - if (out_q != nullptr) - { - // MXFP8 copy of the normed row for the next GEMM, as the CuTe DSL mxfp8 quantize of `normed` (the tree's - // decode quantize, MX-TAC; see mxfp8_ue8m0): four threads hold one 32-column scale block. - float amax = 0.f; - if (active) - { -#pragma unroll - for (int e = 0; e < 8; ++e) - amax = fmaxf(amax, fabsf(s[e])); - } - amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); - amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); - if (active) - { - uint32_t const code = mxfp8_ue8m0(amax); - float const inv_scale = mxfp8_inv_scale(code); - uint32_t packed[2]; -#pragma unroll - for (int h = 0; h < 2; ++h) - { - packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) - | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); - } - *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); - if ((threadIdx.x & 3) == 0) - { - int const kb = c0 / 32; - out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 - + ((t % 128) / 32) * 4 + kb % 4] - = static_cast(code); - } - } - } - if constexpr ((OPTS & 4) == 0) - { - __syncthreads(); - if (threadIdx.x == 0) - flags[blockIdx.x] = parity ^ 1; - } -#elif defined(__CUDA_ARCH__) - __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op -#endif -} - -// Refuses (raises) to launch a kernel whose loaded binary holds no sm_10x code (PTX and SASS version >= 100): the -// device passes for other architectures compile the kernels to a trap, and a build that left out every sm_10x target -// must not reach the GPU at all. -template -bool checkBinary(Kernel kernel, char const* name) -{ - cudaFuncAttributes attr{}; - TLLM_CUDA_CHECK(cudaFuncGetAttributes(&attr, reinterpret_cast(kernel))); - TLLM_CHECK_WITH_INFO(attr.ptxVersion >= 100 && attr.binaryVersion >= 100, - "%s: the loaded binary holds no sm_10x code (PTX %d, SASS %d): rebuild with an sm_100f / sm_107a target", name, - attr.ptxVersion, attr.binaryVersion); - return true; -} - -template -void launch(MinimaxM3Sm107Tp8McAllreduceNormParams const& p, bool pdl, cudaStream_t stream) -{ - int const M = p.numTokens, H = p.hidden; - TLLM_CHECK_WITH_INFO(H % (CLUSTER * 8) == 0 && H / CLUSTER / 8 <= THREADS, "unsupported hidden size %d", H); - TLLM_CHECK_WITH_INFO(M >= 1 && M <= MAXT, "%d tokens do not fit the %d-token bucket", M, MAXT); - auto kernel = p.capOut != nullptr ? mc_allreduce_norm_kernel - : mc_allreduce_norm_kernel; - static bool const binaryOk - = checkBinary(mc_allreduce_norm_kernel, "mc_allreduce_norm_kernel") - && checkBinary( - mc_allreduce_norm_kernel, "mc_allreduce_norm_kernel (capture)"); - (void) binaryOk; - // Both instantiations (with and without the capture epilogue) get the same attributes, whichever runs first. - static bool const configured = []() - { - auto setup = [](auto k) - { - bool ok = CLUSTER <= 8 - || cudaFuncSetAttribute(k, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) == cudaSuccess; - if ((OPTS & 8) != 0) - ok = ok - && cudaFuncSetAttribute( - k, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared) - == cudaSuccess; - return ok; - }; - bool const ok = setup(mc_allreduce_norm_kernel) - && setup(mc_allreduce_norm_kernel); - if (!ok) - cudaGetLastError(); - return ok; - }(); - TLLM_CHECK_WITH_INFO(configured, "cluster size not supported"); - FinalizeArgs const fin{p.fc2, p.perm, p.scales, p.scaleFp32 ? 1 : 0, p.topk}; - cudaLaunchConfig_t cfg = {}; - // The grid is sized to the call's rows: min(M, MAXT) clusters. Each (row, slice) slot keeps its own parity flag, - // read by the producers before they push and rotated only by the slot's CTA after it polled: a slot that a call - // does not use keeps its parity, and every rank runs the same calls, so the parities stay consistent when M - // changes between calls. - cfg.gridDim = dim3((M < MAXT ? M : MAXT) * CLUSTER); - cfg.blockDim = dim3(THREADS); - cfg.stream = stream; - if (p.capOut != nullptr) - { - TLLM_CHECK_WITH_INFO((OPTS & 16) == 0 && H == kCapHidden && 2 * THREADS == kCapThreads, - "the capture epilogue takes H %d on the 128-thread cluster kernel, not the drafter boundary", kCapHidden); - cfg.dynamicSmemBytes = static_cast(H / CLUSTER / 8) * sizeof(uint4); - } - cudaLaunchAttribute attr[2]; - attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - attr[1].id = cudaLaunchAttributeClusterDimension; - attr[1].val.clusterDim.x = CLUSTER; - attr[1].val.clusterDim.y = 1; - attr[1].val.clusterDim.z = 1; - cfg.attrs = attr; - cfg.numAttrs = 2; - cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, fin, M, H, p.residual, p.normWeight, p.eps, - p.wsUc, p.wsMc, p.flags, p.rank, p.world, p.outNormed, p.outResidual, p.outQuant, p.outSf, (H / 32 + 3) / 4, - p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); - if (err != cudaSuccess) - { - cudaGetLastError(); - TLLM_CUDA_CHECK(err); - } -} - -// ws_uc / ws_mc are this rank's unicast and multicast views of ONE exchange buffer (they alias): no __restrict__. -// prepushed >= 2 contract: the producer (m3_mx_gemv's o_proj push, the m3_moe engine's finalize push) passes its own -// grid dependency, pushes its rows into the exchange, and only then releases this grid (griddepcontrol -// .launch_dependents); this kernel therefore skips griddepcontrol.wait and reads only the pushed (Lamport-polled) -// words, its parity flags and the residual / norm weight, which were written two or more grids earlier. -// -// The TP8 row kernel: one CTA of H / 8 = 768 threads per token row, used for every call but the MoE finalize. Its -// residual add + RMSNorm is TRT-LLM's rms_norm_kernel (customAllReduceKernels.cu), the stock fused norm after the -// all-reduce on this target, operation for operation at H 6144: thread j owns columns [8j, 8j + 8), its sum of squares -// is acc += v * v over them in order, then the warp butterfly 16..1, the 24 warp sums through shared memory into a -// second 32-lane butterfly (lanes >= 24 zero), rsqrtf(acc / H + eps) and (v * denom) * w. With the BF16 rank sum -// before the residual add, the normed rows, the updated residual and the MXFP8 copy are bitwise the stock ones. The -// exchange slots and the per-(row, slice) parity flags keep the cluster kernel's layout (slice = 96 threads), so -// producers that push (the o_proj GEMV) are unchanged. -constexpr int kSlices = 8; -constexpr int kSliceThreads = 96; -constexpr int kRowThreads = 768; - -// grid = M CTAs (one per row), THREADS = H / 8 threads. WORLD ranks (compile time). DRAFTER (opts bit 16): the Eagle3 -// drafter boundary, whose stock is a plain BF16 all-reduce then FlashInfer fused_add_rmsnorm: the RMS statistic and -// the normalization take the unrounded FP32 sum bf16(rank sum) + residual (only the stored residual is rounded). -// ws_uc / ws_mc alias (one exchange buffer): no __restrict__. -template -__global__ void __launch_bounds__(THREADS, 1) mc_allreduce_norm_row_kernel(__nv_bfloat16 const* __restrict__ x, - __nv_bfloat16 const* __restrict__ x2, int H, __nv_bfloat16 const* __restrict__ residual, - __nv_bfloat16 const* __restrict__ norm_w, float eps, int* ws_uc, int* ws_mc, int* __restrict__ flags, int rank, - int maxt, __nv_bfloat16* __restrict__ out_normed, __nv_bfloat16* __restrict__ out_residual, - uint8_t* __restrict__ out_q, uint8_t* __restrict__ out_sf, int sf_k_tiles, int prepushed, - __nv_bfloat16* __restrict__ cap_out, int64_t cap_stride, __nv_bfloat16 const* __restrict__ cap_w, float cap_eps) -{ -#if MINIMAX_M3_KERNELS_ENABLED - __shared__ float warp_sq[THREADS / 32]; - __shared__ float total_sh; - extern __shared__ __align__(16) uint4 cap_row[]; // [THREADS] the updated row's chunks (cap_out only) - int const t = blockIdx.x; - int const j = threadIdx.x; - int const c0 = j * 8; - int const slice = j / kSliceThreads; - int64_t const half_words = static_cast(maxt) * WORLD * (H / 2); - if (prepushed < 2) - asm volatile("griddepcontrol.wait;" ::: "memory"); - int const parity = *reinterpret_cast(flags + t * kSlices + slice) & 1; - uint4 const rraw = *reinterpret_cast(residual + static_cast(t) * H + c0); - uint4 const wraw = *reinterpret_cast(norm_w + c0); - if (!prepushed) - { - uint4 raw = *reinterpret_cast(x + static_cast(t) * H + c0); - if (x2 != nullptr) - { - uint4 const raw2 = *reinterpret_cast(x2 + static_cast(t) * H + c0); - __nv_bfloat162* a = reinterpret_cast<__nv_bfloat162*>(&raw); - __nv_bfloat162 const* b = reinterpret_cast<__nv_bfloat162 const*>(&raw2); -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const fa = __bfloat1622float2(a[q]); - float2 const fb = __bfloat1622float2(b[q]); - a[q] = __floats2bfloat162_rn(fa.x + fb.x, fa.y + fb.y); - } - } - int64_t const slot = parity * half_words + (static_cast(t) * WORLD + rank) * (H / 2) + c0 / 2; - *reinterpret_cast(ws_mc + slot) - = make_int4(static_cast(canon_bf16x2(raw.x)), static_cast(canon_bf16x2(raw.y)), - static_cast(canon_bf16x2(raw.z)), static_cast(canon_bf16x2(raw.w))); - } - asm volatile("griddepcontrol.launch_dependents;"); - float s[8] = {}; - float sq = 0.f; - { - int* const p0 = ws_uc + parity * half_words + static_cast(t) * WORLD * (H / 2) + c0 / 2; - int* const p1 = p0 + 2 * half_words; // set B - bool const two = prepushed == 3; - int4 v[WORLD]; - int4 v2[WORLD]; - uint32_t pending = (1u << WORLD) - 1u; - uint32_t pending2 = two ? pending : 0u; - long long spins = 0; - while ((pending | pending2) != 0u) - { -#pragma unroll - for (int r = 0; r < WORLD; ++r) - { - if ((pending >> r) & 1u) - v[r] = ld_volatile_v4(p0 + static_cast(r) * (H / 2)); - if ((pending2 >> r) & 1u) - v2[r] = ld_volatile_v4(p1 + static_cast(r) * (H / 2)); - } -#pragma unroll - for (int r = 0; r < WORLD; ++r) - { - if (((pending >> r) & 1u) && v[r].x != kEmpty && v[r].y != kEmpty && v[r].z != kEmpty - && v[r].w != kEmpty) - pending &= ~(1u << r); - if (((pending2 >> r) & 1u) && v2[r].x != kEmpty && v2[r].y != kEmpty && v2[r].z != kEmpty - && v2[r].w != kEmpty) - pending2 &= ~(1u << r); - } - if (++spins > (1LL << 31)) - __trap(); - } -#pragma unroll - for (int r = 0; r < WORLD; ++r) - { - *reinterpret_cast(p0 + static_cast(r) * (H / 2)) - = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); - uint32_t ww[4] = {static_cast(v[r].x), static_cast(v[r].y), - static_cast(v[r].z), static_cast(v[r].w)}; - if (two) - { - *reinterpret_cast(p1 + static_cast(r) * (H / 2)) - = make_int4(kEmpty, kEmpty, kEmpty, kEmpty); - uint32_t const w2[4] = {static_cast(v2[r].x), static_cast(v2[r].y), - static_cast(v2[r].z), static_cast(v2[r].w)}; -#pragma unroll - for (int q = 0; q < 4; ++q) - { - __nv_bfloat162 const u - = __floats2bfloat162_rn(bf16_lo(ww[q]) + bf16_lo(w2[q]), bf16_hi(ww[q]) + bf16_hi(w2[q])); - ww[q] = *reinterpret_cast(&u); - } - } -#pragma unroll - for (int q = 0; q < 4; ++q) - { - s[2 * q] += bf16_lo(ww[q]); - s[2 * q + 1] += bf16_hi(ww[q]); - } - } - // Stock rounding points: the rank sum is a BF16 all-reduce output before the residual add. -#pragma unroll - for (int i = 0; i < 8; ++i) - s[i] = __bfloat162float(__float2bfloat16_rn(s[i])); - __nv_bfloat162 const* rr = reinterpret_cast<__nv_bfloat162 const*>(&rraw); - uint32_t upd[4]; -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const rf = __bfloat1622float2(rr[q]); - __nv_bfloat162 const u = __floats2bfloat162_rn(s[2 * q] + rf.x, s[2 * q + 1] + rf.y); - upd[q] = *reinterpret_cast(&u); - if constexpr (DRAFTER) - { - s[2 * q] += rf.x; - s[2 * q + 1] += rf.y; - } - else - { - float2 const uf = __bfloat1622float2(u); - s[2 * q] = uf.x; - s[2 * q + 1] = uf.y; - } - } -#pragma unroll - for (int e = 0; e < 8; ++e) - sq += s[e] * s[e]; - *reinterpret_cast(out_residual + static_cast(t) * H + c0) - = make_uint4(upd[0], upd[1], upd[2], upd[3]); - if constexpr (!DRAFTER) - { - if constexpr (CAP) - cap_row[j] = make_uint4(upd[0], upd[1], upd[2], upd[3]); - } - } - // TRT-LLM reduce_fusion::block_reduce_sum. - sq += __shfl_xor_sync(~0, sq, 16); - sq += __shfl_xor_sync(~0, sq, 8); - sq += __shfl_xor_sync(~0, sq, 4); - sq += __shfl_xor_sync(~0, sq, 2); - sq += __shfl_xor_sync(~0, sq, 1); - if ((j & 31) == 0) - warp_sq[j / 32] = sq; - __syncthreads(); - { - float val = (j & 31) < THREADS / 32 ? warp_sq[j & 31] : 0.f; - val += __shfl_xor_sync(~0, val, 16); - val += __shfl_xor_sync(~0, val, 8); - val += __shfl_xor_sync(~0, val, 4); - val += __shfl_xor_sync(~0, val, 2); - val += __shfl_xor_sync(~0, val, 1); - if (j == 0) - total_sh = val; - } - // The parity of this row's slices rotates once every thread has read the exchange (later grids only). - if (j % kSliceThreads == 0) - flags[t * kSlices + slice] = parity ^ 1; - __syncthreads(); - if constexpr (!DRAFTER) - { - if constexpr (CAP) - { - // The whole row is in this CTA: threads 0..255 run capture_norm's chains over the staged chunks. - __shared__ float cap_red[kCapWarps]; - if (j < kCapThreads) - { - float css = 0.f; -#pragma unroll - for (int it = 0; it < kCapIters; ++it) - css = cap_chunk_sq(cap_row[it * kCapThreads + j], css); - css = cap_warp_sum(css); - if ((j & 31) == 0) - cap_red[j / 32] = css; - } - __syncthreads(); - float tot = 0.f; -#pragma unroll - for (int w = 0; w < kCapWarps; ++w) - tot += cap_red[w]; - float const rs = rsqrtf(tot / H + cap_eps); - cap_store(s, rs, cap_w + c0, cap_out + static_cast(t) * cap_stride + c0); - } - } - float const inv = rsqrtf(total_sh / H + eps); - __nv_bfloat162 const* wv = reinterpret_cast<__nv_bfloat162 const*>(&wraw); - uint32_t o[4]; -#pragma unroll - for (int q = 0; q < 4; ++q) - { - float2 const wf = __bfloat1622float2(wv[q]); - __nv_bfloat162 const n = __floats2bfloat162_rn(s[2 * q] * inv * wf.x, s[2 * q + 1] * inv * wf.y); - o[q] = *reinterpret_cast(&n); - float2 const nf = __bfloat1622float2(n); - s[2 * q] = nf.x; - s[2 * q + 1] = nf.y; - } - *reinterpret_cast(out_normed + static_cast(t) * H + c0) = make_uint4(o[0], o[1], o[2], o[3]); - if (out_q != nullptr) - { - float amax = 0.f; -#pragma unroll - for (int e = 0; e < 8; ++e) - amax = fmaxf(amax, fabsf(s[e])); - amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); - amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2)); - uint32_t const code = mxfp8_ue8m0(amax); - float const inv_scale = mxfp8_inv_scale(code); - uint32_t packed[2]; -#pragma unroll - for (int h = 0; h < 2; ++h) - { - packed[h] = mxfp8_e4m3x2(s[4 * h], s[4 * h + 1], inv_scale) - | (mxfp8_e4m3x2(s[4 * h + 2], s[4 * h + 3], inv_scale) << 16); - } - *reinterpret_cast(out_q + static_cast(t) * H + c0) = make_uint2(packed[0], packed[1]); - if ((j & 3) == 0) - { - int const kb = c0 / 32; - out_sf[static_cast(t / 128) * sf_k_tiles * 512 + (kb / 4) * 512 + (t % 32) * 16 - + ((t % 128) / 32) * 4 + kb % 4] - = static_cast(code); - } - } -#elif defined(__CUDA_ARCH__) - __trap(); // compiled for an architecture without the sm_10x features: never a silent no-op -#endif -} - -void launchRow( - MinimaxM3Sm107Tp8McAllreduceNormParams const& p, int32_t maxTokens, bool drafter, bool pdl, cudaStream_t stream) -{ - TLLM_CHECK_WITH_INFO(p.hidden == kRowThreads * 8 && p.world == 2, "the row kernel takes H 6144 at 2 ranks"); - static bool const binaryOk - = checkBinary(mc_allreduce_norm_row_kernel<2, kRowThreads, false, false>, "mc_allreduce_norm_row_kernel") - && checkBinary( - mc_allreduce_norm_row_kernel<2, kRowThreads, true, false>, "mc_allreduce_norm_row_kernel (drafter)") - && checkBinary( - mc_allreduce_norm_row_kernel<2, kRowThreads, false, true>, "mc_allreduce_norm_row_kernel (capture)"); - (void) binaryOk; - TLLM_CHECK_WITH_INFO(p.numTokens >= 1 && p.numTokens <= maxTokens, "%d tokens do not fit the %d-token bucket", - p.numTokens, maxTokens); - cudaLaunchConfig_t cfg = {}; - cfg.gridDim = dim3(p.numTokens); - cfg.blockDim = dim3(kRowThreads); - cfg.stream = stream; - if (p.capOut != nullptr) - { - TLLM_CHECK_WITH_INFO( - !drafter && p.hidden == kCapHidden, "the capture epilogue takes H %d, not the drafter", kCapHidden); - cfg.dynamicSmemBytes = static_cast(kRowThreads) * sizeof(uint4); - } - cudaLaunchAttribute attr[1]; - attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - cfg.attrs = attr; - cfg.numAttrs = 1; - // The capture epilogue never runs on the drafter boundary (checked above), so there are three instantiations. - auto kernel = drafter ? mc_allreduce_norm_row_kernel<2, kRowThreads, true, false> - : p.capOut != nullptr ? mc_allreduce_norm_row_kernel<2, kRowThreads, false, true> - : mc_allreduce_norm_row_kernel<2, kRowThreads, false, false>; - cudaError_t const err = cudaLaunchKernelEx(&cfg, kernel, p.x, p.x2, p.hidden, p.residual, p.normWeight, p.eps, - p.wsUc, p.wsMc, p.flags, p.rank, maxTokens, p.outNormed, p.outResidual, p.outQuant, p.outSf, - (p.hidden / 32 + 3) / 4, p.prepushed, p.capOut, p.capStride, p.capWeight, p.capEps); - if (err != cudaSuccess) - { - cudaGetLastError(); - TLLM_CUDA_CHECK(err); - } -} - -constexpr int kCluster = 8; -constexpr int kThreads = 128; - -} // namespace - -bool isMinimaxM3Sm107Tp8McAllreduceNormSupported(int32_t maxTokens, int32_t opts) -{ - return (maxTokens == 8 && (opts == 7 || opts == 15 || opts == 23 || opts == 31)) || (maxTokens == 16 && opts == 15) - || (maxTokens == 48 && opts == 15); -} - -void invokeMinimaxM3Sm107Tp8McAllreduceNorm(MinimaxM3Sm107Tp8McAllreduceNormParams const& params, int32_t maxTokens, - int32_t opts, bool pdl, cudaStream_t stream) -{ - if (params.fc2 == nullptr && params.world == 2 && params.hidden == kRowThreads * 8) - { - launchRow(params, maxTokens, (opts & 16) != 0, pdl, stream); - return; - } - if (maxTokens == 8 && opts == 7) - { - launch(params, pdl, stream); - } - else if (maxTokens == 8 && opts == 15) - { - launch(params, pdl, stream); - } - else if (maxTokens == 8 && opts == 23) - { - launch(params, pdl, stream); - } - else if (maxTokens == 8 && opts == 31) - { - launch(params, pdl, stream); - } - else if (maxTokens == 16 && opts == 15) - { - launch(params, pdl, stream); - } - else if (maxTokens == 48 && opts == 15) - { - launch(params, pdl, stream); - } - else - { - TLLM_THROW("unsupported (bucket %d, opts %d)", maxTokens, opts); - } -} - -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu deleted file mode 100644 index 3062080fac2c..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.cu +++ /dev/null @@ -1,340 +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/minimaxM3Sm107Tp8/producerPw8.h" - -#include -#include -#include - -// MiniMax-M3 NVFP4 QKV + indexer norm/RoPE/KV-insert producer with an early PDL trigger, at the SM107 TP8 heads. -// -// The SM107 TP8 target's copy of the SM107 TP2 producerPw2.cu (itself a copy of minimaxM3/producerPw.cu, which the -// sm_103 targets keep), for its pre-wait step inputs. Same math and -// outputs as minimaxM3Nvfp4QKVIndexerNormRopeKVInsertKernel (fusedQKNormRopeKernel.cu, built with --use_fast_math like -// the rest of the library), specialized to 8 query, 1 KV and 1 index head per rank (one index-K head), head_dim 128, -// rotary_dim 64, page 128, NVFP4 main KV and FP8 index-K. One warp per (token, packed head); K and V rows go to their -// head's block of the HND pools (the native kernel's head strides). -// -// The decode producer is latency bound: after griddepcontrol.wait the native kernel runs a chain of dependent global -// loads (packed row -> norm weight -> position id -> RoPE row -> cache slot -> quant scale). This kernel reads the -// constant norm weights and KV quant scale before griddepcontrol.wait and issues the packed row, position id and cache -// slot loads together after it, followed by the RoPE row; with `prewaitStepInputs` the position id, the slot and the -// RoPE row are read before the wait too, leaving only the packed row after it. With `early` it also calls -// griddepcontrol.launch_dependents on entry: consumers read its outputs only after their own griddepcontrol.wait, which -// waits for this grid to complete, so the early trigger only lets the next kernel launch and run its pre-wait prologue -// while this one runs. - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ -namespace producerPw8 -{ -namespace -{ - -constexpr int kHeadDim = 128; -constexpr int kRotaryDim = 64; -constexpr int kPageSize = 128; -constexpr int kElems = kHeadDim / 32; -constexpr int kBlock = 256; - -__device__ __forceinline__ float rcp_approx_ftz(float a) -{ - float b; - asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); - return b; -} - -__device__ __forceinline__ uint32_t sf_group_mask4() -{ - uint32_t laneId; - asm("mov.u32 %0, %%laneid;" : "=r"(laneId)); - return 0xFu << (laneId & ~3u); -} - -__device__ __forceinline__ uint32_t fp32x8_to_e2m1(float (&a)[8]) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - uint32_t val; - asm volatile( - "{\n" - ".reg .b8 byte0;\n" - ".reg .b8 byte1;\n" - ".reg .b8 byte2;\n" - ".reg .b8 byte3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" - "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" - "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" - "}" - : "=r"(val) - : "f"(a[0]), "f"(a[1]), "f"(a[2]), "f"(a[3]), "f"(a[4]), "f"(a[5]), "f"(a[6]), "f"(a[7])); - return val; -#else - return 0; -#endif -} - -__device__ __forceinline__ void store_fp8x4(__nv_fp8_e4m3* out, int64_t off, float const (&e)[kElems]) -{ - __nv_fp8x2_e4m3 const low(make_float2(e[0], e[1])); - __nv_fp8x2_e4m3 const high(make_float2(e[2], e[3])); - *reinterpret_cast(out + off) = static_cast(low.__x) | (static_cast(high.__x) << 16); -} - -__device__ __forceinline__ uint16_t quantize_fp4x4(float const (&e)[4], float secondLevelScale, uint8_t* scaleOut) -{ - __nv_bfloat16 rounded[4]; - float localMax = 0.0F; -#pragma unroll - for (int i = 0; i < 4; ++i) - { - rounded[i] = __float2bfloat16_rn(e[i]); - localMax = fmaxf(localMax, fabsf(__bfloat162float(rounded[i]))); - } - uint32_t const groupMask = sf_group_mask4(); - localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 1)); - localMax = fmaxf(localMax, __shfl_xor_sync(groupMask, localMax, 2)); - float const scaleValue = secondLevelScale * (localMax * rcp_approx_ftz(6.0F)); - __nv_fp8_e4m3 const encodedScale(scaleValue); - if (scaleOut != nullptr) - *scaleOut = encodedScale.__x; - float const narrowedScale = static_cast(encodedScale); - float const outputScale - = localMax != 0.0F ? rcp_approx_ftz(narrowedScale * rcp_approx_ftz(secondLevelScale)) : 0.0F; - float q[8] = {}; -#pragma unroll - for (int i = 0; i < 4; ++i) - q[i] = __bfloat162float(rounded[i]) * outputScale; - return static_cast(fp32x8_to_e2m1(q)); -} - -template -__global__ void __launch_bounds__(kBlock) m3_producer_kernel(Params p) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - if constexpr (kEarly) - { - asm volatile("griddepcontrol.launch_dependents;"); - } -#endif - int const warpId = threadIdx.x / 32; - int const laneId = threadIdx.x % 32; - int const globalWarp = blockIdx.x * (kBlock / 32) + warpId; - int const tokenIdx = globalWarp / kTotalHeads; - int const localHead = globalWarp % kTotalHeads; - if (tokenIdx >= p.numTokens) - return; - - constexpr int kBegin = kHeadsQ; - constexpr int vBegin = kBegin + kHeadsKV; - constexpr int indexQBegin = vBegin + kHeadsKV; - constexpr int indexKHead = indexQBegin + kHeadsIndex; - bool const isQ = localHead < kBegin; - bool const isK = localHead >= kBegin && localHead < vBegin; - bool const isV = localHead >= vBegin && localHead < indexQBegin; - bool const isIndexQ = localHead >= indexQBegin && localHead < indexKHead; - bool const isIndexK = localHead == indexKHead; - - // Constant data read before the wait: norm weights and the KV quant scale. - float w[kElems] = {}; - float quantScale = 0.0F; - __nv_bfloat16 const* normWeight - = isQ ? p.qWeight : (isK ? p.kWeight : (isIndexQ ? p.indexQWeight : p.indexKWeight)); - if (!isV) - { - uint2 const wv = *reinterpret_cast(normWeight + laneId * kElems); - float2 const w01 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.x)); - float2 const w23 = __bfloat1622float2(*reinterpret_cast<__nv_bfloat162 const*>(&wv.y)); - w[0] = w01.x; - w[1] = w01.y; - w[2] = w23.x; - w[3] = w23.y; - } - if (isK || isV) - { - quantScale = p.kvQuantScale[isV ? 2 : 1]; - } - - int64_t const inputOffset = (static_cast(tokenIdx) * kTotalHeads + localHead) * kHeadDim + laneId * kElems; - bool const rotates = !isV && laneId * kElems < kRotaryDim; - bool const slotted = isK || isV || isIndexK; - int const coefficient = (laneId * kElems) % (kRotaryDim / 2); - // kPrewait: the position id, the cache slot and the RoPE row are read before the wait. They are per-step metadata - // written before the replay (host-to-device) or by plain-launched kernels earlier in the step (position_ids by the - // step's input preprocessing, the slots by msa_kv_update or its torch fallback), so they are final before this - // grid can start whatever its predecessors trigger; the RoPE table is constant. - int positionId = 0; - int slot = 0; - float4 c = make_float4(0.0F, 0.0F, 0.0F, 0.0F); - float4 s = c; - auto const loadStepInputs = [&]() - { - positionId = rotates ? p.positionIds[tokenIdx] : 0; - slot = slotted ? p.outCacheLoc[tokenIdx] : 0; - if (rotates) - { - int64_t const ropeRow = static_cast(positionId) * kRotaryDim; - c = *reinterpret_cast(p.rotaryCosSin + ropeRow + coefficient); - s = *reinterpret_cast(p.rotaryCosSin + ropeRow + kRotaryDim / 2 + coefficient); - } - }; - if constexpr (kPrewait) - { - loadStepInputs(); - } -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.wait;" ::: "memory"); -#endif - // The packed row comes from the predecessor grid; without kPrewait the step inputs are issued with it, so the - // native kernel's serial chain of dependent loads becomes two round trips. - uint2 const packed = *reinterpret_cast(p.packed + inputOffset); - if constexpr (!kPrewait) - { - loadStepInputs(); - } - float cosine[kElems] = {c.x, c.y, c.z, c.w}; - float sine[kElems] = {s.x, s.y, s.z, s.w}; - - float e[kElems]; - float sumSquares = 0.0F; -#pragma unroll - for (int pair = 0; pair < 2; ++pair) - { - float2 const v = __bfloat1622float2( - *reinterpret_cast<__nv_bfloat162 const*>(reinterpret_cast(&packed) + pair)); - e[2 * pair] = v.x; - e[2 * pair + 1] = v.y; - if (!isV) - sumSquares += v.x * v.x + v.y * v.y; - } - - if (!isV) - { -#pragma unroll - for (int mask = 16; mask > 0; mask >>= 1) - sumSquares += __shfl_xor_sync(0xffffffff, sumSquares, mask, 32); - float const rmsReciprocal = rsqrtf(sumSquares / static_cast(kHeadDim) + p.eps); -#pragma unroll - for (int i = 0; i < kElems; ++i) - { - e[i] *= rmsReciprocal * (1.0F + w[i]); - } - constexpr int kPairOffset = (kRotaryDim / 2) / kElems; -#pragma unroll - for (int i = 0; i < kElems; ++i) - { - int const dim = laneId * kElems + i; - float paired = __shfl_xor_sync(0xffffffff, e[i], kPairOffset); - if (dim < kRotaryDim) - { - if (dim < kRotaryDim / 2) - paired = -paired; - // The native kernel's contraction of e * cos + paired * sin; spelled out so this build matches it. - e[i] = __fmaf_rn(paired, sine[i], __fmul_rn(e[i], cosine[i])); - } - } - } - - if (isQ) - { - store_fp8x4(p.qOut, (static_cast(tokenIdx) * kHeadsQ + localHead) * kHeadDim + laneId * kElems, e); - return; - } - if (isIndexQ) - { - store_fp8x4(p.indexQOut, - (static_cast(tokenIdx) * kHeadsIndex + localHead - indexQBegin) * kHeadDim + laneId * kElems, e); - return; - } - - if (slot < 0) - return; - int const page = slot / kPageSize; - if (page >= p.numPages) - return; - int const withinPage = slot & (kPageSize - 1); - if (isIndexK) - { - store_fp8x4(p.indexK, - static_cast(page) * p.indexPageStride + static_cast(withinPage) * p.indexTokenStride - + laneId * kElems, - e); - return; - } - - int const head = isK ? localHead - kBegin : localHead - vBegin; - int const plane = isV ? 1 : 0; - int const scaleColumn = laneId / 4; - int64_t const scaleWithinPage = isV - ? static_cast(withinPage / 4) * 32 + static_cast(scaleColumn) * 4 + withinPage % 4 - : static_cast(withinPage) * 8 + scaleColumn; - int64_t const scaleOffset = static_cast(page) * p.scalePageStride - + static_cast(plane) * p.scalePlaneStride + static_cast(head) * p.scaleHeadStride - + scaleWithinPage; - uint8_t* scaleOut = laneId % 4 == 0 ? p.kvScale + scaleOffset : nullptr; - uint16_t const packedFp4 = quantize_fp4x4(e, quantScale, scaleOut); - int64_t const dataOffset = static_cast(page) * p.dataPageStride - + static_cast(plane) * p.dataPlaneStride + static_cast(head) * p.dataHeadStride - + static_cast(withinPage) * p.dataTokenStride + laneId * 2; - *reinterpret_cast(p.kvData + dataOffset) = packedFp4; -#endif -} - -template -cudaError_t launch(Params const& p, int numTokens, bool pdl, cudaStream_t stream) -{ - constexpr int kTotalHeads = kHeadsQ + 2 * kHeadsKV + kHeadsIndex + 1; - int const warps = numTokens * kTotalHeads; - cudaLaunchConfig_t cfg = {}; - cfg.gridDim = dim3((warps + kBlock / 32 - 1) / (kBlock / 32)); - cfg.blockDim = dim3(kBlock); - cfg.stream = stream; - cudaLaunchAttribute attr[1]; - attr[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; - cfg.attrs = attr; - cfg.numAttrs = 1; - return cudaLaunchKernelEx(&cfg, m3_producer_kernel, p); -} - -} // namespace - -cudaError_t invokeProducerPw8( - Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream) -{ - if (numTokens <= 0) - return cudaSuccess; - if (prewaitStepInputs) - { - return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); - } - return early ? launch(p, numTokens, pdl, stream) : launch(p, numTokens, pdl, stream); -} - -} // namespace producerPw8 -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h deleted file mode 100644 index 5b269b94587c..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h +++ /dev/null @@ -1,76 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ -//! The MiniMax-M3 NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the SM107 TP8 heads (8 query, 1 KV and -//! 1 index head per rank) with its constant loads before the PDL grid-dependency wait and an optional early dependents -//! trigger. Same outputs, bit for bit, as minimaxM3Nvfp4QKVIndexerNormRopeKVInsert. See producerPw8.cu. -namespace producerPw8 -{ - -struct Params -{ - __nv_bfloat16 const* packed; - __nv_fp8_e4m3* qOut; - __nv_fp8_e4m3* indexQOut; - uint8_t* kvData; - uint8_t* kvScale; - __nv_fp8_e4m3* indexK; - int const* outCacheLoc; - float const* kvQuantScale; - int64_t dataPageStride, dataPlaneStride, dataHeadStride, dataTokenStride; - int64_t scalePageStride, scalePlaneStride, scaleHeadStride; - int64_t indexPageStride, indexTokenStride, numPages; - int numTokens; - float eps; - __nv_bfloat16 const* qWeight; - __nv_bfloat16 const* kWeight; - __nv_bfloat16 const* indexQWeight; - __nv_bfloat16 const* indexKWeight; - float const* rotaryCosSin; - int const* positionIds; -}; - -//! The heads per rank this producer is built for. -constexpr int kHeadsQ = 8; -constexpr int kHeadsKV = 1; -constexpr int kHeadsIndex = 1; - -//! Launch over `numTokens` tokens, one warp per (token, packed head). `prewaitStepInputs` reads the position ids, the -//! cache slots and the RoPE rows before the grid-dependency wait (they must not be written by a PDL predecessor). -cudaError_t invokeProducerPw8( - Params const& p, int numTokens, bool pdl, bool early, bool prewaitStepInputs, cudaStream_t stream); - -} // namespace producerPw8 -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu index 1229b05b029f..70e5dc8d5914 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.cu @@ -17,53 +17,6 @@ #include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h" -#include "tensorrt_llm/common/assert.h" - -#include -#include - -#include -#include - -// Fused sampled verification for the MiniMax-M3 Eagle3 target head (TP-sharded vocabulary, decode rows <= 48). -// -// sm_107 (Rubin) TP8-owned copy of kernels/minimaxM3Sm107Tp2/sampledAccept.cu (bit for bit the same kernel: the -// seam for the TP8 SA-1R round-0 decision), itself the TP2 copy of kernels/minimaxM3/sampledAccept.cu: identical -// algorithm and arithmetic; the vocab shard limit is raised to the TP2 shard (200064 / 2 = 100032 -> kMaxVs 100352), so -// a CTA stages up to kMaxVs / kSplits BF16 logits (25 KB at 8 splits, 50 KB at 4) and the launch opts in to that -// dynamic shared memory. -// -// Replaces, per decode step, the full-vocabulary all-gather of the target logits, the fused sampling kernel and the -// strict-acceptance ops for rows sampled with temperature and top-p (the AgentX contract). Rows with top-k or min-p -// active must stay on the generic path; the caller decides. -// -// Sampling is top-p rejection sampling, distributed over the TP ranks. With w = exp(logit / T - max) over the FULL -// vocabulary and Z = sum(w), a token is in the nucleus iff the mass strictly above its weight is at most top_p * Z -// (the fused sampling kernel's kept set, ties included). A candidate is drawn proportionally to w and accepted iff it -// is in the nucleus; a rejected candidate lies below the nucleus, so the next candidate is drawn over the tokens -// heavier than it. The accepted token is distributed as w restricted to the nucleus. The first candidate lands in the -// nucleus with probability >= top_p, so most rows decide on their second exchange. -// -// Per verify row, 8 CTAs (rows <= 32) or 4 CTAs (more rows) on every rank each own a slice of the rank's logits shard, -// staged in shared memory once. A candidate is drawn in two stages without per-token noise: each CTA draws a local -// candidate by inverse CDF over its slice, and the CTA whose log-mass plus Gumbel noise is largest supplies the -// candidate (CTA j wins with probability T_j / sum T, its token with w / T_j). Uniforms come from a counter-based -// Philox keyed by the request's seed and offset, the row, the attempt and the CTA. Each round every CTA multicasts one -// 4-word record (key, candidate logit, and max and mass in round 0 or the mass above the current candidate later) and -// reads the records of all CTAs of the row on all ranks. Every CTA takes the same decision from the same records in the -// same order, so all ranks agree without any intra-GPU handshake. Words carry the call's tag: a slot needs no clearing -// and two call halves keep a rank one call ahead from overwriting a slot still being read. -// -// SA-1R (Args.one_round_k 9 / 82; MoE prototype moe/sa1r/6c753632ec, exact on 39,968 rows, eager + graph): each CTA -// also multicasts its largest bf16 logits with the round-0 record (3 per extra word). After round 0 every CTA knows M, -// the candidate logit cv and wc = __expf(cv - M); round 1 would accept iff S <= target, S = the fp32 sum (its tree) of -// w = __expf(__fmul_rn(x, tInv) - M) over every token with w > wc. x -> w is monotone (__fmul_rn by tInv > 0, fsub, -// __expf: checked exhaustively over all non-positive fp32 inputs), so if no CTA's unlisted value can beat wc (82: the -// bound tau; 9: the 9th listed), the listed w > wc are exactly round 1's summands, bitwise. Two fp32 sums of the same -// m non-negative terms differ by < m * 2^-23 * S, so accepting iff S_1r + m * 2^-21 * S_1r (round up) <= target -// implies round 1 accepts; m <= 1 is exact. Otherwise round 1 runs as before. An accepted row outputs the round-0 -// candidate, as the two-round kernel does, and consumes no Philox draw it would not. - TRTLLM_NAMESPACE_BEGIN namespace kernels @@ -72,888 +25,15 @@ namespace minimaxM3Sm107Tp8 { namespace sampledAccept { -namespace -{ -constexpr int kThreads = 256; -constexpr int kWarps = kThreads / 32; -constexpr int kMaxRec = kMaxWorld * kMaxSplits; -constexpr int kMaxAttempts = 16; // exchange slots per row, reused cyclically by the rounds (no round cap) -constexpr int kTagBits = 64 - 32 - kIdBits; -constexpr float kTempEpsilon = 1e-6f; // the fused sampling kernel's 1/T guard -// Record polls trap after this many spins (the multicast all-reduce's bound): a missing writer becomes a CUDA error, -// never a hang. -constexpr long long kMaxSpins = 1LL << 31; - -// SA-1R variant kTopK: 0 none; K exact top-K list; 10 * K + P pooled (each thread keeps its top P plus its (P + 1)-th -// as a bound; the CTA lists the pool's top K plus tau = max(K-th listed, max thread bound)). -__host__ __device__ constexpr int mode_k(int m) -{ - return m >= 10 ? m / 10 : m; -} - -__host__ __device__ constexpr int mode_pool(int m) -{ - return m >= 10 ? m % 10 : 0; -} - -__host__ __device__ constexpr int mode_vals(int m) -{ - return mode_k(m) + (mode_pool(m) ? 1 : 0); -} - -// Words per record: key, logit of key, f2, f3, then the round-0 list (3 bf16 + a 16-bit tag per word). -template -struct Rec -{ - static constexpr int kXw = (mode_vals(kTopK) + 2) / 3; - static constexpr int kEx = 4 + kXw; -}; - -// Multicast exchange, per call half: [kMaxAttempts][kMaxRows][world * kMaxSplits][kEx] u64. -template -__host__ __device__ inline long long half_words(int world) -{ - return static_cast(kMaxAttempts) * kMaxRows * world * kMaxSplits * Rec::kEx; -} - -__device__ __forceinline__ uint32_t ordered(float v) -{ - uint32_t const b = __float_as_uint(v); - return (b & 0x80000000u) ? ~b : (b | 0x80000000u); -} - -// bf16 bits <-> an order-preserving 16-bit key. -__device__ __forceinline__ uint32_t ord16(uint32_t b) -{ - return (b & 0x8000u) ? (~b & 0xFFFFu) : (b | 0x8000u); -} - -__device__ __forceinline__ uint32_t unord16(uint32_t k) -{ - return (k & 0x8000u) ? (k & 0x7FFFu) : (~k & 0xFFFFu); -} - -// Philox4x32-10. -__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) -{ -#pragma unroll - for (int i = 0; i < 10; ++i) - { - uint32_t const hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; - uint32_t const hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; - c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); - k.x += 0x9E3779B9u; - k.y += 0xBB67AE85u; - } - return c; -} - -// u in (0, 1): (k + 1/2) / 2^23, exact in fp32. -__device__ __forceinline__ float uniform(uint32_t bits) -{ - return (static_cast(bits >> 9) + 0.5f) * (1.0f / 8388608.0f); -} - -// Key on the wire: perturbed log-mass, then the candidate token id; 0 = no candidate. Tag bits are added on the wire. -__device__ __forceinline__ unsigned long long gkey(float lg, int id) -{ - return (static_cast(ordered(lg)) << kIdBits) - | static_cast((1u << kIdBits) - 1u - static_cast(id)); -} - -__device__ __forceinline__ int key_token(unsigned long long k) -{ - return static_cast((1u << kIdBits) - 1u - static_cast(k & ((1u << kIdBits) - 1u))); -} - -__device__ __forceinline__ void st_sys(unsigned long long* p, unsigned long long v) -{ - asm volatile("st.global.relaxed.sys.u64 [%0], %1;" ::"l"(p), "l"(v) : "memory"); -} - -__device__ __forceinline__ unsigned long long ld_sys(unsigned long long const* p) -{ - unsigned long long v; - asm volatile("ld.global.relaxed.sys.u64 %0, [%1];" : "=l"(v) : "l"(p) : "memory"); - return v; -} - -template -struct Shared -{ - __nv_bfloat16 row[kMaxVs / kSplits]; - float redf[kWarps]; - unsigned long long redk[kWarps]; - unsigned long long rk[kMaxRec]; - float rv[kMaxRec], r2[kMaxRec], r3[kMaxRec]; - float wsum[kWarps]; - int owner; - int pick; - float pickx; - bool last; - float dS[kWarps]; // SA-1R: per-warp partial mass above the candidate - int dM[kWarps]; // SA-1R: per-warp count | incomplete << 16 - uint32_t wtop[kWarps][kTopK > 0 ? mode_k(kTopK) : 1]; // SA-1R: per-warp top-K keys - uint32_t wbound[kWarps]; // SA-1R pooled: per-warp max thread bound - uint32_t ctop[kTopK > 0 ? mode_vals(kTopK) : 1]; // SA-1R: the CTA's top-K keys (descending), by rank - unsigned long long rx[kMaxRec][Rec::kXw > 0 ? Rec::kXw : 1]; -}; - -template -__device__ __forceinline__ float block_max(float v, S& s) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o)); - if ((threadIdx.x & 31) == 0) - s.redf[threadIdx.x >> 5] = v; - __syncthreads(); - float r = s.redf[0]; - for (int i = 1; i < kWarps; ++i) - r = fmaxf(r, s.redf[i]); - __syncthreads(); - return r; -} - -// Fixed-order block sum: identical result for identical inputs. -template -__device__ __forceinline__ float block_sum(float v, S& s) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v += __shfl_xor_sync(0xffffffffu, v, o); - if ((threadIdx.x & 31) == 0) - s.redf[threadIdx.x >> 5] = v; - __syncthreads(); - float r = 0.f; - for (int i = 0; i < kWarps; ++i) - r += s.redf[i]; - __syncthreads(); - return r; -} - -template -__device__ __forceinline__ unsigned long long block_max_u64(unsigned long long v, S& s) -{ -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - { - unsigned long long const w = __shfl_xor_sync(0xffffffffu, v, o); - v = w > v ? w : v; - } - if ((threadIdx.x & 31) == 0) - s.redk[threadIdx.x >> 5] = v; - __syncthreads(); - unsigned long long r = s.redk[0]; - for (int i = 1; i < kWarps; ++i) - r = s.redk[i] > r ? s.redk[i] : r; - __syncthreads(); - return r; -} - -// Exclusive prefix of v over the threads in thread order, and the total. -template -__device__ __forceinline__ float block_excl_scan(float v, float& total, S& s) -{ - int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; - float inc = v; -#pragma unroll - for (int o = 1; o < 32; o <<= 1) - { - float const t = __shfl_up_sync(0xffffffffu, inc, o); - if (lane >= o) - inc += t; - } - if (lane == 31) - s.wsum[warp] = inc; - __syncthreads(); - float base = 0.f; - total = 0.f; - for (int i = 0; i < kWarps; ++i) - { - base += i < warp ? s.wsum[i] : 0.f; - total += s.wsum[i]; - } - __syncthreads(); - return base + inc - v; -} - -// Warp-wide integer max / sum / or (the redux instructions from sm_80; a butterfly otherwise, same results). -__device__ __forceinline__ uint32_t warp_max_u32(uint32_t v) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - return __reduce_max_sync(0xffffffffu, v); -#else -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - { - uint32_t const w = __shfl_xor_sync(0xffffffffu, v, o); - v = w > v ? w : v; - } - return v; -#endif -} - -__device__ __forceinline__ uint32_t warp_add_u32(uint32_t v) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - return __reduce_add_sync(0xffffffffu, v); -#else -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v += __shfl_xor_sync(0xffffffffu, v, o); - return v; -#endif -} - -__device__ __forceinline__ uint32_t warp_or_u32(uint32_t v) -{ -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - return __reduce_or_sync(0xffffffffu, v); -#else -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - v |= __shfl_xor_sync(0xffffffffu, v, o); - return v; -#endif -} - -// Insert key into a descending register list of K (fully unrolled: no local memory). -template -__device__ __forceinline__ void topk_insert(uint32_t (&a)[K], uint32_t key) -{ -#pragma unroll - for (int j = 0; j < K; ++j) - { - uint32_t const hi = key > a[j] ? key : a[j]; - key = key > a[j] ? a[j] : key; - a[j] = hi; - } -} - -template -__device__ __forceinline__ void topk_pop(uint32_t (&a)[K]) -{ -#pragma unroll - for (int j = 0; j + 1 < K; ++j) - a[j] = a[j + 1]; - a[K - 1] = 0u; -} - -// Per-thread insertion with an early reject: after the first few elements almost every key fails the K-th test. -template -__device__ __forceinline__ void thread_insert(uint32_t (&a)[K], uint32_t key) -{ - if (key > a[K - 1]) - topk_insert(a, key); -} - -// The warp's K largest keys (descending) into s.wtop[warp]; published by the caller's next block barrier. -template -__device__ __forceinline__ void warp_topk(uint32_t (&a)[K], S& s) -{ - int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; -#pragma unroll - for (int j = 0; j < K; ++j) - { - uint32_t const m = warp_max_u32(a[0]); - if (a[0] == m && m != 0u) - topk_pop(a); - if (lane == 0) - s.wtop[warp][j] = m; - } -} - -// Pooled: K rounds over each lane's P pool keys (descending) and the warp's max thread bound. -template -__device__ __forceinline__ void warp_pool(uint32_t (&a)[P + 1], S& s) -{ - int const lane = threadIdx.x & 31, warp = threadIdx.x >> 5; - uint32_t const bound = warp_max_u32(a[P]); - uint32_t pl[P]; -#pragma unroll - for (int j = 0; j < P; ++j) - pl[j] = a[j]; -#pragma unroll - for (int j = 0; j < K; ++j) - { - uint32_t const m = warp_max_u32(pl[0]); - if (pl[0] == m && m != 0u) - topk_pop

(pl); - if (lane == 0) - s.wtop[warp][j] = m; - } - if (lane == 0) - s.wbound[warp] = bound; -} - -// A row's kSplits CTAs wait for each other's records, so they launch as one thread-block cluster: the hardware -// schedules a cluster's CTAs together, whatever other streams hold. Rows never wait on other rows. -template -__global__ void __launch_bounds__(kThreads, 2) sampled_accept_kernel(Args const a) -{ - using R = Rec; - extern __shared__ __align__(16) unsigned char smem[]; - Shared& s = *reinterpret_cast*>(smem); - float constexpr kNegInf = -__builtin_huge_valf(); - int const tid = threadIdx.x; - int const row = blockIdx.x / kSplits; - int const split = blockIdx.x % kSplits; - int const world = a.world; - int const nrec = world * kSplits; - // The row's sampling parameters are step inputs (the spec metadata, written before the target forward), not the - // predecessor head GEMM's output, so with prewait_params they load before the dependency wait and overlap the - // head's tail; the same values either way. - float temperature = 0.f, topP = 0.f; - long long seed_in = 0, offset_in = 0; - auto load_params = [&]() - { - temperature = a.temperatures[row]; - topP = a.top_ps[row]; - seed_in = a.seeds[row]; - offset_in = a.offsets[row]; - }; - if (a.prewait_params) - load_params(); -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("griddepcontrol.wait;" ::: "memory"); -#endif - if (!a.prewait_params) - load_params(); -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - // early_trigger: the dependent grid (the first draft pass's input) may launch now and park at its own - // griddepcontrol.wait, which still returns only once this grid has completed and flushed; it skips the launch - // latency after this kernel ends. Otherwise dependents are released at the acceptance epilogue (the CTAs of - // splits 1.. exit without releasing). - if (a.early_trigger) - asm volatile("griddepcontrol.launch_dependents;"); -#endif - float const tempInv = 1.0f / (temperature + kTempEpsilon); - unsigned long long const seed = static_cast(seed_in); - unsigned long long const offset = static_cast(offset_in); - unsigned int const call = *reinterpret_cast(a.counters + 1); - // Wire tags: never 0 (fresh memory), different between calls of the same half and between consecutive laps of - // the slot ring (round r uses slot r % kMaxAttempts and lap parity (r / kMaxAttempts) & 1). A rank decides round - // r only after reading every record of round r, so no rank is more than one round ahead and a slot is rewritten - // only after every reader has consumed it; the lap bit tells the new record from the previous lap's. - auto ktag_of = [&](int r) -> unsigned long long - { - unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; - return (((call >> 1) << 1) | lap) % ((1u << kTagBits) - 1u) + 1u; - }; - auto ftag_of = [&](int r) -> unsigned long long - { - unsigned int const lap = static_cast(r / kMaxAttempts) & 1u; - return static_cast(((((call >> 1) << 1) | lap) & 0x7FFFFFFFu) | 0x80000000u) << 32; - }; - // SA-1R round-0 list words: a 16-bit tag, never 0. - unsigned long long const xtag = (((call >> 1) << 1) & 0x7FFFu) | 0x8000u; - long long const rec_stride = static_cast(world) * kMaxSplits * R::kEx; - unsigned long long* const xu = a.ws_uc + static_cast(call & 1u) * half_words(world) - + static_cast(row) * rec_stride; - unsigned long long* const xm = a.ws_mc + static_cast(call & 1u) * half_words(world) - + static_cast(row) * rec_stride + (a.rank * kMaxSplits + split) * R::kEx; - long long const round_stride = static_cast(kMaxRows) * rec_stride; - - // ---- stage this CTA's slice and its max ------------------------------------------------------------------------- - __nv_bfloat16 const* src = a.logits + row * a.ld; - int const sl_len = ((a.Vs + kSplits - 1) / kSplits + 7) / 8 * 8; - int const s0 = min(split * sl_len, a.Vs); - int const n = min(s0 + sl_len, a.Vs) - s0; - int const id0 = a.vocab_offset + s0; - float tmax = kNegInf; - // SA-1R: each thread's largest (ord16(bf16) << 16 | ~index) keys while staging (only rows with a nucleus cut). - bool const want_x = kTopK > 0 && topP < 1.0f; - constexpr int kK = mode_k(kTopK), kPool = mode_pool(kTopK), kVals = mode_vals(kTopK); - constexpr int kTk = kPool ? kPool + 1 : (kK > 0 ? kK : 1); - uint32_t tk[kTk]; -#pragma unroll - for (int j = 0; j < kTk; ++j) - tk[j] = 0u; - bool const vec = (a.ld % 8) == 0 && (reinterpret_cast(a.logits) & 15) == 0; - int const nvec = vec ? n / 8 : 0; - for (int v = tid; v < nvec; v += kThreads) - { - uint4 const w = *reinterpret_cast(src + s0 + 8 * v); - *reinterpret_cast(&s.row[8 * v]) = w; - __nv_bfloat16 const* e = reinterpret_cast<__nv_bfloat16 const*>(&w); -#pragma unroll - for (int j = 0; j < 8; ++j) - tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(e[j]), tempInv)); - if constexpr (kTopK > 0) - { - if (want_x) - { -#pragma unroll - for (int j = 0; j < 8; ++j) - { - uint32_t const key - = (ord16(__bfloat16_as_ushort(e[j])) << 16) | (0xFFFFu - static_cast(8 * v + j)); - if constexpr (kPool > 0) - topk_insert(tk, key); // branchless: P + 1 compare-swaps - else - thread_insert(tk, key); - } - } - } - } - for (int i = nvec * 8 + tid; i < n; i += kThreads) - { - s.row[i] = src[s0 + i]; - tmax = fmaxf(tmax, __fmul_rn(__bfloat162float(src[s0 + i]), tempInv)); - if constexpr (kTopK > 0) - { - if (want_x) - topk_insert( - tk, (ord16(__bfloat16_as_ushort(src[s0 + i])) << 16) | (0xFFFFu - static_cast(i))); - } - } - if constexpr (kTopK > 0) - { - if (tid < kVals) - s.ctop[tid] = 0u; // published with s.wtop by block_max's barrier - if (want_x) - { - if constexpr (kPool > 0) - warp_pool(tk, s); // s.wtop / s.wbound published by block_max's barrier - else - warp_topk(tk, s); // s.wtop published by block_max's barrier - } - } - float const smax = block_max(tmax, s); // the staged slice is visible to all threads after this barrier - if constexpr (kTopK > 0) - { - // SA-1R: the CTA's top K of the warp lists by rank (keys are unique): a lane quad per key, no serial merge; - // s.ctop is read after the round-0 draw's first barrier - constexpr int kN = kWarps * kK; - constexpr int kSp = kN * 4 <= kThreads ? 4 : 1; // threads per key - if (want_x && tid < kN * kSp) - { - uint32_t const* const wl = &s.wtop[0][0]; - int const part = tid % kSp; - uint32_t const key = wl[tid / kSp]; - int rank = 0; -#pragma unroll - for (int j = part; j < kN; j += kSp) - rank += wl[j] > key ? 1 : 0; - if constexpr (kSp == 4) - { - rank += __shfl_xor_sync(0xffffffffu, rank, 1); - rank += __shfl_xor_sync(0xffffffffu, rank, 2); - } - if (part == 0 && key != 0u && rank < kK) - s.ctop[rank] = key; - } - } - float M = 0.f, wc = 0.f, target = 0.f, cv = 0.f; // global max, candidate weight and logit, top_p * Z - unsigned long long ckey = 0ULL; - int const gid = a.rank * kMaxSplits + split; - for (int r = 0;; ++r) - { - // ---- local candidate by inverse CDF over the slice's tokens heavier than the candidate (w > 0 in round 0) - float const ref = r == 0 ? smax : M; - float const wmin = r == 0 ? 0.f : wc; - float ts = 0.f; - for (int i = tid; i < n; i += kThreads) - { - float const w = __expf(__fmul_rn(__bfloat162float(s.row[i]), tempInv) - ref); - ts += w > wmin ? w : 0.f; - } - float total; - float const excl = block_excl_scan(ts, total, s); - if constexpr (kTopK > 0) - { - // SA-1R round 0: the list words depend only on the staged slice, so they go out now (s.ctop / s.wbound are - // published by the barrier above), one lane per word, ahead of the main record: off the readers' path - if (r == 0 && want_x && tid < R::kXw) - { - uint32_t tau = 0u; - if constexpr (kPool > 0) - { - tau = s.ctop[kK - 1]; // the largest unlisted pool key is <= the K-th listed -#pragma unroll - for (int wq = 0; wq < kWarps; ++wq) - tau = s.wbound[wq] > tau ? s.wbound[wq] : tau; - } - unsigned long long w = xtag; -#pragma unroll - for (int t = 0; t < 3; ++t) - { - int const e = 3 * tid + t; - uint32_t const kv = e < kK ? s.ctop[e] : (kPool > 0 && e == kK ? tau : 0u); - // missing entries: -inf (bf16 0xFF80), w = 0 - uint32_t const b = kv != 0u ? unord16(kv >> 16) : 0xFF80u; - w |= static_cast(b) << (16 * (3 - t)); - } - st_sys(xm + 4 + tid, w); // round 0 = slot 0 - } - } - uint4 const rnd = philox(make_uint4(static_cast(gid), static_cast(row + 64 * r), - static_cast(offset), static_cast(offset >> 32)), - make_uint2(static_cast(seed), static_cast(seed >> 32))); - float const tgt = uniform(rnd.x) * total; - if (tid == 0) - s.owner = -1; - __syncthreads(); - if (ts > 0.f && excl <= tgt) - atomicMax(&s.owner, tid); - __syncthreads(); - if (tid == s.owner) - { - float acc = excl; - int pick = -1; - float px = 0.f; - if constexpr (kIlpWalk) - { - // SA-PW: the owner's strided tokens 8 at a time. The 8 x and w (the same expressions) are - // independent, so their MUFU work overlaps; the serial sum then adds them in the same order with - // the same early exit: the same acc sequence, pick and px as the one-token walk, bit for bit. - constexpr int kW = 8; - bool done = false; - for (int i0 = tid; i0 < n && !done; i0 += kW * kThreads) - { - float xs[kW], wv[kW]; -#pragma unroll - for (int k = 0; k < kW; ++k) - { - int const i = i0 + k * kThreads; - xs[k] = i < n ? __fmul_rn(__bfloat162float(s.row[i]), tempInv) : 0.f; - wv[k] = __expf(xs[k] - ref); - } -#pragma unroll - for (int k = 0; k < kW; ++k) - { - int const i = i0 + k * kThreads; - if (!done && i < n && wv[k] > wmin) - { - acc += wv[k]; - pick = i; - px = xs[k]; - if (acc > tgt) - done = true; - } - } - } - } - else - { - for (int i = tid; i < n; i += kThreads) - { - float const x = __fmul_rn(__bfloat162float(s.row[i]), tempInv); - float const w = __expf(x - ref); - if (w > wmin) - { - acc += w; - pick = i; - px = x; - if (acc > tgt) - break; - } - } - } - s.pick = pick; - s.pickx = px; - } - __syncthreads(); - unsigned long long sb = 0ULL; - if (total > 0.f && s.owner >= 0) - sb = gkey(__logf(total) + ref - logf(-logf(uniform(rnd.y))), id0 + s.pick); - float const smab = total; - unsigned long long const ktag = ktag_of(r); - unsigned long long const ftag = ftag_of(r); - unsigned long long* const xr = xu + (r % kMaxAttempts) * round_stride; - bool const with_x = want_x && r == 0; - if (tid == 0) - { - unsigned long long* const xw = xm + (r % kMaxAttempts) * round_stride; - st_sys(xw + 0, (sb << kTagBits) | ktag); - st_sys(xw + 1, ftag | __float_as_uint(sb != 0ULL ? s.pickx : 0.f)); - st_sys(xw + 2, ftag | __float_as_uint(r == 0 ? smax : smab)); - st_sys(xw + 3, ftag | __float_as_uint(r == 0 ? smab : 0.f)); - } - // ---- every record of the row, every rank --------------------------------------------------------------- - int const nw = with_x ? R::kEx : 4; - for (int t = tid; t < nrec * nw; t += kThreads) - { - int const q = t / nw, f = t % nw; - unsigned long long const* const p = xr + (q / kSplits * kMaxSplits + q % kSplits) * R::kEx + f; - unsigned long long v; - long long spins = 0; - if (f == 0) - { - do - { - v = ld_sys(p); - if (++spins > kMaxSpins) - __trap(); - } while ((v & ((1ULL << kTagBits) - 1ULL)) != ktag); - s.rk[q] = v >> kTagBits; - } - else if (f < 4) - { - do - { - v = ld_sys(p); - if (++spins > kMaxSpins) - __trap(); - } while ((v & 0xFFFFFFFF00000000ULL) != ftag); - float const fv = __uint_as_float(static_cast(v)); - (f == 1 ? s.rv : f == 2 ? s.r2 : s.r3)[q] = fv; - } - else - { - do - { - v = ld_sys(p); - if (++spins > kMaxSpins) - __trap(); - } while ((v & 0xFFFFULL) != xtag); - s.rx[q][f - 4] = v; - } - } - __syncthreads(); - // ---- the decision, identical in every CTA of the row on every rank: fixed record order ----------------- - unsigned long long gk = 0ULL; - float gv = 0.f; - for (int q = 0; q < nrec; ++q) - { - if (s.rk[q] > gk) - { - gk = s.rk[q]; - gv = s.rv[q]; - } - } - bool accept; - if (r == 0) - { - M = kNegInf; - for (int q = 0; q < nrec; ++q) - M = fmaxf(M, s.r2[q]); - float Z = 0.f; - for (int q = 0; q < nrec; ++q) - Z += s.r2[q] == kNegInf ? 0.f : s.r3[q] * __expf(s.r2[q] - M); - target = topP * Z; - accept = !(topP < 1.0f); // no nucleus cut: the first draw stands - ckey = gk; - cv = gv; - if constexpr (kTopK > 0) - { - // SA-1R: every thread takes (record, entry) items; the per-warp partials are published by the round's - // barrier below and every thread combines them in warp order after it (no new barrier) - if (with_x) - { - float const wcand = __expf(cv - M); // = round 1's wmin - float part = 0.f; - int cnt = 0, inc = 0; - for (int t = tid; t < nrec * kVals; t += kThreads) - { - int const q = t / kVals, e = t % kVals; - uint32_t const b = static_cast(s.rx[q][e / 3] >> (16 * (3 - e % 3))) & 0xFFFFu; - float const w = __expf( - __fmul_rn(__bfloat162float(__ushort_as_bfloat16(static_cast(b))), tempInv) - - M); // w bitwise as round 1 computes it - if (w > wcand) - { - if (kPool > 0 && e == kK) - inc = 1; // tau (the largest unlisted value) is heavier than the candidate - else - { - part += w; - ++cnt; - if (kPool == 0 && e == kK - 1) - inc = 1; // the K-th is heavier: unlisted heavier tokens may exist - } - } - } -#pragma unroll - for (int o = 16; o > 0; o >>= 1) - part += __shfl_xor_sync(0xffffffffu, part, o); // any order: the margin covers it - cnt = static_cast(warp_add_u32(static_cast(cnt))); - inc = static_cast(warp_or_u32(static_cast(inc))); - if ((tid & 31) == 0) - { - s.dS[tid >> 5] = part; - s.dM[tid >> 5] = cnt | (inc << 16); - } - } - } - } - else - { - float S = 0.f; // mass strictly above the current candidate - for (int q = 0; q < nrec; ++q) - S += s.r2[q]; - accept = S <= target || gk == 0ULL; - if (!accept) - { - ckey = gk; - cv = gv; - } - } - wc = __expf(cv - M); - __syncthreads(); // records consumed before the next round's writes to s - bool decided = false; - if constexpr (kTopK > 0) - { - if (r == 0 && with_x) - { - float S1 = 0.f; - int m = 0, inc = 0; - for (int wq = 0; wq < kWarps; ++wq) - { - S1 += s.dS[wq]; - m += s.dM[wq] & 0xFFFF; - inc |= s.dM[wq] >> 16; - } - if (!inc) - { - float const margin = m <= 1 ? 0.f : S1 * static_cast(m) * (1.0f / 2097152.0f); - decided = __fadd_ru(S1, margin) <= target; - accept = accept || decided; - } - } - } - if (accept) - { - if (split != 0) - return; - if (tid == 0) - { - a.tokens[row] = key_token(ckey); - if (a.extra_attempts != nullptr && r > 1) - atomicAdd(a.extra_attempts, r - 1); - if (a.decided != nullptr && decided) - atomicAdd(a.decided, 1); - } - break; - } - } - - // ---- split 0: the last row to finish runs the strict acceptance ---------------------------------------------- -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - if (!a.early_trigger) - asm volatile("griddepcontrol.launch_dependents;"); -#endif - __threadfence(); - __syncthreads(); - if (tid == 0) - s.last = atomicAdd(a.counters, 1u) == static_cast(a.rows) - 1; - __syncthreads(); - if (!s.last) - return; - __threadfence(); - if (tid == 0) - { - a.counters[0] = 0u; - a.counters[1] = call + 1u; - } - if (tid < a.nreq) - { - int const* tok = a.tokens; - int const b = tid; - int const k1 = a.draft_len + 1; - int nacc = 1; - bool run = true; - for (int k = 0; k < a.draft_len; ++k) - { - int const tt = *reinterpret_cast(tok + b * k1 + k); - a.acc_tokens[b * k1 + k] = tt; - run = run && a.draft[b * a.draft_len + k] == tt; - nacc += run ? 1 : 0; - } - a.acc_tokens[b * k1 + a.draft_len] = *reinterpret_cast(tok + b * k1 + a.draft_len); - if (a.force_vec != nullptr) - nacc = a.force_vec[b]; - else if (a.force_total > 0) - nacc = a.force_total; - a.num_acc[b] = nacc; - a.gather_ids[b] = static_cast(a.batch_idx[b]) * k1 + nacc - 1; - if (a.next_new != nullptr) - a.next_new[b * a.next_stride] - = *reinterpret_cast(tok + a.batch_idx[b] * k1 + nacc - 1); - } -} - -template -cudaError_t launch(Args const& a, bool use_pdl, cudaStream_t stream) -{ - int device; - cudaError_t const deviceStatus = cudaGetDevice(&device); - if (deviceStatus != cudaSuccess) - return deviceStatus; - static std::mutex configMutex; - static std::unordered_map deviceClusters; - int max_clusters = 0; - cudaLaunchConfig_t cfg = {}; - cfg.gridDim = dim3(a.rows * kSplits); - cfg.blockDim = dim3(kThreads); - cfg.dynamicSmemBytes = sizeof(Shared); - cfg.stream = stream; - cudaLaunchAttribute attrs[2]; - attrs[0].id = cudaLaunchAttributeClusterDimension; - attrs[0].val.clusterDim.x = kSplits; - attrs[0].val.clusterDim.y = 1; - attrs[0].val.clusterDim.z = 1; - attrs[1].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[1].val.programmaticStreamSerializationAllowed = use_pdl ? 1 : 0; - cfg.attrs = attrs; - cfg.numAttrs = 2; - { - std::lock_guard lock(configMutex); - max_clusters = deviceClusters[device]; - if (max_clusters == 0) - { - cudaError_t err = cudaFuncSetAttribute(sampled_accept_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(sizeof(Shared))); - if (err == cudaSuccess) - { - err = cudaOccupancyMaxActiveClusters( - &max_clusters, sampled_accept_kernel, &cfg); - } - if (err != cudaSuccess) - { - return err; - } - } - deviceClusters[device] = max_clusters; - } - TLLM_CHECK_WITH_INFO(max_clusters >= 1, "m3_sampled_accept: a %d-CTA cluster does not fit on one GPC", kSplits); - return cudaLaunchKernelEx(&cfg, sampled_accept_kernel, a); -} - -template -cudaError_t launchRows(Args const& a, bool use_pdl, cudaStream_t stream) -{ - // TP4 caps 2 / 4 / 8 (8 / 16 / 32 rows) and TP8 cap 1 (4 rows): 8 CTAs per row, at most 256 CTAs (2 per SM); - // more rows: 4 CTAs per row. - if (a.ilp_walk) - { - return a.rows <= 32 ? launch<8, kTopK, true>(a, use_pdl, stream) : launch<4, kTopK, true>(a, use_pdl, stream); - } - return a.rows <= 32 ? launch<8, kTopK, false>(a, use_pdl, stream) : launch<4, kTopK, false>(a, use_pdl, stream); -} - -} // namespace - long long sampledAcceptWsWords(int world, int oneRoundK) { - TLLM_CHECK_WITH_INFO( - oneRoundK == 0 || oneRoundK == 9 || oneRoundK == 82, "SA-1R variant %d is not built", oneRoundK); - return 2 * (oneRoundK == 9 ? half_words<9>(world) : oneRoundK == 82 ? half_words<82>(world) : half_words<0>(world)); + return minimaxM3Sm107::sampledAccept::sampledAcceptWsWords(world, oneRoundK); } cudaError_t invokeSampledAccept(Args const& a, bool usePdl, cudaStream_t stream) { - switch (a.one_round_k) - { - case 0: return launchRows<0>(a, usePdl, stream); - case 9: return launchRows<9>(a, usePdl, stream); - case 82: return launchRows<82>(a, usePdl, stream); - default: TLLM_CHECK_WITH_INFO(false, "SA-1R variant %d is not built", a.one_round_k); - } - return cudaErrorInvalidValue; + return minimaxM3Sm107::sampledAccept::invokeSampledAccept(a, usePdl, stream); } - } // namespace sampledAccept } // namespace minimaxM3Sm107Tp8 } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h index aa52cf4a0559..92f00bf92df5 100644 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h +++ b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/sampledAccept.h @@ -18,6 +18,7 @@ #pragma once #include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/sampledAccept.h" #include #include @@ -40,39 +41,7 @@ constexpr int kMaxWorld = 8; // TP ranks constexpr int kMaxVs = 100352; // vocab shard width (TP2: 200064 / 2 = 100032, padded to a multiple of 512) constexpr int kIdBits = 18; // token ids < 2^18 -struct Args -{ - __nv_bfloat16 const* logits; // [rows, Vs] this rank's shard - long long ld; - int rows, Vs, vocab_offset; - float const* temperatures; - float const* top_ps; - long long const* seeds; - long long const* offsets; - unsigned long long* ws_uc; - unsigned long long* ws_mc; - unsigned int* counters; // [0] finished rows, [1] call - int* tokens; // [rows] sampled token per row - int* extra_attempts; // optional [1]: rounds beyond the second, summed over rows - int rank, world; - int const* draft; - int const* batch_idx; - int* acc_tokens; - int* num_acc; - long long* gather_ids; - int draft_len, nreq, force_total; - int* next_new; - int next_stride; - int const* force_vec; - int prewait_params; // nonzero: the row's sampling parameters load before the PDL dependency wait - int early_trigger; // nonzero: dependents are released right after the dependency wait, not at the epilogue - // SA-1R (one-round decision): 0 the two-round kernel; 9 each CTA's top-9 logits ride its round-0 record; 82 pooled - // (each thread's top 2, the CTA lists the pool's top 8 plus the bound of everything unlisted). A row whose first - // candidate is provably in the nucleus then decides in round 0, with the same token as the two-round kernel. - int one_round_k; - int* decided; // optional [1]: rows decided in round 0 (SA-1R), summed - int ilp_walk; // nonzero: the owner walk takes its tokens 8 at a time (same weights, same summation order) -}; +using Args = minimaxM3Sm107::sampledAccept::Args; //! Words (uint64) of the multicast exchange workspace for `world` ranks and SA-1R variant `oneRoundK`: two call halves. long long sampledAcceptWsWords(int world, int oneRoundK = 0); diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu deleted file mode 100644 index 8cbf5d90175c..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.cu +++ /dev/null @@ -1,89 +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/minimaxM3Sm107Tp8/siluAndMul.h" - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ - -namespace -{ - -constexpr int kThreads = 128; - -// One packed bf16 pair of gate and up values -> its two bf16 results (the low half from the low halves). The -// instruction sequence of the Triton kernel's PTX at its defaults (host/silu/silu_and_mul_plain.ptx): g * -log2(e) as -// one product, ex2.approx, + 1, div.full 1 / that, g * s, (u + 0.0) as an IEEE add, the product, one rounding to bf16. -__device__ __forceinline__ uint32_t siluMulBf16x2(uint32_t gate, uint32_t up) -{ - uint32_t r; - asm("{\n\t.reg .b32 gl, gh, ul, uh, t0, t1, e0, e1, s0, s1, p0, p1, q0, q1, r0, r1, one, zero;\n\t" - ".reg .b16 hl, hh;\n\t" - "mov.b32 {hl, hh}, %1;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" - "mov.b32 {hl, hh}, %2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" - "mul.rn.f32 t0, gl, 0fBFB8AA3B;\n\tmul.rn.f32 t1, gh, 0fBFB8AA3B;\n\t" - "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" - "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" - "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" - "mul.rn.f32 p0, gl, s0;\n\tmul.rn.f32 p1, gh, s1;\n\t" - "mov.b32 zero, 0f00000000;\n\tadd.rn.f32 q0, ul, zero;\n\tadd.rn.f32 q1, uh, zero;\n\t" - "mul.rn.f32 r0, q0, p0;\n\tmul.rn.f32 r1, q1, p1;\n\t" - "cvt.rn.bf16x2.f32 %0, r1, r0;\n\t}" - : "=r"(r) - : "r"(gate), "r"(up)); - return r; -} - -// Eight outputs per thread: 16-byte loads of the gate and the up half, one 16-byte store. -__global__ void __launch_bounds__(kThreads) - siluAndMulKernel(uint4 const* __restrict__ x, uint4* __restrict__ out, int64_t rows, int64_t vecs) -{ - int64_t const row = blockIdx.y; - int64_t const v = static_cast(blockIdx.x) * kThreads + threadIdx.x; - if (row >= rows || v >= vecs) - { - return; - } - uint4 const g = x[row * 2 * vecs + v]; - uint4 const u = x[row * 2 * vecs + vecs + v]; - uint4 o; - o.x = siluMulBf16x2(g.x, u.x); - o.y = siluMulBf16x2(g.y, u.y); - o.z = siluMulBf16x2(g.z, u.z); - o.w = siluMulBf16x2(g.w, u.w); - out[row * vecs + v] = o; -} - -} // namespace - -cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream) -{ - int64_t const vecs = d / 8; - dim3 const grid(static_cast((vecs + kThreads - 1) / kThreads), static_cast(rows)); - siluAndMulKernel<<>>( - reinterpret_cast(x), reinterpret_cast(out), rows, vecs); - return cudaGetLastError(); -} - -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h deleted file mode 100644 index 669e33eb5cf5..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/siluAndMul.h +++ /dev/null @@ -1,40 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ - -//! out[m, j] = bf16(silu(g) * (u + 0)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both -//! contiguous, d a multiple of 8, 16-byte aligned rows): bit for bit the Triton silu_and_mul kernel at its defaults -//! (no limit, alpha 1, beta 0), the plain SwiGLU of the Eagle3 drafter MLP. -cudaError_t invokeSiluAndMul(uint16_t const* x, uint16_t* out, int64_t rows, int64_t d, cudaStream_t stream); - -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h deleted file mode 100644 index 3cfd05b73369..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/specOffsets.h +++ /dev/null @@ -1,41 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ - -//! The overlap scheduler's speculative step correction in one launch: pos[i] += posOff[i] for i < nPos and -//! kv[j] += kvOff[j] for j < nKv (int32, wrapping as the ATen adds do). A plain launch: it is a step-input writer that -//! pre-wait readers rely on (rubin-20261015/PREWAIT_AUDIT.md item b). -cudaError_t invokeSpecOffsets(int32_t* pos, int32_t const* posOff, int64_t nPos, int32_t* kv, int32_t const* kvOff, - int64_t nKv, cudaStream_t stream); - -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h deleted file mode 100644 index c125dede4361..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h +++ /dev/null @@ -1,47 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ - -//! The decode graph's step-input prologue in one launch (step_inputs.StepInputPack with early_copy): the work of the -//! two parallel scatter launches and the header clear that follow them, in that order. -//! -//! The window is the pack's staging layout (kernels/stepInputsScatter.h): an int64 header `[numHostCopies, -//! numDeviceCopies]`, then one `(dst, src, bytes)` int64 triple per copy, host copies first; a host copy's `src` is a -//! byte offset into the window, a device copy's a device address. Up to `maxDesc` copies of each kind are written (the -//! slots the parallel scatter's grid covers), every host copy before any device copy, so a device copy may read a -//! destination a host copy filled; the copies of one kind must not overlap (the pack issues the window only then). -//! Afterwards the header is zeroed, so a replay with no new upload writes nothing. With `epoch` non-null, it is -//! incremented once. -cudaError_t invokeStepPrologue(uint8_t* window, int64_t* epoch, int64_t maxDesc, cudaStream_t stream); - -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h b/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h deleted file mode 100644 index 78033e018997..000000000000 --- a/cpp/tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h +++ /dev/null @@ -1,43 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/config.h" - -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace kernels -{ -namespace minimaxM3Sm107Tp8 -{ - -//! The dense MLP's swigluoai activation: out[m, j] = bf16(min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + -//! 1)) with g = x[m, j], u = x[m, d + j] (x bf16 [rows, 2d], out bf16 [rows, d], both contiguous, d even): bit for bit -//! the Triton silu_and_mul kernel at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With `pdl` the launch allows -//! programmatic serialization and the grid waits on its dependency, then releases its dependents before it computes -//! (the down GEMV stages its weights during the activation and reads it only after its own grid wait). -cudaError_t invokeSwigluOai( - uint32_t const* x, uint32_t* out, int64_t rows, int64_t halfD, bool pdl, cudaStream_t stream); - -} // namespace minimaxM3Sm107Tp8 -} // namespace kernels - -TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SiluAndMulOp.cpp similarity index 51% rename from cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SiluAndMulOp.cpp index b312cf79a715..b4dd15bbe7ed 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SiluAndMulOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SiluAndMulOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/siluAndMul.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/siluAndMul.h" #include #include @@ -26,35 +26,49 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp4/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, +//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, //! d]. -void minimaxM3Sm107Tp4SiluAndMul(torch::Tensor const& x, torch::Tensor& out) +template +void minimaxM3Sm107SiluAndMul(torch::Tensor const& x, torch::Tensor& out) { - TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), - "minimax_m3_sm107_tp4_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); + constexpr char const* kOpName = kTp == 2 ? "minimax_m3_sm107_tp2_silu_and_mul" + : kTp == 4 ? "minimax_m3_sm107_tp4_silu_and_mul" + : "minimax_m3_sm107_tp8_silu_and_mul"; + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), kOpName, + ": x must be a contiguous 2-D bfloat16 CUDA tensor"); int64_t const rows = x.size(0); int64_t const d = x.size(1) / 2; - TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, - "minimax_m3_sm107_tp4_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); + TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, kOpName, + ": x rows must be 16-byte aligned with a half width a multiple of 8"); TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, - "minimax_m3_sm107_tp4_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); - TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp4_silu_and_mul: at most 65535 rows"); + kOpName, ": out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); + TORCH_CHECK(rows <= 65535, kOpName, ": at most 65535 rows"); if (rows == 0) { return; } c10::cuda::CUDAGuard const deviceGuard{x.device()}; auto const err - = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), + = tensorrt_llm::kernels::minimaxM3Sm107::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_silu_and_mul launch failed: ", cudaGetErrorString(err)); + TORCH_CHECK(err == cudaSuccess, kOpName, " launch failed: ", cudaGetErrorString(err)); } } // namespace torch_ext TRTLLM_NAMESPACE_END +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107SiluAndMul<2>); +} + TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def("minimax_m3_sm107_tp4_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); @@ -62,5 +76,15 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { - m.impl("minimax_m3_sm107_tp4_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SiluAndMul); + m.impl("minimax_m3_sm107_tp4_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107SiluAndMul<4>); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107SiluAndMul<8>); } diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SpecOffsetsOp.cpp similarity index 52% rename from cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SpecOffsetsOp.cpp index 2142baa18599..3c47a51dd7b8 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SpecOffsetsOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SpecOffsetsOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/specOffsets.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/specOffsets.h" #include #include @@ -29,40 +29,54 @@ namespace torch_ext namespace { -void check(torch::Tensor const& t, char const* name) +void check(torch::Tensor const& t, char const* name, char const* opName) { - TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, - "minimax_m3_sm107_tp4_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); + TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, opName, ": ", name, + " must be a contiguous 1-D int32 CUDA tensor"); } } // namespace //! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch -//! (kernels/minimaxM3Sm107Tp4/specOffsets.h). -void minimaxM3Sm107Tp4SpecOffsets( +//! (kernels/minimaxM3Sm107/specOffsets.h). +template +void minimaxM3Sm107SpecOffsets( torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) { - check(pos, "pos"); - check(posOff, "pos_off"); - check(kv, "kv"); - check(kvOff, "kv_off"); - TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), - "minimax_m3_sm107_tp4_spec_offsets: each tensor and its offsets must have the same length"); + constexpr char const* kOpName = kTp == 2 ? "minimax_m3_sm107_tp2_spec_offsets" + : kTp == 4 ? "minimax_m3_sm107_tp4_spec_offsets" + : "minimax_m3_sm107_tp8_spec_offsets"; + check(pos, "pos", kOpName); + check(posOff, "pos_off", kOpName); + check(kv, "kv", kOpName); + check(kvOff, "kv_off", kOpName); + TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), kOpName, + ": each tensor and its offsets must have the same length"); if (pos.numel() == 0 && kv.numel() == 0) { return; } c10::cuda::CUDAGuard const deviceGuard{pos.device()}; - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSpecOffsets(pos.data_ptr(), + auto const err = tensorrt_llm::kernels::minimaxM3Sm107::invokeSpecOffsets(pos.data_ptr(), posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), at::cuda::getCurrentCUDAStream(pos.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_spec_offsets launch failed: ", cudaGetErrorString(err)); + TORCH_CHECK(err == cudaSuccess, kOpName, " launch failed: ", cudaGetErrorString(err)); } } // namespace torch_ext TRTLLM_NAMESPACE_END +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107SpecOffsets<2>); +} + TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def("minimax_m3_sm107_tp4_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); @@ -70,5 +84,15 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { - m.impl("minimax_m3_sm107_tp4_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SpecOffsets); + m.impl("minimax_m3_sm107_tp4_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107SpecOffsets<4>); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp8_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp8_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107SpecOffsets<8>); } diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107StepPrologueOp.cpp similarity index 57% rename from cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107StepPrologueOp.cpp index f340b50dc7dd..0eec259cd689 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8StepPrologueOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107StepPrologueOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/stepPrologue.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/stepPrologue.h" #include #include @@ -27,36 +27,60 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp8/stepPrologue.h). +//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107/stepPrologue.h). //! //! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the //! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses //! the caller guarantees are live. //! @param epoch Optional one-element int64 step counter, incremented once. -void minimaxM3Sm107Tp8StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) +template +void minimaxM3Sm107StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) { - TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), - "minimax_m3_sm107_tp8_step_prologue: window must be a contiguous uint8 CUDA tensor"); - TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, - "minimax_m3_sm107_tp8_step_prologue: window is shorter than its header and max_desc descriptors"); + constexpr char const* kOpName = kTp == 2 ? "minimax_m3_sm107_tp2_step_prologue" + : kTp == 4 ? "minimax_m3_sm107_tp4_step_prologue" + : "minimax_m3_sm107_tp8_step_prologue"; + TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), kOpName, + ": window must be a contiguous uint8 CUDA tensor"); + TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, kOpName, + ": window is shorter than its header and max_desc descriptors"); c10::cuda::CUDAGuard const deviceGuard{window.device()}; int64_t* epochPtr = nullptr; if (epoch.has_value()) { TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 && epoch->device() == window.device(), - "minimax_m3_sm107_tp8_step_prologue: epoch must be a one-element int64 tensor on the window's device"); + kOpName, ": epoch must be a one-element int64 tensor on the window's device"); epochPtr = epoch->data_ptr(); } - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeStepPrologue( + auto const err = tensorrt_llm::kernels::minimaxM3Sm107::invokeStepPrologue( window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_step_prologue launch failed: ", cudaGetErrorString(err)); + TORCH_CHECK(err == cudaSuccess, kOpName, " launch failed: ", cudaGetErrorString(err)); } } // namespace torch_ext TRTLLM_NAMESPACE_END +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107StepPrologue<2>); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107StepPrologue<4>); +} + TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def("minimax_m3_sm107_tp8_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); @@ -64,5 +88,5 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { - m.impl("minimax_m3_sm107_tp8_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8StepPrologue); + m.impl("minimax_m3_sm107_tp8_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107StepPrologue<8>); } diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SwigluOaiOp.cpp similarity index 56% rename from cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp rename to cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SwigluOaiOp.cpp index 8b5e59751f24..6dd0034b0a70 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SwigluOaiOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107SwigluOaiOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/swigluOai.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/swigluOai.h" #include #include @@ -26,34 +26,58 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp8/swigluOai.h): x bf16 [rows, 2d] -> out bf16 +//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107/swigluOai.h): x bf16 [rows, 2d] -> out bf16 //! [rows, d], PDL-launched with an early release when `pdl`. -void minimaxM3Sm107Tp8SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) +template +void minimaxM3Sm107SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) { - TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), - "minimax_m3_sm107_tp8_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); + constexpr char const* kOpName = kTp == 2 ? "minimax_m3_sm107_tp2_swiglu_oai" + : kTp == 4 ? "minimax_m3_sm107_tp4_swiglu_oai" + : "minimax_m3_sm107_tp8_swiglu_oai"; + TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), kOpName, + ": x must be a contiguous 2-D bfloat16 CUDA tensor"); int64_t const rows = x.size(0); int64_t const d = x.size(1) / 2; - TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, - "minimax_m3_sm107_tp8_swiglu_oai: x rows must be 4-byte aligned with an even half width"); + TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, kOpName, + ": x rows must be 4-byte aligned with an even half width"); TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, - "minimax_m3_sm107_tp8_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); + kOpName, ": out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); if (rows == 0) { return; } c10::cuda::CUDAGuard const deviceGuard{x.device()}; - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSwigluOai( + auto const err = tensorrt_llm::kernels::minimaxM3Sm107::invokeSwigluOai( reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, at::cuda::getCurrentCUDAStream(x.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_swiglu_oai launch failed: ", cudaGetErrorString(err)); + TORCH_CHECK(err == cudaSuccess, kOpName, " launch failed: ", cudaGetErrorString(err)); } } // namespace torch_ext TRTLLM_NAMESPACE_END +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp2_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp2_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107SwigluOai<2>); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def("minimax_m3_sm107_tp4_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("minimax_m3_sm107_tp4_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107SwigluOai<4>); +} + TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def("minimax_m3_sm107_tp8_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); @@ -61,5 +85,5 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { - m.impl("minimax_m3_sm107_tp8_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SwigluOai); + m.impl("minimax_m3_sm107_tp8_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107SwigluOai<8>); } diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp index 6b154d20d3f2..2a1dc45daf85 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2DrafterRopeFp8KvOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.h" #include #include @@ -28,9 +28,9 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp2::drafterRopeFp8Kv; +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107::drafterRopeFp8Kv; -//! The MiniMax-M3 TP2 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp2/drafterRopeFp8Kv.h): RoPE on +//! The MiniMax-M3 TP2 drafter layer's generation preprocess (kernels/minimaxM3Sm107/drafterRopeFp8Kv.h): RoPE on //! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages //! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. void minimaxM3Sm107Tp2DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp index a0f27263a152..9ff32293c563 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2McAllreduceNormOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/mcAllreduceNormTp2.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.h" #include #include @@ -70,13 +70,13 @@ std::vector minimaxM3Sm107Tp2McAllreduceNorm(torch::Tensor const& int const sm = tensorrt_llm::common::getSMVersion(); TORCH_CHECK(sm == 107, "minimax_m3_sm107_tp2_mc_allreduce_norm requires SM107, got SM", sm); TORCH_CHECK(cluster == kCluster, "minimax_m3_sm107_tp2_mc_allreduce_norm: only cluster 8 is built, got ", cluster); - TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107Tp2McAllreduceNormSupported( + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107McAllreduceNormSupported( static_cast(maxt), static_cast(opts)), "minimax_m3_sm107_tp2_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); TORCH_CHECK( world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); - tensorrt_llm::kernels::MinimaxM3Sm107Tp2McAllreduceNormParams p; + tensorrt_llm::kernels::MinimaxM3Sm107McAllreduceNormParams p; if (x2.has_value()) { TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), @@ -153,7 +153,7 @@ std::vector minimaxM3Sm107Tp2McAllreduceNorm(torch::Tensor const& p.capWeight = reinterpret_cast<__nv_bfloat16 const*>(capW->data_ptr()); p.capEps = static_cast(capEps); } - tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp2McAllreduceNorm(p, static_cast(maxt), + tensorrt_llm::kernels::invokeMinimaxM3Sm107McAllreduceNorm(p, static_cast(maxt), static_cast(opts), pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); if (quant) { diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp index ed50e6d10ba2..4ddd4eb207c8 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2ProducerPwOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp2/producerPw2.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/producerPw.h" #include #include @@ -33,7 +33,7 @@ namespace torch_ext namespace { -namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp2::producerPw2; +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107::producerPw; constexpr int64_t kHeadDim = 128; constexpr int64_t kPageSize = 128; @@ -70,9 +70,9 @@ std::vector minimaxM3Sm107Tp2ProducerPw(torch::Tensor const& pack c10::cuda::CUDAGuard const deviceGuard{packed.device()}; static int const sSmVersion = tensorrt_llm::common::getSMVersion(); TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp2_producer_pw requires SM100 or newer"); - int64_t const headsQ = pw::kHeadsQ; - int64_t const headsKV = pw::kHeadsKV; - int64_t const headsIndex = pw::kHeadsIndex; + int64_t const headsQ = pw::kHeadsQ<2>; + int64_t const headsKV = pw::kHeadsKV<2>; + int64_t const headsIndex = pw::kHeadsIndex<2>; int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim && packed.is_contiguous(), @@ -156,7 +156,7 @@ std::vector minimaxM3Sm107Tp2ProducerPw(torch::Tensor const& pack p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); p.rotaryCosSin = rotaryCosSin.data_ptr(); p.positionIds = positionIds.data_ptr(); - cudaError_t const err = pw::invokeProducerPw2(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + cudaError_t const err = pw::invokeProducerPw<2>(p, static_cast(numTokens), pdl, early, prewaitStepInputs, at::cuda::getCurrentCUDAStream(packed.get_device())); if (err != cudaSuccess) { diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp deleted file mode 100644 index 52f60de1b7ca..000000000000 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SiluAndMulOp.cpp +++ /dev/null @@ -1,66 +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/minimaxM3Sm107Tp2/siluAndMul.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp2/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, -//! d]. -void minimaxM3Sm107Tp2SiluAndMul(torch::Tensor const& x, torch::Tensor& out) -{ - TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), - "minimax_m3_sm107_tp2_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); - int64_t const rows = x.size(0); - int64_t const d = x.size(1) / 2; - TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, - "minimax_m3_sm107_tp2_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); - TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() - && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, - "minimax_m3_sm107_tp2_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); - TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp2_silu_and_mul: at most 65535 rows"); - if (rows == 0) - { - return; - } - c10::cuda::CUDAGuard const deviceGuard{x.device()}; - auto const err - = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), - reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_silu_and_mul launch failed: ", cudaGetErrorString(err)); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def("minimax_m3_sm107_tp2_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("minimax_m3_sm107_tp2_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SiluAndMul); -} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp deleted file mode 100644 index 87b6ae599742..000000000000 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SpecOffsetsOp.cpp +++ /dev/null @@ -1,74 +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/minimaxM3Sm107Tp2/specOffsets.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -namespace -{ - -void check(torch::Tensor const& t, char const* name) -{ - TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, - "minimax_m3_sm107_tp2_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); -} - -} // namespace - -//! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch -//! (kernels/minimaxM3Sm107Tp2/specOffsets.h). -void minimaxM3Sm107Tp2SpecOffsets( - torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) -{ - check(pos, "pos"); - check(posOff, "pos_off"); - check(kv, "kv"); - check(kvOff, "kv_off"); - TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), - "minimax_m3_sm107_tp2_spec_offsets: each tensor and its offsets must have the same length"); - if (pos.numel() == 0 && kv.numel() == 0) - { - return; - } - c10::cuda::CUDAGuard const deviceGuard{pos.device()}; - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSpecOffsets(pos.data_ptr(), - posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), - at::cuda::getCurrentCUDAStream(pos.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_spec_offsets launch failed: ", cudaGetErrorString(err)); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def("minimax_m3_sm107_tp2_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("minimax_m3_sm107_tp2_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SpecOffsets); -} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp deleted file mode 100644 index 9e8229ba9e16..000000000000 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2StepPrologueOp.cpp +++ /dev/null @@ -1,68 +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/minimaxM3Sm107Tp2/stepPrologue.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp2/stepPrologue.h). -//! -//! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the -//! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses -//! the caller guarantees are live. -//! @param epoch Optional one-element int64 step counter, incremented once. -void minimaxM3Sm107Tp2StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) -{ - TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), - "minimax_m3_sm107_tp2_step_prologue: window must be a contiguous uint8 CUDA tensor"); - TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, - "minimax_m3_sm107_tp2_step_prologue: window is shorter than its header and max_desc descriptors"); - c10::cuda::CUDAGuard const deviceGuard{window.device()}; - int64_t* epochPtr = nullptr; - if (epoch.has_value()) - { - TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 - && epoch->device() == window.device(), - "minimax_m3_sm107_tp2_step_prologue: epoch must be a one-element int64 tensor on the window's device"); - epochPtr = epoch->data_ptr(); - } - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeStepPrologue( - window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_step_prologue launch failed: ", cudaGetErrorString(err)); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def("minimax_m3_sm107_tp2_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("minimax_m3_sm107_tp2_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2StepPrologue); -} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp deleted file mode 100644 index 2839f6036ffc..000000000000 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp2SwigluOaiOp.cpp +++ /dev/null @@ -1,65 +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/minimaxM3Sm107Tp2/swigluOai.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp2/swigluOai.h): x bf16 [rows, 2d] -> out bf16 -//! [rows, d], PDL-launched with an early release when `pdl`. -void minimaxM3Sm107Tp2SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) -{ - TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), - "minimax_m3_sm107_tp2_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); - int64_t const rows = x.size(0); - int64_t const d = x.size(1) / 2; - TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, - "minimax_m3_sm107_tp2_swiglu_oai: x rows must be 4-byte aligned with an even half width"); - TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() - && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, - "minimax_m3_sm107_tp2_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); - if (rows == 0) - { - return; - } - c10::cuda::CUDAGuard const deviceGuard{x.device()}; - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp2::invokeSwigluOai( - reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, - at::cuda::getCurrentCUDAStream(x.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp2_swiglu_oai launch failed: ", cudaGetErrorString(err)); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def("minimax_m3_sm107_tp2_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("minimax_m3_sm107_tp2_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp2SwigluOai); -} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp index 034521d29796..7b5cb937bc17 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4DrafterRopeFp8KvOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.h" #include #include @@ -28,9 +28,9 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp4::drafterRopeFp8Kv; +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107::drafterRopeFp8Kv; -//! The MiniMax-M3 TP4 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp4/drafterRopeFp8Kv.h): RoPE on +//! The MiniMax-M3 TP4 drafter layer's generation preprocess (kernels/minimaxM3Sm107/drafterRopeFp8Kv.h): RoPE on //! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages //! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. void minimaxM3Sm107Tp4DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp index 3df36817814d..ca8f944851a4 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h" #include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h" #include @@ -39,8 +39,8 @@ static void foldTailImpl(bool oneWarp, torch::Tensor const& ctrl, torch::Tensor srcDraft, dstDraft, srcLens, dstLens, ids, draftIds, posOff, kvOff, tokensPerStep); auto const stream = at::cuda::getCurrentCUDAStream(ctrl.get_device()); auto const err = oneWarp - ? tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeFoldTailOneWarp(p, stream) - : tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeFoldTail(p, static_cast(rows), stream); + ? tensorrt_llm::kernels::minimaxM3Sm107::invokeFoldTailOneWarp(p, stream) + : tensorrt_llm::kernels::minimaxM3Sm107::invokeFoldTail(p, static_cast(rows), stream); TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_fold_tail launch failed: ", cudaGetErrorString(err)); } diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h index 7c25f9093405..a874d42e871f 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FoldTailParams.h @@ -17,7 +17,7 @@ #pragma once -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/foldTail.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h" #include #include @@ -45,7 +45,7 @@ inline void check1d(torch::Tensor const& t, char const* name) //! The fold tail's operands checked and packed (the standalone fold tail op and any op that runs the tail in-kernel, //! e.g. the drafter's last-pass draft_step, build them identically). -inline tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams buildParams(torch::Tensor const& ctrl, +inline tensorrt_llm::kernels::minimaxM3Sm107::FoldTailParams buildParams(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, torch::Tensor& ids, torch::Tensor& draftIds, @@ -78,7 +78,7 @@ inline tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams buildParams(torc TORCH_CHECK(ids.numel() >= rows * tokensPerStep && posOff.numel() >= rows * tokensPerStep && draftIds.numel() >= rows * (tokensPerStep - 1), "minimax_m3_sm107_tp4_fold_tail: staging outputs are too small for rows * tokens_per_step"); - tensorrt_llm::kernels::minimaxM3Sm107Tp4::FoldTailParams p{}; + tensorrt_llm::kernels::minimaxM3Sm107::FoldTailParams p{}; p.maxRows = static_cast(rows); p.ctrl = ctrl.data_ptr(); p.slots = slots.data_ptr(); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp index 6664be44858a..2b8bb84e17d1 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4FullCtxReduceOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/fullCtxReduce.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.h" #include #include @@ -29,7 +29,7 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107Tp4::fullCtxReduce; +namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107::fullCtxReduce; //! Merges the PrimTS split-KV decode scratch (partial_o [B, Hkv, S, R, 128] 16-bit, partial_stats [B, Hkv, S, R] //! fp32 log2-LSE) of the MiniMax-M3 TP4 full-context decode into out (BF16 [B * SQ, Hkv * hR, 128]); with q8 / sf diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp index 3ae68e7b4681..af7a88001999 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4ProducerPwOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp4/producerPw4.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/producerPw.h" #include #include @@ -33,7 +33,7 @@ namespace torch_ext namespace { -namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp4::producerPw4; +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107::producerPw; constexpr int64_t kHeadDim = 128; constexpr int64_t kPageSize = 128; @@ -70,9 +70,9 @@ std::vector minimaxM3Sm107Tp4ProducerPw(torch::Tensor const& pack c10::cuda::CUDAGuard const deviceGuard{packed.device()}; static int const sSmVersion = tensorrt_llm::common::getSMVersion(); TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp4_producer_pw requires SM100 or newer"); - int64_t const headsQ = pw::kHeadsQ; - int64_t const headsKV = pw::kHeadsKV; - int64_t const headsIndex = pw::kHeadsIndex; + int64_t const headsQ = pw::kHeadsQ<4>; + int64_t const headsKV = pw::kHeadsKV<4>; + int64_t const headsIndex = pw::kHeadsIndex<4>; int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim && packed.is_contiguous(), @@ -156,7 +156,7 @@ std::vector minimaxM3Sm107Tp4ProducerPw(torch::Tensor const& pack p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); p.rotaryCosSin = rotaryCosSin.data_ptr(); p.positionIds = positionIds.data_ptr(); - cudaError_t const err = pw::invokeProducerPw4(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + cudaError_t const err = pw::invokeProducerPw<4>(p, static_cast(numTokens), pdl, early, prewaitStepInputs, at::cuda::getCurrentCUDAStream(packed.get_device())); if (err != cudaSuccess) { diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp deleted file mode 100644 index 59009e461d33..000000000000 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4StepPrologueOp.cpp +++ /dev/null @@ -1,68 +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/minimaxM3Sm107Tp4/stepPrologue.h" - -#include -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -//! The decode graph's step-input prologue in one launch (see kernels/minimaxM3Sm107Tp4/stepPrologue.h). -//! -//! @param window The prologue's device window: contiguous uint8, at least the header and `maxDesc` descriptors (the -//! pack's descriptor area; the header's two counts add up to at most that). Its destinations are raw device addresses -//! the caller guarantees are live. -//! @param epoch Optional one-element int64 step counter, incremented once. -void minimaxM3Sm107Tp4StepPrologue(torch::Tensor& window, std::optional epoch, int64_t maxDesc) -{ - TORCH_CHECK(window.is_cuda() && window.scalar_type() == torch::kUInt8 && window.is_contiguous(), - "minimax_m3_sm107_tp4_step_prologue: window must be a contiguous uint8 CUDA tensor"); - TORCH_CHECK(maxDesc > 0 && window.numel() >= 16 + 24 * maxDesc, - "minimax_m3_sm107_tp4_step_prologue: window is shorter than its header and max_desc descriptors"); - c10::cuda::CUDAGuard const deviceGuard{window.device()}; - int64_t* epochPtr = nullptr; - if (epoch.has_value()) - { - TORCH_CHECK(epoch->is_cuda() && epoch->scalar_type() == torch::kInt64 && epoch->numel() == 1 - && epoch->device() == window.device(), - "minimax_m3_sm107_tp4_step_prologue: epoch must be a one-element int64 tensor on the window's device"); - epochPtr = epoch->data_ptr(); - } - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeStepPrologue( - window.data_ptr(), epochPtr, maxDesc, at::cuda::getCurrentCUDAStream(window.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_step_prologue launch failed: ", cudaGetErrorString(err)); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def("minimax_m3_sm107_tp4_step_prologue(Tensor(a!) window, Tensor(b!)? epoch, int max_desc) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("minimax_m3_sm107_tp4_step_prologue", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4StepPrologue); -} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp deleted file mode 100644 index 1172d809a817..000000000000 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp4SwigluOaiOp.cpp +++ /dev/null @@ -1,65 +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/minimaxM3Sm107Tp4/swigluOai.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -//! The dense MLP's swigluoai activation (see kernels/minimaxM3Sm107Tp4/swigluOai.h): x bf16 [rows, 2d] -> out bf16 -//! [rows, d], PDL-launched with an early release when `pdl`. -void minimaxM3Sm107Tp4SwigluOai(torch::Tensor const& x, torch::Tensor& out, bool pdl) -{ - TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), - "minimax_m3_sm107_tp4_swiglu_oai: x must be a contiguous 2-D bfloat16 CUDA tensor"); - int64_t const rows = x.size(0); - int64_t const d = x.size(1) / 2; - TORCH_CHECK(x.size(1) % 4 == 0 && reinterpret_cast(x.data_ptr()) % 4 == 0, - "minimax_m3_sm107_tp4_swiglu_oai: x rows must be 4-byte aligned with an even half width"); - TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() - && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 4 == 0, - "minimax_m3_sm107_tp4_swiglu_oai: out must be a contiguous, 4-byte aligned bfloat16 [rows, d]"); - if (rows == 0) - { - return; - } - c10::cuda::CUDAGuard const deviceGuard{x.device()}; - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp4::invokeSwigluOai( - reinterpret_cast(x.data_ptr()), reinterpret_cast(out.data_ptr()), rows, d / 2, pdl, - at::cuda::getCurrentCUDAStream(x.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp4_swiglu_oai launch failed: ", cudaGetErrorString(err)); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def("minimax_m3_sm107_tp4_swiglu_oai(Tensor x, Tensor(a!) out, bool pdl) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("minimax_m3_sm107_tp4_swiglu_oai", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp4SwigluOai); -} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp index 7b64d911263f..72c6c223a507 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8DrafterRopeFp8KvOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/drafterRopeFp8Kv.h" #include #include @@ -28,9 +28,9 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -namespace drk = tensorrt_llm::kernels::minimaxM3Sm107Tp8::drafterRopeFp8Kv; +namespace drk = tensorrt_llm::kernels::minimaxM3Sm107::drafterRopeFp8Kv; -//! The MiniMax-M3 TP8 drafter layer's generation preprocess (kernels/minimaxM3Sm107Tp8/drafterRopeFp8Kv.h): RoPE on +//! The MiniMax-M3 TP8 drafter layer's generation preprocess (kernels/minimaxM3Sm107/drafterRopeFp8Kv.h): RoPE on //! the packed BF16 qkv rows, the E4M3 query into q_out, and the step's K / V rows into the paged E4M3 cache kv_pages //! ([sub-pages, Hkv, tokens_per_block, 128]) at the sub-pages block_tables ([requests, 2, max blocks]) names. void minimaxM3Sm107Tp8DrafterRopeFp8Kv(torch::Tensor const& qkv, torch::Tensor const& cosSin, diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp index e4d5ca2f6b54..c1021f2c7b0d 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailOp.cpp @@ -15,7 +15,7 @@ * limitations under the License. */ -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h" #include "tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h" #include @@ -39,8 +39,8 @@ static void foldTailImpl(bool oneWarp, torch::Tensor const& ctrl, torch::Tensor srcDraft, dstDraft, srcLens, dstLens, ids, draftIds, posOff, kvOff, tokensPerStep); auto const stream = at::cuda::getCurrentCUDAStream(ctrl.get_device()); auto const err = oneWarp - ? tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeFoldTailOneWarp(p, stream) - : tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeFoldTail(p, static_cast(rows), stream); + ? tensorrt_llm::kernels::minimaxM3Sm107::invokeFoldTailOneWarp(p, stream) + : tensorrt_llm::kernels::minimaxM3Sm107::invokeFoldTail(p, static_cast(rows), stream); TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_fold_tail launch failed: ", cudaGetErrorString(err)); } diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h index fd8bcc99f10b..508b5cad1837 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FoldTailParams.h @@ -17,7 +17,7 @@ #pragma once -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/foldTail.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/foldTail.h" #include #include @@ -45,7 +45,7 @@ inline void check1d(torch::Tensor const& t, char const* name) //! The fold tail's operands checked and packed (the standalone fold tail op and any op that runs the tail in-kernel, //! e.g. the drafter's last-pass draft_step, build them identically). -inline tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams buildParams(torch::Tensor const& ctrl, +inline tensorrt_llm::kernels::minimaxM3Sm107::FoldTailParams buildParams(torch::Tensor const& ctrl, torch::Tensor const& slots, int64_t rows, torch::Tensor const& srcNew, torch::Tensor& dstNew, torch::Tensor const& srcNext, torch::Tensor& dstNext, torch::Tensor const& srcDraft, torch::Tensor& dstDraft, torch::Tensor const& srcLens, torch::Tensor& dstLens, torch::Tensor& ids, torch::Tensor& draftIds, @@ -78,7 +78,7 @@ inline tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams buildParams(torc TORCH_CHECK(ids.numel() >= rows * tokensPerStep && posOff.numel() >= rows * tokensPerStep && draftIds.numel() >= rows * (tokensPerStep - 1), "minimax_m3_sm107_tp8_fold_tail: staging outputs are too small for rows * tokens_per_step"); - tensorrt_llm::kernels::minimaxM3Sm107Tp8::FoldTailParams p{}; + tensorrt_llm::kernels::minimaxM3Sm107::FoldTailParams p{}; p.maxRows = static_cast(rows); p.ctrl = ctrl.data_ptr(); p.slots = slots.data_ptr(); diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp index b7b06bf2749f..45fb4fae1b34 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8FullCtxReduceOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/fullCtxReduce.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/fullCtxReduce.h" #include #include @@ -29,7 +29,7 @@ TRTLLM_NAMESPACE_BEGIN namespace torch_ext { -namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107Tp8::fullCtxReduce; +namespace fcr = tensorrt_llm::kernels::minimaxM3Sm107::fullCtxReduce; //! Merges the PrimTS split-KV decode scratch (partial_o [B, Hkv, S, R, 128] 16-bit, partial_stats [B, Hkv, S, R] //! fp32 log2-LSE) of the MiniMax-M3 TP8 full-context decode into out (BF16 [B * SQ, Hkv * hR, 128]); with q8 / sf diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp index eb3759b28a58..d6a946dd405f 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8McAllreduceNormOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/mcAllreduceNormTp8.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/mcAllreduceNorm.h" #include #include @@ -70,13 +70,13 @@ std::vector minimaxM3Sm107Tp8McAllreduceNorm(torch::Tensor const& int const sm = tensorrt_llm::common::getSMVersion(); TORCH_CHECK(sm == 107, "minimax_m3_sm107_tp8_mc_allreduce_norm requires SM107, got SM", sm); TORCH_CHECK(cluster == kCluster, "minimax_m3_sm107_tp8_mc_allreduce_norm: only cluster 8 is built, got ", cluster); - TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107Tp8McAllreduceNormSupported( + TORCH_CHECK(tensorrt_llm::kernels::isMinimaxM3Sm107McAllreduceNormSupported( static_cast(maxt), static_cast(opts)), "minimax_m3_sm107_tp8_mc_allreduce_norm: unsupported (bucket ", maxt, ", opts ", opts, ")"); TORCH_CHECK( world >= 1 && world <= kMaxWorld && rank >= 0 && rank < world, "invalid rank ", rank, " / world ", world); - tensorrt_llm::kernels::MinimaxM3Sm107Tp8McAllreduceNormParams p; + tensorrt_llm::kernels::MinimaxM3Sm107McAllreduceNormParams p; if (x2.has_value()) { TORCH_CHECK(x2->scalar_type() == at::kBFloat16 && x2->is_contiguous() && x2->sizes() == residual.sizes(), @@ -153,7 +153,7 @@ std::vector minimaxM3Sm107Tp8McAllreduceNorm(torch::Tensor const& p.capWeight = reinterpret_cast<__nv_bfloat16 const*>(capW->data_ptr()); p.capEps = static_cast(capEps); } - tensorrt_llm::kernels::invokeMinimaxM3Sm107Tp8McAllreduceNorm(p, static_cast(maxt), + tensorrt_llm::kernels::invokeMinimaxM3Sm107McAllreduceNorm(p, static_cast(maxt), static_cast(opts), pdl, at::cuda::getCurrentCUDAStream(residual.get_device())); if (quant) { diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp index acd7061eadbb..7df6f5eac0ef 100644 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp +++ b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8ProducerPwOp.cpp @@ -16,7 +16,7 @@ */ #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/minimaxM3Sm107Tp8/producerPw8.h" +#include "tensorrt_llm/kernels/minimaxM3Sm107/producerPw.h" #include #include @@ -33,7 +33,7 @@ namespace torch_ext namespace { -namespace pw = tensorrt_llm::kernels::minimaxM3Sm107Tp8::producerPw8; +namespace pw = tensorrt_llm::kernels::minimaxM3Sm107::producerPw; constexpr int64_t kHeadDim = 128; constexpr int64_t kPageSize = 128; @@ -70,9 +70,9 @@ std::vector minimaxM3Sm107Tp8ProducerPw(torch::Tensor const& pack c10::cuda::CUDAGuard const deviceGuard{packed.device()}; static int const sSmVersion = tensorrt_llm::common::getSMVersion(); TORCH_CHECK(sSmVersion >= 100, "minimax_m3_sm107_tp8_producer_pw requires SM100 or newer"); - int64_t const headsQ = pw::kHeadsQ; - int64_t const headsKV = pw::kHeadsKV; - int64_t const headsIndex = pw::kHeadsIndex; + int64_t const headsQ = pw::kHeadsQ<8>; + int64_t const headsKV = pw::kHeadsKV<8>; + int64_t const headsIndex = pw::kHeadsIndex<8>; int64_t const totalHeads = headsQ + 2 * headsKV + headsIndex + 1; TORCH_CHECK(packed.scalar_type() == at::kBFloat16 && packed.dim() == 2 && packed.size(1) == totalHeads * kHeadDim && packed.is_contiguous(), @@ -156,7 +156,7 @@ std::vector minimaxM3Sm107Tp8ProducerPw(torch::Tensor const& pack p.indexKWeight = reinterpret_cast<__nv_bfloat16 const*>(indexKWeight.data_ptr()); p.rotaryCosSin = rotaryCosSin.data_ptr(); p.positionIds = positionIds.data_ptr(); - cudaError_t const err = pw::invokeProducerPw8(p, static_cast(numTokens), pdl, early, prewaitStepInputs, + cudaError_t const err = pw::invokeProducerPw<8>(p, static_cast(numTokens), pdl, early, prewaitStepInputs, at::cuda::getCurrentCUDAStream(packed.get_device())); if (err != cudaSuccess) { diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp deleted file mode 100644 index a6cda645ba7f..000000000000 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SiluAndMulOp.cpp +++ /dev/null @@ -1,66 +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/minimaxM3Sm107Tp8/siluAndMul.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -//! The drafter MLP's plain SwiGLU (see kernels/minimaxM3Sm107Tp8/siluAndMul.h): x bf16 [rows, 2d] -> out bf16 [rows, -//! d]. -void minimaxM3Sm107Tp8SiluAndMul(torch::Tensor const& x, torch::Tensor& out) -{ - TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kBFloat16 && x.dim() == 2 && x.is_contiguous(), - "minimax_m3_sm107_tp8_silu_and_mul: x must be a contiguous 2-D bfloat16 CUDA tensor"); - int64_t const rows = x.size(0); - int64_t const d = x.size(1) / 2; - TORCH_CHECK(x.size(1) % 16 == 0 && reinterpret_cast(x.data_ptr()) % 16 == 0, - "minimax_m3_sm107_tp8_silu_and_mul: x rows must be 16-byte aligned with a half width a multiple of 8"); - TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kBFloat16 && out.dim() == 2 && out.is_contiguous() - && out.size(0) == rows && out.size(1) == d && reinterpret_cast(out.data_ptr()) % 16 == 0, - "minimax_m3_sm107_tp8_silu_and_mul: out must be a contiguous, 16-byte aligned bfloat16 [rows, d]"); - TORCH_CHECK(rows <= 65535, "minimax_m3_sm107_tp8_silu_and_mul: at most 65535 rows"); - if (rows == 0) - { - return; - } - c10::cuda::CUDAGuard const deviceGuard{x.device()}; - auto const err - = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSiluAndMul(reinterpret_cast(x.data_ptr()), - reinterpret_cast(out.data_ptr()), rows, d, at::cuda::getCurrentCUDAStream(x.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_silu_and_mul launch failed: ", cudaGetErrorString(err)); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def("minimax_m3_sm107_tp8_silu_and_mul(Tensor x, Tensor(a!) out) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("minimax_m3_sm107_tp8_silu_and_mul", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SiluAndMul); -} diff --git a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp b/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp deleted file mode 100644 index 401aa5cc61fd..000000000000 --- a/cpp/tensorrt_llm/thop/minimaxM3/minimaxM3Sm107Tp8SpecOffsetsOp.cpp +++ /dev/null @@ -1,74 +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/minimaxM3Sm107Tp8/specOffsets.h" - -#include -#include -#include - -TRTLLM_NAMESPACE_BEGIN - -namespace torch_ext -{ - -namespace -{ - -void check(torch::Tensor const& t, char const* name) -{ - TORCH_CHECK(t.is_cuda() && t.scalar_type() == torch::kInt32 && t.dim() == 1 && t.stride(0) == 1, - "minimax_m3_sm107_tp8_spec_offsets: ", name, " must be a contiguous 1-D int32 CUDA tensor"); -} - -} // namespace - -//! pos += pos_off and kv += kv_off (equal lengths per pair) in one plain launch -//! (kernels/minimaxM3Sm107Tp8/specOffsets.h). -void minimaxM3Sm107Tp8SpecOffsets( - torch::Tensor& pos, torch::Tensor const& posOff, torch::Tensor& kv, torch::Tensor const& kvOff) -{ - check(pos, "pos"); - check(posOff, "pos_off"); - check(kv, "kv"); - check(kvOff, "kv_off"); - TORCH_CHECK(pos.numel() == posOff.numel() && kv.numel() == kvOff.numel(), - "minimax_m3_sm107_tp8_spec_offsets: each tensor and its offsets must have the same length"); - if (pos.numel() == 0 && kv.numel() == 0) - { - return; - } - c10::cuda::CUDAGuard const deviceGuard{pos.device()}; - auto const err = tensorrt_llm::kernels::minimaxM3Sm107Tp8::invokeSpecOffsets(pos.data_ptr(), - posOff.data_ptr(), pos.numel(), kv.data_ptr(), kvOff.data_ptr(), kv.numel(), - at::cuda::getCurrentCUDAStream(pos.get_device())); - TORCH_CHECK(err == cudaSuccess, "minimax_m3_sm107_tp8_spec_offsets launch failed: ", cudaGetErrorString(err)); -} - -} // namespace torch_ext - -TRTLLM_NAMESPACE_END - -TORCH_LIBRARY_FRAGMENT(trtllm, m) -{ - m.def("minimax_m3_sm107_tp8_spec_offsets(Tensor(a!) pos, Tensor pos_off, Tensor(b!) kv, Tensor kv_off) -> ()"); -} - -TORCH_LIBRARY_IMPL(trtllm, CUDA, m) -{ - m.impl("minimax_m3_sm107_tp8_spec_offsets", &tensorrt_llm::torch_ext::minimaxM3Sm107Tp8SpecOffsets); -} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/config.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/config.py new file mode 100644 index 000000000000..0d494c303466 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/config.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Configuration interpretation shared by the SM107 MiniMax-M3 targets.""" + +from __future__ import annotations + +import copy +import dataclasses +from typing import List + +from transformers import PretrainedConfig + +from tensorrt_llm._torch.model_config import ModelConfig + + +def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: + model_type = getattr(pretrained_config, "model_type", None) + if model_type == "minimax_m3_vl": + return True + if model_type == "minimax_m3": + return False + return any( + "ForConditionalGeneration" in a + for a in (getattr(pretrained_config, "architectures", None) or []) + ) + + +def _text_model_config(model_config: ModelConfig) -> ModelConfig: + """The model config with the language model's ``text_config`` as its pretrained config. + + The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed + sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and + embedding tie carried over. A text config is returned unchanged. + """ + if not _is_vl_config(model_config.pretrained_config): + return model_config + cfg = copy.deepcopy(model_config) + text = getattr(cfg.pretrained_config, "text_config", None) + if text is None: + raise ValueError("MiniMax-M3 VL config has no `text_config`") + if isinstance(text, dict): + out = PretrainedConfig() + for k, v in text.items(): + setattr(out, k, v) + text = out + if ( + getattr(text, "torch_dtype", None) is None + and getattr(cfg.pretrained_config, "torch_dtype", None) is not None + ): + text.torch_dtype = cfg.pretrained_config.torch_dtype + if ( + getattr(text, "tie_word_embeddings", None) is None + and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None + ): + text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings + return dataclasses.replace(cfg, pretrained_config=text) + + +def _validate_runtime(model_config: ModelConfig) -> None: + """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing + sees only its KV cache algorithm; this check at construction refuses every other runtime. + + The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index + projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. + """ + sparse = model_config.sparse_attention_config + if ( + sparse is None + or getattr(sparse, "algorithm", None) != "minimax_m3" + or sparse.implementation != "msa" + ): + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" + ) + if getattr(sparse, "indexer_kv_dtype", None) != "fp8": + raise ValueError( + "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" + ) + mode = model_config.quant_config.quant_mode + if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): + raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") + + +def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = cfg.sparse_attention_config["sparse_attention_freq"] + return [i for i, f in enumerate(freq) if int(f) != 0] + + +def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: + freq = getattr(cfg, "moe_layer_freq", None) + if freq is None: + return list(range(cfg.num_hidden_layers)) + return [i for i, f in enumerate(freq) if int(f) != 0] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/__init__.py new file mode 100644 index 000000000000..38b4908f0eb4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/index_score_runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/index_score_runtime.py new file mode 100644 index 000000000000..2efe3a772e99 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/index_score_runtime.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""Shared scorer compilation; subclasses own scheduling flags and compiled caches.""" + +import cutlass +import torch +from cutlass import cute + +from tensorrt_llm.logger import logger + +from .m3_index_score import IndexDecodeScoreKernel + + +class IndexScoreRunner: + """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, + split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile + serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block + per CTA; ``early_trigger`` is ``INDEX_SCORE_SET`` for every key of the process.""" + + kernel_cache: dict + PRESTAGE_OLD: bool + EARLY_TRIGGER: bool + SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim + SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K + # BLOCK_Q must fit one warp's worth of epilogue lanes. + MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile + EARLY_PREFETCH_DEPTH = 1 + + @classmethod + def is_supported( + cls, + *, + q_dtype: torch.dtype, + num_heads: int, + head_dim: int, + page_size: int, + max_decode_query_len: int, + ) -> bool: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + from tensorrt_llm._utils import is_sm_100f + + return ( + is_sm_100f() + and q_dtype in TORCH_TO_CUTE_DTYPE + and head_dim == cls.SUPPORTED_HEAD_DIM + and page_size == cls.SUPPORTED_PAGE_SIZE + and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q + and max_decode_query_len >= 1 + ) + + @classmethod + def _compile( + cls, + q_dtype: torch.dtype, + num_heads: int, + max_decode_query_len: int, + head_dim: int, + split_k: int, + early: bool, + ) -> None: + from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE + + key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + if key in cls.kernel_cache: + return + cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] + page_size = cls.SUPPORTED_PAGE_SIZE + sym_total_tokens = cute.sym_int() + sym_batch = cute.sym_int() + + # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. + def _sym_stride(): + return cute.sym_int64(divisibility=16) + + q_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (sym_total_tokens, num_heads, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). + k_fake = cute.runtime.make_fake_tensor( + cute_dtype, + (cute.sym_int(), page_size, head_dim), + stride=(_sym_stride(), _sym_stride(), 1), + ) + bt_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) + ) + # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. + score_fake = cute.runtime.make_fake_tensor( + cutlass.Float32, + (num_heads, sym_total_tokens, cute.sym_int()), + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), + ) + sl_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (sym_batch,), stride_order=(0,) + ) + fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = IndexDecodeScoreKernel( + cute_dtype, + num_heads, + max_decode_query_len, + split_k, + head_dim, + prefetch_old=True, + early_issue=early, + prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, + prestage_old=early and cls.PRESTAGE_OLD, + early_trigger=cls.EARLY_TRIGGER, + ) + if cls.EARLY_TRIGGER and not cls.kernel_cache: + logger.info( + "tp8 index scorer: the block selector is released after the scorer's setup (INDEX_SCORE_SET)" + ) + cls.kernel_cache[key] = cute.compile( + kernel, + q_fake, + k_fake, + bt_fake, + score_fake, + sl_fake, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/m3_index_score.py new file mode 100644 index 000000000000..77cf7d70a367 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/m3_index_score.py @@ -0,0 +1,549 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""Shared index-score core; targets retain scheduling policy and compile caches.""" + +import cutlass +from cuda.bindings.driver import CUstream +from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( + EVICT_FIRST, + fp8x4_to_fp16x4, + mma_sync, + simple_tma_copy, +) +from tensorrt_llm._torch.kernels.blackwell.utils import ( + TRTLLM_ENABLE_PDL, + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + + +@dsl_user_op +def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: + """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" + llvm.inline_asm( + res=None, + operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], + asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", + constraints="l,r", + has_side_effects=True, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@cute.jit +def _fp8_to_f16_mma_fragments(src: cute.Tensor): + """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" + src_elems = cute.size(src) + src_u32 = cute.recast_tensor(src, Uint32) + src_f16 = cute.make_rmem_tensor(src_elems, Float16) + src_f16_u32 = cute.recast_tensor(src_f16, Uint32) + # Packed conversion; faster and fewer SASS instructions than + # src.load().to(Float16). + for i in cutlass.range_constexpr(src_elems // 4): + converted = fp8x4_to_fp16x4(src_u32[i]) + src_f16_u32[i * 2] = converted[0] + src_f16_u32[i * 2 + 1] = converted[1] + lower = cute.make_rmem_tensor(src_elems // 2, Float16) + upper = cute.make_rmem_tensor(src_elems // 2, Float16) + + # FP8 ldmatrix yields four consecutive values along K. Split each group + # into the lower two and upper two for the two FP16 MMA k-fragments. + for i in cutlass.range_constexpr(src_elems // 2): + lower[i] = src_f16[(i // 2) * 4 + i % 2] + upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] + return lower, upper + + +class IndexDecodeScoreKernel: + """Per-block max index score for one decode step. + + Grid is (batch, split_k): CTA (b, s) walks KV blocks + s, s + split_k, s + 2 * split_k, ... of request b and writes each block's + score directly, so no cross-CTA reduction is needed. CTAs whose split_id + exceeds the request's block count exit immediately. + """ + + BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint + BAR_MMA = 1 + num_stages = 2 + + def __init__( + self, + dtype: type[cutlass.Numeric], + num_heads: int, + max_decode_query_len: int, + split_k: int, + head_dim: int = 128, + prefetch_old: bool = False, + early_issue: bool = False, + prefetch_depth: int = 0, + prestage_old: bool = False, + early_trigger: bool = False, + ): + # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that + # hold only tokens older than this step (the predecessor writes only this step's tokens). + self.prefetch_old = prefetch_old + # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency + # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. + # The tiles are read after the wait, so their contents are final; the block table is only written between + # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles + # turn out to lie past the request's true length drains them before exiting. + self.early_issue = early_issue + # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at + # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. + self.prefetch_depth = prefetch_depth + # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K + # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's + # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. + self.prestage_old = prestage_old and early_issue + # early_trigger: release the dependent block selector right after this CTA's setup (see the module docstring). + self.early_trigger = early_trigger + self.dtype = dtype + self.num_heads = num_heads + self.max_decode_query_len = max_decode_query_len + self.split_k = split_k + self.head_dim = head_dim + + def _swizzle_elems(self) -> int: + """Elements spanned by one 128-byte swizzle atom.""" + return 128 * 8 // self.dtype.width + + def _sq_layout(self): + """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_q = self.num_heads * self.max_decode_query_len + layout = cute.make_layout( + (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), + stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + def _sk_layout(self): + """Composed SMEM layout for the pipelined K tiles.""" + elems = self._swizzle_elems() + head_dim = self.head_dim + block_k = self.BLOCK_K + layout = cute.make_layout( + (1, block_k, (elems, head_dim // elems), self.num_stages), + stride=(0, elems, (1, block_k * elems), block_k * head_dim), + ) + return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) + + @cute.jit + def __call__( + self, + gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] + gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] + block_table: cute.Tensor, # [bs, max_pages] + score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] + seq_lens: cute.Tensor, # [bs] + stream: CUstream, + ): + num_heads = self.num_heads + head_dim = self.head_dim + MAX_DQL = self.max_decode_query_len + assert num_heads * MAX_DQL <= 32 + + batch = seq_lens.shape[0] + decode_query_len = gQ.shape[0] // batch + grid = (batch, self.split_k, 1) + block = (32 * 5, 1, 1) + + tma_g2s = cpasync.CopyBulkTensorTileG2SOp() + elems = self._swizzle_elems() + + q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gQ, (None, None, elems)), + self._sq_layout(), + (MAX_DQL, num_heads, head_dim), + ) + k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( + tma_g2s, + cute.logical_divide(gK_cache, (None, None, elems)), + self._sk_layout(), + (1, self.BLOCK_K, head_dim), + ) + + self.kernel( + q_tma_atom, + q_tma_tensor, + k_tma_atom, + k_tma_tensor, + gK_cache, + block_table, + score, + seq_lens, + decode_query_len, + ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) + + @cute.kernel + def kernel( + self, + q_tma_atom: cute.CopyAtom, + q_tma_tensor: cute.Tensor, + k_tma_atom: cute.CopyAtom, + k_tma_tensor: cute.Tensor, + gK_cache: cute.Tensor, + block_table: cute.Tensor, + score: cute.Tensor, + seq_lens: cute.Tensor, + decode_query_len, + ): + batch_id, split_id, _ = cute.arch.block_idx() + _, split_k, _ = cute.arch.grid_dim() + warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = cute.arch.lane_idx() + + NUM_HEADS = self.num_heads + MAX_DQL = self.max_decode_query_len + BLOCK_Q = NUM_HEADS * MAX_DQL + BLOCK_K = self.BLOCK_K + head_dim = self.head_dim + dtype = self.dtype + MMA_N = 8 + num_stages = self.num_stages + Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) + EPI_Q = Q_TILES * MMA_N + + sq_layout = self._sq_layout() + sk_layout = self._sk_layout() + + smem = SmemAllocator() + sK = smem.allocate_tensor( + dtype, + sk_layout.outer, + byte_alignment=128, + swizzle=sk_layout.inner, + )[0, None, None, None] + # sQ aliases the first K stage: Q is consumed into registers before the + # first K tile is needed, so the two never overlap in time. + sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) + # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q + # column mode. The target profile keeps the rank-2 view even for + # degenerate shapes such as MAX_DQL == 1. + q_tma_elems = self._swizzle_elems() + sQ = cute.coalesce( + cute.group_modes(sQ_tma, 0, 2), + target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), + ) + epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) + + tma_full_mbar = smem.allocate_array(Int64, num_stages) + tma_empty_mbar = smem.allocate_array(Int64, num_stages) + + # Barrier init and descriptor prefetch do not depend on the predecessor + # grid, so they are the work PDL can overlap. Every CTA runs them, even + # the ones that own no blocks, since ownership is not known until + # seq_lens is read below. + if warp_id == 0: + with cute.arch.elect_one(): + for i in cutlass.range_constexpr(num_stages): + cute.arch.mbarrier_init(tma_full_mbar + i, 1) + cute.arch.mbarrier_init(tma_empty_mbar + i, 128) + cute.arch.mbarrier_init_fence() + elif warp_id == 1: + cpasync.prefetch_descriptor(q_tma_atom) + cpasync.prefetch_descriptor(k_tma_atom) + cute.arch.sync_threads() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + + if cutlass.const_expr(self.prefetch_old): + if warp_id == 4: + with cute.arch.elect_one(): + seqlen0 = seq_lens[batch_id] + num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) + stable = (seqlen0 - decode_query_len) // BLOCK_K + if stable > num_blocks0: + stable = num_blocks0 + if cutlass.const_expr(self.prefetch_depth > 0): + if stable > split_id + self.prefetch_depth * split_k: + stable = split_id + self.prefetch_depth * split_k + for block_id in range(split_id, stable, split_k): + page = block_table[batch_id, block_id] + _l2_prefetch_bulk( + gK_cache[page, None, None].iterator.toint().to(Int64), + Int32(BLOCK_K * head_dim * (dtype.width // 8)), + ) + + # Without early_trigger the dependent block selector is released as the CTAs exit. + early_q = False + early_k = False + early_page = Int32(0) + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + seqlen_hint = seq_lens[batch_id] + hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) + if hint_blocks > block_table.shape[1]: + hint_blocks = block_table.shape[1] + if split_id < hint_blocks: + early_q = True + early_page = block_table[batch_id, split_id] + if cutlass.const_expr(self.prestage_old): + # Stage the first K tile now when every token of the block predates this step. + if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: + early_k = True + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + griddepcontrol_wait() + + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + gQ_early = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) + ) + simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) + if not early_k: + cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + tma_full_mbar + 1 % num_stages, + BLOCK_K * head_dim * (dtype.width // 8), + ) + simple_tma_copy( + k_tma_atom, + k_tma_tensor[early_page, None, None], + sK[None, None, 1 % num_stages], + tma_full_mbar + 1 % num_stages, + cache_policy=EVICT_FIRST, + ) + + # The length that bounds this grid's work is read after the wait (the + # pre-wait reads above are prefetch / early-issue hints only, clamped to + # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens + # patches every step, and num_blocks bounds an unmasked block_table load + # and an unmasked score store, so a stale length is an out-of-bounds write + # rather than just a wrong answer. + seqlen = seq_lens[batch_id] + num_blocks = cute.ceil_div(seqlen, BLOCK_K) + + if split_id < num_blocks: + if warp_id == 4: + # TMA warp + tma_stage = 0 + tma_parity = 1 + first_block = split_id + + if not early_q: + gQ_tile = cute.local_tile( + cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), + tiler=(MAX_DQL, NUM_HEADS, head_dim), + coord=(0, 0, 0), + ) + cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) + with cute.arch.elect_one(): + Q_size = BLOCK_Q * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) + # TMA bounds-checks rows when the runtime decode_query_len is + # below MAX_DQL; padded Q columns are masked before the stores. + simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + if early_q: + # The first K tile was issued right after the wait. + first_block = split_id + split_k + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + for block_id in range(first_block, num_blocks, split_k): + page_id = block_table[batch_id, block_id] + gK_tile = k_tma_tensor[page_id, None, None] + k_mbar = tma_full_mbar + tma_stage + + cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) + with cute.arch.elect_one(): + K_size = BLOCK_K * head_dim * (dtype.width // 8) + cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) + simple_tma_copy( + k_tma_atom, + gK_tile, + sK[None, None, tma_stage], + k_mbar, + cache_policy=EVICT_FIRST, + ) + + tma_stage = (tma_stage + 1) % num_stages + if tma_stage == 0: + tma_parity ^= 1 + + else: + # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T + sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) + q_start = seqlen - decode_query_len + + elems = 128 // dtype.width # 16B + MMA_K = 32 * 8 // dtype.width # 32B + + # Pre-compute ldmatrix addresses. + # sK loads a [16 x 16B] tile: + # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) + # sQ loads an [8 x 32B] tile: + # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) + sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) + + # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) + # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) + sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] + sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] + + ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) + ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) + + rQ = cute.make_rmem_tensor( + ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype + ) + rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) + rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + for q in cutlass.range_constexpr(Q_TILES): + cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) + cute.arch.mbarrier_arrive(tma_empty_mbar) + + tma_stage = 1 % self.num_stages + tma_parity = 0 + if tma_stage == 0: + tma_parity ^= 1 + + if cutlass.const_expr(dtype is Float8E4M3FN): + rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) + q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) + rQ_f16[None, None, None, 0].store(q_lower.load()) + rQ_f16[None, None, None, 1].store(q_upper.load()) + + for block_id in range(split_id, num_blocks, split_k): + rC.fill(0.0) + + if warp_id == 0: + cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + for k in cutlass.range_constexpr(head_dim // MMA_K): + cute.copy( + ldsm_atom, + sK_ldsm[None, (None, k, tma_stage)], + rK[None, None, k], + ) + for m in cutlass.range_constexpr(2): + if cutlass.const_expr(dtype is Float8E4M3FN): + rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK_lower, + rQ_f16[None, k, n, 0], + rC[None, m, n], + ) + rC[None, m, n] = mma_sync( + rK_upper, + rQ_f16[None, k, n, 1], + rC[None, m, n], + ) + else: + for n in cutlass.range_constexpr(Q_TILES): + rC[None, m, n] = mma_sync( + rK[None, m, k], + rQ[(None, k % 2), k // 2, n], + rC[None, m, n], + ) + + cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) + + k_start = block_id * BLOCK_K + warp_id * 32 + + for q in cutlass.range_constexpr(Q_TILES): + for i in cutlass.range_constexpr(4): + for j in cutlass.range_constexpr(2): + col = q * 8 + (lane_id % 4) * 2 + j + q_local_pos = col % MAX_DQL + q_pos = q_start + q_local_pos + k_pos = k_start + i * 8 + lane_id // 4 + rC[q * 8 + i * 2 + j] = ( + rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") + ) + + for q in cutlass.range_constexpr(Q_TILES): + # Thread-local reduction along the BLOCK_K dim. + rScore = cute.make_rmem_tensor(2, Float32) + rScore.fill(float("-inf")) + for i in cutlass.range_constexpr(4): + rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) + rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) + + # Warp reduction among lanes 0, 4, 8, 12, ... + for i in cutlass.range_constexpr(3): + offset = 4 << i + other0 = cute.arch.shuffle_sync_bfly( + rScore[0], offset=offset, mask=-1, mask_and_clamp=31 + ) + other1 = cute.arch.shuffle_sync_bfly( + rScore[1], offset=offset, mask=-1, mask_and_clamp=31 + ) + rScore[0] = cute.arch.fmax(rScore[0], other0) + rScore[1] = cute.arch.fmax(rScore[1], other1) + + if lane_id * 2 < MMA_N: + epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] + epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] + cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) + + head_id = lane_id // MAX_DQL + q_local_pos = lane_id - head_id * MAX_DQL + valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len + if lane_id < BLOCK_Q and valid_q: + final_score = epi_buffer[lane_id, 0] + for i in cutlass.range_constexpr(1, 4): + final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) + + t = batch_id * decode_query_len + q_local_pos + score[head_id, t, block_id] = final_score + + tma_stage = (tma_stage + 1) % self.num_stages + if tma_stage == 0: + tma_parity ^= 1 + else: + if cutlass.const_expr(self.early_issue): + if warp_id == 4: + if early_q: + # Early tiles past the request's true length: let them land before the CTA exits. + cute.arch.mbarrier_wait(tma_full_mbar, 0) + cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/m3_select_long_dsl.py new file mode 100644 index 000000000000..10a84c896135 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/m3_select_long_dsl.py @@ -0,0 +1,285 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). + +Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long``, +one 1024-thread CTA per row), from a CTA of 384 threads: + +* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and + mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's + valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned + key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped + block and masks the result, so a thread's loads issue back to back. +* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of + the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. +* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by + that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate + against the whole row instead. +* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. + +The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait +(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are +compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and +loads them before the executor's warmup, so a decode step never compiles one. +""" + +import cutlass +import cutlass.cute as cute +from cutlass.utils.smem_allocator import SmemAllocator + +from tensorrt_llm._torch.kernels.blackwell.utils import ( + griddepcontrol_launch_dependents, + griddepcontrol_wait, +) + +SELECT_THREADS = 384 +SELECT_WARPS = SELECT_THREADS // 32 +# 1M tokens + draft tokens at 128 tokens per block. +MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row +MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging +TOPK = 16 # capacity: fixed: MSA's top-k +NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) +# The key of -inf (the reference's 0xff800000): never emitted. +NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size +INIT_SCORE = 1.0e30 +LOCAL_SCORE = 1.0e29 +# (keys per thread, rows up to this many valid blocks) +ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) + +# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare +# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group +# minima [16], per-warp candidate counts [SELECT_WARPS]. +_CAND = 0 +_PICK_KEY = _CAND + 2 * MAX_CAND +_PICK_BLOCK = _PICK_KEY + TOPK +_WARP_MIN = _PICK_BLOCK + TOPK +_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 +_WARP_CNT = _GROUP_MIN + 16 +SMEM_WORDS = _WARP_CNT + SELECT_WARPS +SMEM_BYTES = 4 * SMEM_WORDS + + +def _key(score): + """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" + bits = cutlass.Float32(score).bitcast(cutlass.Int32) + return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) + + +def _imin(a, b): + return cutlass.select_(b < a, b, a) + + +def _imax(a, b): + return cutlass.select_(b > a, b, a) + + +def _forced_key(scores, head, block, query, init_blocks, local_start): + raw = scores[head, block, query] + s = cutlass.select_( + block >= local_start, + cutlass.Float32(LOCAL_SCORE), + cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), + ) + return _key(s) + + +@cute.jit +def _select_row( + row, + scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides + n_valid: cute.Tensor, # int32 [>= total_q] + out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] + init_blocks, + local_blocks, + num_kv_heads, + smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory + bar_id: cutlass.Constexpr, + loaded_trigger: cutlass.Constexpr, + raw_prewait, + n_valid_prewait: cutlass.Constexpr, +): + """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier + ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the + row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by + the caller before its grid wait; otherwise it is read here.""" + tid, _, _ = cute.arch.thread_idx() + lane = tid % cutlass.Int32(32) + warp = tid // cutlass.Int32(32) + query = row // num_kv_heads + head = row % num_kv_heads + if cutlass.const_expr(n_valid_prewait): + raw = raw_prewait + else: + raw = n_valid[query] + valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) + local_start = _imax(raw - local_blocks, cutlass.Int32(0)) + cand = cute.make_tensor( + cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) + ) + + if tid < TOPK: + smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) + smem[_PICK_BLOCK + tid] = cutlass.Int32(0) + + if valid <= TOPK: + # Every valid block is picked (thread t holds block t). + if tid < valid: + smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) + smem[_PICK_BLOCK + tid] = tid + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): + items = ITEM_CLASSES[ci][0] + lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] + hi = ITEM_CLASSES[ci][1] + if (valid > lo) & (valid <= hi): + key = cute.make_rmem_tensor(items, cutlass.Int32) + local_min = cutlass.Int32(NO_KEY) + for s in cutlass.range_constexpr(items): + # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. + block = tid + cutlass.Int32(s * SELECT_THREADS) + k = _forced_key( + scores, head, _imin(block, valid - 1), query, init_blocks, local_start + ) + key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) + local_min = _imin(local_min, key[s]) + local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) + if lane < 16: + smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + if cutlass.const_expr(loaded_trigger): + griddepcontrol_launch_dependents() + if tid < 16: + g = cutlass.Int32(NO_KEY) + for w in cutlass.range_constexpr(SELECT_WARPS): + g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) + smem[_GROUP_MIN + tid] = g + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + bound = cutlass.Int32( + -2147483648 + ) # the smallest key: signed keys, unlike the reference's unsigned ones + for g in cutlass.range_constexpr(16): + bound = _imax(bound, smem[_GROUP_MIN + g]) + + # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. + cnt = cutlass.Int32(0) + for s in cutlass.range_constexpr(items): + cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) + incl = cnt + for d in cutlass.range_constexpr(5): + up = cute.arch.shuffle_sync( + incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) + ) + incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) + if lane == 31: + smem[_WARP_CNT + warp] = incl + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + before = cutlass.Int32(0) + total = cutlass.Int32(0) + for w in cutlass.range_constexpr(SELECT_WARPS): + c = smem[_WARP_CNT + w] + before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) + total += c + slot = before + incl - cnt + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + if slot < MAX_CAND: + cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( + tid + cutlass.Int32(s * SELECT_THREADS) + ) + slot += 1 + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if total <= MAX_CAND: + for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): + mine = cand[i] + rank = cutlass.Int32(0) + for j in cutlass.range(total, unroll=8): + rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) + smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) + else: + # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. + for s in cutlass.range_constexpr(items): + if key[s] <= bound: + block = tid + cutlass.Int32(s * SELECT_THREADS) + rank = cutlass.Int32(0) + for j in cutlass.range(valid, unroll=1): + kj = _forced_key(scores, head, j, query, init_blocks, local_start) + less = (kj < key[s]) | ((kj == key[s]) & (j < block)) + rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) + if rank < TOPK: + smem[_PICK_KEY + rank] = key[s] + smem[_PICK_BLOCK + rank] = block + cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) + + if warp == 0: + # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. + v = cutlass.Int32(NO_KEY) + if lane < TOPK: + pk = smem[_PICK_KEY + lane] + if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): + v = smem[_PICK_BLOCK + lane] + for size_log in cutlass.range_constexpr(1, 6): + for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): + other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) + ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) + lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) + v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) + if lane < TOPK: + out[query, head, lane] = cutlass.select_( + v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v + ) + + +class _SelectLongDslKernel: + """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count + before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because + ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's + first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it + (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" + + def __init__( + self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool + ): + self.use_pdl = use_pdl + self.early_trigger = early_trigger + self.loaded_trigger = (not early_trigger) and trigger_phase == 1 + self.n_valid_prewait = n_valid_prewait + + @cute.kernel + def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): + row, _, _ = cute.arch.block_idx() + smem = SmemAllocator() + words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) + raw_prewait = cutlass.Int32(0) + if cutlass.const_expr(self.n_valid_prewait): + raw_prewait = n_valid[row // num_kv_heads] + griddepcontrol_wait() + if cutlass.const_expr(self.early_trigger): + griddepcontrol_launch_dependents() + _select_row( + row, + scores, + n_valid, + out, + init_blocks, + local_blocks, + num_kv_heads, + words, + bar_id=1, + loaded_trigger=self.loaded_trigger, + raw_prewait=raw_prewait, + n_valid_prewait=self.n_valid_prewait, + ) + + @cute.jit + def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): + num_kv_heads = scores.shape[0] + self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( + grid=(scores.shape[2] * num_kv_heads, 1, 1), + block=(SELECT_THREADS, 1, 1), + stream=stream, + use_pdl=self.use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/producer_tp8.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_layer/phase_producer.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/producer_tp8.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/triton_sparse_decode.py new file mode 100644 index 000000000000..5effe83c9824 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/triton_sparse_decode.py @@ -0,0 +1,892 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import triton +import triton.language as tl + + +@triton.jit +def _e2m1x2_to_f16x2(packed): + """Convert one packed E2M1 byte into its low/high FP16 values.""" + return tl.inline_asm_elementwise( + asm="""{ + .reg .b8 src; + .reg .b32 pair; + mov.b32 {src, _, _, _}, $2; + cvt.rn.f16x2.e2m1x2 pair, src; + mov.b32 {$0, $1}, pair; + }""", + constraints="=h,=h,r", + args=(packed,), + dtype=(tl.float16, tl.float16), + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_nvfp4_rows( + packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte + scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements + BLOCK_N: tl.constexpr, + D: tl.constexpr, + SCALE_COLS: tl.constexpr, + SF_VEC: tl.constexpr, +): + """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. + + The caller gathers the block scales in whichever layout its role stores + them, so this handles only the value math and the scale broadcast. The + per-tensor scale is applied by the caller. + """ + # Blackwell converts both E2M1 nibbles with one instruction. Joining the + # low/high results along a new innermost axis and flattening it rebuilds + # the original element order without scalar exponent/sign reconstruction. + value_lo, value_hi = _e2m1x2_to_f16x2(packed) + values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) + + scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) + scales = tl.reshape( + tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) + ) + return values * scales + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary +# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused + counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + stride_out_n, + stride_out_h, + stride_out_d, + num_kv_heads, + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place + EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query token, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + pid_b = pid_bc % total_q + pid_c = pid_bc // total_q + req_id = pid_b // decode_query_len + q_offset = pid_b - req_id * decode_query_len + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + if EARLY_TRIGGER: + # Dependents still wait for this grid to complete; triggering here only + # lets them start their input-independent prologue (e.g. weight loads). + tl.extra.cuda.gdc_launch_dependents() + + seq_len = tl.load(seq_lens + req_id) + query_pos = seq_len - decode_query_len + q_offset + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + + # Bound the walk by the token's own valid block count rather than a + # sentinel: the selector emits ascending ids and only pads with -1 past + # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk = tl.minimum(max_topk, num_blocks) + chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + off_h = tl.arange(0, BLOCK_SIZE_H) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + h_mask = off_h < gqa_group_size + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + pid_b * stride_qn + + (pid_h + off_h[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result + # matches a caller that widened q itself before the call. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + for it in tl.range(chunk_start, chunk_end): + if IDENTITY: + # Slot `it` of the identity selection, bounded by this row's own kv extent above. + blk = it + else: + blk = tl.load(cur_idx_ptr) + cur_idx_ptr = cur_idx_ptr + stride_tk + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, + # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K block scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) + p = tl.exp2(qk - m_ij[:, None]) + l_ij = tl.sum(p, axis=1) + if USE_LINEAR_SOFTMAX: + alpha = tl.exp2(m_i - m_ij) + acc_o = acc_o * alpha[:, None] + else: + acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] + if KV_NVFP4: + # V block scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = l_i * alpha + l_ij + else: + lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h + tl.store( + o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h + tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) + + if FUSED_MERGE: + # Split-K serial reduction: every chunk publishes its partial, and the + # chunk that arrives last for this (token, KV head) merges all of them + # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. + tl.debug_barrier() + counter = counter_ptr + pid_b * num_kv_heads + pid_kh + arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") + if arrived == NUM_TOPK_CHUNKS - 1: + tl.debug_barrier() + o_rows = ( + o_ptr + + pid_b * stride_o_b + + (pid_h + off_h[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d + ) + l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h + lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + lse_max = tl.maximum(lse_max, lse_c) + # All chunks empty (a CUDA-graph padding row): zero, never NaN. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + denom += tl.exp2(lse_c - lse_max) + denom = tl.where(denom > 0, denom, 1.0) + merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + for c in tl.static_range(NUM_TOPK_CHUNKS): + lse_c = tl.load( + l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" + ) + w_c = tl.exp2(lse_c - lse_max) / denom + o_c = tl.load( + o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" + ).to(tl.float32) + merged += o_c * w_c[:, None] + tl.store( + out_ptr + + pid_b * stride_out_n + + (pid_h + off_h[:, None]) * stride_out_h + + off_d[None, :] * stride_out_d, + merged.to(out_ptr.dtype.element_ty), + mask=hd_mask, + ) + tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.heuristics( + { + "BLOCK_SIZE_H": lambda args: max( + 16, + triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), + ), + "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), + } +) +@triton.jit( + do_not_specialize=["decode_query_len", "total_q"], + do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], +) +def _gqa_sparse_decode_pair_kernel( + q_ptr, # [total_q, num_heads, head_dim] + k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] + v_ptr, # paged V: same layout as K + kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False + k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] + v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order + k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 + v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 + t_ptr, # topk_idx: [num_kv_heads, total_q, topk] + o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT + lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT + block_table_ptr, # [num_reqs, max_blocks] + seq_lens, # [num_reqs] + total_q, + gqa_group_size, + head_dim, + max_topk, + sm_scale, + decode_query_len, + stride_qn, + stride_qh, + stride_qd, + stride_k_blk, + stride_k_h, + stride_k_pos, + stride_k_d, + stride_v_blk, + stride_v_h, + stride_v_pos, + stride_v_d, + stride_ksc_blk, + stride_ksc_h, + stride_vsc_blk, + stride_vsc_h, + stride_th, + stride_tn, + stride_tk, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_bt_b, + num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request + num_pages, # KV pool pages: a page id outside [0, num_pages) is not read + BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache + WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register + COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math + USE_PDL: tl.constexpr, + KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales + USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence + QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set + MAX_TOPK: tl.constexpr, + SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only + DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge + IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) +): + sm_scale_log2e = sm_scale * 1.4426950409 + if KV_NVFP4: + # Both per-tensor dequant scales factor out of the inner loop: QK is + # linear in K, and V's scale commutes with the softmax weights. Reading + # them from device memory keeps the launch free of a host sync, so + # CUDA-graph capture stays unconstrained. + sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) + v_global_scale = tl.load(v_gs_ptr) + # Split-K over the topk dimension: pid(0) folds (query group, chunk). + pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) + num_query_groups = total_q // QUERY_GROUP_SIZE + pid_g = pid_bc % num_query_groups + pid_c = pid_bc // num_query_groups + groups_per_req = decode_query_len // QUERY_GROUP_SIZE + req_id = pid_g // groups_per_req + q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE + pid_b = req_id * decode_query_len + q_group_offset + pid_h = pid_kh * gqa_group_size + chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_start = pid_c * chunk_size_topk + + if USE_PDL: + tl.extra.cuda.gdc_wait() + + seq_len = tl.load(seq_lens + req_id) + # Query rows are flattened as [query-in-group, GQA head]. + off_h = tl.arange(0, BLOCK_SIZE_H) + q_lane = off_h // gqa_group_size + h_lane = off_h - q_lane * gqa_group_size + h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE + q_token = pid_b + q_lane + query_pos = seq_len - decode_query_len + q_group_offset + q_lane + # Full-CG padding uses zero-length request rows. Clamp to an empty + # attention range instead of letting padded rows produce negative lengths. + kv_len = tl.maximum(query_pos + 1, 0) + kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) + + # Each query can have a different valid-block count at a page boundary. + # The selector puts those valid entries first and pads only the tail. + kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) + num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_0 = tl.minimum(max_topk, num_blocks_0) + chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) + if QUERY_GROUP_SIZE == 2: + kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) + num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K + real_topk_1 = tl.minimum(max_topk, num_blocks_1) + chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) + + off_n = tl.arange(0, BLOCK_SIZE_K) + off_d = tl.arange(0, BLOCK_SIZE_D) + if KV_NVFP4: + # Packed byte column and block-scale column. The NVFP4 path runs only at + # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. + off_dp = tl.arange(0, BLOCK_SIZE_D // 2) + off_s = tl.arange(0, SCALE_COLS) + d_mask = off_d < head_dim + hd_mask = h_mask[:, None] & d_mask[None, :] + bt_row = block_table_ptr + req_id * stride_bt_b + + m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + if USE_LINEAR_SOFTMAX: + # Keep the linear softmax denominator and reuse the producer's alpha. + # This removes one exp2 plus one log2 from every visited KV page; only + # the final LSE written for split-K merging still needs a logarithm. + l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) + else: + lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) + acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) + q = tl.load( + q_ptr + + q_token[:, None] * stride_qn + + (pid_h + h_lane[:, None]) * stride_qh + + off_d[None, :] * stride_qd, + mask=hd_mask, + other=0.0, + ) + # Widen before anything reads q.dtype: K, V and p are all cast to it below, + # so leaving q narrow would silently run the whole attention in FP8 and + # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this + # reproduces the values the caller's standalone widening kernel produced. + if WIDEN_Q: + q = q.to(COMPUTE_DTYPE) + kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 + + idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk + if QUERY_GROUP_SIZE == 2: + idx_ptr_1 = idx_ptr_0 + stride_tn + idx_0 = chunk_start + idx_1 = chunk_start + count_0 = chunk_end_0 - chunk_start + count_1 = chunk_end_1 - chunk_start + chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS + chunk_offsets = tl.arange(0, chunk_capacity) + sentinel = 2147483647 + if IDENTITY: + blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) + blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) + else: + blocks_0 = tl.load( + idx_ptr_0 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_0, + other=sentinel, + ) + blocks_1 = tl.load( + idx_ptr_1 + chunk_offsets * stride_tk, + mask=chunk_offsets < count_1, + other=sentinel, + ) + chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 + # Equal chunks take the measured shared-page loop. A divergent chunk + # walks the sorted union, bounded by the sum of both row lengths. + num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) + else: + num_candidates = chunk_end_0 - chunk_start + + for cand in tl.range(0, num_candidates): + if QUERY_GROUP_SIZE == 2: + if chunk_equal: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + else: + valid_0 = idx_0 < chunk_end_0 + valid_1 = idx_1 < chunk_end_1 + if IDENTITY: + blk_0 = tl.where(valid_0, idx_0, sentinel) + blk_1 = tl.where(valid_1, idx_1, sentinel) + else: + blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) + blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) + use_0 = valid_0 & (blk_0 <= blk_1) + use_1 = valid_1 & (blk_1 <= blk_0) + has_block = use_0 | use_1 + blk = tl.minimum(blk_0, blk_1) + idx_0 += use_0.to(tl.int32) + idx_1 += use_1.to(tl.int32) + idx_ptr_0 += use_0.to(tl.int64) * stride_tk + idx_ptr_1 += use_1.to(tl.int64) * stride_tk + query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) + else: + has_block = True + if IDENTITY: + blk = chunk_start + cand + else: + blk = tl.load(idx_ptr_0) + idx_ptr_0 += stride_tk + query_active = h_mask + + if has_block: + # int64 page offsets: a large cache overflows int32 well before the + # per-page block offsets do. + # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the + # pool, reads nothing. + blk_ok = (blk >= 0) & (blk < num_bt_cols) + page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) + blk_ok = blk_ok & (page >= 0) & (page < num_pages) + page = tl.where(blk_ok, page, 0) + pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max + q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] + # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block + # would make exp2(-inf - -inf) = NaN). + if blk_ok: + if KV_NVFP4: + # K scales are token-major linear within the page-head region; + # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. + k_deq = _dequant_nvfp4_rows( + tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[:, None] * stride_k_pos + + off_dp[None, :] * stride_k_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + k_sc_ptr + + page * stride_ksc_blk + + pid_kh * stride_ksc_h + + off_n[:, None] * SCALE_COLS + + off_s[None, :], + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ) + k = tl.trans(k_deq).to(q.dtype) + else: + k = tl.load( + k_ptr + + page * stride_k_blk + + pid_kh * stride_k_h + + off_n[None, :] * stride_k_pos + + off_d[:, None] * stride_k_d, + mask=d_mask[:, None] & pos_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + k = (k * kv_scale).to(q.dtype) + active_pos = query_active[:, None] & q_pos_mask + qk = tl.where(active_pos, 0.0, float("-inf")) + qk += tl.dot(q, k) * sm_scale_log2e + m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) + p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) + l_ij = tl.sum(p, axis=1) + alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) + acc_o = acc_o * alpha[:, None] + if KV_NVFP4: + # V scales use vLLM's 4x4 token-quad order rather than K's + # linear one: [token // 4, scale_col, token % 4]. + v = _dequant_nvfp4_rows( + tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_dp[None, :] * stride_v_d, + mask=pos_mask[:, None], + other=0, + ), + tl.load( + v_sc_ptr + + page * stride_vsc_blk + + pid_kh * stride_vsc_h + + (off_n[:, None] // 4) * (4 * SCALE_COLS) + + 4 * off_s[None, :] + + (off_n[:, None] % 4), + mask=pos_mask[:, None], + other=0, + ), + BLOCK_N=BLOCK_SIZE_K, + D=BLOCK_SIZE_D, + SCALE_COLS=SCALE_COLS, + SF_VEC=BLOCK_SIZE_D // SCALE_COLS, + ).to(q.dtype) + else: + v = tl.load( + v_ptr + + page * stride_v_blk + + pid_kh * stride_v_h + + off_n[:, None] * stride_v_pos + + off_d[None, :] * stride_v_d, + mask=pos_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(q.dtype) + if USE_SCALE: + v = (v * kv_scale).to(q.dtype) + acc_o += tl.dot(p.to(v.dtype), v) + if USE_LINEAR_SOFTMAX: + l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) + else: + lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) + lse_i = tl.where(query_active, lse_ij, lse_i) + m_i = m_ij + + # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. + if USE_LINEAR_SOFTMAX: + scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) + lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) + else: + scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) + lse_out = lse_i + if KV_NVFP4: + # Applying V's scale per partial is exact: the merge is a weighted + # average whose weights sum to one, so a constant survives it. + scale = scale * v_global_scale + acc_o = acc_o * scale[:, None] + tl.store( + o_ptr + + pid_c * stride_o_c + + q_token[:, None] * stride_o_b + + (pid_h + h_lane[:, None]) * stride_o_h + + off_d[None, :] * stride_o_d, + acc_o.to(o_ptr.dtype.element_ty), + mask=hd_mask, + ) + if not DIRECT_OUTPUT: + tl.store( + lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, + lse_out, + mask=h_mask, + ) + + # After the stores, never before: this releases either the merge grid or, + # for direct output, the next PDL-enabled consumer of the final tensor. + if USE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def _rcp_approx_ftz(x): + return tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + +@triton.jit +def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): + """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as + torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, + scales in the 128x4-swizzled layout.""" + NUM_SF: tl.constexpr = HEAD_DIM // 32 + xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) + amax = tl.max(tl.abs(xr), axis=1) + sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) + bits = sfv.to(tl.int32, bitcast=True) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + e = tl.minimum(e, 254) + scale = ((254 - e) << 23).to(tl.float32, bitcast=True) + scale = tl.where(amax != 0, scale, 0.0) + q = (xr * scale[:, None]).to(tl.float8e4nv) + off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) + tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) + k_sf = head * NUM_SF + tl.arange(0, NUM_SF) + off_sf = ( + (row // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + k_sf % 4 + ) + tl.store(sf_ptr + off_sf, e.to(tl.uint8)) + + +@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) +@triton.jit +def _merge_topk_attn_out_kernel( + o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] + lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] + out_ptr, # merged out: [total_q, num_heads, head_dim] + head_dim, + stride_o_c, + stride_o_b, + stride_o_h, + stride_o_d, + stride_l_c, + stride_l_b, + stride_l_h, + stride_out_n, + stride_out_h, + stride_out_d, + NUM_TOPK_CHUNKS: tl.constexpr, + BLOCK_SIZE_D: tl.constexpr, + USE_PDL: tl.constexpr, + EARLY_TRIGGER: tl.constexpr = False, + q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 + sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled + num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) + MXFP8_OUT: tl.constexpr = False, +): + pid_b, pid_h = tl.program_id(0), tl.program_id(1) + + if USE_PDL: + if EARLY_TRIGGER: + # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight + # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits + # for this grid (and therefore the producer) to complete. Because this trigger precedes this grid's own + # wait, it says nothing about what upstream has finished: a dependent may read only step-invariant data + # (weights) before its griddepcontrol.wait, never this step's attention output, q, top-k table or lengths. + tl.extra.cuda.gdc_launch_dependents() + tl.extra.cuda.gdc_wait() + + off_c = tl.arange(0, NUM_TOPK_CHUNKS) + off_d = tl.arange(0, BLOCK_SIZE_D) + o = tl.load( + o_ptr + + pid_b * stride_o_b + + pid_h * stride_o_h + + off_c[:, None] * stride_o_c + + off_d[None, :] * stride_o_d, + mask=off_d[None, :] < head_dim, + other=0.0, + ).to(tl.float32) + # Empty chunks contribute -inf, hence weight 0. + lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) + lse_max = tl.max(lse, axis=0) + # A row whose every chunk is empty (a CUDA-graph padding row) would give + # -inf - -inf here. Zero it instead: the padded output is discarded, but a + # NaN would survive into the residual stream and the all-reduce. + lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) + weights = tl.exp2(lse - lse_max) + denom = tl.sum(weights, axis=0) + weights = weights / tl.where(denom > 0, denom, 1.0) + o_merged = tl.sum(o * weights[:, None], axis=0) + out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d + tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) + if MXFP8_OUT: + _store_mxfp8_head( + o_merged.to(out_ptr.dtype.element_ty), + q8_ptr, + sf_ptr, + pid_b, + pid_h, + num_sf_k_tiles, + BLOCK_SIZE_D, + ) + + if USE_PDL and not EARLY_TRIGGER: + tl.extra.cuda.gdc_launch_dependents() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/triton_sparse_decode_mq.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode_mq.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/attention/triton_sparse_decode_mq.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/embedding_firstnorm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/embedding_firstnorm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/embedding_firstnorm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/embedding_firstnorm/m3_embedding_firstnorm_kernel.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/m3_lowm_splitk.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_lowm_splitk.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/m3_lowm_splitk.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/m3_mx_gemv_kernel.py similarity index 95% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/m3_mx_gemv_kernel.py index 3ba6fb0d7c8c..9dcd5eb2ccbf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/m3_mx_gemv_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/gemm/m3_mx_gemv_kernel.py @@ -46,7 +46,11 @@ from cutlass import dsl_user_op from cutlass.experimental import primitives as prims -from ..m3_layer.phase_producer import ProducerArgs, bf16_bits_to_f32, producer_head +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.producer_tp8 import ( + ProducerArgs, + bf16_bits_to_f32, + producer_head, +) CTA_M = 128 # capacity: fixed: a kernel tile constant (weight rows per CTA = the MMA's M) MMA_N = 8 # token columns @@ -1373,7 +1377,7 @@ def _atom_map(t, atoms): @cute.jit -def m3_mx_gemv( +def m3_mx_gemv_producer( w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 @@ -1470,3 +1474,82 @@ def m3_mx_gemv( stream=stream, use_pdl=use_pdl, ) # fmt: skip + + +@cute.jit +def m3_mx_gemv( + w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous + wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 + x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 + xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 + y: cute.Tensor, # [M * N] bf16 + x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) + timers: cute.Tensor, # [grid * 8] int64, written when `timed` + num_tokens: cutlass.Int32, + n_out: cutlass.Constexpr[int], + k_in: cutlass.Constexpr[int], + split: cutlass.Constexpr[int], + ring: cutlass.Constexpr[int], + trigger_early: cutlass.Constexpr[bool], + contig: cutlass.Constexpr[bool], + timed: cutlass.Constexpr[bool], + nacc: cutlass.Constexpr[int], + a_tmem: cutlass.Constexpr[bool], + st_async: cutlass.Constexpr[bool], + x_bf16: cutlass.Constexpr[bool], + ar_ws: cutlass.Int64, + ar_flags: cutlass.Int64, + ar_rank: cutlass.Int32, + ar_push: cutlass.Constexpr[int], + ar_world: cutlass.Constexpr[int], + ar_maxt: cutlass.Constexpr[int], + mma_n: cutlass.Constexpr[int], + use_pdl: cutlass.Constexpr[bool], + stream: cuda_driver.CUstream, + din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) + din_table: cute.Tensor, + din_hidden: cute.Tensor, + din_w: cute.Tensor, + din_f: cute.Tensor, + din_vocab: cutlass.Int32, + din_g: cute.Tensor, + din_res: cute.Tensor, + din_flags: cutlass.Int32, + rope_cs: cute.Tensor, + rope_i: cute.Tensor, + rope_bt: cute.Tensor, + rope_q: cute.Tensor, + rope_pool: cute.Tensor, + rope_scale: cute.Tensor, + rope_dims: cute.Tensor, + din: cutlass.Constexpr[bool], + rope: cutlass.Constexpr[bool], + x_swiglu: cutlass.Constexpr[bool] = False, +) -> None: + # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. + assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( + "x_swiglu needs x_bf16 and excludes din" + ) + k_tiles = num_k_tiles(k_in) + tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) + tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) + tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) + tma_desc_xsf = _atom_map(xsf, k_tiles) + m3_mx_gemv_kernel( + tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, + trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, + ar_maxt, mma_n, + din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, + rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, x_swiglu, + cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), + cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), + cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), + cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), + cutlass.Int32(0), cutlass.Float32(0), False, + ).launch( + grid=((n_out // CTA_M) * split, 1, 1), + block=(THREADS, 1, 1), + cluster=(split, 1, 1), + stream=stream, + use_pdl=use_pdl, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/moe/m3_moe_kernel.py similarity index 98% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/moe/m3_moe_kernel.py index 084f54951209..48405c3b37db 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/m3_moe_kernel.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/moe/m3_moe_kernel.py @@ -12,7 +12,7 @@ # 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. -"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP2 rank (routed and shared intermediate 1536): +"""MiniMax-M3 routed experts for up to 8 decode tokens on one sm_107 TP2/TP4/TP8 rank: ``m3_moe`` (group-phased weight stream). One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), @@ -54,6 +54,8 @@ from cutlass.experimental import primitives as prims _CFG = globals().get("M3_CONFIG") or {} +TP_SIZE = int(_CFG["tp_size"]) +assert TP_SIZE in (2, 4, 8), TP_SIZE def _cfg(key: str, default): @@ -69,7 +71,7 @@ def _cfg(key: str, default): # MMAs / scale atoms are issued; no padding, no extra bytes) I_TP = _cfg("i_tp", 768) TWO_I = 2 * I_TP -assert I_TP % 128 == 0 and 128 <= I_TP <= 1536, I_TP +assert I_TP % 128 == 0 and 128 <= I_TP <= (768 if TP_SIZE == 8 else 1536), I_TP E_LOCAL = 128 TOP_K = 4 # capacity: fixed: the checkpoint's top-k ROUTE_SCALE = 2.0 @@ -108,19 +110,19 @@ def _cfg(key: str, default): # TAIL_TRIGGER (with LATE_TRIGGER, no all-reduce push): the CTAs that own more FC2 blocks than the rest (the first # H / 32 mod NCTA, which end last) release the dependents right after the grid dependency, so the dependents launch # when the other CTAs tear down, onto the SMs they leave, instead of after the last of the heavy CTAs. -TAIL_TRIGGER = bool(_cfg("tail_trigger", 1)) +TAIL_TRIGGER = TP_SIZE != 8 and bool(_cfg("tail_trigger", 1)) NHEAVY = (H // 32) % NCTA if NCTA < H // 32 else 0 # SH_TAKE: the last NHEAVY shared-expert CTAs, after their shared work, each run the routed body as an FC2-only CTA for # one FC2 output block (the second block of routed CTA c < NHEAVY, which then owns one). The plan carries their rows # (PLAN_COLS columns). Each block keeps its tiles, scales and finalize, so the output is unchanged. -SH_TAKE = bool(_cfg("sh_take", 0)) and SHARED and 0 < NHEAVY <= 24 +SH_TAKE = TP_SIZE == 2 and bool(_cfg("sh_take", 0)) and SHARED and 0 < NHEAVY <= 24 SH_HELPER0 = 24 - NHEAVY # first helper shared index (NSH = 24) PLAN_COLS = NCTA + 24 if SH_TAKE else NCTA # EARLY_AR (scheduling only): when the finalize pushes into the multicast all-reduce, EVERY CTA releases the # dependents right after the grid dependency, so the all-reduce (which reads only the pushed tagged words, without a # grid wait) and, through its own early release, the next layer's QKV GEMV become resident on each SM this kernel # frees, instead of after the last CTA's teardown. The arithmetic is unchanged. -EARLY_AR = bool(_cfg("early_ar", 1)) +EARLY_AR = TP_SIZE != 4 and bool(_cfg("early_ar", 1)) # AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's # FC2 runs (not two dependent global loads per push on the kernel's tail). AR_PRE = bool(_cfg("ar_pre", 1)) @@ -151,7 +153,7 @@ def _cfg(key: str, default): _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) ) # (in-chain: a win at M <= 2, not 4..8) if SHARED: - assert SH_I in (384, 768, 1536), SH_I + assert SH_I in ((384, 768) if TP_SIZE == 8 else (384, 768, 1536)), SH_I assert SH_PRETILE or SH_SPLIT == 2, ( "the un-pretiled shared stream is the 32-row (TP4) layout only" ) @@ -256,10 +258,13 @@ def _nkb2(t2: int) -> int: S_WARP = 6 # capacity: fixed: a warp role index MMA_WARP = 7 # capacity: fixed: a warp role index # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path). -E0_WARP = MMA_WARP # capacity: fixed: a warp role index +E0_M4 = ( + TP_SIZE != 2 and bool(_cfg("e0_m4", 1)) and NRW0 == 4 and RW_EXTRA == 0 and THREADS // 32 > 8 +) +E0_WARP = 8 if E0_M4 else MMA_WARP E0_BAR = 8 E1_BAR = 9 -E0_OFFLOAD = EARLY0 and NRW0 <= 3 and bool(_cfg("e0_off", 1)) +E0_OFFLOAD = EARLY0 and (NRW0 <= 3 or E0_M4) and bool(_cfg("e0_off", 1)) # B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing # warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the # two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. @@ -300,11 +305,11 @@ def _nkb2(t2: int) -> int: assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) # TP2 at 8 tokens: the 98 KB of hidden states fit no 5-stage ring; the router GEMV then reads them from global (the # kernel's non-XSTAGE path, same math) -if XSTAGE and M_MAX * H * 2 > STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA): +if TP_SIZE != 8 and XSTAGE and M_MAX * H * 2 > STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA): XSTAGE = False -assert not XSTAGE or M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( - "x staging must fit the ring" -) +assert (TP_SIZE != 8 and not XSTAGE) or M_MAX * H * 2 <= STAGES * ( + NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA +), "x staging must fit the ring" # one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) _OFF_SFAR = STAGES * NUM_BYTES_A _OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE @@ -345,9 +350,9 @@ def _nkb2(t2: int) -> int: SH_SFB_GU = SH_SFA + 4 * 8 SH_SFB_DN = SH_SFB_GU + 4 * SH_GT # pre-tiled stages beyond the ring that warp 0 prefetches into L2, per lane (TP2: 72 stages) -SH_PF = max(2, (SH_STAGES - SH_NST + 31) // 32) if SHARED else 2 +SH_PF = max(2, (SH_STAGES - SH_NST + 31) // 32) if SHARED and TP_SIZE != 8 else 2 # TP2 (SH_GT 48 gate_up k-tiles, SH_DKT 12) needs more than 256 columns: allocate the next power of two -SH_TMEM = 256 if SH_SFB_DN + 4 * SH_DKT <= 256 else 512 +SH_TMEM = 256 if TP_SIZE == 8 or SH_SFB_DN + 4 * SH_DKT <= 256 else 512 # SH_TAKE helpers hand the shared allocation to the routed body, whose teardown frees TMEM_COLS assert not SH_TAKE or SH_TMEM == TMEM_COLS, ( "SH_TAKE needs the shared TMEM allocation to equal the routed one" @@ -355,7 +360,11 @@ def _nkb2(t2: int) -> int: if SHARED: assert FUSED_ROUTER, "shared expert: fused router only" assert SH_STAGES - SH_NST <= 32 * SH_PF, SH_STAGES - assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 4 * SH_DKT <= SH_TMEM, ( + assert ( + SH_NST >= 3 + and _SB_END <= BIG + and SH_SFB_DN + (24 if TP_SIZE == 8 else 4 * SH_DKT) <= SH_TMEM + ), ( SH_NST, _SB_END, BIG, @@ -1883,9 +1892,11 @@ def _shared_cta( base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) hv = [cutlass.Int32(0) for _ in range(32)] pend = cutlass.Int32(1) - t0h = cutlass.Int64(0) + if cutlass.const_expr(TP_SIZE != 8): + t0h = cutlass.Int64(0) while pend != cutlass.Int32(0): - t0h = _hang_tick(t0h) + if cutlass.const_expr(TP_SIZE != 8): + t0h = _hang_tick(t0h) for q4 in cutlass.range_constexpr(8): v4 = prims.load_ext( hx.subview(base + cutlass.Int32(4 * q4)), @@ -2180,9 +2191,10 @@ def m3_moe_kernel( if cutlass.const_expr(XSTAGE): prims.mbarrier_init(x_bar.subview(0), 1) prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) - if (warp == MMA_WARP) & ( - bx < cutlass.Int32(NCTA) - ): # SH_TAKE helpers reuse the shared allocation + allocate_tmem = warp == MMA_WARP + if cutlass.const_expr(TP_SIZE == 2): + allocate_tmem = allocate_tmem & (bx < cutlass.Int32(NCTA)) + if allocate_tmem: prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) prims.tcgen05_relinquish_alloc_permit() if warp == W_WARP: @@ -2382,9 +2394,11 @@ def m3_moe_kernel( ) done = cutlass.Int32(0) res = [cutlass.Int64(0) for _ in range(4)] - t0h = cutlass.Int64(0) + if cutlass.const_expr(TP_SIZE != 8): + t0h = cutlass.Int64(0) while done == cutlass.Int32(0): - t0h = _hang_tick(t0h) + if cutlass.const_expr(TP_SIZE != 8): + t0h = _hang_tick(t0h) ok = cutlass.Boolean(True) for z in cutlass.range_constexpr(4): res[z] = _ld_relaxed_u64( @@ -2963,9 +2977,11 @@ def m3_moe_kernel( caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() if lane == 0: cv = _ld_acquire_u32(caddr) - t0h = cutlass.Int64(0) + if cutlass.const_expr(TP_SIZE != 8): + t0h = cutlass.Int64(0) while (cv - cnt_target) < cutlass.Int32(0): - t0h = _hang_tick(t0h) + if cutlass.const_expr(TP_SIZE != 8): + t0h = _hang_tick(t0h) _nanosleep(B2IMG_NS) cv = _ld_acquire_u32(caddr) _fence_proxy_async_global() @@ -3092,9 +3108,11 @@ def m3_moe_kernel( _stamp(tdbg, bx, 29) npoll = cutlass.Int32(0) got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] - t0h = cutlass.Int64(0) + if cutlass.const_expr(TP_SIZE != 8): + t0h = cutlass.Int64(0) while done == cutlass.Int32(0): - t0h = _hang_tick(t0h) + if cutlass.const_expr(TP_SIZE != 8): + t0h = _hang_tick(t0h) ok = cutlass.Boolean(True) for z in cutlass.range_constexpr(NZ): qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi @@ -3504,9 +3522,11 @@ def m3_moe_kernel( else: pbits = _ld_relaxed_u32(paddr) spins = cutlass.Int32(0) - t0h = cutlass.Int64(0) + if cutlass.const_expr(TP_SIZE != 8): + t0h = cutlass.Int64(0) while pbits == cutlass.Int32(PARTIAL_EMPTY): - t0h = _hang_tick(t0h) + if cutlass.const_expr(TP_SIZE != 8): + t0h = _hang_tick(t0h) pbits = _ld_relaxed_u32(paddr) spins = spins + cutlass.Int32(1) _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) @@ -3758,9 +3778,11 @@ def m3_moe_kernel( if cutlass.const_expr(SH_ADD): # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) wy = wys[bt] - t0h = cutlass.Int64(0) + if cutlass.const_expr(TP_SIZE != 8): + t0h = cutlass.Int64(0) while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: - t0h = _hang_tick(t0h) + if cutlass.const_expr(TP_SIZE != 8): + t0h = _hang_tick(t0h) wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) acc_o = _add_rn( (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/m3_fold_tail.py new file mode 100644 index 000000000000..38c1a339f4aa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/m3_fold_tail.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared integer store/staging kernel; targets own launch policy.""" + +import triton +import triton.language as tl + + +@triton.jit +def _fold_tail_kernel( + ctrl, + slots, + src_new, + sn_s0, + sn_s1, + SN_W: tl.constexpr, + dst_new, + dn_s0, + dn_s1, + DN_W: tl.constexpr, + src_next, + sx_s0, + sx_s1, + SX_W: tl.constexpr, + dst_next, + dx_s0, + dx_s1, + DX_W: tl.constexpr, + src_draft, + sd_s0, + sd_s1, + SD_W: tl.constexpr, + dst_draft, + dd_s0, + dd_s1, + DD_W: tl.constexpr, + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + T: tl.constexpr, + D: tl.constexpr, + BLOCK: tl.constexpr, +): + i = tl.program_id(0) + n = tl.load(ctrl) + if i < n: + s = tl.load(slots + i).to(tl.int64) + c = tl.arange(0, BLOCK) + new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) + tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) + nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) + tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) + drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) + tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) + ln = tl.load(src_lens + i) + tl.store(dst_lens + s, ln) + # The next step's inputs read the store's padded rows: zero past each source width. + tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) + tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) + tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) + tl.store(kv_off + i, ln - T) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/m3_force_vec.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_force_vec.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/kernels/spec/m3_force_vec.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py index 4418d2908b33..15926f7e5259 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/modeling.py @@ -63,9 +63,9 @@ from __future__ import annotations import atexit -import copy import dataclasses import os +import sys import threading import time from dataclasses import dataclass @@ -81,6 +81,7 @@ Mapping, Optional, Tuple, + cast, ) import torch @@ -108,6 +109,98 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import ( + modeling_support as _modeling_support, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import reporting as _reporting +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _is_vl_config as _is_vl_config, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _moe_layer_ids as _moe_layer_ids, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _sparse_layer_ids as _sparse_layer_ids, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _text_model_config as _text_model_config, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _validate_runtime as _validate_runtime, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.activation.m3_silu_and_mul import ( + m3_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + msa_kv_update_hook as _msa_kv_update_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_fp8_indexer_qk_norm_rope import ( + m3_fp8_indexer_qk_norm_rope as m3_fp8_indexer_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qk_norm_rope_fp8 import ( + m3_qk_norm_rope_fp8, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qkv_insert import ( + m3_qkv_insert, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_select_blocks import ( + m3_select_blocks, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( + m3_allreduce, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.m3_embedding import ( + m3_embedding, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_bf16_linear import ( + m3_bf16_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + ALPHA as DENSE_ALPHA, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + BETA as DENSE_BETA, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + LIMIT as DENSE_LIMIT, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + m3_dense_swiglu, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_fp32_linear import ( + m3_fp32_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( + m3_mxfp8_gemm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_linear import ( + m3_mxfp8_linear as m3_mxfp8_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + MxMode as MxMode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + _mx_mode as _mx_mode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + _Mxfp8GemmState as _Mxfp8GemmState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_tc import ( + m3_router_tc, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( + m3_capture_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_chunk_norm import ( + m3_chunk_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_draft_input import ( + m3_draft_input, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( + m3_force_vec, + m3_force_vec_req, +) from tensorrt_llm._torch.attention.backends import AttentionMetadata from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams @@ -121,22 +214,27 @@ ) from tensorrt_llm._torch.attention.backends.utils import create_attention from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( - FlashInferMXFP8GemmRunner, - MXFP8QuantizeRunner, + FlashInferMXFP8GemmRunner as FlashInferMXFP8GemmRunner, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + MXFP8QuantizeRunner as MXFP8QuantizeRunner, ) from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_speculative import ( _ensure_draft_vocab_size, - _get_requested_draft_moe_backend, - _set_draft_kv_cache_quant_algo, _slice_spec_position_ids, ) +from tensorrt_llm._torch.models.modeling_speculative import ( + _get_requested_draft_moe_backend as _get_requested_draft_moe_backend, +) +from tensorrt_llm._torch.models.modeling_speculative import ( + _set_draft_kv_cache_quant_algo as _set_draft_kv_cache_quant_algo, +) from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM from tensorrt_llm._torch.modules.embedding import LMHead from tensorrt_llm._torch.modules.linear import ( - MXFP8LinearMethod, TensorParallelMode, flashinfer_mxfp8_autotune_active, flashinfer_mxfp8_decode_graph_capture_active, @@ -157,7 +255,9 @@ _FORCE_ACCEPT_RNG_COUNTER_STRIDE, _FORCE_ACCEPT_RNG_POOL_SIZE, _FORCE_ACCEPT_RNG_SLOT_STRIDE, - should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.speculative.interface import ( + should_use_separate_draft_kv_cache as should_use_separate_draft_kv_cache, ) from tensorrt_llm._torch.utils import ActType_TrtllmGen, get_model_extra_attrs from tensorrt_llm._utils import mpi_rank, prefer_pinned @@ -179,35 +279,21 @@ from .ops import m3_index_score as _index_score_kernel from .ops import m3_sampled_accept as m3_sampled_accept_kernel from .ops import m3_shard_accept as m3_shard_accept_kernel -from .ops import msa_kv_update as _msa_kv_update_kernel from .ops import triton_sparse_decode as _sparse_decode_kernel -from .ops.activation.m3_silu_and_mul import m3_silu_and_mul from .ops.attention.m3_draft_attention import m3_draft_attention -from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope from .ops.attention.m3_fullctx_decode import m3_fullctx_decode from .ops.attention.m3_index_score import m3_index_score from .ops.attention.m3_msa_attention import m3_msa_attention from .ops.attention.m3_producer import m3_producer -from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 -from .ops.attention.m3_qkv_insert import m3_qkv_insert -from .ops.attention.m3_select_blocks import m3_select_blocks from .ops.attention.m3_select_long import m3_select_long from .ops.attention.m3_select_long_dsl import m3_select_long_dsl from .ops.attention.m3_sparse_decode import m3_sparse_decode from .ops.attention.m3_sparse_decode_lc import m3_sparse_decode_lc -from .ops.comm.m3_allreduce import m3_allreduce from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm -from .ops.embedding.m3_embedding import m3_embedding from .ops.gemm import m3_ugpu_head as m3_ugpu_head_op -from .ops.gemm.m3_bf16_linear import m3_bf16_linear -from .ops.gemm.m3_dense_swiglu import ALPHA as DENSE_ALPHA -from .ops.gemm.m3_dense_swiglu import BETA as DENSE_BETA -from .ops.gemm.m3_dense_swiglu import LIMIT as DENSE_LIMIT -from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS from .ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for -from .ops.gemm.m3_fp32_linear import m3_fp32_linear from .ops.gemm.m3_mx_gemv import ( DENSE_DOWN, DENSE_GATE_UP, @@ -220,21 +306,14 @@ served_max_rows, ) from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 -from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm -from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear from .ops.gemm.m3_ugpu_head import m3_ugpu_head from .ops.long_select import long_select_blocks from .ops.m3_moe_finalize import m3_moe_finalize from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine -from .ops.moe.m3_router_tc import m3_router_tc from .ops.mx_gemv import op as mx_gemv from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm -from .ops.spec.m3_capture_norm import m3_capture_norm -from .ops.spec.m3_chunk_norm import m3_chunk_norm -from .ops.spec.m3_draft_input import m3_draft_input from .ops.spec.m3_draft_silu import m3_draft_silu from .ops.spec.m3_draft_step import m3_draft_step -from .ops.spec.m3_force_vec import m3_force_vec, m3_force_vec_req from .ops.spec.m3_sampled_accept import m3_sampled_accept from .ops.spec.m3_shard_accept import m3_shard_accept from .ops.spec.m3_spec_offsets import m3_spec_offsets @@ -564,7 +643,7 @@ class DecodePlan: # (The drafter epilogue, opts bit 16, is built for the 8-token bucket only: a B4 spec-tail first pass of 16 / 32 rows # takes the built-in all-reduce + RMSNorm; its later passes of up to 8 rows stay multicast.) DRAFT_MC_AR_MAX_ROWS = 8 # capacity: fixed: the drafter epilogue's 8-token bucket -# The drafter's o_proj and down projection, on the copied split-K kernel (DRAFT_LOWM_OWN_SPLITK), also push their rows +# The drafter's o_proj and down projection, on the shared split-K kernel (DRAFT_LOWM_OWN_SPLITK), push their rows # into the drafter boundary all-reduce's exchange (the canonical words its own push would store, at the call parity), # which then runs prepushed 2 instead of pushing x itself; the copy releases its dependents only after its own grid # wait and the pushes (the prepushed >= 2 contract). Bitwise. TP2: the 4-row first pass only (its 1-row linears run @@ -608,7 +687,7 @@ class DecodePlan: # keep the in-tree cuBLAS GEMM. 0:. Rubin TP2 (mclk locked, 4 rotating sets) at 1 row: qkv 13.18 -> 9.71, # o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the certified low-M rows -# The split-K ('s') tactics run this target's copy of FlashInfer's kernel (ops/m3_lowm_splitk): bit for bit the +# The split-K ('s') tactics run the shared FlashInfer kernel (kernels/gemm/m3_lowm_splitk): bit for bit the # same output at every certified cell (spec/test_modeling_v2_m3_lowm_splitk.py); the seam for the drafter-boundary push # epilogue and the down projection's SwiGLU prologue. False: (FlashInfer's kernel). DRAFT_LOWM_OWN_SPLITK = True @@ -801,7 +880,9 @@ def covers(self, key: Key) -> bool: return key in self.keys def certified(self) -> FrozenSet[Key]: - return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + cells_by_tp = getattr(self.entry, "CELLS_BY_TP", None) + cells = self.entry.CELLS if cells_by_tp is None else cells_by_tp[2] + return frozenset(k for cell in cells for k in self.certifies(cell.spec)) def _rows(lo: int, hi: int) -> FrozenSet[int]: @@ -1142,7 +1223,7 @@ def moe_path(num_tokens: int) -> Optional[str]: MX_GEMV_PATH = MX_GEMV.path MX_GEMV_BF16_PATH = MX_GEMV_BF16.path # The dense layers' swigluoai activation as this target's C++ op (ops/gemm/m3_dense_swiglu, -# minimaxM3Sm107Tp2/swigluOai.cu, LongAttn c267035d2f: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, +# minimaxM3Sm107/swigluOai.cu, LongAttn c267035d2f: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, # beta 1), PDL-launched with an early release so the down GEMV stages its weights during it; the down GEMV reads the # activation only after its own grid wait. Indexer kdev17 783285 (TP2 chain): -0.69 us per dense layer vs the Triton # kernel; LongAttn gate a 783511 (oracle 23/23). False runs the Triton kernel. @@ -1183,58 +1264,36 @@ class FastPathUnavailable(RuntimeError): _lock = threading.Lock() _require = [False] _logged: set = set() +_POLICY_TARGET = cast(_reporting.ReportingTarget, sys.modules[__name__]) def set_require(require: bool) -> None: """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" - _require[0] = bool(require) + return _reporting.set_require(_POLICY_TARGET, require) def require_mode() -> bool: - return _require[0] + return _reporting.require_mode(_POLICY_TARGET) def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: - key = (path, reason) - with _lock: - if key in _logged: - return - _logged.add(key) - level(message) + return _reporting._log_once(_POLICY_TARGET, level, path, reason, message) def missed(path: str, rows: Optional[int], reason: str) -> None: """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" - at = "" if rows is None else f" at {rows} rows" - if _require[0]: - raise FastPathUnavailable( - f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " - f"{reason}" - ) - _log_once( - logger.warning, - path, - reason, - f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", - ) + return _reporting.missed(_POLICY_TARGET, path, rows, reason) def capture_missed(path: str, rows: Optional[int], what: str) -> None: """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" - missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + return _reporting.capture_missed(_POLICY_TARGET, path, rows, what) def taken(path: str, rows: Optional[int], what: str = "") -> None: """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" - at = "" if rows is None else f" at {rows} rows" - detail = f" ({what})" if what else "" - _log_once( - logger.info, - path, - f"taken: {rows}" + (f": {what}" if what else ""), - f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", - ) + return _reporting.taken(_POLICY_TARGET, path, rows, what) def long_select_covers(num_blocks: int) -> bool: @@ -1275,13 +1334,12 @@ def long_select_fallback(num_blocks: int, rows: int) -> None: def uncovered(path: str, reason: str) -> None: """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" - _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + return _reporting.uncovered(_POLICY_TARGET, path, reason) def reset_logged() -> None: """Forget which (path, reason) pairs were logged.""" - with _lock: - _logged.clear() + return _reporting.reset_logged(_POLICY_TARGET) # Kernel modules report from inside their calls through ops/m3_report: the same policy. @@ -1339,132 +1397,6 @@ def _print_values() -> None: atexit.register(_print_values) -def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: - model_type = getattr(pretrained_config, "model_type", None) - if model_type == "minimax_m3_vl": - return True - if model_type == "minimax_m3": - return False - return any( - "ForConditionalGeneration" in a - for a in (getattr(pretrained_config, "architectures", None) or []) - ) - - -def _text_model_config(model_config: ModelConfig) -> ModelConfig: - """The model config with the language model's ``text_config`` as its pretrained config. - - The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed - sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and - embedding tie carried over. A text config is returned unchanged. - """ - if not _is_vl_config(model_config.pretrained_config): - return model_config - cfg = copy.deepcopy(model_config) - text = getattr(cfg.pretrained_config, "text_config", None) - if text is None: - raise ValueError("MiniMax-M3 VL config has no `text_config`") - if isinstance(text, dict): - out = PretrainedConfig() - for k, v in text.items(): - setattr(out, k, v) - text = out - if ( - getattr(text, "torch_dtype", None) is None - and getattr(cfg.pretrained_config, "torch_dtype", None) is not None - ): - text.torch_dtype = cfg.pretrained_config.torch_dtype - if ( - getattr(text, "tie_word_embeddings", None) is None - and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None - ): - text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings - return dataclasses.replace(cfg, pretrained_config=text) - - -def _validate_runtime(model_config: ModelConfig) -> None: - """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing - sees only its KV cache algorithm; this check at construction refuses every other runtime. - - The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index - projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. - """ - sparse = model_config.sparse_attention_config - if ( - sparse is None - or getattr(sparse, "algorithm", None) != "minimax_m3" - or sparse.implementation != "msa" - ): - raise ValueError( - "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" - ) - if getattr(sparse, "indexer_kv_dtype", None) != "fp8": - raise ValueError( - "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" - ) - mode = model_config.quant_config.quant_mode - if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): - raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") - - -def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: - freq = cfg.sparse_attention_config["sparse_attention_freq"] - return [i for i, f in enumerate(freq) if int(f) != 0] - - -def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: - freq = getattr(cfg, "moe_layer_freq", None) - if freq is None: - return list(range(cfg.num_hidden_layers)) - return [i for i, f in enumerate(freq) if int(f) != 0] - - -class _Mxfp8GemmState(nn.Module): - """The MXFP8 GEMM backend state the engine's warmup drives. - - The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM - tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of - this target dispatches identically, so one method object carries that state for all of them; the forward reads it - at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. - """ - - def __init__(self) -> None: - super().__init__() - self.quant_method = MXFP8LinearMethod() - - -@dataclass(frozen=True) -class MxMode: - """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" - - decode_graph: bool # inside a decode CUDA-graph warmup or capture - autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass - tuned_quantize_gemm: ( - bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) - ) - flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation - native_tuning: ( - bool # the native GEMM through the autotuner (warmup, before its tactics are registered) - ) - - -def _mx_mode(state: _Mxfp8GemmState) -> MxMode: - method = state.quant_method - decode_graph = flashinfer_mxfp8_decode_graph_capture_active() - autotune = flashinfer_mxfp8_autotune_active() - return MxMode( - decode_graph=decode_graph, - autotune=autotune, - tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, - flashinfer=method.backend == "flashinfer" - or ( - method.backend == "auto" - and (autotune or (method._flashinfer_autotuned and decode_graph)) - ), - native_tuning=method.needs_native_autotune, - ) - - # -------------------------------------------------------------------------------------------------------------------- # Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs # -------------------------------------------------------------------------------------------------------------------- @@ -1474,50 +1406,27 @@ def _mxfp8_linear( core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode ) -> torch.Tensor: """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" - _record( - "m3_mxfp8_linear", - tuple(w.shape), - mode.tuned_quantize_gemm, - mode.flashinfer, - mode.native_tuning, - ) - return m3_mxfp8_linear( - x.contiguous(), - w, - s, - core.mx_alpha, - torch.bfloat16, - mode.tuned_quantize_gemm, - mode.autotune, - mode.flashinfer, - mode.native_tuning, - ) + return _modeling_support._mxfp8_linear(_SUPPORT_TARGET, core, x, w, s, mode) _PINNED_RUNNERS: dict = {} def _pinned_runner(cls, dtype): - runner = _PINNED_RUNNERS.get((cls, dtype)) - if runner is None: - runner = _PINNED_RUNNERS[(cls, dtype)] = cls(dtype) - return runner + return _modeling_support._pinned_runner(_SUPPORT_TARGET, cls, dtype) def _pinned_mxfp8_gemm(act, act_scale, weight, weight_scale, out_dtype): """The decode-graph MXFP8 GEMM at the target's fixed tactic (``MXFP8_DECODE_GEMM_TACTIC``), no AutoTuner: the same choice on every rank and every start.""" - _record("pinned_mxfp8_gemm", tuple(weight.shape), act.shape[0], MXFP8_DECODE_GEMM_TACTIC) - runner = _pinned_runner(FlashInferMXFP8GemmRunner, out_dtype) - return runner.forward([act, act_scale, weight, weight_scale], tactic=MXFP8_DECODE_GEMM_TACTIC) + return _modeling_support._pinned_mxfp8_gemm( + _SUPPORT_TARGET, act, act_scale, weight, weight_scale, out_dtype + ) def _pinned_mxfp8_quantize(x): """The decode-graph MXFP8 activation quantize at the target's fixed tactic (``MXFP8_DECODE_QUANT_TACTIC``).""" - _record("pinned_mxfp8_quantize", x.shape[0], MXFP8_DECODE_QUANT_TACTIC) - return _pinned_runner(MXFP8QuantizeRunner, x.dtype).forward( - [x], tactic=MXFP8_DECODE_QUANT_TACTIC - ) + return _modeling_support._pinned_mxfp8_quantize(_SUPPORT_TARGET, x) _MX_TAC_ENGAGED: set = set() @@ -1882,8 +1791,7 @@ def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): def _split_chunk(rows: int, num_chunks: int) -> List[int]: """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" - base, extra = divmod(rows, num_chunks) - return [base + 1] * extra + [base] * (num_chunks - extra) + return _modeling_support._split_chunk(_SUPPORT_TARGET, rows, num_chunks) # -------------------------------------------------------------------------------------------------------------------- @@ -1892,7 +1800,7 @@ def _split_chunk(rows: int, num_chunks: int) -> List[int]: def _nvfp4_layer(manager, i: int) -> bool: - return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + return _modeling_support._nvfp4_layer(_SUPPORT_TARGET, manager, i) def _producer_covered(core) -> bool: @@ -1919,31 +1827,7 @@ def _producer_layout(kv_data, kv_scale) -> bool: def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: """Reject invalid HND pools; select the fused producer only for its packed layout.""" - if ( - data.dim() != 5 - or scales.dim() != 5 - or data.shape[0] <= 0 - or data.shape[0] != scales.shape[0] - or tuple(data.shape[1:]) != (2, num_heads, 128, 64) - or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) - or data.dtype not in (torch.int8, torch.uint8) - or scales.dtype != torch.uint8 - ): - raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") - for cache in (data, scales): - if ( - cache.stride(4) != 1 - or cache.stride(3) != cache.shape[4] - or cache.stride(2) < 128 * cache.stride(3) - or cache.stride(1) < num_heads * cache.stride(2) - or cache.stride(0) < 2 * cache.stride(1) - ): - raise ValueError( - "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" - ) - if scales.stride(2) != 128 * 8: - raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") - return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + return _modeling_support._nvfp4_fused_write_layout(_SUPPORT_TARGET, data, scales, num_heads) #: The step phases ("decode" / "prefill") whose producer engagement has been logged, once per process. @@ -2191,14 +2075,14 @@ def _dense_qk_norm_rope_fp8( def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: - """The dense MLP's SwiGLU-OAI activation: this target's C++ op (``m3_dense_swiglu``, minimaxM3Sm107Tp2/ + """The dense MLP's SwiGLU-OAI activation: this target's C++ op (``m3_dense_swiglu``, minimaxM3Sm107/ swigluOai.cu: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, beta 1, PDL-launched with an early release so the down GEMV stages its weights during it) at the rows its cells certify, in any graph or eager step; otherwise the Triton kernel (``m3_silu_and_mul``). Indexer kdev17 (TP2 chain): -0.69 us per dense layer vs the Triton kernel. Every decline is logged once.""" path = DENSE_SWIGLU_CPP_PATH rows = gu.shape[0] - max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS_BY_TP[2]) params = (core.swiglu_alpha, core.swiglu_limit, core.swiglu_beta) act = None if rows > max_rows: @@ -2207,7 +2091,12 @@ def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: uncovered(path, f"activation parameters {params}") else: act = m3_dense_swiglu( - gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + gu, + core.swiglu_limit, + core.swiglu_alpha, + core.swiglu_beta, + get_env_enable_pdl(), + tp_size=2, ) if act is None: uncovered(path, "the gate_up layout") @@ -2225,31 +2114,9 @@ def _fp8_index_qk_norm_rope( """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` are the calling target's ``step_args`` entries.""" - if position_ids is None or idx_qk.dtype != torch.bfloat16: - raise NotImplementedError( - "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." - ) - index_k_cache = msa_idx_k_cache(i) - if index_k_cache.dtype != torch.float8_e4m3fn: - raise ValueError( - "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " - f"got {index_k_cache.dtype}." - ) - rows = int(idx_qk.shape[0]) - _record("m3_fp8_indexer_qk_norm_rope", rows) - return m3_fp8_indexer_qk_norm_rope( - idx_qk.contiguous(), - index_k_cache, - msa_out_cache_loc[:rows], - core.index_heads, - core.index_dim, - core.rotary_dim, - core.rms_eps, - core.w[f"l{i}_idx_q_norm"], - core.w[f"l{i}_idx_k_norm"], - core.rope_theta, - position_ids.reshape(-1).contiguous().to(torch.int32), - ).flatten(1) + return _modeling_support._fp8_index_qk_norm_rope( + _SUPPORT_TARGET, core, i, idx_qk, position_ids, msa_idx_k_cache, msa_out_cache_loc + ) # -------------------------------------------------------------------------------------------------------------------- @@ -4225,6 +4092,7 @@ def _index_scores(self, i: int, idx_q, step, num_tokens: int, qlen: int): qlen, split_k, early, + runner=_index_score_kernel.IndexScoreRunner, ) taken(INDEX_SCORE_PATH, rows) else: @@ -4627,8 +4495,8 @@ class M3Eagle3Draft(nn.Module): (``ops/gemm/m3_draft_lowm``, the TP8 port) at the rows ``DRAFT_LOWM`` covers, compiled before the warmup (``prebuild_linears``); every other call is the in-tree cuBLAS GEMM. Rubin TP2 (mclk locked, 4 rotating sets) at 1 row: qkv 13.18 -> 9.71, o 7.17 -> 5.33, gate_up 22.27 -> 17.70, down 12.52 -> 9.76, fc 21.50 -> 17.77 us. The - split-K ('s') tactics run this target's copy of FlashInfer's kernel (``ops/m3_lowm_splitk``, bit for bit the - same output); with ``DRAFT_AR_PUSH`` the o_proj and down projection of the fused draft loop also push their rows + split-K ('s') tactics run the shared FlashInfer kernel (``kernels/gemm/m3_lowm_splitk``, bit for bit the + same output); with ``DRAFT_AR_PUSH`` the o_proj and down projection of the fused draft loop push their rows into the drafter boundary all-reduce's exchange (``arm_push`` / ``take_pushed``), as do the direct ('d') 1-row linears through this target's copy of the direct kernel (``ops/m3_lowm_direct``; o_proj 9.22 -> 8.09 us, down 13.81 -> 12.72 us per 1-row boundary). @@ -4803,8 +4671,11 @@ def _lowm_linear(self, weight: torch.Tensor, x: torch.Tensor) -> Optional[torch. def prebuild_linears(self) -> int: """Compile (and run once) every covered row count of every drafter projection, outside capture, on the current stream, before the executor's warmup. Returns the number of calls compiled.""" + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_lowm_splitk import ( + compile_push, + ) + from .ops.m3_lowm_direct import compile_direct_push - from .ops.m3_lowm_splitk import compile_push use_pdl = get_env_enable_pdl() built = 0 @@ -4900,7 +4771,7 @@ def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: def mlp(self, h: torch.Tensor, reduce: bool, cpp_silu: bool = False) -> torch.Tensor: """gate_up, the plain SwiGLU, down (all-reduced unless ``reduce`` is False). ``cpp_silu`` (the fused draft - loop): the SwiGLU on this target's C++ op (``ops/spec/m3_draft_silu``, minimaxM3Sm107Tp2/siluAndMul.cu), bit + loop): the SwiGLU on this target's C++ op (``ops/spec/m3_draft_silu``, minimaxM3Sm107/siluAndMul.cu), bit for bit the in-tree Triton silu_and_mul in one launch either way (no Triton on the GEN decode path).""" gu = self.linear(h, "gu") act = None @@ -4968,41 +4839,14 @@ def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, captu in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; otherwise it is the in-tree capture. """ - if capture_norm is None: - spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) - return - for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): - if captured_layer_id == layer_id: - assert spec_metadata.hidden_states is not None - num_tokens = spec_metadata.num_capture_tokens - if num_tokens > spec_metadata.hidden_states.shape[0]: - raise RuntimeError( - "EAGLE3 hidden-state capture token count exceeds buffer capacity: " - f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" - ) - if num_tokens > hidden_states.shape[0]: - raise RuntimeError( - "EAGLE3 hidden-state capture token count exceeds available hidden states: " - f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" - ) - hidden_size = spec_metadata.hidden_size - m3_spec_tail.capture_norm( - hidden_states[:num_tokens], - residual[:num_tokens] if residual is not None else None, - capture_norm[0][i], - capture_norm[1], - spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], - ) - break + return _modeling_support._eagle3_capture( + _SUPPORT_TARGET, spec_metadata, layer_id, hidden_states, residual, capture_norm + ) def _sharded_head(lm_head: nn.Module) -> bool: """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" - return ( - getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN - and getattr(lm_head, "padding_size", 0) == 0 - and getattr(lm_head, "tp_size", 1) > 1 - ) + return _modeling_support._sharded_head(_SUPPORT_TARGET, lm_head) class _M3ShardDraftLogits(nn.Module): @@ -5219,6 +5063,7 @@ def _forward_impl( max_draft_len=runtime_draft_len, ) + next_new_tokens = None if self._uses_external_shared_target_kv: next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( position_ids=position_ids, @@ -5232,6 +5077,8 @@ def _forward_impl( batch_size=batch_size, ) else: + # Retain the pre-draft tensor before metadata is replaced, as in the parent worker. + pre_draft_seq_lens = attn_metadata._seq_lens_cuda # Save the old attn_metadata and spec_metadata self._prepare_attn_metadata_for_spec_dec(attn_metadata) @@ -5254,7 +5101,7 @@ def _forward_impl( # Get the draft KV cache manager if using separate layouts draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - next_draft_tokens = self._forward_draft_loop( + next_draft_tokens, next_new_tokens = self._forward_draft_loop( inputs, attn_metadata, spec_metadata, @@ -5266,6 +5113,8 @@ def _forward_impl( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ) # restore attn_metadata to support cuda graph self._restore_attn_metadata_from_spec_dec(attn_metadata) @@ -5277,7 +5126,7 @@ def _forward_impl( if self._m3_next_new_tokens is not None: # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). next_new_tokens = self._m3_next_new_tokens - else: + elif next_new_tokens is None: next_new_tokens = self._prepare_next_new_tokens( accepted_tokens, next_draft_tokens, @@ -5958,11 +5807,13 @@ def _forward_draft_loop( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ): """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop (W2b: drafting from the head's local shard where ``_m3_wide_draft_logits_ok``).""" if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): - return self._m3_fused_draft_loop( + next_draft_tokens = self._m3_fused_draft_loop( inputs, attn_metadata, spec_metadata, @@ -5971,6 +5822,7 @@ def _forward_draft_loop( batch_size, num_accepted_tokens, ) + return next_draft_tokens, None # W2b: on a step the fused loop declines, the drafter head returns this rank's BF16 vocab shard (not the # gathered fp32 logits), so the greedy top-1 runs the fused shard argmax (greedy_sample_draft_with_tp_gather). swap = self._m3_wide_draft_logits_ok(spec_metadata, draft_model) @@ -5990,6 +5842,8 @@ def _forward_draft_loop( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ) finally: if swap: @@ -6167,40 +6021,7 @@ def _m3_fused_draft_loop( def _build_speculation(shell, model_config: ModelConfig): """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for the llama3 drafter architecture).""" - spec = model_config.spec_config - if ( - not spec.spec_dec_mode.use_one_engine() - or not spec.spec_dec_mode.is_eagle3_one_model() - or spec.eagle3_model_arch != "llama3" - or getattr(spec, "use_dynamic_tree", False) - ): - raise ValueError( - "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" - ) - draft_config = ModelConfig.from_pretrained( - spec.speculative_model, - trust_remote_code=True, - attn_backend=model_config.attn_backend, - moe_backend=_get_requested_draft_moe_backend(model_config, spec), - mapping=model_config.mapping, - spec_config=model_config.spec_config, - max_num_tokens=model_config.max_num_tokens, - moe_max_num_tokens=model_config.moe_max_num_tokens, - ) - _set_draft_kv_cache_quant_algo(draft_config, model_config) - draft_config.extra_attrs = model_config.extra_attrs - use_separate = should_use_separate_draft_kv_cache(spec) - draft = M3Eagle3Draft( - shell.model, - draft_config, - model_config.pretrained_config.num_hidden_layers, - spec.num_capture_layers, - ) - worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) - worker.set_draft_model(draft) - shell.epilogue.append(draft) - shell.epilogue.append(worker) - return draft_config, draft, worker, use_separate + return _modeling_support._build_speculation(_SUPPORT_TARGET, shell, model_config) # -------------------------------------------------------------------------------------------------------------------- @@ -6451,3 +6272,6 @@ def forward( # at its tail (one launch inside the step graph). spec_fold.store_in_forward(outputs) return outputs + + +_SUPPORT_TARGET = cast(_modeling_support.PinnedModelingSupport, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py index c2b23aa18750..2989104e372c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_draft_attention.py @@ -4,17 +4,24 @@ from __future__ import annotations -import torch +import torch as torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.draft_attention import ( + GroupedDraftAttentionBase, +) +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs as AttentionForwardArgs, +) from tensorrt_llm._torch.attention.backends.interface import ( - AttentionForwardArgs, - PredefinedAttentionMask, + PredefinedAttentionMask as PredefinedAttentionMask, ) -from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.functional import PositionEmbeddingType as PositionEmbeddingType +from tensorrt_llm.functional import RotaryScalingType as RotaryScalingType -class M3DraftAttention(OpWrapper): +class M3DraftAttention(GroupedDraftAttentionBase): """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` @@ -38,94 +45,5 @@ class M3DraftAttention(OpWrapper): kernel's input precision, its P rounding and accumulation order), which the test bounds. """ - def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: - num_tokens = metadata.num_tokens - out = backend.forward( - qkv[:num_tokens, :], - None, - None, - metadata, - forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), - ) - return out[0] if isinstance(out, tuple) else out - - def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: - """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. - - Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and - attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from - the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated - after the call (whose roped, quantized append it then reads back; the test checks that append itself). - Mirrors grouped KV heads (query head h reads KV head h // (heads / kv heads); two per rank at tp2) and - un-scaled NeoX RoPE, the drafter's configuration. - """ - manager = metadata.kv_cache_manager - layer = backend.layer_idx - heads, dim = backend.num_heads, backend.head_dim - kv_heads = int(backend.num_kv_heads) - assert heads % kv_heads == 0 - group = heads // kv_heads - rope = backend.rope_params - assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) - assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 - page = int(manager.tokens_per_block) - sm_scale = dim**-0.5 / float(backend.q_scaling) - kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) - num_seqs = int(metadata.num_seqs) - qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] - cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] - total = int(metadata.num_tokens) - assert sum(qlens) == total - device = qkv.device - - rot = int(rope.dim) - inv_freq = float(rope.theta) ** ( - -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot - ) - pool = manager.get_buffers(layer, "HND") - q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) - out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) - first = 0 - for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): - n, length = qlens[row], cached[row] + qlens[row] - pages = torch.tensor( - manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ - : (length + page - 1) // page - ], - dtype=torch.int64, - device=device, - ) - # [kv heads, tokens, dim]: each KV head's rows in token order across the pages. - keys = ( - pool[pages, 0] - .to(torch.float64) - .transpose(0, 1) - .reshape(kv_heads, -1, dim)[:, :length] - * kv_scale - ) - values = ( - pool[pages, 1] - .to(torch.float64) - .transpose(0, 1) - .reshape(kv_heads, -1, dim)[:, :length] - * kv_scale - ) - keys = keys.repeat_interleave(group, dim=0) # [heads, tokens, dim] - values = values.repeat_interleave(group, dim=0) - qpos = cached[row] + torch.arange(n, device=device) - angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] - cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] - rows = q[first : first + n] - x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] - rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) - allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] - logits = torch.einsum("thd,hkd->thk", rows, keys) * sm_scale - logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) - out[first : first + n] = torch.einsum( - "thk,hkd->thd", torch.softmax(logits, dim=-1), values - ) - first += n - return out.reshape(total, heads * dim).to(qkv.dtype) - m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py index bc3cd23b289f..1d3530215ed6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fullctx_decode.py @@ -8,12 +8,15 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_dense_decode import ( + dense_attention_reference, + gather_paged_kv, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( fp8_e4m3_saturate, ) from .. import m3_fullctx_decode as kernel -from .m3_dense_decode import dense_attention_reference, gather_paged_kv __all__ = ["m3_fullctx_decode"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py index be1c89b97261..b46fd83da452 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_index_score.py @@ -8,13 +8,9 @@ from __future__ import annotations -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_index_score import ( + M3IndexScoreBase, ) __all__ = ["m3_index_score"] @@ -42,17 +38,7 @@ def _wide_lens(requests: int) -> tuple: HEAD_DIM = 128 -def compiled_scorer(key: tuple): - """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss - (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" - from ..m3_index_score import IndexScoreRunner - - if key not in IndexScoreRunner.kernel_cache: - IndexScoreRunner._compile(*key) - return IndexScoreRunner.kernel_cache[key] - - -class M3IndexScore(OpWrapper): +class M3IndexScore(M3IndexScoreBase): """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is @@ -162,78 +148,5 @@ class M3IndexScore(OpWrapper): differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. """ - def __call__( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> torch.Tensor: - _, num_heads, head_dim = idx_q.shape - key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) - compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) - return score - - def reference( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> torch.Tensor: - """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and - ``early`` do not change them.""" - total_q, heads, _ = idx_q.shape - qlen = max_decode_query_len - out = torch.full( - (heads, total_q, score.shape[2]), - -float("inf"), - dtype=torch.float32, - device=idx_q.device, - ) - with true_fp32_matmul(): - for b, seq in enumerate(seq_lens.tolist()): - nblocks = (seq + PAGE - 1) // PAGE - if nblocks == 0: - continue - t0 = b * qlen - q = idx_q[t0 : t0 + qlen].float() - k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() - s = torch.einsum("qhd,kd->hqk", q, k) - q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) - k_pos = torch.arange(k.shape[0], device=idx_q.device) - s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) - out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) - return out - - def is_valid( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> None: - # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. - nblocks = (seq_lens + PAGE - 1) // PAGE - assert bool((nblocks <= block_table.shape[1]).all()), ( - "a request runs past its block-table row" - ) - assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" - assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( - "idx_q must hold max_decode_query_len rows per request" - ) - m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py index 6ec67e73edd4..c45b8595e0e1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_msa_attention.py @@ -4,20 +4,29 @@ from __future__ import annotations -from typing import Optional +from typing import Optional as Optional import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa import ( + GroupedMsaAttentionBase, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa import ( + _selected_blocks as _selected_blocks, +) from tensorrt_llm._torch.attention.backends.interface import ( - AttentionForwardArgs, - SparseBackendForwardArgs, + AttentionForwardArgs as AttentionForwardArgs, +) +from tensorrt_llm._torch.attention.backends.interface import ( + SparseBackendForwardArgs as SparseBackendForwardArgs, ) -from .m3_sparse_decode_lc import dequant_nvfp4_pages +from .m3_sparse_decode_lc import dequant_nvfp4_pages as dequant_nvfp4_pages -class M3MsaAttention(OpWrapper): +class M3MsaAttention(GroupedMsaAttentionBase): """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` [T, heads * 128] (bf16, the step's live rows). @@ -31,6 +40,14 @@ class M3MsaAttention(OpWrapper): objects for this layer and step. """ + @staticmethod + def _dequant_pages(*args, **kwargs) -> torch.Tensor: + return dequant_nvfp4_pages(*args, **kwargs) + + @staticmethod + def _select_blocks(*args, **kwargs) -> torch.Tensor: + return _selected_blocks(*args, **kwargs) + ARCHS = frozenset({Arch.SM_107}) CELLS: tuple[Cell, ...] = ( @@ -62,188 +79,5 @@ class M3MsaAttention(OpWrapper): rounding and accumulation order), which the test bounds. """ - def __call__( - self, - backend, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - idx_q: Optional[torch.Tensor], - metadata, - output: torch.Tensor, - sparse: bool, - kv_scale_orig_quant: Optional[torch.Tensor], - kv_scale_quant_orig: Optional[torch.Tensor], - ) -> torch.Tensor: - if k is not None: - backend.write_layer_caches( - k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant - ) - if sparse: - blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) - args = AttentionForwardArgs( - output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) - ) - else: - args = AttentionForwardArgs(output=output) - args.kv_scale_orig_quant = kv_scale_orig_quant - args.kv_scale_quant_orig = kv_scale_quant_orig - backend.forward(q, None, None, metadata, forward_args=args) - return output - - def reference( - self, - backend, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - idx_q: Optional[torch.Tensor], - metadata, - output: torch.Tensor, - sparse: bool, - kv_scale_orig_quant: Optional[torch.Tensor], - kv_scale_quant_orig: Optional[torch.Tensor], - ) -> torch.Tensor: - """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. - - Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the - metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of - ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the - test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and - attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors grouped KV heads (query - head h reads KV head h // (heads / kv heads); two per rank at tp2), each KV head selecting its own top-k - blocks from its own index heads against the rank's single index-K head. - """ - manager = metadata.kv_cache_manager - layer = backend.layer_idx - heads, dim = backend.num_heads, backend.head_dim - kv_heads = int(backend.num_kv_heads) - assert heads % kv_heads == 0 - group = heads // kv_heads - page = int(manager.tokens_per_block) - sm_scale = dim**-0.5 / float(backend.q_scaling) - num_seqs = int(metadata.num_seqs) - qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] - cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] - total = sum(qlens) - assert q.shape[0] == total and output.shape[0] == total - - pool = manager.get_buffers(layer, "HND") - nvfp4 = manager.is_nvfp4_layer(layer) - if nvfp4: - scale_pool = manager.get_block_scale_buffers(layer, "HND") - k_global = float(kv_scale_quant_orig.reshape(-1)[1]) - v_global = float(kv_scale_quant_orig.reshape(-1)[2]) - q_rows = q.reshape(total, heads, dim) - if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: - # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. - q_rows = q_rows.to(torch.float8_e4m3fn) - q64 = q_rows.to(torch.float64) - if sparse: - config = backend.m3_config - idx_cache = metadata.msa_idx_k_cache(layer) - idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( - torch.float64 - ) - assert config.num_index_heads % kv_heads == 0 - idx_group = config.num_index_heads // kv_heads - - out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) - first = 0 - for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): - n, length = qlens[row], cached[row] + qlens[row] - blocks = (length + page - 1) // page - pages = torch.tensor( - manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], - dtype=torch.int64, - device=q.device, - ) - qpos = cached[row] + torch.arange(n, device=q.device) - causal = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] - for g in range(kv_heads): - if nvfp4: - keys = dequant_nvfp4_pages( - pool[:, 0].view(torch.uint8), - scale_pool[:, 0], - pages, - quad_scales=False, - kv_head=g, - ) - values = dequant_nvfp4_pages( - pool[:, 1].view(torch.uint8), - scale_pool[:, 1], - pages, - quad_scales=True, - kv_head=g, - ) - keys, values = keys * k_global, values * v_global - else: - keys = pool[pages, 0, g].to(torch.float64) - values = pool[pages, 1, g].to(torch.float64) - keys = keys.reshape(-1, dim)[:length] - values = values.reshape(-1, dim)[:length] - allowed = causal - if sparse: - # One index-K head per rank, shared by every index head (the cache is [pages, 1, page, dim]). - idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] - chosen = _selected_blocks( - idx64[first : first + n, g * idx_group : (g + 1) * idx_group], - idx_keys, - causal, - qpos, - page, - config, - ) - allowed = causal & chosen.repeat_interleave(page, dim=1)[:, :length] - h0, h1 = g * group, (g + 1) * group - for lo in range(0, n, 256): - hi = min(n, lo + 256) - logits = ( - torch.einsum("thd,kd->thk", q64[first + lo : first + hi, h0:h1], keys) - * sm_scale - ) - logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) - out[first + lo : first + hi, h0:h1] = torch.softmax(logits, dim=-1) @ values - first += n - return out.reshape(total, heads * dim).to(output.dtype) - - -def _selected_blocks( - idx_q: torch.Tensor, - idx_keys: torch.Tensor, - allowed: torch.Tensor, - qpos: torch.Tensor, - page: int, - config, -) -> torch.Tensor: - """The blocks each query token of one request attends, [n, blocks] bool. - - A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of - ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last - ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the - backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the - ranking, so none is applied. - """ - n, length = allowed.shape - blocks = (length + page - 1) // page - scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) - scores = scores.masked_fill(~allowed, float("-inf")) - padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) - padded[:, :length] = scores - block_scores = padded.view(n, blocks, page).amax(dim=-1) - ids = torch.arange(blocks, device=scores.device)[None, :] - n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] - block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) - forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( - (ids >= n_valid - config.local_blocks) & (ids < n_valid) - ) - block_scores = block_scores.masked_fill(forced, float("inf")) - order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] - keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid - chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) - rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) - chosen[rows[keep], order[keep]] = True - return chosen - m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py index ff9a05c0233b..c819ed326a31 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_producer.py @@ -2,14 +2,20 @@ # SPDX-License-Identifier: Apache-2.0 """The sparse layers' NVFP4 QKV + indexer RMSNorm / RoPE / KV-insert producer at the TP2 heads, with its constant loads staged before the PDL grid-dependency wait and an early dependents trigger -(``kernels/minimaxM3Sm107Tp2/producerPw2.cu``).""" +(``kernels/minimaxM3Sm107/producerPw.cu``).""" from __future__ import annotations +import sys +from typing import cast + import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + producer_reference as _producer_shared, +) __all__ = ["m3_producer"] @@ -32,29 +38,20 @@ def _norm_rope( ): """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. x [T, 128] fp32; cos / sin [T, 32].""" - y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) - lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] - return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + return _producer_shared._norm_rope(_REFERENCE_TARGET, x, weight, cos, sin, eps) def decode_nvfp4_token( data: torch.Tensor, scales: torch.Tensor, quant_scale: float ) -> torch.Tensor: """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" - table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) - nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() - vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) - s = scales.view(torch.float8_e4m3fn).float() / quant_scale - return vals * s.repeat_interleave(FP4_BLOCK) + return _producer_shared.decode_nvfp4_token(_REFERENCE_TARGET, data, scales, quant_scale) def scale_index(within_page: int, plane: int) -> torch.Tensor: """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token swizzle for V.""" - cols = torch.arange(8) - if plane == 0: - return within_page * 8 + cols - return (within_page // 4) * 32 + cols * 4 + within_page % 4 + return _producer_shared.scale_index(_REFERENCE_TARGET, within_page, plane) class M3Producer(OpWrapper): @@ -158,27 +155,27 @@ def reference( """fp32 values before their quantization: "q" [T, 32, 128], "index_q" [T, 2, 128], per token and KV head "k" and "v" [T, 2, 128], and "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl``, ``early`` and ``prewait_step_inputs`` do not change them.""" - t = packed.shape[0] - x = packed.float().view(t, HEADS_PACKED, HEAD_DIM) - rope = rotary_cos_sin[position_ids[:t].long()] - cos, sin = rope[:, 0], rope[:, 1] - - def heads(first, count, weight): - return torch.stack( - [_norm_rope(x[:, h], weight, cos, sin, eps) for h in range(first, first + count)], - dim=1, - ) - - k0 = HEADS_Q - v0 = k0 + HEADS_KV - iq0 = v0 + HEADS_KV - return { - "q": heads(0, HEADS_Q, q_weight), - "k": heads(k0, HEADS_KV, k_weight), - "v": x[:, v0 : v0 + HEADS_KV], - "index_q": heads(iq0, HEADS_INDEX, index_q_weight), - "index_k": _norm_rope(x[:, iq0 + HEADS_INDEX], index_k_weight, cos, sin, eps), - } + return _producer_shared.producer_reference_multi_kv( + self, + _REFERENCE_TARGET, + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + prewait_step_inputs, + ) def is_valid( self, @@ -201,16 +198,29 @@ def is_valid( prewait_step_inputs: bool = False, ) -> None: assert heads_q == HEADS_Q, f"the producer is built for {HEADS_Q} query heads, not {heads_q}" - t = packed.shape[0] - # The RoPE row is read at the position with no bound. - pos = position_ids[:t] - assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( - "position_ids must index rotary_cos_sin" + return _producer_shared.validate_producer( + self, + _REFERENCE_TARGET, + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + prewait_step_inputs, ) - # Two tokens writing one slot race. - slots = out_cache_loc[:t] - live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] - assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" m3_producer = M3Producer() + +_REFERENCE_TARGET = cast(_producer_shared.ProducerTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py deleted file mode 100644 index e253c70a53aa..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qk_norm_rope_fp8.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( - fused_qk_norm_rope, -) - - -class M3QkNormRopeFp8(OpWrapper): - """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded - once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place - entry. - - The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps - of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="one decode token, tp2 heads (32 q, 2 kv)", spec=dict(m=1, q=32, kv=2)), - Cell(why="an 8-row decode verify", spec=dict(m=8, q=32, kv=2)), - Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=32, kv=2)), - ) - - note = """ - `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from - fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. - """ - - def __call__( - self, - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) -> torch.Tensor: - return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) - - def reference( - self, - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) -> torch.Tensor: - """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" - out = fused_qk_norm_rope( - qkv.clone(), - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) - return out.to(torch.float8_e4m3fn) - - -m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py index 66f7f33e8452..da14065a93b2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_long.py @@ -4,10 +4,19 @@ from __future__ import annotations +import sys +from typing import cast + import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.select_reference import ( + SelectReferenceTarget, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.select_reference import ( + select_reference as _shared_select_reference, +) #: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). MAX_BLOCKS = 9216 @@ -134,32 +143,18 @@ def reference( (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and ``trigger_phase`` only say where and how the op writes; they do not enter the result. """ - num_kv_heads, num_blocks, total_q = scores.shape - device = scores.device - blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) - raw = n_valid_blocks[:total_q].to(torch.int64) - valid = raw.clamp(0, num_blocks) - local_start = (raw - local_blocks).clamp_min(0) - result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) - for head in range(num_kv_heads): - forced = scores[head].float().t().clone() # [total_q, num_blocks] - forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) - forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) - bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF - key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) - in_row = blocks[None, :] < valid[:, None] - packed = torch.where( - in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max - ) - k = min(TOPK, num_blocks) - picked = torch.topk(packed, k, dim=1, largest=False).values - picked_key = picked >> 16 - keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) - ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) - ids = ids.sort(dim=1).values - ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) - result[:, head, :k] = ids.to(torch.int32) - return result + return _shared_select_reference( + self, + _SELECT_REFERENCE_TARGET, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) def is_valid( self, @@ -190,3 +185,5 @@ def is_valid( m3_select_long = M3SelectLong() + +_SELECT_REFERENCE_TARGET = cast(SelectReferenceTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py index 98f04538e784..ae91d98714cb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode.py @@ -9,6 +9,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_reference import ( + validate_sparse_selection, +) from ..triton_sparse_decode import minimax_m3_sparse_attn_decode from .m3_sparse_decode_lc import PAGE, sparse_attention_reference @@ -175,33 +178,9 @@ def is_valid( num_topk_chunks: Optional[int] = None, identity_blocks: bool = False, ) -> None: - """The inputs this call takes and answers wrongly. - - Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` - table entries and indexes the block table with them unchecked, so they - must be real blocks of the request. ``identity_blocks`` is a promise the - multi-query producer relies on: every token of a request reads the same - blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the - host: a test-only cost. - """ - total_q = q.shape[0] - seq = seq_lens.tolist() - max_topk = topk_idx.shape[-1] - for tok in range(total_q): - req, offset = divmod(tok, decode_query_len) - kv_len = max(seq[req] - decode_query_len + offset + 1, 0) - nblocks = (kv_len + PAGE - 1) // PAGE - ids = topk_idx[:, tok, : min(max_topk, nblocks)] - assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( - f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " - "the block table is indexed with them unchecked" - ) - if identity_blocks: - assert nblocks <= max_topk and bool( - (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() - ), ( - f"token {tok}: identity_blocks promises the table holds every block of the request in order" - ) + validate_sparse_selection( + q, topk_idx, block_table, seq_lens, decode_query_len, identity_blocks, PAGE + ) m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py index 94b160e3f4d4..c1dbb671eaa1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_sparse_decode_lc.py @@ -16,6 +16,13 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from ....ops.attention.sparse_reference import _E2M1 as _E2M1 +from ....ops.attention.sparse_reference import dequant_nvfp4_pages as _dequant_nvfp4_pages +from ....ops.attention.sparse_reference import mxfp8_handoff_reference as _mxfp8_handoff_reference +from ....ops.attention.sparse_reference import ( + sparse_attention_reference as _sparse_attention_reference, +) + #: The TP2 geometry the kernel is compiled for: 16 query heads on each of 2 KV heads per rank. HEADS_PER_KV = 16 KV_HEADS = 2 @@ -29,8 +36,6 @@ MX_BLOCK = 32 WIDTH = NUM_HEADS * HEAD_DIM -_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) - def dequant_nvfp4_pages( packed: torch.Tensor, @@ -39,26 +44,8 @@ def dequant_nvfp4_pages( quad_scales: bool, kv_head: int = 0, ) -> torch.Tensor: - """Dequantize KV head ``kv_head`` of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. - - ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low - nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, - HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token - ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at - ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor - global scale is not applied. Every product is exact in float64. - """ - table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) - rows = packed[pages, kv_head].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] - nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) - cols = HEAD_DIM // SF_VEC - region = block_scale[pages, kv_head].reshape(rows.shape[0], PAGE * cols) - t = torch.arange(PAGE, device=packed.device)[:, None] - c = torch.arange(cols, device=packed.device)[None, :] - index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c - scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) - scales = scales.view(torch.float8_e4m3fn).to(torch.float64) - return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + """Use the shared reference with this target's geometry.""" + return _dequant_nvfp4_pages(packed, block_scale, pages, quad_scales, kv_head) def sparse_attention_reference( @@ -76,91 +63,29 @@ def sparse_attention_reference( sm_scale: float, kv_head: int = 0, ) -> torch.Tensor: - """Block-sparse decode attention over KV head ``kv_head`` of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. - - Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset - ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / - PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at - ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that - attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or - its byte view); softmax and both products are float64. - """ - q64 = ( - q.view(torch.float8_e4m3fn).to(torch.float64) - if q.element_size() == 1 - else q.to(torch.float64) + """Use the shared reference with this target's geometry.""" + return _sparse_attention_reference( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, + kv_head, ) - total_q, num_heads, _ = q64.shape - kgs = float(k_global_scale.reshape(-1)[0]) - vgs = float(v_global_scale.reshape(-1)[0]) - out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) - seq = seq_lens.tolist() - positions = torch.arange(PAGE, device=q.device) - for tok in range(total_q): - req, offset = divmod(tok, decode_query_len) - kv_len = max(seq[req] - decode_query_len + offset + 1, 0) - real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) - blocks = topk[tok, :real_topk] - blocks = blocks[blocks >= 0].to(torch.int64) - if blocks.numel() == 0: - continue - pages = block_table[req, blocks].to(torch.int64) - keys = ( - dequant_nvfp4_pages( - k, k_block_scale, pages, quad_scales=False, kv_head=kv_head - ).reshape(-1, HEAD_DIM) - * kgs - ) - values = ( - dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True, kv_head=kv_head).reshape( - -1, HEAD_DIM - ) - * vgs - ) - valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) - logits = (q64[tok] @ keys.T) * sm_scale - logits = logits.masked_fill(~valid[None, :], float("-inf")) - values = values.masked_fill(~valid[:, None], 0) - out[tok] = torch.softmax(logits, dim=-1) @ values - return out def mxfp8_handoff_reference( out: torch.Tensor, num_rows_alloc: int ) -> Tuple[torch.Tensor, torch.Tensor]: - """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, - sf [num_rows_alloc * WIDTH // 32] uint8)``. - - Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` - rounded up when its mantissa is non-zero (capped at 254), the elements are - scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands - in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` - for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a - comparison against it is exact. - """ - total_q = out.shape[0] - x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) - amax = x.abs().amax(dim=-1) - sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) - bits = sfv.view(torch.int32) - e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) - e = e.clamp(max=254) - qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) - xq = (x * qs[..., None]).clamp(-448.0, 448.0) - q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) - num_sf_k_tiles = WIDTH // 128 - sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) - tok = torch.arange(total_q, device=out.device)[:, None] - k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] - offset = ( - (tok // 128) * num_sf_k_tiles * 512 - + (k_sf // 4) * 512 - + (tok % 32) * 16 - + ((tok % 128) // 32) * 4 - + k_sf % 4 - ) - sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) - return q8, sf + """Use the shared reference with this target's geometry.""" + return _mxfp8_handoff_reference(out, num_rows_alloc, width=WIDTH) class M3SparseDecodeLc(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py deleted file mode 100644 index 7f480855c170..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_allreduce.py +++ /dev/null @@ -1,61 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams - - -class M3AllReduce(OpWrapper): - """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the - plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: - ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; - returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the - communicator's per call. - - The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps - beyond the multicast kernel's rows. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="tp2, a prefill chunk", spec=dict(world=2, m=2048, hidden=6144)), - Cell(why="tp2, an 8-row generation step", spec=dict(world=2, m=8, hidden=6144)), - Cell( - why="tp2, plain sum of eight drafter rows", - spec=dict(world=2, m=8, hidden=6144, plain=True), - ), - ) - - note = """ - The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then - adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum - to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. - """ - - def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): - return comm(x, all_reduce_params=params) - - def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): - """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" - total = torch.stack([x.float() for x in xs]).sum(0) - if params is None: - return total.to(xs[0].dtype) - updated = (params.residual.float() + total).to(xs[0].dtype) - h = updated.float() - normed = ( - h - * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) - * params.norm_weight.float() - ) - return normed.to(xs[0].dtype), updated - - -m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py index b624ae64e59e..858aa4e9c76e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_allreduce_norm.py @@ -4,85 +4,23 @@ from __future__ import annotations -from typing import Optional, Sequence +from typing import Optional import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper -_E4M3_MAX = 448.0 +from ....ops.comm.reference import _E4M3_MAX as _E4M3_MAX +from ....ops.comm.reference import McAllreduceNormReference +from ....ops.comm.reference import boundary as boundary +from ....ops.comm.reference import mxfp8_rows as mxfp8_rows +from ....ops.comm.reference import sf_written as sf_written +from ....ops.comm.reference import swizzle_sf as swizzle_sf -def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. - - One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), - data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the - op leaves them (the comparison masks them out). - """ - m, h = normed.shape - blocks = normed.float().view(m, h // 32, 32) - amax = blocks.abs().amax(-1) - bits = (amax / _E4M3_MAX).view(torch.int32) - byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) - scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) - data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) - data = data.view(m, h).to(torch.float8_e4m3fn) - return data, swizzle_sf(byte.to(torch.uint8)) - - -def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: - """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" - rows, cols = byte.shape - padded_rows = (rows + 127) // 128 * 128 - r = torch.arange(rows, device=byte.device).view(-1, 1) - c = torch.arange(cols, device=byte.device).view(1, -1) - index = ( - (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) - ) - out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) - out[index.reshape(-1)] = byte.reshape(-1) - return out - - -def sf_written(rows: int, cols: int, device) -> torch.Tensor: - """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" - return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() - - -def boundary( - partials: Sequence[torch.Tensor], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - drafter: bool = False, -) -> tuple[torch.Tensor, ...]: - """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). - - ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer - fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; - only the stored residual is rounded to BF16. - - The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are - summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic - and the normalization are done in float64 and rounded to BF16 once. - """ - total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) - for p in partials: - total = total + p.float() - pre = total.to(torch.bfloat16).float() + residual.float() - updated = pre.to(torch.bfloat16) - u = pre.double() if drafter else updated.double() - inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) - normed = (u * inv * norm_w.double()).to(torch.bfloat16) - if not quant: - return normed, updated - return (normed, updated, *mxfp8_rows(normed)) - - -class M3McAllreduceNorm(OpWrapper): +class M3McAllreduceNorm(McAllreduceNormReference): """One TP group's boundary: ``updated = bf16(residual + bf16(sum_r bf16(x_r + x2_r)))``, ``normed = RMSNorm(updated)``: bitwise the stock all-reduce + fused residual RMSNorm (+ mxfp8_quantize) of SM107 TP2. @@ -252,29 +190,6 @@ def __call__( cap_eps, ) - def reference( - self, - per_rank_x: Sequence[torch.Tensor], - per_rank_x2: Optional[Sequence[torch.Tensor]], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - drafter: bool = False, - ) -> tuple[torch.Tensor, ...]: - """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). - - The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling - rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. - """ - partials = [] - for r, x in enumerate(per_rank_x): - if per_rank_x2 is None: - partials.append(x) - else: - partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) - return boundary(partials, residual, norm_w, eps, quant, drafter) - def is_valid( self, x: torch.Tensor, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py index 9861ed2f0513..ec1b419454e6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -4,40 +4,20 @@ from __future__ import annotations -from typing import Optional, Sequence +from typing import Optional import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper -from .m3_mc_allreduce_norm import boundary +from ....ops.comm.reference import McFinalizeAllreduceNormReference +from ....ops.comm.reference import boundary as boundary +from ....ops.comm.reference import finalize_partial as finalize_partial -def finalize_partial( - fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] -) -> torch.Tensor: - """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k - order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). - - The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two - FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. - """ - m, topk = perm.shape - acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) - sc = scales.reshape(m, topk).float() - for k in range(topk): - rows = perm[:, k].long() - keep = (rows >= 0).unsqueeze(-1) - gathered = fc2[rows.clamp_min(0)].double() - step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() - acc = torch.where(keep, step, acc) - if shared is not None: - acc = acc.to(torch.bfloat16).float() + shared.float() - return acc.to(torch.bfloat16) - - -class M3McFinalizeAllreduceNorm(OpWrapper): +class M3McFinalizeAllreduceNorm(McFinalizeAllreduceNormReference): """The MoE boundary with the routed experts' finalize folded in. Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, @@ -139,30 +119,6 @@ def __call__( cap_eps, ) - def reference( - self, - per_rank_fc2: Sequence[torch.Tensor], - per_rank_perm: Sequence[torch.Tensor], - per_rank_scales: Sequence[torch.Tensor], - per_rank_shared: Optional[Sequence[torch.Tensor]], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - ) -> tuple[torch.Tensor, ...]: - """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the - collective exception, as in ``m3_mc_allreduce_norm.reference``).""" - partials = [ - finalize_partial( - fc2, - per_rank_perm[r], - per_rank_scales[r], - None if per_rank_shared is None else per_rank_shared[r], - ) - for r, fc2 in enumerate(per_rank_fc2) - ] - return boundary(partials, residual, norm_w, eps, quant) - def is_valid( self, fc2: torch.Tensor, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py deleted file mode 100644 index 71fce6b253cb..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding/m3_embedding.py +++ /dev/null @@ -1,48 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops - - -class M3Embedding(OpWrapper): - """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). - - The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode - embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) - ), - Cell( - why="an 8-row decode verify, int64 ids", - spec=dict(rows=8, id_dtype=torch.int64, pad=False), - ), - Cell( - why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) - ), - ) - - note = """ - An exact gather: `reference` must match bit for bit. - """ - - def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: - return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) - - def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: - valid = (ids >= 0) & (ids < table.shape[0]) - rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] - return rows * valid.unsqueeze(-1).to(table.dtype) - - -m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py index e0caa5e350dc..d0d860de8645 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/embedding_firstnorm/op.py @@ -8,15 +8,22 @@ (:class:`DecodeEmbedding` compiles at construction). """ +import sys import threading -from typing import Dict, Tuple +from typing import Dict, Tuple, cast -import cuda.bindings.driver as cuda_driver -import cutlass.cute as cute +import cuda.bindings.driver as cuda_driver # noqa: F401 — used by the bound family implementation +import cutlass.cute as cute # noqa: F401 — used by the bound family implementation import torch -from cutlass.cute.runtime import from_dlpack +from cutlass.cute.runtime import from_dlpack as from_dlpack -from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding import ( + decode_embedding as _embedding_shared, +) + +from ....kernels.embedding_firstnorm.m3_embedding_firstnorm_kernel import ( # noqa: F401 + m3_embed_norm as _m3_embed_norm_jit, +) HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size MAX_ROWS = 32 # capacity: fixed: the kernel's built rows (decode steps) @@ -26,12 +33,7 @@ def arg(tensor): - return from_dlpack( - tensor.detach(), - assumed_align=tensor.element_size() - if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) - else 16, - ) + return _embedding_shared.arg(_EMBED_TARGET, tensor) def _check(ids, table, weight, raw, out) -> None: @@ -58,32 +60,7 @@ def _check(ids, table, weight, raw, out) -> None: def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): - key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) - fn = _compiled.get(key) - if fn is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "m3_embed_norm compiles on its first call per shape: call it once before capture" - ) - with _lock: - fn = _compiled.get(key) - if fn is None: - stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) - fn = _compiled[key] = cute.compile( - _m3_embed_norm_jit, - arg(ids), - arg(table), - arg(weight), - arg(raw), - arg(out), - int(table.shape[0]), - float(eps), - ids.shape[0], - HIDDEN, - bool(use_pdl), - stream, - ) - return fn + return _embedding_shared._compile(_EMBED_TARGET, ids, table, weight, eps, raw, out, use_pdl) def m3_embed_norm( @@ -97,20 +74,9 @@ def m3_embed_norm( ) -> Tuple[torch.Tensor, torch.Tensor]: """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / ``eps`` into ``out``; returns ``(out, raw)``.""" - _check(ids, table, weight, raw, out) - fn = _compile(ids, table, weight, eps, raw, out, use_pdl) - stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) - fn( - arg(ids), - arg(table), - arg(weight), - arg(raw), - arg(out), - int(table.shape[0]), - float(eps), - stream, + return _embedding_shared.m3_embed_norm( + _EMBED_TARGET, ids, table, weight, eps, raw, out, use_pdl ) - return out, raw class DecodeEmbedding: @@ -119,25 +85,12 @@ class DecodeEmbedding: ``eps``. Built (and compiled) before CUDA-graph capture.""" def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): - assert not torch.cuda.is_current_stream_capturing() - assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN - assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) - assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) - self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( - table, - weight, - rows, - id_dtype, - use_pdl, + _embedding_shared.init_decode_embedding( + self, _EMBED_TARGET, table, weight, eps, rows, id_dtype, use_pdl ) - self.eps = float(eps) - self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) - self.output = torch.empty_like(self.raw) - ids = torch.empty(rows, device=table.device, dtype=id_dtype) - _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) def __call__(self, ids): - assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() - return m3_embed_norm( - ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl - ) + return _embedding_shared.call_decode_embedding(self, _EMBED_TARGET, ids) + + +_EMBED_TARGET = cast(_embedding_shared.EmbeddingTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py deleted file mode 100644 index f1b619a13564..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_dense_swiglu.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The dense MLP's swigluoai activation as one C++ launch with PDL and an early release (DENSE_SWIGLU_CPP).""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -# The dense layers' swigluoai parameters, which the kernel hard-codes (modeling_minimaxm3._build_swiglu_oai_dense_mlp). -LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 - - -class M3DenseSwiglu(OpWrapper): - """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, - rounded once to bf16: trtllm::minimax_m3_sm107_tp2_swiglu_oai (minimaxM3Sm107Tp2/swigluOai.cu), bit for bit the - Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the - grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights - during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even - half width; ``__call__`` returns None otherwise or for other activation parameters (the caller runs the Triton - op).""" - - ARCHS = frozenset({Arch.SM_107}) - - # The TP2 dense MLP shard: gate | up 12288 wide (d = 6144), at the decode rows. - CELLS: tuple[Cell, ...] = ( - Cell(why="the TP2 dense MLP at 1 row", spec=dict(rows=1, d=6144)), - Cell(why="the TP2 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=6144)), - Cell(why="the TP2 dense MLP at 8 rows", spec=dict(rows=8, d=6144)), - ) - - note = """ - Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the - alpha-scaled ex2.approx sigmoid, div.full), checked over every bf16 gate and every bf16 up pattern. - """ - - def __call__( - self, - x: torch.Tensor, - limit: Optional[float] = LIMIT, - alpha: Optional[float] = ALPHA, - beta: Optional[float] = BETA, - use_pdl: bool = True, - ) -> Optional[torch.Tensor]: - if ( - (limit, alpha, beta) != (LIMIT, ALPHA, BETA) - or x.dtype != torch.bfloat16 - or x.dim() != 2 - or not x.is_cuda - or not x.is_contiguous() - or x.shape[1] % 4 != 0 - or x.data_ptr() % 4 != 0 - ): - return None - out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) - torch.ops.trtllm.minimax_m3_sm107_tp2_swiglu_oai(x, out, bool(use_pdl)) - return out - - def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: - return torch.ops.trtllm.silu_and_mul( - x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta - ) - - def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: - assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" - - -m3_dense_swiglu = M3DenseSwiglu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py index 94f5e112e592..8afbbdebce86 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_draft_lowm.py @@ -110,9 +110,13 @@ def __call__( run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) elif own_splitk: - # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # The shared split-K kernel (kernels/gemm/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` # it also pushes the rows into a drafter boundary all-reduce (modeling's DRAFT_AR_PUSH). - from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_lowm_splitk import ( + SplitKTactic, + run_splitk_dense, + run_splitk_dense_push, + ) if ar is not None: run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py deleted file mode 100644 index 3bfcbc625016..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_gemm.py +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned - -from .m3_mxfp8_linear import mxfp8_reference - - -class M3Mxfp8Gemm(OpWrapper): - """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight - [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key - and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The decode dispatch hands it the boundary / merge MXFP8 activation on rows the GEMV does not serve: the - # unsplit shared down projection at 17..32 rows. - CELLS: tuple[Cell, ...] = ( - Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=1536)), - Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=1536)), - ) - - note = """ - Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and - multiplies in fp32; the test's band covers the accumulation order. - """ - - def __call__( - self, - act: torch.Tensor, - act_scale: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - output_dtype: torch.dtype, - tune: bool, - ) -> torch.Tensor: - return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) - - def reference( - self, - act: torch.Tensor, - act_scale: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - output_dtype: torch.dtype, - tune: bool, - ) -> torch.Tensor: - """The dequantized product; ``tune`` is ignored.""" - return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) - - -m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py deleted file mode 100644 index 9eadb1e0d20a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_mxfp8_linear.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) -from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned - -_BLOCK = 32 - - -def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: - """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" - pad_r = (rows + 127) // 128 * 128 - pad_c = (cols + 3) // 4 * 4 - s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) - return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] - - -def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: - rows, k = x.shape - s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) - factor = torch.pow(2.0, (s - 127).to(torch.float32)) - return ( - x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) - ).view(rows, k) - - -def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: - """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" - with true_fp32_matmul(): - return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() - - -class M3Mxfp8Linear(OpWrapper): - """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 - out, on the backend the engine selected: - - * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize - and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); - * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or - the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by - ``alpha`` (fp32 [1], 1.0 for these checkpoints). - - Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it - outside the decode-step kernels. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="sparse-layer QKV, a prefill chunk, native GEMM", - spec=dict(m=2048, n=4992, k=6144, path="native"), - ), - Cell( - why="o_proj, a prefill chunk, native GEMM", - spec=dict(m=2048, n=6144, k=4096, path="native"), - ), - Cell( - why="dense down projection, eager decode, native GEMM", - spec=dict(m=4, n=6144, k=6144, path="native"), - ), - Cell( - why="shared gate_up, FlashInfer GEMM", - spec=dict(m=64, n=3072, k=6144, path="flashinfer"), - ), - Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=4096, path="tuned")), - ) - - note = """ - Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` - quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and - multiplies dequantized operands in fp32; the test's band covers the accumulation order. - """ - - def __call__( - self, - x: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - alpha: torch.Tensor, - output_dtype: torch.dtype, - tuned_quantize_gemm: bool, - tune: bool, - flashinfer: bool, - native_tuning: bool, - ) -> torch.Tensor: - if tuned_quantize_gemm: - return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) - act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) - if flashinfer: - return torch.ops.trtllm.flashinfer_mm_mxfp8( - act, act_scale, weight, weight_scale, output_dtype - ) - gemm = ( - torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned - if native_tuning - else torch.ops.trtllm.mxfp8_mxfp8_gemm - ) - return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) - - def reference( - self, - x: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - alpha: torch.Tensor, - output_dtype: torch.dtype, - tuned_quantize_gemm: bool, - tune: bool, - flashinfer: bool, - native_tuning: bool, - ) -> torch.Tensor: - """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are - ignored.""" - m, k = x.shape - blocks = x.float().view(m, k // _BLOCK, _BLOCK) - amax = blocks.abs().amax(-1, keepdim=True) - exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) - scale = torch.pow(2.0, exp) - q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale - with true_fp32_matmul(): - out = q.view(m, k) @ _dequant(weight, weight_scale).t() - return (out * alpha.float()).to(output_dtype) - - -m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py index c61ab4acbf14..75d4576ac261 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_draft_step.py @@ -27,9 +27,12 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + encode_chunked, + round_trip_ok, +) from tensorrt_llm.logger import logger -from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op from . import m3_report diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py index a4c7b5ccb8fd..c8ff1799bf25 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_drafter_attn.py @@ -20,27 +20,50 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Mapping, Optional, Tuple +import sys +from typing import TYPE_CHECKING as TYPE_CHECKING +from typing import Mapping as Mapping +from typing import Optional as Optional +from typing import Tuple as Tuple +from typing import cast import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + drafter as _drafter_shared, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + DrafterFullCtxBase, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + drafter_rope_applies as _shared_rope_applies, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + initialize_drafter_state as _shared_initialize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + install as _shared_install, +) from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( - FlashInferTrtllmGenFmha, + FlashInferTrtllmGenFmha as FlashInferTrtllmGenFmha, ) -from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase -from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase as FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams as FmhaParams from tensorrt_llm._torch.attention.backends.fmha.utils import ( - get_attention_chunk_size, - get_bmm1_scale, + get_attention_chunk_size as get_attention_chunk_size, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import get_bmm1_scale as get_bmm1_scale +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionInputType as AttentionInputType, ) -from tensorrt_llm._torch.attention.backends.interface import AttentionInputType from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( - fp8_e4m3_saturate, + fp8_e4m3_saturate as fp8_e4m3_saturate, ) -from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor -from tensorrt_llm.bindings.internal import thop -from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.quantization.mode import QuantMode +from tensorrt_llm._utils import TensorWrapper as TensorWrapper +from tensorrt_llm._utils import convert_to_torch_tensor as convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop as thop +from tensorrt_llm.functional import PositionEmbeddingType as PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode as QuantMode from . import m3_fullctx_decode as fullctx @@ -60,9 +83,15 @@ def fwd_has_no_mrope(fwd) -> bool: from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention -class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): +class M3DrafterFullCtxFmha(DrafterFullCtxBase): """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + FULLCTX = fullctx + GEOMETRY_ERROR = "not the TP2 FP8 drafter geometry" + + def _launch_rope(self, *args) -> None: + torch.ops.trtllm.minimax_m3_sm107tp2_drafter_rope_fp8_kv(*args, ROPE_CARVEOUT) + def __init__( self, attn: "TrtllmAttention", @@ -77,291 +106,25 @@ def __init__( on_rope_taken=None, ) -> None: super().__init__(attn) - self._recipes = recipes - self._max_kv_len = max_kv_len - self._own_reducer = own_reducer - self._fused_rope_mode = fused_rope_mode - self._pdl = pdl - self._on_rope_taken = on_rope_taken - self._one_request_multi = ROPE_ONE_REQUEST_MULTI - self._on_taken = on_taken - self._on_uncovered = on_uncovered - self._on_missed = on_missed - - def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: - nreq = int(metadata.num_generations) - if nreq <= 0 or int(q.shape[0]) % nreq: - return None - return int(q.shape[0]) // nreq, nreq - - def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: - if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: - return False - if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: - return False - attn = self.attn - if ( - attn.num_heads != fullctx.HEADS - or attn.num_kv_heads != fullctx.KV_HEADS - or attn.head_dim != fullctx.HEAD_DIM - or not QuantMode(attn.quant_mode).has_fp8_kv_cache() - ): - self._on_uncovered("not the TP2 FP8 drafter geometry") - return False - # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's - # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. - if ( - forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant - or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig - ): - self._on_uncovered("per-step KV cache scales") - return False - if int(metadata.max_seq_len) > self._max_kv_len: - self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") - return False - shape = self._shape(q, metadata) - if shape is None or shape not in self._recipes: - self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") - return False - # Compiled before warmup (``install``); the forward never compiles one. - if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: - self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") - return False - return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) - - def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: - attn = params.attn - width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim - rotary_dim = int(attn.rope_params.dim) - return ( - self._fused_rope_mode is not None - and params.qkv_input is not None - and attention_input.dim() == 2 - and attention_input.dtype == torch.bfloat16 - and attention_input.is_contiguous() - and int(attention_input.shape[1]) == width - and not params.is_cross - and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) - and attn.rotary_cos_sin is not None - and attn.rotary_cos_sin.dtype == torch.float32 - and fwd_has_no_mrope(params.fwd) - and 0 < rotary_dim <= attn.head_dim - and rotary_dim % 16 == 0 - and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) - # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo - # runs for several requests only) and reads the workspace's; with ROPE_ONE_REQUEST_MULTI - # the fused launch (which places rows by request = token / q_len, no map) takes it too, else it keeps - # the in-tree path (and its token-map guard). - and ( - self._one_request_multi - or not (self._multi_token(params) and int(params.batch_size) == 1) - ) - ) - - @staticmethod - def _multi_token(params: FmhaParams) -> bool: - return ( - params.spec_decoding_generation_lengths is not None - and int(params.attn.predicted_tokens_per_seq) > 1 + _shared_initialize( + self, + _DRAFTER_BINDINGS, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode, + pdl, + on_rope_taken, + has_one_request_multi=True, ) - def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: - """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" - if not self._fused_rope_applies(params, attention_input): - return False - attn = params.attn - meta = params.meta - qlen, nreq = int(params.input_seq_length), int(params.batch_size) - plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) - if plan is None: - return False - # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary - # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the - # block table is the pool's [requests, 2, blocks] K / V sub-page indices. - pool_index, layer_in_pool = ( - int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] - ) - tpb = int(params.tokens_per_block) - sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim - base = ( - int(meta.host_kv_cache_pool_pointers[pool_index, 0]) - + layer_in_pool * int(params.kv_factor) * sub_page_bytes - ) - pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 - kv_pages = convert_to_torch_tensor( - TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) - ) - tables = meta.kv_cache_block_offsets[ - pool_index, params.seq_offset : params.seq_offset + nreq - ] - multi = self._multi_token(params) - if not torch.cuda.is_current_stream_capturing(): - # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess - # path checks; the op itself writes V at the table's V sub-page. - tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb - cols = torch.arange(tables.shape[-1], device=tables.device) - in_use = cols[None, :] < tpb_held[:, None] - offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() - if offsets.numel() != 1 or int(offsets) != 1: - raise RuntimeError( - f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" - ) - q = torch.empty( - int(params.num_tokens), - attn.num_heads, - attn.head_dim, - dtype=torch.float8_e4m3fn, - device=attention_input.device, - ) - torch.ops.trtllm.minimax_m3_sm107tp2_drafter_rope_fp8_kv( - attention_input, - attn.rotary_cos_sin, - params.sequence_lengths, - params.spec_decoding_position_offsets if multi else None, - kv_pages, - tables, - q, - qlen, - attn.num_heads, - attn.num_kv_heads, - int(attn.rope_params.dim), - tpb, - self._fused_rope_mode, - self._pdl, - ROPE_CARVEOUT, - ) - plan.run( - q, - kv_pages[:-1], - kv_pages[1:], - tables[:, 0], - params.sequence_lengths[:nreq], - get_bmm1_scale(attn), - params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), - ) - if self._on_rope_taken is not None: - self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") - self._on_taken(params.num_tokens) - return True - - def run_generation(self, params: FmhaParams) -> None: - attention_input = params.qkv_input if params.qkv_input is not None else params.query_input - if attention_input is None: - raise RuntimeError("the drafter decode requires QKV or query input.") - attn = params.attn - meta = params.meta - fwd = params.fwd - rope_params = attn.rope_params - if self._run_fused_rope(params, attention_input): - return - # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 - # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). - fp8_context_fmha = False - ( - q_processed, - kv_pool, - block_tables, - _kv_scale_pool, - _bmm1_scale, - _bmm2_scale, - _fmha_workspace, - _cu_seqlens, - _max_q_len, - _max_kv_len, - _window_left, - is_multi_token_gen, - ) = thop.trtllm_gen_generation_preprocess( - attention_input, - params.workspace, - params.sequence_lengths, - params.spec_decoding_generation_lengths, - params.spec_decoding_position_offsets, - meta.kv_cache_block_offsets, - meta.host_kv_cache_pool_pointers, - meta.host_kv_cache_pool_mapping, - fwd.kv_scale_orig_quant, - fwd.kv_scale_quant_orig, - fwd.out_scale, - attn.rotary_inv_freq, - attn.rotary_cos_sin, - fwd.mrope_position_deltas, - attn.local_layer_idx, - params.seq_offset, - attn.num_heads, - attn.num_kv_heads, - attn.head_dim, - params.tokens_per_block, - attn.quant_mode, - params.max_attention_window_size, - params.cyclic_attention_window_size, - params.num_tokens, - params.batch_size, - params.input_seq_length, - params.max_past_kv_length, - rope_params.dim, - rope_params.theta, - int(rope_params.scale_type), - rope_params.scale, - rope_params.max_positions, - attn.position_embedding_type, - get_bmm1_scale(attn), - 1.0, - fp8_context_fmha, - attn.predicted_tokens_per_seq, - get_attention_chunk_size(attn), - self._multi_processor_count, - params.total_num_blocks, - params.kv_factor, - True, - params.is_cross, - ) - qlen, nreq = params.input_seq_length, params.batch_size - plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) - # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the - # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row - # pass is never multi-token. - if plan is None or (is_multi_token_gen and qlen == 1): - raise RuntimeError( - f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" - ) - rows = params.num_tokens * attn.num_heads * attn.head_dim - q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) - q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) - # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 - # of the pool (the two views have the same page count, as the kernel requires). - if not torch.cuda.is_current_stream_capturing(): - # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. - lengths = params.spec_decoding_generation_lengths - if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): - raise RuntimeError( - f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" - ) - # Only the blocks a request holds are laid out; the table's tail entries are unset. - tpb = int(params.tokens_per_block) - held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb - cols = torch.arange(block_tables.shape[-1], device=block_tables.device) - in_use = cols[None, :] < held[:, None] - offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() - if offsets.numel() != 1 or int(offsets) != 1: - raise RuntimeError( - f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" - ) - # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. - # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. - pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 - kv_pages = convert_to_torch_tensor( - TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) - ) - plan.run( - q, - kv_pages[:-1], - kv_pages[1:], - block_tables[:nreq, 0], - params.sequence_lengths[:nreq], - get_bmm1_scale(attn), - params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + return _shared_rope_applies( + self, _DRAFTER_BINDINGS, params, attention_input, has_one_request_multi=True ) - self._on_taken(params.num_tokens) def install( @@ -379,35 +142,20 @@ def install( """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it was installed on.""" - from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention - - count = 0 - device = next(draft_model.parameters()).device - # The backend is not a module: each attention module of the drafter holds it as ``attn``. - for module in draft_model.modules(): - backend = getattr(module, "attn", None) - if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): - continue - if not bool((backend.kv_cache_scaling_factor == 1.0).all()): - on_uncovered("a scaled FP8 KV cache") - continue - fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) - libs = backend._fmha_manager.fmha_libs - if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): - libs.insert( - 0, - M3DrafterFullCtxFmha( - backend, - recipes, - max_kv_len, - own_reducer, - on_taken, - on_uncovered, - on_missed, - fused_rope_mode=fused_rope_mode, - pdl=pdl, - on_rope_taken=on_rope_taken, - ), - ) - count += 1 - return count + return _shared_install( + draft_model=draft_model, + recipes=recipes, + max_kv_len=max_kv_len, + own_reducer=own_reducer, + pdl=pdl, + on_taken=on_taken, + on_uncovered=on_uncovered, + on_missed=on_missed, + fused_rope_mode=fused_rope_mode, + on_rope_taken=on_rope_taken, + prebuild=fullctx.prebuild, + fmha_class=M3DrafterFullCtxFmha, + ) + + +_DRAFTER_BINDINGS = cast(_drafter_shared.MultiRequestDrafterBindings, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py index 1736115f1386..fb091a142a38 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fold_tail.py @@ -18,66 +18,18 @@ """ import torch -import triton -import triton.language as tl +import triton as triton +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.spec.m3_fold_tail import ( + _fold_tail_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.fold_tail import ( + fold_tail as _shared_fold_tail, +) -@triton.jit -def _fold_tail_kernel( - ctrl, - slots, - src_new, - sn_s0, - sn_s1, - SN_W: tl.constexpr, - dst_new, - dn_s0, - dn_s1, - DN_W: tl.constexpr, - src_next, - sx_s0, - sx_s1, - SX_W: tl.constexpr, - dst_next, - dx_s0, - dx_s1, - DX_W: tl.constexpr, - src_draft, - sd_s0, - sd_s1, - SD_W: tl.constexpr, - dst_draft, - dd_s0, - dd_s1, - DD_W: tl.constexpr, - src_lens, - dst_lens, - ids, - draft_ids, - pos_off, - kv_off, - T: tl.constexpr, - D: tl.constexpr, - BLOCK: tl.constexpr, -): - i = tl.program_id(0) - n = tl.load(ctrl) - if i < n: - s = tl.load(slots + i).to(tl.int64) - c = tl.arange(0, BLOCK) - new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) - tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) - nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) - tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) - drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) - tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) - ln = tl.load(src_lens + i) - tl.store(dst_lens + s, ln) - # The next step's inputs read the store's padded rows: zero past each source width. - tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) - tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) - tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) - tl.store(kv_off + i, ln - T) + +def _launch_cpp(*args: object) -> None: + torch.ops.trtllm.minimax_m3_sm107_tp2_fold_tail(*args) def fold_tail( @@ -101,72 +53,24 @@ def fold_tail( ) -> bool: """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width].""" - two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) - one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) - if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): - return False - if ( - any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) - or slots.dtype != torch.long - ): - return False - t_step = tokens_per_step - d_step = t_step - 1 - if ( - rows <= 0 - or slots.numel() < rows - or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) - or ids.numel() < rows * t_step - or pos_off.numel() < rows * t_step - or draft_ids.numel() < rows * d_step - or kv_off.numel() < rows - ): - return False - widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) - widths += (src_draft.shape[1], dst_draft.shape[1], t_step) - if cpp and max(widths) <= 32: - torch.ops.trtllm.minimax_m3_sm107_tp2_fold_tail( - ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, - src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, - ) # fmt: skip - return True - block = triton.next_power_of_2(max(widths)) - _fold_tail_kernel[(rows,)]( - ctrl, - slots, - src_new, - src_new.stride(0), - src_new.stride(1), - src_new.shape[1], - dst_new, - dst_new.stride(0), - dst_new.stride(1), - dst_new.shape[1], - src_next, - src_next.stride(0), - src_next.stride(1), - src_next.shape[1], - dst_next, - dst_next.stride(0), - dst_next.stride(1), - dst_next.shape[1], - src_draft, - src_draft.stride(0), - src_draft.stride(1), - src_draft.shape[1], - dst_draft, - dst_draft.stride(0), - dst_draft.stride(1), - dst_draft.shape[1], - src_lens, - dst_lens, - ids, - draft_ids, - pos_off, - kv_off, - t_step, - d_step, - block, - num_warps=1, + return _shared_fold_tail( + ctrl=ctrl, + slots=slots, + rows=rows, + src_new=src_new, + dst_new=dst_new, + src_next=src_next, + dst_next=dst_next, + src_draft=src_draft, + dst_draft=dst_draft, + src_lens=src_lens, + dst_lens=dst_lens, + ids=ids, + draft_ids=draft_ids, + pos_off=pos_off, + kv_off=kv_off, + tokens_per_step=tokens_per_step, + cpp=cpp, + native_launch=_launch_cpp, + triton_kernel=_fold_tail_kernel, ) - return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py index b6705bdb46b7..542422319210 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_fullctx_decode.py @@ -26,6 +26,18 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + FullCtxGeometry, + FullCtxPlanBase, + check_layout, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + device_index as _index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + kv_page_views as kv_page_views, +) + HEADS = 32 KV_HEADS = 2 # capacity: fixed: the tp2 KV heads per rank HEAD_DIM = 128 # capacity: fixed: the checkpoint's head dim @@ -41,116 +53,10 @@ _lock = threading.Lock() -class FullCtxPlan: +class FullCtxPlan(FullCtxPlanBase): """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" - def __init__( - self, - device: torch.device, - qlen: int, - nreq: int, - max_kv_len: int, - recipe: Mapping, - own_reducer: bool, - pdl: bool, - ): - from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec - from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( - get_max_active_clusters_for_cluster_size, - make_decode_config, - ) - - self._tsdec = tsdec - self.qlen = qlen - self.nreq = nreq - self.pdl = pdl - self.recipe = dict(recipe) - splits = int(self.recipe["splits"]) - fp8 = tsdec._cutlass_dtype("float8_e4m3fn") - # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer - # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents - # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an - # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the - # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own - # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge - # (minimax_m3_sm107_tp2_fullctx_reduce) waits as its first instruction. - cfg = make_decode_config( - headdim=HEAD_DIM, - args={"use_pdl": pdl, **self.recipe.get("config", {})}, - seq_len_q=qlen, - seq_len_kv=min(max_kv_len, POLICY_KV_LEN), - batch_size=nreq, - num_heads_q=HEADS, - num_heads_kv=KV_HEADS, - q_dtype=fp8, - k_dtype=fp8, - v_dtype=fp8, - o_dtype=tsdec._cutlass_dtype("float16"), - qkv_layout="pagedKv", - num_tokens_per_page=PAGE, - storage_tokens_per_page=PAGE, - split_kv_mode="gmem_reduction_with_separate_kernel", - splits_kv=splits, - max_splits_kv=splits, - min_loop_iters_per_split=1, - mask_type="causal" if qlen > 1 else "dense", - auto_tuner=True, - split_kv=True, - ) - _check_layout(cfg, qlen, splits) - with torch.cuda.device(device): - max_active = get_max_active_clusters_for_cluster_size(1) - spec = tsdec._decode_launch_spec_from_config( - cfg, - batch_size=nreq, - num_qo_heads=HEADS, - num_kv_heads=KV_HEADS, - head_dim=HEAD_DIM, - seq_len_q=qlen, - max_active_clusters=int(max_active), - ) - compile_spec = tsdec._make_decode_compile_spec( - spec, - device_index=_index(device), - num_qo_heads=HEADS, - num_kv_heads=KV_HEADS, - head_dim=HEAD_DIM, - page_size=PAGE, - max_kv_len=max_kv_len, - seq_len_q=qlen, - q_dtype_key="float8_e4m3fn", - k_dtype_key="float8_e4m3fn", - v_dtype_key="float8_e4m3fn", - output_dtype_key="float16", - use_packed_q=False, - kv_prefix_mode="dynamic", - kv_lengths_mode="dynamic", - ) - self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) - layout = tsdec._make_decode_workspace_layout( - spec.scratch_shapes, - torch.float16, - use_separate_reduction_kernel=True, - use_split_kv=True, - ) - self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) - self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) - self.workspace.split_kv_counter.zero_() - self.workspace.cu_seqlens_q.zero_() - self.workspace.attention_sinks.zero_() - rows = HEADS_PER_KV * qlen - if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): - raise RuntimeError( - f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " - f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" - ) - self.own_reducer = own_reducer - self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) - self.tile_kv = int(cfg.tile_size_kv) - self.insts_kv = int(cfg.num_insts_kv) - self.causal = qlen > 1 - shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) - self.staging = torch.empty(shape, dtype=torch.float16, device=device) + GEOMETRY = FullCtxGeometry(HEADS, KV_HEADS, HEAD_DIM, PAGE, POLICY_KV_LEN, HEADS_PER_KV) def run( self, @@ -229,32 +135,7 @@ def run( def _check_layout(cfg, qlen: int, splits: int) -> None: - """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. - - Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is - ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding - window, attention sinks or packed Q. - """ - problems = [] - if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): - problems.append("not split-KV with the separate GMEM reducer") - if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: - problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") - if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: - problems.append("cluster-SMEM reduction or a persistent scheduler") - if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: - problems.append("packed Q, a sliding window or attention sinks") - if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): - problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") - if qlen > 1: - if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): - problems.append("multi-token Q without the grouped [token, head] tile layout") - elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): - problems.append( - f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" - ) - if problems: - raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + return check_layout(cfg, qlen, splits, heads_per_kv=HEADS_PER_KV) def has_own_reducer() -> bool: @@ -286,12 +167,3 @@ def prebuild( _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) built += 1 return built - - -def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: - """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" - return pool[0::subpages_per_slot], pool[1::subpages_per_slot] - - -def _index(device: torch.device) -> int: - return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py index 8088d79f014f..db05da4f2be4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_index_score.py @@ -23,22 +23,18 @@ and called through this target's ``m3_index_score`` entry. """ -import cutlass -import torch -from cuda.bindings.driver import CUstream -from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute -from cutlass._mlir.dialects import llvm -from cutlass.cute.nvgpu import cpasync, warp -from cutlass.cutlass_dsl import dsl_user_op -from cutlass.utils.smem_allocator import SmemAllocator - -from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( - EVICT_FIRST, - fp8x4_to_fp16x4, - mma_sync, - simple_tma_copy, +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.index_score_runtime import ( + IndexScoreRunner as SharedIndexScoreRunner, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + IndexDecodeScoreKernel as IndexDecodeScoreKernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + _fp8_to_f16_mma_fragments as _fp8_to_f16_mma_fragments, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + _l2_prefetch_bulk as _l2_prefetch_bulk, ) -from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, griddepcontrol_wait # With the early issue, a CTA whose first block holds only tokens older than this step also TMA-loads that block's # index-K tile before the grid wait (bit-identical scores; replay period -47 us per step on Rubin TP2). @@ -47,624 +43,9 @@ __all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] -@dsl_user_op -def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: - """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" - llvm.inline_asm( - res=None, - operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], - asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", - constraints="l,r", - has_side_effects=True, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@cute.jit -def _fp8_to_f16_mma_fragments(src: cute.Tensor): - """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" - src_elems = cute.size(src) - src_u32 = cute.recast_tensor(src, Uint32) - src_f16 = cute.make_rmem_tensor(src_elems, Float16) - src_f16_u32 = cute.recast_tensor(src_f16, Uint32) - # Packed conversion; faster and fewer SASS instructions than - # src.load().to(Float16). - for i in cutlass.range_constexpr(src_elems // 4): - converted = fp8x4_to_fp16x4(src_u32[i]) - src_f16_u32[i * 2] = converted[0] - src_f16_u32[i * 2 + 1] = converted[1] - lower = cute.make_rmem_tensor(src_elems // 2, Float16) - upper = cute.make_rmem_tensor(src_elems // 2, Float16) - - # FP8 ldmatrix yields four consecutive values along K. Split each group - # into the lower two and upper two for the two FP16 MMA k-fragments. - for i in cutlass.range_constexpr(src_elems // 2): - lower[i] = src_f16[(i // 2) * 4 + i % 2] - upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] - return lower, upper - - -class IndexDecodeScoreKernel: - """Per-block max index score for one decode step. - - Grid is (batch, split_k): CTA (b, s) walks KV blocks - s, s + split_k, s + 2 * split_k, ... of request b and writes each block's - score directly, so no cross-CTA reduction is needed. CTAs whose split_id - exceeds the request's block count exit immediately. - """ - - BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint - BAR_MMA = 1 - num_stages = 2 - - def __init__( - self, - dtype: type[cutlass.Numeric], - num_heads: int, - max_decode_query_len: int, - split_k: int, - head_dim: int = 128, - prefetch_old: bool = False, - early_issue: bool = False, - prefetch_depth: int = 0, - prestage_old: bool = False, - ): - # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that - # hold only tokens older than this step (the predecessor writes only this step's tokens). - self.prefetch_old = prefetch_old - # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency - # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. - # The tiles are read after the wait, so their contents are final; the block table is only written between - # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles - # turn out to lie past the request's true length drains them before exiting. - self.early_issue = early_issue - # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at - # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. - self.prefetch_depth = prefetch_depth - # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K - # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's - # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. - self.prestage_old = prestage_old and early_issue - self.dtype = dtype - self.num_heads = num_heads - self.max_decode_query_len = max_decode_query_len - self.split_k = split_k - self.head_dim = head_dim - - def _swizzle_elems(self) -> int: - """Elements spanned by one 128-byte swizzle atom.""" - return 128 * 8 // self.dtype.width - - def _sq_layout(self): - """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" - elems = self._swizzle_elems() - head_dim = self.head_dim - block_q = self.num_heads * self.max_decode_query_len - layout = cute.make_layout( - (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), - stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), - ) - return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) - - def _sk_layout(self): - """Composed SMEM layout for the pipelined K tiles.""" - elems = self._swizzle_elems() - head_dim = self.head_dim - block_k = self.BLOCK_K - layout = cute.make_layout( - (1, block_k, (elems, head_dim // elems), self.num_stages), - stride=(0, elems, (1, block_k * elems), block_k * head_dim), - ) - return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) - - @cute.jit - def __call__( - self, - gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] - gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] - block_table: cute.Tensor, # [bs, max_pages] - score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] - seq_lens: cute.Tensor, # [bs] - stream: CUstream, - ): - num_heads = self.num_heads - head_dim = self.head_dim - MAX_DQL = self.max_decode_query_len - assert num_heads * MAX_DQL <= 32 - - batch = seq_lens.shape[0] - decode_query_len = gQ.shape[0] // batch - grid = (batch, self.split_k, 1) - block = (32 * 5, 1, 1) - - tma_g2s = cpasync.CopyBulkTensorTileG2SOp() - elems = self._swizzle_elems() - - q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( - tma_g2s, - cute.logical_divide(gQ, (None, None, elems)), - self._sq_layout(), - (MAX_DQL, num_heads, head_dim), - ) - k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( - tma_g2s, - cute.logical_divide(gK_cache, (None, None, elems)), - self._sk_layout(), - (1, self.BLOCK_K, head_dim), - ) - - self.kernel( - q_tma_atom, - q_tma_tensor, - k_tma_atom, - k_tma_tensor, - gK_cache, - block_table, - score, - seq_lens, - decode_query_len, - ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) - - @cute.kernel - def kernel( - self, - q_tma_atom: cute.CopyAtom, - q_tma_tensor: cute.Tensor, - k_tma_atom: cute.CopyAtom, - k_tma_tensor: cute.Tensor, - gK_cache: cute.Tensor, - block_table: cute.Tensor, - score: cute.Tensor, - seq_lens: cute.Tensor, - decode_query_len, - ): - batch_id, split_id, _ = cute.arch.block_idx() - _, split_k, _ = cute.arch.grid_dim() - warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane_id = cute.arch.lane_idx() - - NUM_HEADS = self.num_heads - MAX_DQL = self.max_decode_query_len - BLOCK_Q = NUM_HEADS * MAX_DQL - BLOCK_K = self.BLOCK_K - head_dim = self.head_dim - dtype = self.dtype - MMA_N = 8 - num_stages = self.num_stages - Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) - EPI_Q = Q_TILES * MMA_N - - sq_layout = self._sq_layout() - sk_layout = self._sk_layout() - - smem = SmemAllocator() - sK = smem.allocate_tensor( - dtype, - sk_layout.outer, - byte_alignment=128, - swizzle=sk_layout.inner, - )[0, None, None, None] - # sQ aliases the first K stage: Q is consumed into registers before the - # first K tile is needed, so the two never overlap in time. - sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) - # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q - # column mode. The target profile keeps the rank-2 view even for - # degenerate shapes such as MAX_DQL == 1. - q_tma_elems = self._swizzle_elems() - sQ = cute.coalesce( - cute.group_modes(sQ_tma, 0, 2), - target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), - ) - epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) - - tma_full_mbar = smem.allocate_array(Int64, num_stages) - tma_empty_mbar = smem.allocate_array(Int64, num_stages) - - # Barrier init and descriptor prefetch do not depend on the predecessor - # grid, so they are the work PDL can overlap. Every CTA runs them, even - # the ones that own no blocks, since ownership is not known until - # seq_lens is read below. - if warp_id == 0: - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(num_stages): - cute.arch.mbarrier_init(tma_full_mbar + i, 1) - cute.arch.mbarrier_init(tma_empty_mbar + i, 128) - cute.arch.mbarrier_init_fence() - elif warp_id == 1: - cpasync.prefetch_descriptor(q_tma_atom) - cpasync.prefetch_descriptor(k_tma_atom) - cute.arch.sync_threads() - if cutlass.const_expr(self.prefetch_old): - if warp_id == 4: - with cute.arch.elect_one(): - seqlen0 = seq_lens[batch_id] - num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) - stable = (seqlen0 - decode_query_len) // BLOCK_K - if stable > num_blocks0: - stable = num_blocks0 - if cutlass.const_expr(self.prefetch_depth > 0): - if stable > split_id + self.prefetch_depth * split_k: - stable = split_id + self.prefetch_depth * split_k - for block_id in range(split_id, stable, split_k): - page = block_table[batch_id, block_id] - _l2_prefetch_bulk( - gK_cache[page, None, None].iterator.toint().to(Int64), - Int32(BLOCK_K * head_dim * (dtype.width // 8)), - ) - - # Otherwise the dependent block selector is released as the CTAs exit. - early_q = False - early_k = False - early_page = Int32(0) - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - seqlen_hint = seq_lens[batch_id] - hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) - if hint_blocks > block_table.shape[1]: - hint_blocks = block_table.shape[1] - if split_id < hint_blocks: - early_q = True - early_page = block_table[batch_id, split_id] - if cutlass.const_expr(self.prestage_old): - # Stage the first K tile now when every token of the block predates this step. - if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: - early_k = True - cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - tma_full_mbar + 1 % num_stages, - BLOCK_K * head_dim * (dtype.width // 8), - ) - simple_tma_copy( - k_tma_atom, - k_tma_tensor[early_page, None, None], - sK[None, None, 1 % num_stages], - tma_full_mbar + 1 % num_stages, - cache_policy=EVICT_FIRST, - ) - - griddepcontrol_wait() - - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - if early_q: - gQ_early = cute.local_tile( - cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), - tiler=(MAX_DQL, NUM_HEADS, head_dim), - coord=(0, 0, 0), - ) - cute.arch.mbarrier_wait(tma_empty_mbar, 1) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) - ) - simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) - if not early_k: - cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - tma_full_mbar + 1 % num_stages, - BLOCK_K * head_dim * (dtype.width // 8), - ) - simple_tma_copy( - k_tma_atom, - k_tma_tensor[early_page, None, None], - sK[None, None, 1 % num_stages], - tma_full_mbar + 1 % num_stages, - cache_policy=EVICT_FIRST, - ) - - # The length that bounds this grid's work is read after the wait (the - # pre-wait reads above are prefetch / early-issue hints only, clamped to - # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens - # patches every step, and num_blocks bounds an unmasked block_table load - # and an unmasked score store, so a stale length is an out-of-bounds write - # rather than just a wrong answer. - seqlen = seq_lens[batch_id] - num_blocks = cute.ceil_div(seqlen, BLOCK_K) - - if split_id < num_blocks: - if warp_id == 4: - # TMA warp - tma_stage = 0 - tma_parity = 1 - first_block = split_id - - if not early_q: - gQ_tile = cute.local_tile( - cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), - tiler=(MAX_DQL, NUM_HEADS, head_dim), - coord=(0, 0, 0), - ) - cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) - with cute.arch.elect_one(): - Q_size = BLOCK_Q * head_dim * (dtype.width // 8) - cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) - # TMA bounds-checks rows when the runtime decode_query_len is - # below MAX_DQL; padded Q columns are masked before the stores. - simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) - - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - if early_q: - # The first K tile was issued right after the wait. - first_block = split_id + split_k - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - for block_id in range(first_block, num_blocks, split_k): - page_id = block_table[batch_id, block_id] - gK_tile = k_tma_tensor[page_id, None, None] - k_mbar = tma_full_mbar + tma_stage - - cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) - with cute.arch.elect_one(): - K_size = BLOCK_K * head_dim * (dtype.width // 8) - cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) - simple_tma_copy( - k_tma_atom, - gK_tile, - sK[None, None, tma_stage], - k_mbar, - cache_policy=EVICT_FIRST, - ) - - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - else: - # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T - sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) - q_start = seqlen - decode_query_len - - elems = 128 // dtype.width # 16B - MMA_K = 32 * 8 // dtype.width # 32B - - # Pre-compute ldmatrix addresses. - # sK loads a [16 x 16B] tile: - # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) - # sQ loads an [8 x 32B] tile: - # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) - sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) - sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) - - # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) - # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) - sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] - sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] - - ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) - ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) - - rQ = cute.make_rmem_tensor( - ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype - ) - rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) - rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) - - if warp_id == 0: - cute.arch.mbarrier_wait(tma_full_mbar, 0) - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - for q in cutlass.range_constexpr(Q_TILES): - cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) - cute.arch.mbarrier_arrive(tma_empty_mbar) - - tma_stage = 1 % self.num_stages - tma_parity = 0 - if tma_stage == 0: - tma_parity ^= 1 - - if cutlass.const_expr(dtype is Float8E4M3FN): - rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) - q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) - rQ_f16[None, None, None, 0].store(q_lower.load()) - rQ_f16[None, None, None, 1].store(q_upper.load()) - - for block_id in range(split_id, num_blocks, split_k): - rC.fill(0.0) - - if warp_id == 0: - cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - - for k in cutlass.range_constexpr(head_dim // MMA_K): - cute.copy( - ldsm_atom, - sK_ldsm[None, (None, k, tma_stage)], - rK[None, None, k], - ) - for m in cutlass.range_constexpr(2): - if cutlass.const_expr(dtype is Float8E4M3FN): - rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) - for n in cutlass.range_constexpr(Q_TILES): - rC[None, m, n] = mma_sync( - rK_lower, - rQ_f16[None, k, n, 0], - rC[None, m, n], - ) - rC[None, m, n] = mma_sync( - rK_upper, - rQ_f16[None, k, n, 1], - rC[None, m, n], - ) - else: - for n in cutlass.range_constexpr(Q_TILES): - rC[None, m, n] = mma_sync( - rK[None, m, k], - rQ[(None, k % 2), k // 2, n], - rC[None, m, n], - ) - - cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) - - k_start = block_id * BLOCK_K + warp_id * 32 - - for q in cutlass.range_constexpr(Q_TILES): - for i in cutlass.range_constexpr(4): - for j in cutlass.range_constexpr(2): - col = q * 8 + (lane_id % 4) * 2 + j - q_local_pos = col % MAX_DQL - q_pos = q_start + q_local_pos - k_pos = k_start + i * 8 + lane_id // 4 - rC[q * 8 + i * 2 + j] = ( - rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") - ) - - for q in cutlass.range_constexpr(Q_TILES): - # Thread-local reduction along the BLOCK_K dim. - rScore = cute.make_rmem_tensor(2, Float32) - rScore.fill(float("-inf")) - for i in cutlass.range_constexpr(4): - rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) - rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) - - # Warp reduction among lanes 0, 4, 8, 12, ... - for i in cutlass.range_constexpr(3): - offset = 4 << i - other0 = cute.arch.shuffle_sync_bfly( - rScore[0], offset=offset, mask=-1, mask_and_clamp=31 - ) - other1 = cute.arch.shuffle_sync_bfly( - rScore[1], offset=offset, mask=-1, mask_and_clamp=31 - ) - rScore[0] = cute.arch.fmax(rScore[0], other0) - rScore[1] = cute.arch.fmax(rScore[1], other1) - - if lane_id * 2 < MMA_N: - epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] - epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - - head_id = lane_id // MAX_DQL - q_local_pos = lane_id - head_id * MAX_DQL - valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len - if lane_id < BLOCK_Q and valid_q: - final_score = epi_buffer[lane_id, 0] - for i in cutlass.range_constexpr(1, 4): - final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) - - t = batch_id * decode_query_len + q_local_pos - score[head_id, t, block_id] = final_score - - tma_stage = (tma_stage + 1) % self.num_stages - if tma_stage == 0: - tma_parity ^= 1 - else: - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - if early_q: - # Early tiles past the request's true length: let them land before the CTA exits. - cute.arch.mbarrier_wait(tma_full_mbar, 0) - cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) - - -class IndexScoreRunner: - """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, - split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile - serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block - per CTA.""" +class IndexScoreRunner(SharedIndexScoreRunner): + """TP2 schedule and private compilation cache.""" kernel_cache: dict = {} - SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim - SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K - # BLOCK_Q must fit one warp's worth of epilogue lanes. - MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile - EARLY_PREFETCH_DEPTH = 1 - - @classmethod - def is_supported( - cls, - *, - q_dtype: torch.dtype, - num_heads: int, - head_dim: int, - page_size: int, - max_decode_query_len: int, - ) -> bool: - from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE - from tensorrt_llm._utils import is_sm_100f - - return ( - is_sm_100f() - and q_dtype in TORCH_TO_CUTE_DTYPE - and head_dim == cls.SUPPORTED_HEAD_DIM - and page_size == cls.SUPPORTED_PAGE_SIZE - and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q - and max_decode_query_len >= 1 - ) - - @classmethod - def _compile( - cls, - q_dtype: torch.dtype, - num_heads: int, - max_decode_query_len: int, - head_dim: int, - split_k: int, - early: bool, - ) -> None: - from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE - - key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) - if key in cls.kernel_cache: - return - cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] - page_size = cls.SUPPORTED_PAGE_SIZE - sym_total_tokens = cute.sym_int() - sym_batch = cute.sym_int() - - # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. - def _sym_stride(): - return cute.sym_int64(divisibility=16) - - q_fake = cute.runtime.make_fake_tensor( - cute_dtype, - (sym_total_tokens, num_heads, head_dim), - stride=(_sym_stride(), _sym_stride(), 1), - ) - # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). - k_fake = cute.runtime.make_fake_tensor( - cute_dtype, - (cute.sym_int(), page_size, head_dim), - stride=(_sym_stride(), _sym_stride(), 1), - ) - bt_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) - ) - # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. - score_fake = cute.runtime.make_fake_tensor( - cutlass.Float32, - (num_heads, sym_total_tokens, cute.sym_int()), - stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), - ) - sl_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (sym_batch,), stride_order=(0,) - ) - fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - kernel = IndexDecodeScoreKernel( - cute_dtype, - num_heads, - max_decode_query_len, - split_k, - head_dim, - prefetch_old=True, - early_issue=early, - prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, - prestage_old=early and INDEX_SCORE_PRESTAGE_OLD, - ) - cls.kernel_cache[key] = cute.compile( - kernel, - q_fake, - k_fake, - bt_fake, - score_fake, - sl_fake, - fake_stream, - options="--enable-tvm-ffi", - ) + PRESTAGE_OLD = INDEX_SCORE_PRESTAGE_OLD + EARLY_TRIGGER = False diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py index 3e88c7fb1ce2..6af262dafe82 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_locality.py @@ -1,183 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Locality domains ("uGPU") of a multi-die GPU (Rubin VR200: 2 domains of 100 SMs + 12 remainder SMs) for the -streaming decode kernels of this target: the SM -> domain map, per-domain weight placement, and the CTA -> SM -assignment a persistent kernel uses to read only its own die's memory. - -A kernel launched with one CTA per SM reads ``%smid`` once and takes its logical block index from a host-built -permutation (:meth:`LocalityMap.logical_map`), so every per-CTA structure (work plan row, epochs, roles) follows its -die. Weights that are split by output rows into per-domain halves (:meth:`LocalityMap.split_rows`) live in one virtual -range whose halves are physically in the two domains. Placement only: the arithmetic is unchanged, so results are -bitwise those of the unplaced weights. - -The SM -> domain map is attested by the driver: the SMs are split by locality domain -(``cuDevSmResourceSplit`` with ``CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID``), and a ``%smid`` probe runs on a green -context of each split. Allocations use driver VMM with ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN`` and are checked -with ``CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL``; any failure raises (fail closed).""" - -from __future__ import annotations +"""Target-owned locality-map cache with shared placement implementation.""" import threading -from typing import Dict, List, Sequence, Tuple +from typing import Dict import torch +from ...ops.gemm.locality import LocalityMap as _LocalityMap + _LOCK = threading.Lock() _MAPS: Dict[int, "LocalityMap"] = {} -_PROBE_SRC = r""" -extern "C" __global__ void smid_probe(int* out) { - if (threadIdx.x == 0) { - unsigned s; - asm volatile("mov.u32 %0, %%smid;" : "=r"(s)); - out[blockIdx.x] = (int)s; - } -} -""" - - -def _drv(): - from cuda.bindings import driver - - return driver - - -def _ck(res, what: str): - """cuda-python returns (CUresult, *values): raise on an error, return the values.""" - err = res[0] if isinstance(res, tuple) else res - if int(err) != 0: - raise RuntimeError(f"locality: {what} failed: {err}") - if isinstance(res, tuple): - return res[1] if len(res) == 2 else res[1:] - return None - - -class _Mapping: - """One virtual range backed by per-domain physical allocations, exposed to torch through - ``__cuda_array_interface__``. Lives as long as a tensor that views it (the tensor holds this object).""" - - def __init__(self, device: int, pieces: Sequence[Tuple[int, int]]): - """``pieces``: (bytes, domain) in address order; each a multiple of the allocation granularity.""" - d = _drv() - self.device = device - self.size = sum(n for n, _ in pieces) - self.va = _ck(d.cuMemAddressReserve(self.size, 0, 0, 0), "cuMemAddressReserve") - self.handles = [] - off = 0 - try: - for nbytes, dom in pieces: - prop = _alloc_prop(device, dom) - h = _ck(d.cuMemCreate(nbytes, prop, 0), f"cuMemCreate(domain {dom})") - self.handles.append((h, nbytes)) - _ck(d.cuMemMap(int(self.va) + off, nbytes, 0, h, 0), "cuMemMap") - off += nbytes - acc = d.CUmemAccessDesc() - acc.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - acc.location.id = device - acc.flags = d.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE - _ck(d.cuMemSetAccess(self.va, self.size, [acc], 1), "cuMemSetAccess") - except Exception: - self.close() - raise - self.__cuda_array_interface__ = { - "shape": (self.size,), - "typestr": "|u1", - "data": (int(self.va), False), - "version": 3, - "strides": None, - "stream": None, - } - - def close(self): - d = _drv() - off = 0 - for h, n in self.handles: - d.cuMemUnmap(int(self.va) + off, n) - d.cuMemRelease(h) - off += n - self.handles = [] - if self.va is not None: - d.cuMemAddressFree(self.va, self.size) - self.va = None - - def __del__(self): - try: - self.close() - except Exception: - pass - - -_LIBCUDA = None - - -def _libcuda(): - global _LIBCUDA - if _LIBCUDA is None: - import ctypes - _LIBCUDA = ctypes.CDLL("libcuda.so.1") - return _LIBCUDA - - -def _alloc_prop(device: int, domain: int): - d = _drv() - prop = d.CUmemAllocationProp() - prop.type = d.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED - prop.requestedHandleTypes = d.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE - prop.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN - prop.location.localized.deviceId = device - prop.location.localized.localityDomainId = domain - return prop - - -class LocalityMap: - """The locality domains of one device. Build with :meth:`get` before CUDA-graph capture.""" - - def __init__(self, device: int): - d = _drv() - _ck(d.cuInit(0), "cuInit") - dev = _ck(d.cuDeviceGet(device), "cuDeviceGet") - attr = d.CUdevice_attribute - self.device = device - self.n_domains = int( - _ck( - d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, dev), - "domain count", - ) - ) - self.sms_per_domain = int( - _ck( - d.cuDeviceGetAttribute( - attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_MULTIPROCESSOR_COUNT, dev - ), - "SMs per domain", - ) - ) - self.num_sms = int( - _ck(d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev), "SMs") - ) - if self.n_domains != 2: - raise RuntimeError( - f"locality: {self.n_domains} locality domains (this layer handles 2)" - ) - self._domain_sms = self._probe(d, dev) - dom = [-1] * self.num_sms - for k, sms in enumerate(self._domain_sms): - for s in sms: - if dom[s] != -1: - raise RuntimeError(f"locality: SM {s} probed in two domains") - dom[s] = k - self.sm_domain: List[int] = dom - self.remainder: List[int] = [s for s in range(self.num_sms) if dom[s] < 0] - rt = _alloc_prop(device, 0) - self.granularity = int( - _ck( - d.cuMemGetAllocationGranularity( - rt, d.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM - ), - "granularity", - ) - ) +class LocalityMap(_LocalityMap): @classmethod def get(cls, device) -> "LocalityMap": idx = torch.device(device).index if not isinstance(device, int) else device @@ -191,258 +27,3 @@ def get(cls, device) -> "LocalityMap": ) m = _MAPS[idx] = cls(idx) return m - - def sms(self, domain: int) -> List[int]: - return list(self._domain_sms[domain]) - - @property - def sm_table(self) -> torch.Tensor: - """int32 [num_sms] on the device: (domain << 16) | rank within the domain's sorted SMs; -1 = remainder SM. - For kernels that pick work by %smid themselves.""" - t = getattr(self, "_sm_table", None) - if t is None: - v = [-1] * self.num_sms - for k, sms in enumerate(self._domain_sms): - for i, s in enumerate(sms): - v[s] = (k << 16) | i - t = self._sm_table = torch.tensor(v, dtype=torch.int32, device=f"cuda:{self.device}") - return t - - def _probe(self, d, dev) -> List[List[int]]: - """Split the SMs by locality domain (strict, no backfill) and record %smid of a probe kernel launched on a - green context of each part: the driver's own SM -> domain assignment.""" - from cuda.bindings import nvrtc - - full = _ck( - d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), - "dev resource", - ) - params = [] - for k in range(self.n_domains): - p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() - p.smCount = self.sms_per_domain - p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID - p.localityDomainId = k - params.append(p) - res, _rem = _ck( - d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" - ) - # the probe module, compiled for this device - major = _ck( - d.cuDeviceGetAttribute( - d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev - ), - "cc", - ) - minor = _ck( - d.cuDeviceGetAttribute( - d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev - ), - "cc", - ) - prog = _ck( - nvrtc.nvrtcCreateProgram(_PROBE_SRC.encode(), b"smid_probe.cu", 0, [], []), - "nvrtcCreateProgram", - ) - opts = [f"--gpu-architecture=sm_{major}{minor}".encode()] - err = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)[0] - if int(err) != 0: - raise RuntimeError(f"locality: probe compile failed: {err}") - n = _ck(nvrtc.nvrtcGetCUBINSize(prog), "cubin size") - cubin = b" " * n - _ck(nvrtc.nvrtcGetCUBIN(prog, cubin), "cubin") - out: List[List[int]] = [] - nb = 4 * self.sms_per_domain # more blocks than SMs: every SM of the part shows up - torch.cuda.set_device(self.device) - buf = torch.full((nb,), -1, dtype=torch.int32, device=f"cuda:{self.device}") - for k in range(self.n_domains): - desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") - g = _ck( - d.cuGreenCtxCreate(desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM), - "green ctx", - ) - st = _ck( - d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), - "green stream", - ) - ctx = _ck(d.cuCtxFromGreenCtx(g), "ctx from green ctx") - _ck(d.cuCtxPushCurrent(ctx), "push ctx") - try: - mod = _ck(d.cuModuleLoadData(cubin), "probe module") - fn = _ck(d.cuModuleGetFunction(mod, b"smid_probe"), "probe fn") - buf.fill_(-1) - torch.cuda.synchronize(self.device) - import ctypes - - arg = ctypes.c_void_p(buf.data_ptr()) - _ck( - d.cuLaunchKernel( - fn, nb, 1, 1, 32, 1, 1, 0, st, ((arg,), (ctypes.c_void_p,)), 0 - ), - "probe launch", - ) - _ck(d.cuStreamSynchronize(st), "probe sync") - sms = sorted(set(int(x) for x in buf.tolist() if x >= 0)) - d.cuModuleUnload(mod) - finally: - d.cuCtxPopCurrent() - d.cuStreamDestroy(st) - d.cuGreenCtxDestroy(g) - if len(sms) != self.sms_per_domain: - raise RuntimeError( - f"locality: domain {k} probed {len(sms)} SMs, expected {self.sms_per_domain}" - ) - out.append(sms) - return out - - def logical_map(self, per_domain: int, extra: int) -> Tuple[torch.Tensor, List[int]]: - """The SM -> logical block index permutation for a kernel of ``2 * per_domain + extra`` CTAs, one per SM: - logical blocks [d * per_domain, (d + 1) * per_domain) run on domain d's first ``per_domain`` SMs; logical - blocks 2 * per_domain .. run on the rest (the leftover domain SMs, alternating domains, then the remainder - SMs). Returns (int32 [num_sms] on the device, -1 = no block; the domain of each extra block, -1 remainder).""" - if 2 * per_domain + extra != self.num_sms: - raise ValueError( - f"locality: {2 * per_domain} + {extra} CTAs must cover the {self.num_sms} SMs" - ) - if per_domain > self.sms_per_domain: - raise ValueError(f"locality: {per_domain} CTAs per domain > {self.sms_per_domain} SMs") - lmap = [-1] * self.num_sms - for k in range(2): - for i, s in enumerate(self._domain_sms[k][:per_domain]): - lmap[s] = k * per_domain + i - left = [self._domain_sms[k][per_domain:] for k in range(2)] - rest = [x for pair in zip(*left) for x in pair] + self.remainder - ext_dom = [] - for i, s in enumerate(rest): - lmap[s] = 2 * per_domain + i - ext_dom.append(self.sm_domain[s]) - assert sorted(lmap) == list(range(self.num_sms)) - return torch.tensor(lmap, dtype=torch.int32, device=f"cuda:{self.device}"), ext_dom - - def domain_streams(self) -> List[torch.cuda.ExternalStream]: - """One non-blocking stream per domain on a green context holding exactly that domain's SMs (kept alive with - this map): work launched on stream d runs only on domain d's SMs. Fork / join with events from the caller's - stream; graph capture follows the forks. Built once, before capture.""" - st = getattr(self, "_streams", None) - if st is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("locality.domain_streams: build before CUDA-graph capture") - d = _drv() - dev = _ck(d.cuDeviceGet(self.device), "cuDeviceGet") - full = _ck( - d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), - "dev resource", - ) - params = [] - for k in range(self.n_domains): - p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() - p.smCount = self.sms_per_domain - p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID - p.localityDomainId = k - params.append(p) - res, _rem = _ck( - d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" - ) - self._green = [] - st = [] - for k in range(self.n_domains): - desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") - g = _ck( - d.cuGreenCtxCreate( - desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM - ), - "green", - ) - s_ = _ck( - d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), - "green stream", - ) - self._green.append((g, s_)) - st.append(torch.cuda.ExternalStream(int(s_), device=f"cuda:{self.device}")) - self._streams = st - return st - - def domain_of(self, ptr: int) -> int: - """The locality domain ordinal of a device address (-1: not placed in one domain). Through libcuda - directly: cuda-python's attribute helper does not know this attribute.""" - import ctypes - - d = _drv() - lib = _libcuda() - out = ctypes.c_int(-9) - attr = int(d.CUpointer_attribute.CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL) - err = lib.cuPointerGetAttribute(ctypes.byref(out), ctypes.c_int(attr), ctypes.c_uint64(ptr)) - if err != 0: - raise RuntimeError( - f"locality: cuPointerGetAttribute(LOCALITY_DOMAIN_ORDINAL) failed: {err}" - ) - return int(out.value) - - def place_split(self, t: torch.Tensor, split_row: int) -> torch.Tensor: - """Split-mapped placement: a copy of ``t`` (contiguous, rows along dim 0) with the SAME shape and layout, in one - virtual range whose rows [0, split_row) are physically in domain 0 and rows [split_row, N) in domain 1. The - split must fall on an allocation granule (rows [0, split_row) a multiple of 2 MiB on VR200); the tail is padded - to a granule in domain 1. Stock consumers read it as a plain tensor; no extra memory once ``t`` is dropped.""" - if not t.is_contiguous(): - raise ValueError("locality.place_split: needs a contiguous tensor") - row_bytes = t[0].numel() * t.element_size() if t.dim() > 1 else t.element_size() - head = split_row * row_bytes - total = t.numel() * t.element_size() - if not 0 < head < total or head % self.granularity: - raise ValueError( - f"locality.place_split: rows [0, {split_row}) = {head} B must be a granule multiple " - f"({self.granularity} B) inside the {total} B tensor" - ) - tail = -(-(total - head) // self.granularity) * self.granularity - m = _Mapping(self.device, [(head, 0), (tail, 1)]) - flat = torch.as_tensor(m, device=f"cuda:{self.device}") - out = flat[:total].view(t.dtype).view(t.shape) - out.copy_(t) - out._locality_mapping = m - for q, k in ( - (out.data_ptr(), 0), - (out.data_ptr() + head - 1, 0), - (out.data_ptr() + head, 1), - (out.data_ptr() + total - 1, 1), - ): - if self.domain_of(q) != k: - raise RuntimeError(f"locality.place_split: byte {q:#x} is not in domain {k}") - return out - - def replace_param_(self, param: torch.nn.Parameter, split_row: int) -> torch.nn.Parameter: - """Re-place a parameter in place with :meth:`place_split` (swaps ``param.data``; the old block returns to the - caching allocator, so call ``torch.cuda.empty_cache()`` between large re-placements).""" - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("locality.replace_param_: not during CUDA-graph capture") - param.data = self.place_split(param.data, split_row) - return param - - def split_rows(self, t: torch.Tensor, dim: int = 1) -> torch.Tensor: - """Place ``t`` (contiguous, 1-byte elements, rows along ``dim`` split into two equal halves) as - [2, *t.shape[:dim], t.shape[dim] // 2, *t.shape[dim + 1:]]: half k (rows [k h, (k + 1) h) of every leading - index) in domain k's memory, one virtual range. Each half's byte size must be a multiple of the allocation - granularity (2 MiB on VR200). The copy is exact; the caller drops ``t`` afterwards.""" - if t.element_size() != 1 or not t.is_contiguous(): - raise ValueError("locality.split_rows: needs a contiguous 1-byte tensor") - n = t.shape[dim] - if n % 2: - raise ValueError(f"locality.split_rows: {n} rows don't split in halves") - half_bytes = t.numel() // 2 - if half_bytes % self.granularity: - raise ValueError( - f"locality.split_rows: a half is {half_bytes} B, not a multiple of the {self.granularity} B granularity" - ) - m = _Mapping(self.device, [(half_bytes, 0), (half_bytes, 1)]) - flat = torch.as_tensor(m, device=f"cuda:{self.device}") - shape = (2,) + tuple(t.shape[:dim]) + (n // 2,) + tuple(t.shape[dim + 1 :]) - out = flat.view(torch.uint8).view(shape) - src = t.view(torch.uint8) - for k in range(2): - out[k].copy_(src.narrow(dim, k * (n // 2), n // 2)) - out._locality_mapping = m # the tensor keeps its mapping alive - for k in range(2): - p = out[k].data_ptr() - for q in (p, p + half_bytes - 1): - if self.domain_of(q) != k: - raise RuntimeError(f"locality.split_rows: byte {q:#x} is not in domain {k}") - return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py deleted file mode 100644 index 1cace7ab549a..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_lowm_splitk.py +++ /dev/null @@ -1,1255 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. - -This target's copy of the installed FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the -Eagle3 drafter's low-M linears, modeling's ``DRAFT_LOWM_OWN_SPLITK``): the same kernel and the same split order at every -tactic, bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's -final cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, -the SwiGLU prologue that follow. - -Each cluster rank accumulates an exact K slice in FP32. Peers publish partials -to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public -``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below -use kernel coordinates: kernel-M carries public N and kernel-N carries public M. -""" - -from __future__ import annotations - -import dataclasses -import functools - -import cuda.bindings.driver as _cuda -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.utils.blackwell_helpers as sm100_utils -from cutlass import Int32 -from cutlass._mlir.dialects import llvm -from cutlass.cute import experimental as cute_ext -from cutlass.cute.nvgpu import tcgen05 -from cutlass.cute.runtime import from_dlpack -from cutlass.cutlass_dsl import T, dsl_user_op - -#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. -_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") - -#: K extent of one CTA tile. -_CTA_K = 128 # capacity: fixed: a kernel tile constant - -#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. -_SUPPORTED_MMA_M = (64, 128) - -#: Kernel-N carries public M, which is limited to 32. -_SUPPORTED_MMA_N = (8, 16, 32) - -#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. -_SUPPORTED_SPLIT_K = (1, 2, 3, 4) - -#: Largest public M this low-M policy serves. -_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound - -#: Bytes per FP32 partial exchanged through DSMEM. -_FP32_BYTES = 4 - -#: DSMEM mailbox base alignment, in bytes. -_MAILBOX_ALIGN_BYTES = 128 - -#: Size and alignment of one mbarrier, in bytes. -_MBARRIER_BYTES = 8 - -#: Size and alignment of the TMEM base pointer slot. -_TMEM_POINTER_BYTES = 4 - -#: Bytes per BF16/FP16 element. -_AB_ELEMENT_BYTES = 2 - -#: Alignment of the A/B shared-memory buffers. -_AB_BUFFER_ALIGN_BYTES = 1024 # capacity: fixed: an alignment - -#: A/B pipeline stage bounds. One stage is useful only for a single K tile. -_MIN_AB_STAGES = 1 # capacity: fixed: pipeline stages -_DEFAULT_MIN_AB_STAGES = 2 -_MAX_AB_STAGES = 12 # capacity: fixed: pipeline stages - - -@dataclasses.dataclass(frozen=True) -class SplitKTactic: - """One specialization; mma_m carries public N and mma_n carries public M.""" - - mma_m: int - mma_n: int - split_k: int - ab_stages: int - - -def _align_up(value: int, alignment: int) -> int: - return ((value + alignment - 1) // alignment) * alignment - - -def _smem_bytes( - tactic: SplitKTactic, - ab_stages: int, -) -> int: - """Mirror the device allocator's shared-memory layout.""" - cursor = ( - _align_up( - tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, - _AB_BUFFER_ALIGN_BYTES, - ) - + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages - ) - - cursor = _align_up(cursor, _MBARRIER_BYTES) - cursor += 2 * ab_stages * _MBARRIER_BYTES - cursor += 3 * _MBARRIER_BYTES - cursor = _align_up(cursor, _TMEM_POINTER_BYTES) - cursor += _TMEM_POINTER_BYTES - - if tactic.split_k == 1: - return cursor - - return ( - _align_up( - _align_up(cursor, _MAILBOX_ALIGN_BYTES) - + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, - _MBARRIER_BYTES, - ) - + _MBARRIER_BYTES - ) - - -def _max_ab_stages_for( - tactic: SplitKTactic, - smem_capacity: int, -) -> int: - return next( - ( - stages - for stages in range(_MAX_AB_STAGES, -1, -1) - if _smem_bytes(tactic, stages) <= smem_capacity - ), - 0, - ) - - -def validate_tactic( - tactic: SplitKTactic, - m: int, - n: int, - k: int, - *, - smem_capacity: int = _SMEM_CAPACITY_BYTES, -) -> None: - """Reject a tactic that cannot serve ``(m, n, k)``.""" - if tactic.mma_m not in _SUPPORTED_MMA_M: - raise ValueError(f"unsupported mma_m={tactic.mma_m}") - if tactic.mma_n not in _SUPPORTED_MMA_N: - raise ValueError(f"unsupported mma_n={tactic.mma_n}") - if tactic.split_k not in _SUPPORTED_SPLIT_K: - raise ValueError(f"unsupported split_k={tactic.split_k}") - if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: - raise ValueError( - f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" - ) - if not 1 <= m <= _MAX_M: - raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") - if n <= 0: - raise ValueError(f"N must be positive, got {n}") - if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: - raise ValueError( - f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" - ) - smem_bytes = _smem_bytes(tactic, tactic.ab_stages) - if smem_bytes > smem_capacity: - raise ValueError( - f"tactic {tactic} needs {smem_bytes} B of shared memory but only " - f"{smem_capacity} B are available; max ab_stages is " - f"{_max_ab_stages_for(tactic, smem_capacity)}" - ) - - -def autotune_tactics( - m: int, - n: int, - k: int, - *, - smem_capacity: int = _SMEM_CAPACITY_BYTES, -) -> list[SplitKTactic]: - """Return valid tactics in the shape-specific stage window.""" - tactics: list[SplitKTactic] = [] - for mma_m in _SUPPORTED_MMA_M: - for mma_n in _SUPPORTED_MMA_N: - for split_k in _SUPPORTED_SPLIT_K: - min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES - base = SplitKTactic(mma_m, mma_n, split_k, min_stages) - try: - validate_tactic(base, m, n, k, smem_capacity=smem_capacity) - except ValueError: - continue - max_stages = _max_ab_stages_for(base, smem_capacity) - # Short K favors shallow pipelines; long K stays near the cap. - tactics.extend( - dataclasses.replace(base, ab_stages=ab_stages) - for ab_stages in ( - range(min_stages, min(max_stages, 6) + 1) - if k <= 4 * _CTA_K - else range( - min(max(5, max_stages - 2), max_stages), - max_stages + 1, - ) - ) - ) - return tactics - - -def default_tactic(m: int, n: int, k: int) -> SplitKTactic: - """Choose the default occupancy-oriented tactic.""" - if n <= 512: - mma_m = 64 - mma_n = 16 if n == 512 and m > 24 else 8 - requested_split = 4 - else: - mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 - if n <= 3072: - mma_m = 128 if m <= 16 else 64 - requested_split = 4 if m <= 16 else 2 - elif n < 8192: - mma_m = 64 - requested_split = 2 - else: - mma_m = 128 if k <= 1024 and m <= 24 else 64 - requested_split = 1 - - if k <= 4 * _CTA_K: - requested_split = 1 - split_k = next( - split_k - for split_k in reversed(_SUPPORTED_SPLIT_K) - if split_k <= requested_split and (k // _CTA_K) % split_k == 0 - ) - tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) - max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) - tactic = dataclasses.replace( - tactic, - ab_stages=( - _MIN_AB_STAGES - if k == _CTA_K - else _DEFAULT_MIN_AB_STAGES - if k <= 2 * _CTA_K and m > 8 - else min(max_stages, 6) - if k <= 4 * _CTA_K - else max_stages - ), - ) - validate_tactic(tactic, m, n, k) - return tactic - - -__all__ = [ - "SplitKTactic", - "autotune_tactics", - "default_tactic", - "run_splitk_dense", -] - - -@dsl_user_op -def _map_shared_rank( - smem_ptr: "cute.Pointer", - peer_cta_rank_in_cluster: "Int32", - *, - loc=None, - ip=None, -) -> "Int32": - """Map an SMEM pointer into a peer CTA's address space.""" - return Int32( - llvm.inline_asm( - T.i32(), - [ - smem_ptr.toint(loc=loc, ip=ip).ir_value(), - peer_cta_rank_in_cluster.ir_value(), - ], - "mapa.shared::cluster.u32 $0, $1, $2;", - "=r,r,r", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - ) - - -@dsl_user_op -def _store_shared_remote_v4( - value0, - value1, - value2, - value3, - smem_ptr: "cute.Pointer", - mbar_ptr: "cute.Pointer", - peer_cta_rank_in_cluster: "Int32", - *, - loc=None, - ip=None, -) -> None: - """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" - llvm.inline_asm( - None, - [ - _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), - value0.bitcast(Int32).ir_value(loc=loc, ip=ip), - value1.bitcast(Int32).ir_value(loc=loc, ip=ip), - value2.bitcast(Int32).ir_value(loc=loc, ip=ip), - value3.bitcast(Int32).ir_value(loc=loc, ip=ip), - _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), - ], - "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " - "[$0], {$1, $2, $3, $4}, [$5];", - "r,r,r,r,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - - -#: Rank that gathers partials and stores the output. -OWNER_RANK = 0 # capacity: fixed: a split index, not a size - - -# The drafter-boundary push (``ar_push``, modeling's DRAFT_AR_PUSH): the exchange words of -# ops/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. -@dsl_user_op -def _ld_volatile_s32(addr, *, loc=None, ip=None): - return Int32( - llvm.inline_asm( - T.i32(), - [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], - "ld.volatile.global.s32 $0, [$1];", - "=r,l", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - ) - - -@dsl_user_op -def _st_global_b32(addr, value, *, loc=None, ip=None): - llvm.inline_asm( - None, - [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], - "st.global.b32 [$0], $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - - -@dsl_user_op -def _bf16x2_word(lo, hi, *, loc=None, ip=None): - """bf16(lo) | bf16(hi) << 16 (RN), -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" - return Int32( - llvm.inline_asm( - T.i32(), - [ - cutlass.Float32(hi).ir_value(loc=loc, ip=ip), - cutlass.Float32(lo).ir_value(loc=loc, ip=ip), - ], - "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" - "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" - "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" - "mov.b32 $0, w;\n\t}", - "=r,f,f", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - ) - - -class SplitKDenseGemmKernel: - """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" - - def __init__( - self, - *, - tactic: SplitKTactic, - use_pdl: bool, - has_bias: bool, - ar_push: bool = False, - ) -> None: - self.ar_push = ar_push - self.acc_dtype = cutlass.Float32 - self.cta_m = tactic.mma_m - self.cta_n = tactic.mma_n - self.cta_k = _CTA_K - self.num_ab_stage = tactic.ab_stages - self.split_k = tactic.split_k - self.use_pdl = use_pdl - self.has_bias = has_bias - - self.threads_per_cta = 256 - self.epilog_threads = 128 - self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) - self.cta_group = tcgen05.CtaGroup.ONE - self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD - self.cluster_shape = (1, tactic.split_k, 1) - - values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads - if values_per_thread % 4: - raise ValueError( - f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " - f"{values_per_thread} " - "values per epilogue thread; remote stores require a multiple of 4" - ) - self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread - self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES - - @cute.experimental.jit - def __call__( - self, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - bias: cute.Tensor, - stream: _cuda.CUstream, - ar_ws: cutlass.Int64 = cutlass.Int64(0), - ar_flags: cutlass.Int64 = cutlass.Int64(0), - ar_rank: Int32 = Int32(0), - ar_world: Int32 = Int32(1), - ar_maxt: Int32 = Int32(8), - ar_tokens: Int32 = Int32(0), - ar_n: Int32 = Int32(0), - ): - # Grid-y packs output-N tile and cluster rank. - self.kernel( - a, b, c, bias, ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_tokens, ar_n - ).launch( - grid=( - cute.ceil_div(c.layout.shape[0], self.cta_m), - cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, - c.layout.shape[2], - ), - block=(self.threads_per_cta, 1, 1), - cluster=self.cluster_shape, - smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), - stream=stream, - use_pdl=self.use_pdl, - ) - - @cute.experimental.kernel - def kernel( - self, - mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major - mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major - mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major - mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False - ar_ws: cutlass.Int64, # ar_push: the boundary all-reduce's multicast exchange words (int32) - ar_flags: cutlass.Int64, # ar_push: its per-(token, cluster CTA) call flags (parity = flag & 1) - ar_rank: Int32, - ar_world: Int32, - ar_maxt: Int32, # the bucket's tokens - ar_tokens: Int32, # rows pushed (public M) - ar_n: Int32, # public N (the hidden size): the exchange row width - ): - """Allocate storage and dispatch the specialized warps.""" - stages = self.num_ab_stage - - ab_dtype = mA.element_type - tiled_mma = sm100_utils.make_trivial_tiled_mma( - ab_dtype, - ab_dtype, - utils.LayoutEnum.from_tensor(mA).mma_major_mode(), - utils.LayoutEnum.from_tensor(mB).mma_major_mode(), - self.acc_dtype, - self.cta_group, - self.mma_tiler_mn, - ) - - mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) - block_idx = cute.arch.block_idx() - bidx = block_idx[0] - split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) - n_idx = block_idx[1] // self.split_k - l_idx = block_idx[2] - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - - sA = cute_ext.allocate( - ab_dtype, - cute.AddressSpace.smem, - sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), - alignment=_AB_BUFFER_ALIGN_BYTES, - ) - sB = cute_ext.allocate( - ab_dtype, - cute.AddressSpace.smem, - sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), - alignment=_AB_BUFFER_ALIGN_BYTES, - ) - - acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) - c_tiler_mn = (self.cta_m, self.cta_n) - - bar_full = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(stages), - alignment=_MBARRIER_BYTES, - ).iterator - bar_empty = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(stages), - alignment=_MBARRIER_BYTES, - ).iterator - bar_tma_epilog = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - bar_mma_epilog = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - bar_tmem_alloc = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - tmem_base_ptr = cute_ext.allocate( - cutlass.Int32, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_TMEM_POINTER_BYTES, - ).iterator - # ar_push: the owner's BF16 tile (kernel-M = output column, kernel-N = token), M-major. - sD = cute_ext.allocate( - mC.element_type, - cute.AddressSpace.smem, - cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), - alignment=16, - ) - - if cutlass.const_expr(self.split_k > 1): - mailbox = cute_ext.allocate( - cutlass.Float32, - cute.AddressSpace.smem, - cute.make_layout(self.mailbox_elements), - alignment=_MAILBOX_ALIGN_BYTES, - ) - bar_reduce = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - else: - # Dummy operands for the compile-time-elided reduction. - mailbox = sA - bar_reduce = bar_mma_epilog - - if warp_idx == 0: - with cute.arch.elect_one(): - for i in range(stages): - cute.arch.mbarrier_init(bar_full + i, 2) - cute.arch.mbarrier_init(bar_empty + i, 1) - cute.arch.mbarrier_init(bar_tma_epilog, 32) - cute.arch.mbarrier_init(bar_mma_epilog, 1) - cute.arch.mbarrier_init(bar_tmem_alloc, 160) - - if cutlass.const_expr(self.split_k > 1): - # Owner arrival plus peer transaction-byte credits. - cute.arch.mbarrier_init(bar_reduce, 1) - - cute.arch.mbarrier_init_fence() - if cutlass.const_expr(self.split_k > 1): - # Publish peer barriers before cross-CTA stores. - cute.arch.cluster_arrive_relaxed() - else: - cute.arch.barrier() - - # Host validation guarantees an equal, tail-free K partition. - k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k - k_tile_start = split_rank * k_tile_count - - if cutlass.const_expr(self.split_k > 1): - cute.arch.cluster_wait() - - # Warp 3 is idle; warps 4-7 run the epilogue. - if warp_idx == 0: - self.dma_warp( - bar_full, - bar_empty, - bar_tma_epilog, - cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), - sA, - cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), - k_tile_start, - k_tile_count, - True, - ) - elif warp_idx == 1: - self.dma_warp( - bar_full, - bar_empty, - bar_tma_epilog, - cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), - sB, - cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), - k_tile_start, - k_tile_count, - False, - ) - elif warp_idx == 2: - self.mma_warp( - bar_full, - bar_empty, - bar_mma_epilog, - bar_tmem_alloc, - tiled_mma, - sA, - sB, - tmem_base_ptr, - acc_layout, - self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), - k_tile_count, - ) - elif warp_idx >= 4: - self.epilog_warp( - bar_tma_epilog, - bar_mma_epilog, - bar_tmem_alloc, - tmem_base_ptr, - acc_layout, - cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), - cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), - cute.arch.thread_idx()[0] - 128, - mC.element_type, - utils.LayoutEnum.from_tensor(mC), - mailbox, - bar_reduce, - split_rank, - sD, - bidx, - n_idx * self.cta_n, - ar_ws, - ar_flags, - ar_rank, - ar_world, - ar_maxt, - ar_tokens, - ar_n, - ) - - @cute.experimental.jit - def dma_warp( - self, - bar_full, - bar_empty, - bar_tma_epilog, - g_tile: cute.Tensor, - s_tile: cute.Tensor, - cta_v_map: cute.Layout, - k_tile_start: cutlass.Int32, - k_tile_count: cutlass.Int32, - is_a: cutlass.Constexpr, - ): - stages = self.num_ab_stage - if cutlass.const_expr(not is_a and self.use_pdl): - cute.arch.griddepcontrol_wait() - - empty_phase = cutlass.Int32(1) - for k_tile in cutlass.range(k_tile_count, unroll=1): - stage = k_tile % stages - cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - bar_full + stage, - cute.size_in_bytes( - s_tile.element_type, - cute.slice_(s_tile.layout, (None, None, None, 0)), - ), - ) - cute_ext.tma_load( - g_tile[None, None, k_tile_start + k_tile], - s_tile[None, None, None, stage], - (bar_full + stage).value, - cta_v_map=cta_v_map, - tma_operation_type=self.tma_op, - update_expect_tx=False, - ) - if stage == stages - 1: - empty_phase = empty_phase ^ 1 - - if cutlass.const_expr(is_a and self.use_pdl and not self.ar_push): - # ar_push: the dependent is a prepushed all-reduce, which skips its grid wait: it is released only after - # this grid's own grid wait and its pushes, in the owner's _store_tile (the prepushed >= 2 contract). - cute.arch.griddepcontrol_launch_dependents() - if cutlass.const_expr(not is_a and self.has_bias): - cute.arch.mbarrier_arrive(bar_tma_epilog) - self._drain_producer(bar_empty, empty_phase, k_tile_count) - - @cute.experimental.jit - def _drain_producer( - self, - bar_empty, - empty_phase: cutlass.Int32, - k_tile_count: cutlass.Int32, - ): - stages = self.num_ab_stage - for tail in cutlass.range(stages, unroll=1): - stage = (tail + k_tile_count) % stages - cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) - if stage == stages - 1: - empty_phase = empty_phase ^ 1 - - @cute.experimental.jit - def mma_warp( - self, - bar_full, - bar_empty, - bar_mma_epilog, - bar_tmem_alloc, - tiled_mma: cute.TiledMma, - sA: cute.Tensor, - sB: cute.Tensor, - tmem_base_ptr, - acc_layout: cutlass.Constexpr, - mma_inst_tile_k: cutlass.Constexpr, - k_tile_count: cutlass.Int32, - ): - num_tmem_cols = 256 - 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) - - tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) - accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] - mma_atom = cute.make_mma_atom(tiled_mma.op) - full_phase = cutlass.Int32(0) - for k_tile in cutlass.range(k_tile_count, unroll=1): - stage = k_tile % self.num_ab_stage - cute.arch.mbarrier_wait(bar_full + stage, full_phase) - for k_block in range(mma_inst_tile_k): - if k_block == 0: - mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) - else: - mma_atom.set(tcgen05.Field.ACCUMULATE, True) - cute_ext.dot( - mma_atom, - cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), - cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), - accumulator, - ) - with cute.arch.elect_one(): - tcgen05.commit(bar_empty + stage, None, self.cta_group) - if stage == self.num_ab_stage - 1: - full_phase = full_phase ^ 1 - - with cute.arch.elect_one(): - tcgen05.commit(bar_mma_epilog, None, self.cta_group) - cute.arch.mbarrier_arrive(bar_tmem_alloc) - cute.arch.mbarrier_wait(bar_tmem_alloc, 1) - cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) - - @cute.experimental.jit - def epilog_warp( - self, - bar_tma_epilog, - bar_mma_epilog, - bar_tmem_alloc, - tmem_base_ptr, - acc_layout: cutlass.Constexpr, - gD_tile: cute.Tensor, - gBias_tile: cute.Tensor, - epi_tid: cutlass.Int32, - c_dtype: cutlass.Constexpr, - d_layout: cutlass.Constexpr, - mailbox, - bar_reduce, - split_rank: cutlass.Int32, - sD: cute.Tensor, - bidx: cutlass.Int32, - row_offset: cutlass.Int32, - ar_ws: cutlass.Int64, - ar_flags: cutlass.Int64, - ar_rank: cutlass.Int32, - ar_world: cutlass.Int32, - ar_maxt: cutlass.Int32, - ar_tokens: cutlass.Int32, - ar_n: cutlass.Int32, - ): - # Wait until MMA publishes the TMEM base pointer. - cute.arch.mbarrier_arrive(bar_tmem_alloc) - cute.arch.mbarrier_wait(bar_tmem_alloc, 0) - - acc_view = cute.make_tensor( - cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), - acc_layout, - )[((None, None), 0, 0, 0)] - - epi_tile = (self.cta_m, self.cta_n) - tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( - sm100_utils.get_tmem_load_op( - (self.cta_m, self.cta_n, self.cta_k), - d_layout, - c_dtype, - self.acc_dtype, - epi_tile, - False, - ), - acc_view, - ) - gD_epi = cute.flat_divide(gD_tile, epi_tile) - - # Match each epilogue thread's TMEM partition in RMEM. - rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) - rAcc = cute_ext.allocate( - self.acc_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - rD = cute_ext.allocate( - c_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - thr_t2r = tiled_copy_t2r.get_slice(epi_tid) - - if cutlass.const_expr(self.has_bias): - bias_dtype = gBias_tile.element_type - rBias = cute_ext.allocate( - bias_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - rBiasAcc = cute_ext.allocate( - self.acc_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - if split_rank == OWNER_RANK: - cute.arch.mbarrier_wait(bar_tma_epilog, 0) - cute_ext.partition_and_copy( - cute.make_tiled_copy_D( - cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), - tiled_copy_t2r, - ).get_slice(epi_tid), - cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], - rBias, - ) - rBiasAcc.store(rBias.load().to(self.acc_dtype)) - - cute.arch.mbarrier_wait(bar_mma_epilog, 0) - cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) - # Make tcgen05.ld visible before TMEM release and RMEM use. - cute.arch.fence_view_async_tmem_load() - cute.arch.mbarrier_arrive(bar_tmem_alloc) - - # Peers publish FP32 partials; only rank 0 reduces and stores. - if cutlass.const_expr(self.split_k > 1): - assert cute.size(rmem_layout) == self.mailbox_elements // ( - (self.split_k - 1) * self.epilog_threads - ) - values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) - values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) - if split_rank != OWNER_RANK: - for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): - _store_shared_remote_v4( - rAcc[value_idx], - rAcc[value_idx + 1], - rAcc[value_idx + 2], - rAcc[value_idx + 3], - mailbox.iterator - + (split_rank - Int32(1)) * values_per_peer - + epi_tid * values_per_thread - + value_idx, - bar_reduce, - Int32(OWNER_RANK), - ) - else: - if epi_tid == 0: - cute.arch.mbarrier_arrive_and_expect_tx( - bar_reduce, self.expected_transaction_bytes - ) - cute.arch.mbarrier_wait(bar_reduce, 0) - for peer in cutlass.range_constexpr(self.split_k - 1): - for value_idx in cutlass.range_constexpr(values_per_thread): - rAcc[value_idx] = ( - rAcc[value_idx] - + mailbox[ - peer * values_per_peer + epi_tid * values_per_thread + value_idx - ] - ) - - if split_rank == OWNER_RANK: - if cutlass.const_expr(self.has_bias): - rAcc.store(rAcc.load() + rBiasAcc.load()) - self._store_tile(thr_t2r, rAcc, rD, gD_epi, c_dtype) - if cutlass.const_expr(self.ar_push): - self._push_tile( - thr_t2r, - rD, - sD, - epi_tid, - bidx, - row_offset, - ar_ws, - ar_flags, - ar_rank, - ar_world, - ar_maxt, - ar_tokens, - ar_n, - ) - - # The reduction mbarrier covers remote stores; no cluster barrier needed. - - @cute.experimental.jit - def _store_tile(self, thr_t2r, rAcc, rD, gD_epi, c_dtype: cutlass.Constexpr): - """The owner rank's reduced FP32 tile: one cast to ``c_dtype`` and the store to global memory.""" - rD.store(rAcc.load().to(c_dtype)) - # Preserve TMEM coordinates; the copy predicates output tails. - cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) - - @cute.experimental.jit - def _push_tile( - self, - thr_t2r, - rD, - sD, - epi_tid: cutlass.Int32, - bidx: cutlass.Int32, - row_offset: cutlass.Int32, - ar_ws: cutlass.Int64, - ar_flags: cutlass.Int64, - ar_rank: cutlass.Int32, - ar_world: cutlass.Int32, - ar_maxt: cutlass.Int32, - ar_tokens: cutlass.Int32, - ar_n: cutlass.Int32, - ): - """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output - would store them (ops/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) - cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" - cute_ext.partition_and_copy(thr_t2r, rD, sD) - cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) - if cutlass.const_expr(self.use_pdl): - cute.arch.griddepcontrol_wait() - pairs = self.cta_m // 2 - col0 = bidx * self.cta_m - half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) - for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): - word = epi_tid + w0 - local_t = word // pairs - t = row_offset + local_t - p = word % pairs - if t < ar_tokens: - n = col0 + 2 * p - par = _ld_volatile_s32( - ar_flags + cutlass.Int64(4) * cutlass.Int64(t * 8 + n // (ar_n // 8)) - ) & Int32(1) - off = ( - cutlass.Int64(par) * half - + cutlass.Int64(t * ar_world + ar_rank) * cutlass.Int64(ar_n // 2) - + cutlass.Int64(n // 2) - ) - _st_global_b32( - ar_ws + off * cutlass.Int64(4), - _bf16x2_word( - sD[2 * p, local_t].to(cutlass.Float32), - sD[2 * p + 1, local_t].to(cutlass.Float32), - ), - ) - if cutlass.const_expr(self.use_pdl): - cute.arch.griddepcontrol_launch_dependents() - - -import torch as _torch # noqa: E402 (as in FlashInfer) - -_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) - - -@cute.experimental.jit -def _bmm_no_bias( - gemm_op: cutlass.Constexpr, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - stream: _cuda.CUstream, -): - c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) - gemm_op( - cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), - cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), - c, - cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), - stream, - ) - - -@cute.experimental.jit -def _bmm_push( - gemm_op: cutlass.Constexpr, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - stream: _cuda.CUstream, - ar_ws: cutlass.Int64, - ar_flags: cutlass.Int64, - ar_rank: Int32, - ar_world: Int32, - ar_maxt: Int32, - ar_tokens: Int32, - ar_n: Int32, -): - c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) - gemm_op( - cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), - cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), - c, - cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), - stream, - ar_ws, - ar_flags, - ar_rank, - ar_world, - ar_maxt, - ar_tokens, - ar_n, - ) - - -@cute.experimental.jit -def _bmm_bias( - gemm_op: cutlass.Constexpr, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - bias: cute.Tensor, - stream: _cuda.CUstream, -): - gemm_op( - cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), - cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), - cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), - cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), - stream, - ) - - -def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): - return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( - leading_dim=leading_dim - ) - - -def _detect_leading_dim(tensor: _torch.Tensor) -> int: - # Ignore synthetic batch stride, including 1x1. - for dim, stride in enumerate(tensor.stride()[1:], start=1): - if stride == 1: - return dim - raise ValueError("tensor has no stride-1 dimension") - - -def _make_layout_tensor( - shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int -) -> _torch.Tensor: - permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] - return _torch.empty( - tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" - ).permute([permutation.index(dim) for dim in range(len(shape))]) - - -def _make_compile_repr_tensors( - dtype: _torch.dtype, - has_bias: bool, - a_leading: int, - b_leading: int, - c_leading: int, -): - m, n, k, batch = 64, 8, _CTA_K, 1 - tensors = tuple( - _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) - for shape, leading_dim in ( - ((batch, n, k), a_leading), - ((batch, k, m), b_leading), - ((batch, n, m), c_leading), - ) - ) - if not has_bias: - return (*tensors, None) - return ( - *tensors, - _from_dlpack_dynamic( - _torch.empty((n,), dtype=dtype, device="cuda").as_strided( - size=(batch, n, m), stride=(0, 1, 0) - ), - 1, - 2, - ), - ) - - -def _to_cute_swap(a, b, out, bias): - a_swap = b.unsqueeze(0).transpose(-2, -1) - b_swap = a.unsqueeze(0).transpose(-2, -1) - c_swap = out.unsqueeze(0).transpose(-2, -1) - leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) - cute_tensors = tuple( - _from_dlpack_dynamic(tensor, leading_dim) - for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) - ) - if bias is None: - return (*cute_tensors, None, leading_dims) - return ( - *cute_tensors, - _from_dlpack_dynamic( - bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), - 1, - 2, - ), - leading_dims, - ) - - -@functools.cache -def _get_compiled_splitk_kernel( - dtype, - tactic: SplitKTactic, - use_pdl: bool, - has_bias: bool, - leading_dims: tuple[int, int, int], -): - if dtype not in _SUPPORTED_TORCH_DTYPES: - raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") - - kernel = SplitKDenseGemmKernel( - tactic=tactic, - use_pdl=use_pdl, - has_bias=has_bias, - ) - compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) - stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) - if has_bias: - compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) - else: - compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) - return compiled - - -def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: - tensors = (a, b, out) + ((bias,) if bias is not None else ()) - if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): - raise ValueError("a, b, out, and bias must be torch tensors") - if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: - raise ValueError("split-K dense GEMM accepts only 2D 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, b, out, and bias must share BF16 or FP16 dtype") - if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): - raise ValueError("a, b, and out must be dense row-major or column-major matrices") - if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): - raise ValueError("a, b, and out must be 32-byte aligned") - - m, k = a.shape - if b.shape[0] != k: - raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") - n = b.shape[1] - if out.shape != (m, n): - raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") - if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): - raise ValueError( - f"bias must be contiguous with shape {(n,)}, " - f"got shape {tuple(bias.shape)} and stride {bias.stride()}" - ) - - return m, n, k - - -def run_splitk_dense( - a, - b, - bias, - out, - pdl: bool, - tactic: SplitKTactic, -): - """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" - validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) - has_bias = bias is not None - cute_tensors = _to_cute_swap(a, b, out, bias) - compiled = _get_compiled_splitk_kernel( - dtype=a.dtype, - tactic=tactic, - use_pdl=pdl, - has_bias=has_bias, - leading_dims=cute_tensors[4], - ) - stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) - if has_bias: - compiled(*cute_tensors[:4], stream) - else: - compiled(*cute_tensors[:3], stream) - return out - - -@functools.cache -def _get_compiled_push( - dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] -): - kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, ar_push=True) - compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) - stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) - z64, z32 = cutlass.Int64(0), Int32(0) - return cute_ext.compile( - _bmm_push, - kernel, - *compile_tensors[:3], - stream, - z64, - z64, - z32, - Int32(1), - Int32(8), - z32, - Int32(6144), - ) - - -def compile_push(a, b, out, pdl: bool, tactic) -> None: - """Compile (only) the push variant of ``run_splitk_dense_push`` for these operands' shapes and layouts.""" - cute_tensors = _to_cute_swap(a, b, out, None) - _get_compiled_push(a.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) - - -def run_splitk_dense_push(a, b, out, pdl: bool, tactic: SplitKTactic, ar: tuple): - """``run_splitk_dense`` without bias that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange; - ``ar`` = (multicast words address, flags address, rank, world, bucket tokens) of the 8-token bucket.""" - validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) - cute_tensors = _to_cute_swap(a, b, out, None) - compiled = _get_compiled_push(a.dtype, tactic, pdl, cute_tensors[4]) - stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) - ws, flags, rank, world, maxt = ar - compiled( - *cute_tensors[:3], - stream, - cutlass.Int64(ws), - cutlass.Int64(flags), - Int32(rank), - Int32(world), - Int32(maxt), - Int32(out.shape[0]), - Int32(out.shape[1]), - ) - return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py index 76b73fc9ce91..3d3928610659 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_mc_allreduce.py @@ -28,11 +28,12 @@ epilogue-option choice and the MXFP8 publication. """ -from dataclasses import dataclass -from typing import Any - import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.mc_workspace import ( + McAllReduceWorkspace as McAllReduceWorkspace, +) + from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op from ..ops.comm.m3_mc_finalize_allreduce_norm import ( m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, @@ -73,74 +74,6 @@ def _ar_opts(m: int, bucket: int, finalize: bool = False) -> int: return opts | _OPT_CLUSTER_BARRIER -@dataclass(eq=False) -class McAllReduceWorkspace: - """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call - counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" - - uc: torch.Tensor - mc: torch.Tensor - flags: dict - rank: int - world_size: int - hidden: int - handle: Any - comm: Any - max_tokens: int = 8 - - @classmethod - def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": - from tensorrt_llm._torch.distributed.ops import ( - _make_mnnvl_mcast_buffer, - _mnnvl_workspace_all_succeeded, - ) - from tensorrt_llm._utils import mpi_comm - - session = mpi_comm() - session_group = session.Get_group() - group = session_group.Incl(mapping.tp_group) - session_group.Free() - comm = session.Create_group(group) - group.Free() - words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves - ok = not torch.cuda.is_current_stream_capturing() - if not _mnnvl_workspace_all_succeeded(comm, ok): - comm.Free() - raise RuntimeError( - "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" - ) - handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) - uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) - mc = handle.get_mc_buffer((words,), torch.int32, 0) - uc.fill_(EMPTY_WORD) - torch.cuda.synchronize() - if not _mnnvl_workspace_all_succeeded(comm, True): - raise RuntimeError("McAllReduceWorkspace.create failed on another rank") - return cls( - uc=uc, - mc=mc, - flags={}, - rank=mapping.tp_rank, - world_size=mapping.tp_size, - hidden=hidden, - handle=handle, - comm=comm, - max_tokens=max_tokens, - ) - - def flags_for(self, cluster: int): - """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. - get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) - that None never happens and cannot differ between ranks.""" - flags = self.flags.get(cluster) - if flags is None: - if torch.cuda.is_current_stream_capturing(): - return None - flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) - self.flags[cluster] = flags - return flags - - class McBuckets: """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py index 790194f9a927..b94e01bb15ec 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_moe/op.py @@ -53,6 +53,34 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.phase_plan import ( + MoePlanGeometry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.phase_plan import ( + make_phase_plan as _shared_make_phase_plan, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + M3MoeLayer as M3MoeLayer, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + _view as _view, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + ar_push_desc as ar_push_desc, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + fine_ranges, + pretile_shared, + shared_views, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + is_supported as is_supported, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.state import ( + M3MoeStateBase, + MoeStateBindings, +) + HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size INTER = 1536 # capacity: fixed: the routed experts' width per tp2 rank (the routed intermediate per TP rank (TP2)) NUM_EXPERTS = 128 @@ -129,7 +157,9 @@ def flat_applies(num_tokens: int) -> bool: M3_MOE_LAG = 2 -_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") +_ENGINE_PATH = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../../kernels/moe/m3_moe_kernel.py") +) def _kernel_module(config: dict, path: str): @@ -144,128 +174,13 @@ def _kernel_module(config: dict, path: str): ) spec = importlib.util.spec_from_file_location(name, path) mod = importlib.util.module_from_spec(spec) - mod.M3_CONFIG = dict(config) + mod.M3_CONFIG = dict(config, tp_size=2) sys.modules[name] = mod spec.loader.exec_module(mod) _modules[key] = mod return mod -def _view(t: torch.Tensor, align: int): - from cutlass.cute.runtime import from_dlpack - - return from_dlpack(t.detach(), assumed_align=align) - - -def is_supported(device: Optional[torch.device] = None) -> tuple: - if not torch.cuda.is_available(): - return False, "no CUDA" - major, _ = torch.cuda.get_device_capability(device) - if major != 10: - return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" - try: - import cutlass # noqa: F401 - except ImportError: - return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" - return True, "" - - -def ar_push_desc( - mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device -) -> torch.Tensor: - """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A - of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). - With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old - behaviour.""" - return torch.tensor( - [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device - ) - - -def _shared_views(shared_weights, who): - """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" - if shared_weights is None: - raise ValueError(f"{who} (shared): shared_weights are required") - gu_w, gu_sf, dn_w, dn_sf = shared_weights - want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} - for name, (t, shape) in want.items(): - if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): - raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") - if ( - gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 - or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 - ): - raise ValueError( - f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" - ) - return ( - gu_w.view(torch.uint8).view(-1), - gu_sf.view(torch.uint8).reshape(-1), - dn_w.view(torch.uint8).view(-1), - dn_sf.view(torch.uint8).reshape(-1), - ) - - -def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: - """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB - shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and - its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 - = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row - R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then - down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per - layer.""" - dev = gu8.device - H_, I_ = HIDDEN, SHARED_INTER - P = I_ // n_cta - split = 128 // (2 * P) - gt = H_ // 128 // split - dkt = I_ // 128 - gu = gu8.view(2 * I_, H_) - dn = dn8.view(H_, I_) - gkb, dkb = H_ // 32, I_ // 32 - - def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast - return sf[ - ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) - ] - - r = torch.arange(128, device=dev) - q = r >> 5 - ri = r & 31 - octant = r // P - chunk = torch.arange(8, device=dev) - swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) - out = torch.empty( - n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev - ) # capacity: fixed: a workspace of the kernel's fixed geometry - for s in range(n_cta): - # gate_up: MMA row R -> weight row, slice - nrow = (octant & 1) * I_ + P * s + r % P - sl = octant >> 1 - for g in range(gt): - kt = sl * gt + g # per row - tile = gu[ - nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) - ] # [128, 128] - img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) - out[s, g, :16384] = img.reshape(-1) - kb = torch.arange(4, device=dev) - codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] - atom = torch.empty(512, dtype=torch.uint8, device=dev) - atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) - out[s, g, 16384:] = atom - for gd in range(2 * dkt): - jt, kt = gd // dkt, gd % dkt - row0 = 256 * s + 128 * jt - tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] - img = tile.reshape(128, 8, 16)[r[:, None], swz] - out[s, gt + gd, :16384] = img.reshape(-1) - out[s, gt + gd, 16384:] = dnsf8[ - ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 - ] - return out.view(-1) - - # --------------------------------------------------------------------------------------------------------------------- # m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. # --------------------------------------------------------------------------------------------------------------------- @@ -282,26 +197,6 @@ def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: return hit[1] -def _fine_ranges(weights, total: int): - """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item - counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary - strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit - runs at least to that unit's end. Returns the item counts.""" - n = len(weights) - items = [0] * n - s = 0 - acc = 0.0 - for c in range(n): - acc += weights[c] - e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) - if s % _S1 and e < total: # started inside a unit: run to (at least) its end - e = max(e, min(total, (s // _S1 + 1) * _S1)) - items[c] = e - s - s = e - assert s == total - return items - - def make_phase_plan( num_ctas: int, m_max: int, @@ -318,84 +213,24 @@ def make_phase_plan( group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" - gcap = TOP_K * m_max - nb2 = HIDDEN // 32 - plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) - nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] - b0s = [sum(nbs[:c]) for c in range(num_ctas)] - # (heavy_fc2_only) CTAs with more FC2 blocks than the minimum stream no FC1: they run each group's FC2 as soon as - # its hand-off lands, so their trailing FC2 after the last hand-off matches the others' two lagged groups - min_nb = min(nbs) - no_fc1 = [heavy_fc2_only and nbs[c] > min_nb for c in range(num_ctas)] - BIG_W = 1 << 40 - blk_bytes = 32 * _ROW2_BYTES # one block, one slot - for g in range(TOP_K, gcap + 1): - ng = (g + 3) // 4 - for c in range(num_ctas): - plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) - for p in range(ng): - n_p = min(4, g - 4 * p) - total = n_p * U1 * _S1 - fc2 = [0] * num_ctas - for c in range(num_ctas): - if p >= lag: - fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes - if p == ng - 1 and not ( - early0 and p == 0 and m_max >= 2 - ): # early0: group 0's ranges the same for every G - for q in range(max(0, ng - lag), ng): - fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) - for c in range(num_ctas): - if no_fc1[c]: - fc2[c] += BIG_W - items = [0] * num_ctas - if total >= _S1 * num_ctas: - - def counts_for(limit): - return [ - 0 if no_fc1[c] else max(_S1, (limit - f) // _ITEM_BYTES) - for c, f in enumerate(fc2) - ] - - lo, hi = 0, max(f for c, f in enumerate(fc2) if not no_fc1[c]) + total * _ITEM_BYTES - while lo < hi: - mid = (lo + hi) // 2 - if sum(counts_for(mid)) >= total: - hi = mid - else: - lo = mid + 1 - items = counts_for(lo) - extra = sum(items) - total - while extra > 0: - c = min( - (c for c in range(num_ctas) if items[c] > _S1), - key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), - ) - items[c] -= 1 - extra -= 1 - elif fine: - # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 - lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) - for _ in range(100): - mid = 0.5 * (lo_f + hi_f) - if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: - hi_f = mid - else: - lo_f = mid - wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] - sw = sum(wts) - items = _fine_ranges([w * total / sw for w in wts], total) - else: - chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) - for c in chosen: - items[c] = _S1 - i0 = 0 - for c in range(num_ctas): - if items[c]: - plan[g, c, p] = i0 | ((i0 + items[c]) << 16) - i0 += items[c] - assert i0 == total - return plan, max(nbs) + return _shared_make_phase_plan( + num_ctas=num_ctas, + m_max=m_max, + lag=lag, + tail_w=tail_w, + early0=early0, + fine=fine, + heavy_fc2_only=heavy_fc2_only, + geometry=MoePlanGeometry( + top_k=TOP_K, + hidden=HIDDEN, + u1=U1, + stages=_S1, + row2_bytes=_ROW2_BYTES, + item_bytes=_ITEM_BYTES, + ), + fine_ranges_fn=_fine_ranges, + ) def _shared_takes_fc2_plan(plan: torch.Tensor, num_ctas: int, m_max: int) -> torch.Tensor: @@ -413,7 +248,7 @@ def _shared_takes_fc2_plan(plan: torch.Tensor, num_ctas: int, m_max: int) -> tor return ext -class M3MoeState: +class M3MoeState(M3MoeStateBase): """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated @@ -422,6 +257,15 @@ class M3MoeState: fused_router = True + BINDINGS = MoeStateBindings( + inter=INTER, + shared_inter=SHARED_INTER, + engine_path=_ENGINE_PATH, + compiled=_compiled, + pretiled=_pretiled, + cache_module_config=True, + ) + def __init__( self, device: torch.device, @@ -585,227 +429,14 @@ def __init__( self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) self._fn = None - def layer( - self, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - e_score_correction_bias, - router_weight=None, - fc31_input_scale=None, - shared_weights=None, - ) -> "M3MoeLayer": - """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). - ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [3072, 6144], its swizzled UE8M0 - scales, down e4m3 [6144, 1536], its swizzled UE8M0 scales).""" - layer = M3MoeLayer( - self, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - e_score_correction_bias, - router_weight, - fc31_input_scale, - ) - if self.shared: - layer.shared_weights = _shared_views(shared_weights, "m3_moe") - layer.shared_mxfp8 = tuple(shared_weights) - if self.mod.SH_PRETILE: - # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy - gu8, _, dn8, _ = layer.shared_weights - layer.shared_weights = ( - gu8, - _pretiled(*layer.shared_weights), - dn8, - self._dummy.view(torch.uint8), - ) - return layer - - def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: - m = self.m - if ( - hidden is None - or hidden.dtype != torch.bfloat16 - or tuple(hidden.shape) != (m, HIDDEN) - or not hidden.is_contiguous() - ): - raise ValueError( - f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" - ) - fz = ( - hidden.view(-1).view(torch.int32), - layer.router_weight.view(-1), - layer.fc31_input_scale.view(-1), - self.lgbuf, - ) - lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy - if out is None: - out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) - sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 - args = ( - layer.w3_w1_weight, - layer.w3_w1_weight_scale.view(-1), - layer.w2_weight, - layer.w2_weight_scale.view(-1), - lgv, - layer.bias, - xq, - xs, - layer.escal, - self.plan.view(-1), - self.partials, - self.hbuf32, - self.epochs, - out.view(-1), - self.route_dbg.view(-1), - self.tdbg.view(-1), - *fz, - *sh, - self.hx, - self.dnx, - self._ar_off if arp is None else arp, - layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), - layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), - self.img32, - self.cnt, - ) - if self._fn is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("m3_moe compiles on its first call: call it once before capture") - import cutlass.cute as cute - - aligns = [ - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 4, - 16, - 16, - 16, - 16, - 16, - 4, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - ] - sig = [_view(t, a) for t, a in zip(args, aligns)] - # keyed on the config the kernel module was built with (not self.config), so a module built with any other - # trace-time config can never reuse another build's compiled kernel - key = (_ENGINE_PATH,) + tuple(sorted(self.mod.M3_CONFIG.items())) - fn = _compiled.get(key) - if fn is None: - fn = _compiled[key] = cute.compile( - self.mod.m3_moe, - *sig, - cute.runtime.make_fake_stream(), - options="--enable-tvm-ffi", - ) - self._fn = fn - self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) - return out - - -class M3MoeLayer: - """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in - place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading - weights).""" - - shared_weights = None # the kernel's view of the shared expert (a shared state) - shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given - def __init__( - self, - state, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - bias, - router_weight=None, - fc31_input_scale=None, - ): - e = NUM_EXPERTS - expected = { - "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), - "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), - "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), - "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), - } - for name, (t, shape) in expected.items(): - if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): - raise ValueError( - f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" - ) - self.state = state - self.w3_w1_weight = w3_w1_weight.view(torch.uint8) - self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) - self.w2_weight = w2_weight.view(torch.uint8) - self.w2_weight_scale = w2_weight_scale.view(torch.uint8) - self.escal = torch.stack( - [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] - ).float() - self.escal = self.escal.reshape(6, e).contiguous() - self.bias = bias.float().contiguous() - if ( - router_weight is None - or fc31_input_scale is None - or router_weight.dtype != torch.float32 - or tuple(router_weight.shape) != (e, HIDDEN) - or not router_weight.is_contiguous() - or fc31_input_scale.dtype != torch.float32 - or fc31_input_scale.numel() != 1 - ): - raise ValueError( - "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" - ) - self.router_weight = router_weight - self.fc31_input_scale = fc31_input_scale +def _shared_views(shared_weights, who): + return shared_views(shared_weights, who, shared_inter=SHARED_INTER) - def __call__( - self, - hidden: torch.Tensor, - out: Optional[torch.Tensor] = None, - arp: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output - into an all-reduce exchange.""" - return self.state.run(self, hidden, out, arp=arp) + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24): + return pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta, shared_inter=SHARED_INTER) + + +def _fine_ranges(weights, total: int): + return fine_ranges(weights, total, stages=_S1) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py index 48a7827f59b3..a89f96fe0b10 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_prequant.py @@ -1,111 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Hand-off of an MXFP8-quantized activation from the kernel that produced it to the next MXFP8 GEMM. -A producer that already holds a BF16 row in registers (the multicast all-reduce + RMSNorm, the split-K attention -merge) also writes its MXFP8 form (E4M3 rows + UE8M0 scales in the 128x4-swizzled layout of -``trtllm::mxfp8_quantize(x, True)``) and publishes it here keyed by the BF16 tensor. The MXFP8 linear whose input -is that tensor takes it and skips its own quantize kernel. Every take clears the slot, so a hand-off is only ever -consumed by the first MXFP8 linear that runs after the producer. - -Module state shared across files (one instance each per process): - -* ``SLOT`` (PrequantSlot): published by ``m3_mc_allreduce`` (``quant=True``), the sparse-decode merge and the fused - decode attention (MXFP8 output); cleared by the sparse decode before it runs; taken by the MXFP8 linear dispatch - and by the one-launch QKV + attention + o_proj path (which also peeks ``pending``). -* ``AR_PUSH`` (ArPushSlot): armed by the fast forward around the attention call; the o_proj GEMV takes the target - and marks ``pushed``; the fast forward checks ``take_pushed`` before the attention boundary all-reduce. -* ``MOE_PUSH`` (MoePushSlot): armed by the fast forward before the MoE; the routed-expert engine takes ``engine`` - and the shared expert's down projection takes ``shared``, each marks its pushed output; the fast forward - checks ``take_pushed`` before the MoE boundary all-reduce. -""" - - -class PrequantSlot: - def __init__(self): - self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) - - def publish(self, x, xq, xsf): - self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) - - def clear(self): - self.pending = None - - def take(self, x): - """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" - pending, self.pending = self.pending, None - if ( - pending is None - or x.dim() != 2 - or x.data_ptr() != pending[0] - or x.shape[0] != pending[1] - or x.shape[1] != pending[2] - ): - return None - return pending[3], pending[4] +"""Target-local activation hand-off state.""" +from ...ops.gemm.prequant import ArPushSlot, MoePushSlot, PrequantSlot SLOT = PrequantSlot() - - -class ArPushSlot: - """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. - - The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output - it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" - - def __init__(self): - self.target = None - self.pushed = None # data_ptr of the pushed output - - def arm(self, target): - self.target = target - self.pushed = None - - def take_target(self): - target, self.target = self.target, None - return target - - def take_pushed(self, x): - pushed, self.pushed = self.pushed, None - return pushed is not None and x.data_ptr() == pushed - - AR_PUSH = ArPushSlot() - - -class MoePushSlot: - """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, - ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv - ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" - - def __init__(self): - self.engine = None - self.shared = None - self.engine_pushed = None # data_ptr of the pushed routed output - self.shared_pushed = None # data_ptr of the pushed shared output - - def arm(self, engine, shared): - self.engine, self.shared = engine, shared - self.engine_pushed = self.shared_pushed = None - - def disarm(self): - self.engine = self.shared = None - - def take_engine(self): - engine, self.engine = self.engine, None - return engine - - def take_shared(self): - shared, self.shared = self.shared, None - return shared - - def take_pushed(self, shared, routed): - """(routed pushed, shared pushed) for these outputs; clears the marks.""" - e = routed is not None and self.engine_pushed == routed.data_ptr() - s = shared is not None and self.shared_pushed == shared.data_ptr() - self.engine_pushed = self.shared_pushed = None - return e, s - - MOE_PUSH = MoePushSlot() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py index 30855eb7f05d..d92ac2eb018d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_router_gemv.py @@ -15,9 +15,15 @@ import torch -from ..ops.moe.m3_router_gemv import m3_router_gemv as _router_gemv_op -from ..ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 as _router_gemv_fp4_op -from ..ops.moe.m3_router_tc import m3_router_tc as _router_tc_op +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv import ( + m3_router_gemv as _router_gemv_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv_fp4 import ( + m3_router_gemv_fp4 as _router_gemv_fp4_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_tc import ( + m3_router_tc as _router_tc_op, +) # Opt-in dynamic shared memory left after the kernel's static reduction buffer (SM100: 227 KiB). _MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py index 6e6e23e0c0b8..d4e58a983bef 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_select_long_dsl.py @@ -1,297 +1,47 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). +"""Target-specific select policy and independent compilation cache.""" -Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long``, -one 1024-thread CTA per row), from a CTA of 384 threads: +import sys +from typing import Dict, Tuple, cast -* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and - mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's - valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned - key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped - block and masks the result, so a thread's loads issue back to back. -* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of - the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. -* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by - that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate - against the whole row instead. -* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. - -The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait -(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are -compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and -loads them before the executor's warmup, so a decode step never compiles one. -""" - -from typing import Dict, Tuple - -import cutlass -import cutlass.cute as cute +import cutlass as cutlass +import cutlass.cute as cute # noqa: F401 — used by the bound family implementation import torch -from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor -from cutlass.utils.smem_allocator import SmemAllocator +from cutlass.cute.runtime import make_fake_compact_tensor as make_fake_compact_tensor +from cutlass.cute.runtime import make_fake_stream as make_fake_stream +from cutlass.cute.runtime import make_fake_tensor as make_fake_tensor -from tensorrt_llm._torch.kernels.blackwell.utils import ( - griddepcontrol_launch_dependents, - griddepcontrol_wait, +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + long_select_host as _select_shared, ) -# The 384-thread select reads its row's valid-block count before its grid wait: off (its TP2 in-model pairs did not -# pass); the selection is the same either way. -SELECT_N_VALID_PREWAIT = False - -SELECT_THREADS = 384 -SELECT_WARPS = SELECT_THREADS // 32 -# 1M tokens + draft tokens at 128 tokens per block. -MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row -MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging -TOPK = 16 # capacity: fixed: MSA's top-k -NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) -# The key of -inf (the reference's 0xff800000): never emitted. -NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size -INIT_SCORE = 1.0e30 -LOCAL_SCORE = 1.0e29 -# (keys per thread, rows up to this many valid blocks) -ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) - -# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare -# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group -# minima [16], per-warp candidate counts [SELECT_WARPS]. -_CAND = 0 -_PICK_KEY = _CAND + 2 * MAX_CAND -_PICK_BLOCK = _PICK_KEY + TOPK -_WARP_MIN = _PICK_BLOCK + TOPK -_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 -_WARP_CNT = _GROUP_MIN + 16 -SMEM_WORDS = _WARP_CNT + SELECT_WARPS -SMEM_BYTES = 4 * SMEM_WORDS - - -def _key(score): - """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" - bits = cutlass.Float32(score).bitcast(cutlass.Int32) - return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) - - -def _imin(a, b): - return cutlass.select_(b < a, b, a) - - -def _imax(a, b): - return cutlass.select_(b > a, b, a) - - -def _forced_key(scores, head, block, query, init_blocks, local_start): - raw = scores[head, block, query] - s = cutlass.select_( - block >= local_start, - cutlass.Float32(LOCAL_SCORE), - cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), - ) - return _key(s) - - -@cute.jit -def _select_row( - row, - scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides - n_valid: cute.Tensor, # int32 [>= total_q] - out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] - init_blocks, - local_blocks, - num_kv_heads, - smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory - bar_id: cutlass.Constexpr, - loaded_trigger: cutlass.Constexpr, - raw_prewait, - n_valid_prewait: cutlass.Constexpr, -): - """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier - ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the - row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by - the caller before its grid wait; otherwise it is read here.""" - tid, _, _ = cute.arch.thread_idx() - lane = tid % cutlass.Int32(32) - warp = tid // cutlass.Int32(32) - query = row // num_kv_heads - head = row % num_kv_heads - if cutlass.const_expr(n_valid_prewait): - raw = raw_prewait - else: - raw = n_valid[query] - valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) - local_start = _imax(raw - local_blocks, cutlass.Int32(0)) - cand = cute.make_tensor( - cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) - ) - - if tid < TOPK: - smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) - smem[_PICK_BLOCK + tid] = cutlass.Int32(0) - - if valid <= TOPK: - # Every valid block is picked (thread t holds block t). - if tid < valid: - smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) - smem[_PICK_BLOCK + tid] = tid - if cutlass.const_expr(loaded_trigger): - griddepcontrol_launch_dependents() - for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): - items = ITEM_CLASSES[ci][0] - lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] - hi = ITEM_CLASSES[ci][1] - if (valid > lo) & (valid <= hi): - key = cute.make_rmem_tensor(items, cutlass.Int32) - local_min = cutlass.Int32(NO_KEY) - for s in cutlass.range_constexpr(items): - # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. - block = tid + cutlass.Int32(s * SELECT_THREADS) - k = _forced_key( - scores, head, _imin(block, valid - 1), query, init_blocks, local_start - ) - key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) - local_min = _imin(local_min, key[s]) - local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) - if lane < 16: - smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - if cutlass.const_expr(loaded_trigger): - griddepcontrol_launch_dependents() - if tid < 16: - g = cutlass.Int32(NO_KEY) - for w in cutlass.range_constexpr(SELECT_WARPS): - g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) - smem[_GROUP_MIN + tid] = g - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - bound = cutlass.Int32( - -2147483648 - ) # the smallest key: signed keys, unlike the reference's unsigned ones - for g in cutlass.range_constexpr(16): - bound = _imax(bound, smem[_GROUP_MIN + g]) - - # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. - cnt = cutlass.Int32(0) - for s in cutlass.range_constexpr(items): - cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) - incl = cnt - for d in cutlass.range_constexpr(5): - up = cute.arch.shuffle_sync( - incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) - ) - incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) - if lane == 31: - smem[_WARP_CNT + warp] = incl - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - before = cutlass.Int32(0) - total = cutlass.Int32(0) - for w in cutlass.range_constexpr(SELECT_WARPS): - c = smem[_WARP_CNT + w] - before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) - total += c - slot = before + incl - cnt - for s in cutlass.range_constexpr(items): - if key[s] <= bound: - if slot < MAX_CAND: - cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( - tid + cutlass.Int32(s * SELECT_THREADS) - ) - slot += 1 - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - - if total <= MAX_CAND: - for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): - mine = cand[i] - rank = cutlass.Int32(0) - for j in cutlass.range(total, unroll=8): - rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) - if rank < TOPK: - smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) - smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) - else: - # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. - for s in cutlass.range_constexpr(items): - if key[s] <= bound: - block = tid + cutlass.Int32(s * SELECT_THREADS) - rank = cutlass.Int32(0) - for j in cutlass.range(valid, unroll=1): - kj = _forced_key(scores, head, j, query, init_blocks, local_start) - less = (kj < key[s]) | ((kj == key[s]) & (j < block)) - rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) - if rank < TOPK: - smem[_PICK_KEY + rank] = key[s] - smem[_PICK_BLOCK + rank] = block - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - - if warp == 0: - # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. - v = cutlass.Int32(NO_KEY) - if lane < TOPK: - pk = smem[_PICK_KEY + lane] - if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): - v = smem[_PICK_BLOCK + lane] - for size_log in cutlass.range_constexpr(1, 6): - for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): - other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) - ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) - lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) - v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) - if lane < TOPK: - out[query, head, lane] = cutlass.select_( - v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v - ) - - -class _SelectLongDslKernel: - """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count - before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because - ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's - first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it - (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" - - def __init__( - self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool - ): - self.use_pdl = use_pdl - self.early_trigger = early_trigger - self.loaded_trigger = (not early_trigger) and trigger_phase == 1 - self.n_valid_prewait = n_valid_prewait - - @cute.kernel - def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): - row, _, _ = cute.arch.block_idx() - smem = SmemAllocator() - words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) - raw_prewait = cutlass.Int32(0) - if cutlass.const_expr(self.n_valid_prewait): - raw_prewait = n_valid[row // num_kv_heads] - griddepcontrol_wait() - if cutlass.const_expr(self.early_trigger): - griddepcontrol_launch_dependents() - _select_row( - row, - scores, - n_valid, - out, - init_blocks, - local_blocks, - num_kv_heads, - words, - bar_id=1, - loaded_trigger=self.loaded_trigger, - raw_prewait=raw_prewait, - n_valid_prewait=self.n_valid_prewait, - ) - - @cute.jit - def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): - num_kv_heads = scores.shape[0] - self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( - grid=(scores.shape[2] * num_kv_heads, 1, 1), - block=(SELECT_THREADS, 1, 1), - stream=stream, - use_pdl=self.use_pdl, - ) +from ...kernels.attention.m3_select_long_dsl import _CAND as _CAND +from ...kernels.attention.m3_select_long_dsl import _GROUP_MIN as _GROUP_MIN +from ...kernels.attention.m3_select_long_dsl import _PICK_BLOCK as _PICK_BLOCK +from ...kernels.attention.m3_select_long_dsl import _PICK_KEY as _PICK_KEY +from ...kernels.attention.m3_select_long_dsl import _WARP_CNT as _WARP_CNT +from ...kernels.attention.m3_select_long_dsl import _WARP_MIN as _WARP_MIN +from ...kernels.attention.m3_select_long_dsl import INIT_SCORE as INIT_SCORE +from ...kernels.attention.m3_select_long_dsl import ITEM_CLASSES as ITEM_CLASSES +from ...kernels.attention.m3_select_long_dsl import LOCAL_SCORE as LOCAL_SCORE +from ...kernels.attention.m3_select_long_dsl import MAX_BLOCKS as MAX_BLOCKS +from ...kernels.attention.m3_select_long_dsl import MAX_CAND as MAX_CAND +from ...kernels.attention.m3_select_long_dsl import NEG_INF_KEY as NEG_INF_KEY +from ...kernels.attention.m3_select_long_dsl import NO_KEY as NO_KEY +from ...kernels.attention.m3_select_long_dsl import SELECT_THREADS as SELECT_THREADS +from ...kernels.attention.m3_select_long_dsl import SELECT_WARPS as SELECT_WARPS +from ...kernels.attention.m3_select_long_dsl import SMEM_BYTES as SMEM_BYTES +from ...kernels.attention.m3_select_long_dsl import SMEM_WORDS as SMEM_WORDS +from ...kernels.attention.m3_select_long_dsl import TOPK as TOPK +from ...kernels.attention.m3_select_long_dsl import _forced_key as _forced_key +from ...kernels.attention.m3_select_long_dsl import _imax as _imax +from ...kernels.attention.m3_select_long_dsl import _imin as _imin +from ...kernels.attention.m3_select_long_dsl import _key as _key +from ...kernels.attention.m3_select_long_dsl import _select_row as _select_row +from ...kernels.attention.m3_select_long_dsl import _SelectLongDslKernel as _SelectLongDslKernel +SELECT_N_VALID_PREWAIT = False #: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} @@ -311,28 +61,7 @@ def kernel_key( def _compile(key: Tuple[bool, bool, int, bool]): - if key not in kernel_cache: - heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() - scores = make_fake_tensor( - cutlass.Float32, - (heads, blocks, rows), - stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), - ) - n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) - out = make_fake_tensor( - cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) - ) - kernel_cache[key] = cute.compile( - _SelectLongDslKernel(*key), - scores, - n_valid, - out, - cutlass.Int32(0), - cutlass.Int32(0), - make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) - return kernel_cache[key] + return _select_shared._compile(_SELECT_TARGET, key) def select_long_dsl( @@ -346,21 +75,22 @@ def select_long_dsl( trigger_phase: int, ) -> None: """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" - _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( - scores, n_valid_blocks, out, init_blocks, local_blocks + return _select_shared.select_long_dsl( + _SELECT_TARGET, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, ) def prebuild(keys, device) -> list: """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" - scores = torch.zeros( - 1, 2048, 4, dtype=torch.float32, device=device - ) # capacity: fixed: a dummy prebuild row, not a size - n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) - out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) - built = [] - for key in keys: - key = kernel_key(*key) - _compile(key)(scores, n_valid, out, 0, 1) - built.append(key) - return built + return _select_shared.prebuild(_SELECT_TARGET, keys, device) + + +_SELECT_TARGET = cast(_select_shared.LongSelectTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py index 8382d02fb031..58294bfd48c6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_shard_accept.py @@ -21,6 +21,10 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.accept_workspace import ( + create_accept_workspace, +) + from . import m3_report #: Verify rows per call, and exchange rows per call half of the workspace. @@ -60,49 +64,7 @@ def _num_sms() -> int: def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: - """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the - local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call - outside CUDA-graph capture, in the same order on every rank.""" - words = 2 * slot_rows * max(mapping.tp_size, 1) - if mapping.tp_size > 1: - from tensorrt_llm._torch.distributed.ops import ( - _make_mnnvl_mcast_buffer, - _mnnvl_workspace_all_succeeded, - ) - from tensorrt_llm._utils import mpi_comm - - session = mpi_comm() - session_group = session.Get_group() - group = session_group.Incl(mapping.tp_group) - session_group.Free() - comm = session.Create_group(group) - group.Free() - if not _mnnvl_workspace_all_succeeded(comm, True): - comm.Free() - raise RuntimeError("exchange workspace creation failed on another rank") - handle = _make_mnnvl_mcast_buffer( - comm, max(words * 8, 4096), mapping, mapping.is_multi_node() - ) - uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) - mc = handle.get_mc_buffer((words,), torch.int64, 0) - uc.zero_() - torch.cuda.synchronize() - if not _mnnvl_workspace_all_succeeded(comm, True): - raise RuntimeError("exchange workspace creation failed on another rank") - mc_ptr, rank = mc.data_ptr(), mapping.tp_rank - else: - uc = torch.zeros(words, dtype=torch.int64, device=device) - mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None - return AcceptWorkspace( - uc=uc, - mc_ptr=mc_ptr, - keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), - counters=torch.zeros(2, dtype=torch.int32, device=uc.device), - rank=rank, - world_size=max(mapping.tp_size, 1), - handle=handle, - comm=comm, - ) + return create_accept_workspace(mapping, device, slot_rows, AcceptWorkspace) def _workspace(mapping, device) -> Optional[AcceptWorkspace]: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py index f5a1b229c777..03cae9c855c9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_target_head.py @@ -12,13 +12,30 @@ from __future__ import annotations +import sys from dataclasses import dataclass -from typing import Optional +from typing import Optional, cast import torch -from ..ops.gemm.m3_target_head import CTA_WARPS, MAX_ROWS, encode_chunked, round_trip_ok -from ..ops.gemm.m3_target_head import m3_target_head as _target_head_op +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + target_head_state as _head_shared, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + CTA_WARPS as CTA_WARPS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + MAX_ROWS as MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + encode_chunked as encode_chunked, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + m3_target_head as _target_head_op, # noqa: F401 — used by the bound family implementation +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + round_trip_ok as round_trip_ok, +) @dataclass(eq=False) @@ -43,71 +60,24 @@ class TargetHead: def _num_sms() -> int: - global _NUM_SMS - if _NUM_SMS is None: - _NUM_SMS = torch.cuda.get_device_properties( - torch.cuda.current_device() - ).multi_processor_count - return _NUM_SMS + return _head_shared._num_sms(_HEAD_TARGET) def head_supported(lm_head) -> bool: """A bias-free 2-D BF16 head of whole 16-row tiles whose K the coding covers (a multiple of 64).""" - weight = getattr(lm_head, "weight", None) - if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: - return False - return ( - getattr(lm_head, "bias", None) is None - and weight.shape[0] % 16 == 0 - and weight.shape[1] % 64 == 0 - ) + return _head_shared.head_supported(_HEAD_TARGET, lm_head) def prepare_head(lm_head, grid: Optional[int] = None) -> TargetHead: """Code ``lm_head``'s local shard and check it decodes to the BF16 weight bit for bit (fails closed: a shard that does not is returned with ``coded=False`` and nothing kept). Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_target_head``).""" - weight = lm_head.weight - with torch.no_grad(): - w = weight.detach() - packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) - if not round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): - return TargetHead(weight=weight, coded=False) - if grid is None: - grid = _num_sms() - n = weight.shape[0] - return TargetHead( - weight=weight, - coded=True, - w=packed, - base=base, - exc_off=exc_off, - exc_pos=exc_pos, - exc_w=exc_w.view(-1), - part=torch.zeros(grid * CTA_WARPS * 2 * 4 * 32, dtype=torch.float32, device=weight.device), - counter=torch.zeros((n + 15) // 16, dtype=torch.int32, device=weight.device), - grid=grid, - ) + return _head_shared.prepare_head(_HEAD_TARGET, lm_head, grid) def logits(head: TargetHead, x: torch.Tensor, use_pdl: bool = True) -> torch.Tensor: """This rank's BF16 logits shard [rows, n] of the rows x [rows <= MAX_ROWS, K] (BF16) from the coded shard.""" - assert head.coded and 1 <= x.shape[0] <= MAX_ROWS, (head.coded, x.shape) - if x.stride(1) != 1 or x.stride(0) % 8 or x.data_ptr() % 16: - x = x.contiguous() - n = head.weight.shape[0] - out = torch.empty((x.shape[0], n), dtype=torch.bfloat16, device=x.device) - return _target_head_op( - head.w, - x, - out, - head.part, - head.counter, - head.exc_off, - head.exc_pos, - head.exc_w, - head.base, - n, - head.grid, - use_pdl, - ) + return _head_shared.logits(_HEAD_TARGET, head, x, use_pdl) + + +_HEAD_TARGET = cast(_head_shared.HeadTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py index db0077117fe3..5d99aab162d7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/m3_tokmap_guard.py @@ -17,61 +17,20 @@ from __future__ import annotations -from typing import Callable, Dict, Optional, Tuple +import sys +from typing import Callable, Dict, Optional, Tuple, cast import torch -from tensorrt_llm.bindings.internal import thop -from tensorrt_llm.logger import logger +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + token_map as _shared, +) +from tensorrt_llm.bindings.internal import thop as thop +from tensorrt_llm.logger import logger as logger #: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). -_ARG_NAMES = ( - "qkv_input", - "workspace", - "sequence_lengths", - "spec_decoding_generation_lengths", - "spec_decoding_position_offsets", - "kv_cache_block_offsets", - "host_kv_cache_pool_pointers", - "host_kv_cache_pool_mapping", - "kv_scale_orig_quant", - "kv_scale_quant_orig", - "attention_output_orig_quant", - "rotary_inv_freq", - "rotary_cos_sin", - "mrope_position_deltas", - "layer_idx", - "seq_offset", - "num_heads", - "num_kv_heads", - "head_size", - "tokens_per_block", - "kv_cache_quant_mode", - "max_attention_window_size", - "cyclic_attention_window_size", - "num_tokens", - "batch_beam", - "input_seq_length", - "max_past_kv_length", - "rotary_embedding_dim", - "rotary_embedding_base", - "rotary_embedding_scale_type", - "rotary_embedding_scale", - "rotary_embedding_max_positions", - "position_embedding_type", - "bmm1_scale", - "bmm2_scale", - "fp8_context_fmha", - "predicted_tokens_per_seq", - "attention_chunk_size", - "multi_processor_count", - "total_num_blocks", - "kv_factor", - "need_build_kv_cache_metadata", - "cross_attention", - "skip_fmha_workspace", -) -_INT2_BYTES = 8 # capacity: fixed: an element size +_ARG_NAMES = _shared._ARG_NAMES +_INT2_BYTES = _shared._INT2_BYTES _original: Optional[Callable] = None _max_rows = 0 @@ -91,70 +50,20 @@ def _no_missed(rows: int, reason: str) -> None: _on_missed: Callable[[int, str], None] = _no_missed -def _indexed(device: torch.device) -> torch.device: - """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" - if device.type == "cuda" and device.index is None: - return torch.device("cuda", torch.cuda.current_device()) - return device +_indexed = _shared._indexed def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" - device = _indexed(device) - m = _maps.get(device) - if m is None or m.shape[0] < rows: - reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" - _on_missed(rows, reason) - raise RuntimeError(f"m3_tokmap_guard: {reason}") - return m[:rows] + return _shared._one_request_map(_TARGET, device, rows) def _map_offset(args: dict, rows: int) -> int: - x = args["qkv_input"] - skip = bool(args.get("skip_fmha_workspace", False)) - key = ( - x.dtype, - rows, - int(args["num_heads"]), - int(args["head_size"]), - int(args["rotary_embedding_dim"]), - int(args["num_kv_heads"]), - skip, - ) - off = _offsets.get(key) - if off is None: - layout = thop.get_trtllm_gen_generation_workspace_layout( - x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip - ) - off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) - if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: - raise RuntimeError( - f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " - f"(offset {off}, {size} bytes): the layout changed under this guard" - ) - _offsets[key] = off - return off + return _shared._map_offset(_TARGET, args, rows) def _guarded(*args, **kwargs): - a = dict(zip(_ARG_NAMES, args)) - a.update(kwargs) - rows = int(a["num_tokens"]) - if ( - a.get("spec_decoding_generation_lengths") is not None - and int(a["predicted_tokens_per_seq"]) > 1 - and int(a["batch_beam"]) == 1 - and rows > 1 - and not bool(a.get("cross_attention", False)) - ): - ws = a["workspace"] - src = _one_request_map(ws.device, rows) - off = _map_offset(a, rows) - # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. - dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] - dst.view(torch.int32).view(rows, 2).copy_(src) - _on_taken(rows) - return _original(*args, **kwargs) + return _shared._guarded(_TARGET, *args, **kwargs) def install( @@ -166,20 +75,7 @@ def install( """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this call installed the wrapper.""" - global _original, _max_rows, _on_taken, _on_missed - device = _indexed(device) - if device not in _maps or _maps[device].shape[0] < max_rows: - t = torch.arange(max_rows, dtype=torch.int32, device=device) - _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() - _max_rows = max(_max_rows, max_rows) - _on_taken = on_taken - _on_missed = on_missed - if _original is not None: - return False - _original = thop.trtllm_gen_generation_preprocess - thop.trtllm_gen_generation_preprocess = _guarded - logger.info( - "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " - f"(up to {max_rows} rows)" - ) - return True + return _shared.install(_TARGET, device, max_rows, on_taken, on_missed) + + +_TARGET = cast(_shared.TokenMapTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py index 6e3feee8882a..d811527d4fec 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_moe_engine.py @@ -12,21 +12,34 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm.logger import logger -from tensorrt_llm.quantization.utils.fp4_utils import ( - get_reorder_rows_for_gated_act_gemm_row_indices, - get_shuffle_matrix_a_row_indices, -) -from ..m3_moe.op import ( - HIDDEN, - INTER, - SHARED_INTER, - TOP_K, - M3MoeLayer, - M3MoeState, - ar_push_desc, - is_supported, -) +from ....ops.moe.reference import _E2M1 as _E2M1 +from ....ops.moe.reference import _F32_MAX as _F32_MAX +from ....ops.moe.reference import _LOG2E as _LOG2E +from ....ops.moe.reference import _RCP6 as _RCP6 +from ....ops.moe.reference import ROUTED_SCALING as ROUTED_SCALING +from ....ops.moe.reference import SHARED_ALPHA as SHARED_ALPHA +from ....ops.moe.reference import SHARED_LIMIT as SHARED_LIMIT +from ....ops.moe.reference import _bf16 as _bf16 +from ....ops.moe.reference import _e2m1_decode as _e2m1_decode +from ....ops.moe.reference import _e2m1_round as _e2m1_round +from ....ops.moe.reference import _e4m3_round as _e4m3_round +from ....ops.moe.reference import _f32 as _f32 +from ....ops.moe.reference import _fc1_epilogue as _fc1_epilogue +from ....ops.moe.reference import _gemm as _gemm +from ....ops.moe.reference import _nvfp4_act as _nvfp4_act +from ....ops.moe.reference import _nvfp4_expert as _nvfp4_expert +from ....ops.moe.reference import _row_perm as _row_perm +from ....ops.moe.reference import _shared_expert as _shared_expert_reference +from ....ops.moe.reference import _ue8m0_up as _ue8m0_up +from ....ops.moe.reference import moe_reference as _moe_reference +from ....ops.moe.reference import moe_route as moe_route +from ....ops.moe.reference import mxfp8_dequant as mxfp8_dequant +from ....ops.moe.reference import mxfp8_quant_dequant as mxfp8_quant_dequant +from ....ops.moe.reference import unswizzle_sf as unswizzle_sf +from ..m3_moe.op import HIDDEN as HIDDEN +from ..m3_moe.op import INTER, SHARED_INTER, M3MoeLayer, M3MoeState, ar_push_desc, is_supported +from ..m3_moe.op import TOP_K as TOP_K __all__ = [ "M3MoeLayer", @@ -41,204 +54,16 @@ "unswizzle_sf", ] -# The model's constants the kernels are built with: routed_scaling_factor, the shared expert's SwiGLU-OAI -# (swiglu_alpha, swiglu_limit, the (up + 1) offset). -ROUTED_SCALING = 2.0 -SHARED_ALPHA = 1.702 -SHARED_LIMIT = 7.0 -_LOG2E = 1.4426950216293334961 -_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum -_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) -_F32_MAX = 3.40282346638528859812e38 - - -def _f32(x: float) -> torch.Tensor: - return torch.tensor(x, dtype=torch.float32) - - -def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: - """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" - lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) - lo = lut[(packed & 0xF).long()] - hi = lut[(packed >> 4).long()] - return torch.stack((lo, hi), dim=-1).flatten(-2) - - -def _e2m1_round(x: torch.Tensor) -> torch.Tensor: - """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" - a = x.abs() - q = torch.full_like(a, 6.0) - for hi, value, inclusive in ( - (5.0, 4.0, True), - (3.5, 3.0, False), - (2.5, 2.0, True), - (1.75, 1.5, False), - (1.25, 1.0, True), - (0.75, 0.5, False), - (0.25, 0.0, True), - ): - q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) - return torch.where(x < 0, -q, q) - - -def _e4m3_round(x: torch.Tensor) -> torch.Tensor: - """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" - return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() - - -def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: - """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" - bits = x.contiguous().view(torch.int32) - code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) - return code.clamp(0, 254) - - -def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: - """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" - flat = sf.reshape(-1).view(torch.uint8) - r = torch.arange(rows, device=flat.device).view(-1, 1) - b = torch.arange(blocks, device=flat.device).view(1, -1) - padded_blocks = (blocks + 3) // 4 * 4 - idx = ( - (r // 128) * (128 * padded_blocks) - + (b // 4) * 512 - + (r % 32) * 16 - + ((r % 128) // 32) * 4 - + (b % 4) - ) - return flat[idx] - - -def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: - """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" - rows, k = q.shape - codes = unswizzle_sf(sf, rows, k // 32).float() - return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( - 32, dim=1 - ) - - -def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: - """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to - fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" - rows, k = x.shape - blocks = x.float().view(rows, k // 32, 32) - code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) - scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) - return (_e4m3_round(blocks / scale) * scale).view(rows, k) - - -def _bf16(x: torch.Tensor) -> torch.Tensor: - return x.to(torch.bfloat16).float() - - -def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" - return (a.double() @ b.double().t()).float() - - -def _row_perm(rows: int, gated: bool) -> torch.Tensor: - """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle - of epilogue tile 128).""" - probe = torch.empty(rows, 1) - perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) - if gated: - perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] - return perm - - -def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: - """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row - order (the global scale is not applied: it is the per-expert alpha).""" - rows, half = w.shape - codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() - vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) - return vals[perm_inv] - - -def moe_route(x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor): - """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" - logits = _gemm(x, router_weight) - sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() - keys = sig + bias - ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] - sk = keys.gather(1, ids) - bias[ids] - total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) - w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) - return ids, _bf16(w) - - -def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: - """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" - rows, k = x.shape - blocks = x.view(rows, k // 16, 16) - vmax = blocks.abs().amax(dim=-1, keepdim=True) - sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) - # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element - # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). - osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) - return (_e2m1_round(blocks * osc) * sff).view(rows, k) - - -def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: - """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 - input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with - act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" - a, c, alpha, beta, lim = (esc[i] for i in range(5)) - gc = torch.where(gate > lim, lim, gate) - gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) - uc = torch.minimum(torch.maximum(up, -lim), lim) - c1 = a * (alpha * _f32(_LOG2E).to(up.device)) - sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) - h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) - blocks = h.view(-1, 16) - amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) - sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) - osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) - return (_e2m1_round((blocks * osc) * c) * sff).view(-1) - - -def _shared_expert(x: torch.Tensor, shared_mxfp8) -> torch.Tensor: - """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the - model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" - gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 - gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) - gate, up = gu[:, :SHARED_INTER], gu[:, SHARED_INTER:] - a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) - ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 - sg = 1.0 / ( - torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 - ) - h = _bf16(ub * (sg * a)) - return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) +def _shared_expert( + x: torch.Tensor, + shared_mxfp8: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], +) -> torch.Tensor: + return _shared_expert_reference(x, shared_mxfp8, shared_inter=SHARED_INTER) def moe_reference(hidden: torch.Tensor, layer: M3MoeLayer, shared: bool) -> torch.Tensor: - """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" - x = hidden.float() - m = x.shape[0] - ids, weights = moe_route(x, layer.router_weight, layer.bias) - xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) - inv31 = torch.argsort(_row_perm(2 * INTER, gated=True)).to(x.device) - inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) - y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) - for e in torch.unique(ids).tolist(): - w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) - w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) - esc = layer.escal[:, e] - for t, k in (ids == e).nonzero().tolist(): - acc1 = _gemm(xr[t : t + 1], w31)[0] - # The buffer holds [w3 (up); w1 (gate)]. - h = _fc1_epilogue(acc1[:INTER], acc1[INTER:], esc) - y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) - acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) - for k in range(TOP_K): - acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() - routed = acc.float() - if not shared: - return routed.to(torch.bfloat16) - return (_shared_expert(x, layer.shared_mxfp8) + _bf16(routed)).to(torch.bfloat16) + return _moe_reference(hidden, layer, shared, inter=INTER, shared_inter=SHARED_INTER) _ENGAGED: set = set() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py deleted file mode 100644 index 0306210adaab..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv_fp4.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""fp32 router GEMV plus the routed experts' NVFP4 activation, one token.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .m3_router_gemv import fp32_logits, nvfp4_quantize - - -class M3RouterGemvFp4(OpWrapper): - """``[logits fp32 [M, E], x_fp4 [M, K / 2] uint8, x_sf [M * K / 16] uint8]`` from one kernel. - - The logits are ``m3_router_gemv``'s; the NVFP4 activation is ``fp4_quantize(x, sf_scale, 16, False, False)`` - (16-element blocks, E4M3 scales in the linear layout, per-tensor ``sf_scale``) for the routed experts' FC1. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # The fused-router decode step at one token: 128 experts, hidden 6144, the experts' input scale. - CELLS: tuple[Cell, ...] = ( - Cell( - why="one decode token: logits for 128 experts and the NVFP4 activation of hidden 6144", - spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 8.0), - ), - Cell( - why="a global scale that puts block scales near the E4M3 top, where the scale cast saturates", - spec=dict(m=1, e=128, k=6144, sf_scale=448.0 * 6.0 / 2.0), - ), - ) - - note = """ - The e2m1 codes are exact except within `TIE_WINDOW` of a rounding midpoint, where the kernel's approximate - reciprocals decide the direction. Up to 16 rows are accepted; the target sends one. - """ - - def __call__( - self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool - ) -> list[torch.Tensor]: - return torch.ops.trtllm.minimax_m3_router_gemv_fp4(x, w, sf_scale, pdl) - - def reference( - self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """True fp32 logits and the NVFP4 recipe; ``pdl`` does not enter the arithmetic.""" - packed, sf, _ = nvfp4_quantize(x, sf_scale) - return fp32_logits(x, w), packed, sf - - def is_valid(self, x: torch.Tensor, w: torch.Tensor, sf_scale: torch.Tensor, pdl: bool) -> None: - assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" - # A zero or non-finite global scale makes every block scale 0 or NaN and the data codes garbage; the op - # reads the scalar on the device and cannot refuse it. - value = float(sf_scale.reshape(())) - assert value > 0.0 and value != float("inf"), ( - f"sf_scale must be positive and finite, got {value}" - ) - - -m3_router_gemv_fp4 = M3RouterGemvFp4() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py deleted file mode 100644 index 8c2a67356d3f..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/msa_kv_update.py +++ /dev/null @@ -1,89 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""One kernel for MSA's device-side KV-length correction (MsaAttentionMetadata.on_update_kv_lens). - -The torch formulation is ~23 one-CTA kernels (min, cumsum, gathers, div / clamp / remainder, the block-table lookup, -the valid-block counts) and runs at the head of every captured decode step, twice. This kernel computes the same -int32 results in one launch: - - kv_true[b] = min(kv_lens[b], staged[b]) -> seq_lens_out[b], cu_kv[b + 1] = prefix sum (cu_kv[0] = 0) - pos[t] = kv_true[r] - qo[r] + intra[t] (r = q_batch_row[t]) - out_loc[t] = table[r, min(max(pos, 0) // page, ncols - 1)] * page + max(pos, 0) % page - n_valid[t] = (max(pos + 1, 0) + page - 1) // page - -``msa_kv_update`` is the target's ``MsaKvUpdate`` (``tensorrt_llm._torch.pyexecutor.runtime_hooks``): the MSA -metadata calls it with its own buffers and falls back to the torch formulation when it returns False. The kernel is -``torch.ops.trtllm.minimax_m3_msa_kv_update``, called through this target's ``msa_kv_update`` entry. -""" - -from typing import Optional - -import torch - -from ..ops.attention.msa_kv_update import msa_kv_update as _kv_update_op - - -def msa_kv_update( - kv_lens: torch.Tensor, - kv_lens_staged: torch.Tensor, - seq_lens: torch.Tensor, - q_batch_row: torch.Tensor, - q_intra: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - batch: int, - total_q: int, - msa_seq_lens: torch.Tensor, - cu_kv_lens: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - n_valid_blocks: Optional[torch.Tensor], -) -> bool: - """The fused correction (see module doc); False, with nothing written, when an operand is outside what the - kernel covers. Arguments in ``MsaKvUpdate`` order; ``seq_lens`` are the per-request query lengths (``qo``).""" - ts = [ - kv_lens, - kv_lens_staged, - seq_lens, - q_batch_row, - q_intra, - block_table, - msa_seq_lens, - out_cache_loc, - ] - ts += [t for t in (cu_kv_lens, n_valid_blocks) if t is not None] - if ( - any(t.dtype != torch.int32 or not t.is_cuda for t in ts) - or block_table.dim() != 2 - or block_table.stride(1) != 1 - ): - return False - if any(not t.is_contiguous() for t in ts if t is not block_table): - return False - if ( - kv_lens.numel() < batch - or kv_lens_staged.numel() < batch - or seq_lens.numel() < batch - or msa_seq_lens.numel() < batch - or q_batch_row.numel() < total_q - or q_intra.numel() < total_q - or out_cache_loc.numel() < total_q - or (cu_kv_lens is not None and cu_kv_lens.numel() < batch + 1) - or (n_valid_blocks is not None and n_valid_blocks.numel() < total_q) - ): - return False - _kv_update_op( - kv_lens, - kv_lens_staged, - seq_lens, - q_batch_row, - q_intra, - block_table, - tokens_per_block, - batch, - total_q, - msa_seq_lens, - cu_kv_lens, - out_cache_loc, - n_valid_blocks, - ) - return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py deleted file mode 100644 index 3f74be245ee0..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/m3_mx_gemv_kernel.py +++ /dev/null @@ -1,1237 +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. -"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. - -The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), -fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 -scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a -``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at -most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. - -A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... -through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale -atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in -one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 -partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the -SPLIT partials in rank order and stores bf16. - -Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i -reads TMEM lanes 32 i .. 32 i + 31). -""" - -from __future__ import annotations - -import cuda.bindings.driver as cuda_driver -import cutlass -import cutlass.cute as cute -import cutlass.experimental.cuda as cuda -from cutlass import dsl_user_op -from cutlass.experimental import primitives as prims - -CTA_M = 128 # capacity: fixed: a kernel tile constant (weight rows per CTA = the MMA's M) -MMA_N = 8 # token columns -CTA_K = 128 # capacity: fixed: a kernel tile constant (one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row) -MMA_INST_K = 32 -# 4 MMAs per k-tile, one per scale byte of the atom. -K_BLOCKS = CTA_K // MMA_INST_K # capacity: fixed: a kernel tile constant -THREADS = 256 -SF_ATOM_BYTES = 512 -SF_ATOM_U16 = SF_ATOM_BYTES // 2 -SF_COLS = 4 # TMEM columns of one scale atom -W_TILE_BYTES = CTA_M * CTA_K -X_TILE_BYTES = MMA_N * CTA_K -EVICT_FIRST = 0x12F0000000000000 -A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A - -a_dtype = cutlass.Float8E4M3FN -b_dtype = cutlass.Float8E4M3FN - - -def num_k_tiles(k_in: int) -> int: - return k_in // CTA_K - - -def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: - # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 - # MMAs do not form one serial dependency chain; the epilogue adds them. - return max(16, nacc * mma_n) - - -def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: - # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] - return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 - - -def tmem_cols( - ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N -) -> int: - need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS - if a_tmem: - need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA - cols = 32 - while cols < need: - cols *= 2 - return cols - - -@dsl_user_op -def read_globaltimer(*, loc=None, ip=None): - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int64( - _llvm.inline_asm( - _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): - """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], - "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): - """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" - from cutlass._mlir.dialects import llvm as _llvm - - _llvm.inline_asm( - None, - [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), - cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), - cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], - "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", - has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) # fmt: skip - - -@dsl_user_op -def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): - """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - out = _llvm.inline_asm( - _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], - "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) # fmt: skip - return [ - cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) - ] - - -@dsl_user_op -def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): - """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - done = cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" - "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - return done != cutlass.Int32(0) - - -@dsl_user_op -def _absmax(acc, v, *, loc=None, ip=None): - """max(acc, |v|).""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Float32( - _llvm.inline_asm( - _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _rcp_approx_ftz(x, *, loc=None, ip=None): - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Float32( - _llvm.inline_asm( - _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): - """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], - "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _ld_volatile_s32(addr, *, loc=None, ip=None): - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", - has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _st_global_b32(addr, value, *, loc=None, ip=None): - from cutlass._mlir.dialects import llvm as _llvm - - _llvm.inline_asm( - None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], - "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) # fmt: skip - - -@dsl_user_op -def _bf16x2_word(lo, hi, *, loc=None, ip=None): - """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" - "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" - "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" - "mov.b32 $0, w;\n\t}", "=r,f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _bf16_bits(x, *, loc=None, ip=None): - """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -def _bf16_round(x): - return _bf16_bits(x).bitcast(cutlass.Float32) - - -@dsl_user_op -def _rope_bits(c, a, s, b, *, loc=None, ip=None): - """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], - "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" - "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -def _din_load4(table, hidden, hidden_half, word_idx): - """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform - base-address select instead of a branch.""" - addr = cutlass.select_( - hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() - ) - ptr = cutlass.inttoptr( - addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 - ) - return ptr.load(count=4, alignment=16) - - -def _stamp(timers, timed: bool, bx, slot: int): - # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). - if timed: - timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) - - -def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): - # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... - if contig: - return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) - return rank + cutlass.Int32(i * split) - - -@cute.kernel -def m3_mx_gemv_kernel( - tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 - tma_desc_wsf: cutlass.GridConstant[ - cuda.TensorMap - ], # W scale atoms [N / 128 * K / 128, 256] u16 - tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 - tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 - y: cutlass.Array, # [M * N] bf16, token-major - x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words - timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` - num_tokens: cutlass.Int32, - n_out: cutlass.Constexpr[int], - k_in: cutlass.Constexpr[int], - split: cutlass.Constexpr[int], - ring: cutlass.Constexpr[int], - trigger_early: cutlass.Constexpr[bool], - contig: cutlass.Constexpr[bool], - timed: cutlass.Constexpr[bool], - nacc: cutlass.Constexpr[int], - a_tmem: cutlass.Constexpr[bool], - st_async: cutlass.Constexpr[bool], - x_bf16: cutlass.Constexpr[bool], - ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) - ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) - ar_rank: cutlass.Int32, - ar_push: cutlass.Constexpr[ - int - ], # 1: push the output rows into the next all-reduce; 2: and trigger after it - ar_world: cutlass.Constexpr[int], - ar_maxt: cutlass.Constexpr[int], - mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n - din_i: cutlass.Array, # din: int32 token ids [M] - din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words - din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words - din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words - din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] - din_vocab: cutlass.Int32, - din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words - din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words - din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res - rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] - rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] - rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] - rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 - rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 - rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale - rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] - din: cutlass.Constexpr[bool], - rope: cutlass.Constexpr[bool], -): - k_tiles = num_k_tiles(k_in) - x_tile_bytes = mma_n * CTA_K - my_tiles = k_tiles // split - rows_per_rank = CTA_M // split - tx, _, _ = cute.arch.thread_idx() - bx, _, _ = cute.arch.block_idx() - warp_id = cute.arch.warp_idx() - rank = cutlass.Int32(0) - if cutlass.const_expr(split > 1): - rank = cute.arch.block_idx_in_cluster() - tma_ptr_w = tma_desc_w.get_ptr() - tma_ptr_wsf = tma_desc_wsf.get_ptr() - tma_ptr_x = tma_desc_x.get_ptr() - tma_ptr_xsf = tma_desc_xsf.get_ptr() - n_tile = bx // cutlass.Int32(split) - m_offset = n_tile * cutlass.Int32(CTA_M) - - smem_a = cutlass.Array( - cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 - ) - smem_b = cutlass.Array( - cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 - ) - smem_sfa = cutlass.Array( - cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 - ) - smem_sfb = cutlass.Array( - cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 - ) - w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) - mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) - act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) - acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) - reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) - tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) - # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded - # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a - # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). - mb_stride = mma_n + 4 - mailbox = cutlass.Array( - cutlass.Float32, - split * rows_per_rank * mb_stride, - space=cutlass.AddressSpace.smem, - alignment=16, - ) - # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. - din_rs = None - if cutlass.const_expr(din): - din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) - - if warp_id == 0: - prims.prefetch_tensormap(tma_ptr_w) - prims.prefetch_tensormap(tma_ptr_wsf) - prims.prefetch_tensormap(tma_ptr_x) - prims.prefetch_tensormap(tma_ptr_xsf) - if prims.elect_sync(): - for s in cutlass.range_constexpr(ring): - prims.mbarrier_init(w_full.subview(s), 1) - prims.mbarrier_init(mma_done.subview(s), 1) - prims.mbarrier_init(act_full, 1) - prims.mbarrier_init(acc_done, 1) - if cutlass.const_expr(split > 1 and st_async): - # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. - prims.mbarrier_init(reduce_ready, 1) - prims.mbarrier_arrive_expect_tx( - reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 - ) - elif cutlass.const_expr(split > 1): - prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) - if warp_id == 2: - prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) - prims.tcgen05_relinquish_alloc_permit() - prims.fence_mbarrier_init() - if cutlass.const_expr(split > 1): - prims.barrier_cluster_arrive_relaxed() - prims.barrier_cluster_wait() - prims.barrier_cta_sync(0) - if tx == 0: - _stamp(timers, timed, bx, 0) - - if warp_id == 0: - if prims.elect_sync(): - for i in cutlass.range_constexpr(ring, my_tiles): - k_global = _k_tile(rank, i, split, my_tiles, contig) - prims.cp_async_bulk_tensor_prefetch( - tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] - ) - if cutlass.const_expr((n_out // CTA_M) * split >= 80): - # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which - # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's - # (bx + grid / 2) weights into L2. - if bx < cutlass.Int32((n_out // CTA_M) * split // 2): - pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) - p_rank = pb % cutlass.Int32(split) - p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) - for i in cutlass.range_constexpr(my_tiles): - p_k = _k_tile(p_rank, i, split, my_tiles, contig) - prims.cp_async_bulk_tensor_prefetch( - tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] - ) - for i in cutlass.range_constexpr(my_tiles): - stage = i % ring - if cutlass.const_expr(i >= ring): - while not cute.arch.mbarrier_try_wait( - mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 - ): - pass - k_global = _k_tile(rank, i, split, my_tiles, contig) - prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) - prims.cp_async_bulk_tensor_shared_cta_global( - smem_a.subview(stage * W_TILE_BYTES), - tma_ptr_w, - (k_global * cutlass.Int32(CTA_K), m_offset), - w_full.subview(stage), - l2_cache_hint=EVICT_FIRST, - ) - prims.cp_async_bulk_tensor_shared_cta_global( - smem_sfa.subview(stage * SF_ATOM_BYTES), - tma_ptr_wsf, - (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), - w_full.subview(stage), - l2_cache_hint=EVICT_FIRST, - ) - elif warp_id == 1: - prims.griddepcontrol(prims.GridDepAction.WAIT) - if cutlass.const_expr(trigger_early): - # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read - # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its - # wait, so the chain of releases never runs ahead of a producer still in flight. - prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) - if tx == 32: - _stamp(timers, timed, bx, 1) - if cutlass.const_expr(not x_bf16): - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) - for i in cutlass.range_constexpr(my_tiles): - k_global = _k_tile(rank, i, split, my_tiles, contig) - prims.cp_async_bulk_tensor_shared_cta_global( - smem_b.subview(i * x_tile_bytes), - tma_ptr_x, - (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), - act_full, - ) - prims.cp_async_bulk_tensor_shared_cta_global( - smem_sfb.subview(i * SF_ATOM_BYTES), - tma_ptr_xsf, - (cutlass.Int32(0), k_global), - act_full, - ) - elif warp_id == 2: - tmem_raw = tmem_ptr_i32.load() - base_col = tmem_raw & 0xFFFF - base_row = tmem_raw >> 16 - sfa_col0 = base_col + acc_cols(nacc, mma_n) - sfb_col0 = sfa_col0 + ring * SF_COLS - idesc = prims.Tcgen05MxInstrDesc.build( - a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M - ) - s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 - if cutlass.const_expr(a_tmem): - # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they - # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, - # its scale atoms and the MMAs, A read from TMEM. - base_row = tmem_raw >> 16 - base_col = tmem_raw & 0xFFFF - sfa_col0 = base_col + acc_cols(nacc, mma_n) - sfb_col0 = sfa_col0 + my_tiles * SF_COLS - a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) - s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 - for i in cutlass.range_constexpr(my_tiles): - stage = i % ring - while not cute.arch.mbarrier_try_wait( - w_full.subview(stage).data_ptr(), (i // ring) % 2 - ): - pass - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - desc_a = prims.Tcgen05SmemDesc.build( - smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, - base_offset=0, layout=2, - ) # fmt: skip - desc_sfa = prims.Tcgen05SmemDesc.build( - smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, - base_offset=0, layout=0, - ) # fmt: skip - if prims.elect_sync(): - for kb in cutlass.range_constexpr(K_BLOCKS): - prims.tcgen05_cp( - prims.Tcgen05CpShape.SHAPE_128X256B, - cutlass.inttoptr( - a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), - 6, - cutlass.Int32, - ), - desc_a + ((MMA_INST_K >> 4) * kb), - group=prims.CTAGroup.CTA_1, - ) - prims.tcgen05_cp( - s2t_shape, - cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), - desc_sfa, multicast=s2t_multicast, - ) # fmt: skip - if cutlass.const_expr(i + ring < my_tiles): - prims.tcgen05_commit(mma_done.subview(stage)) - if cutlass.const_expr(x_bf16): - # Completed by a thread's arrive: a suspended try_wait wakes late on those. - while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): - pass - else: - while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): - pass - if tx == 64: - _stamp(timers, timed, bx, 2) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - if prims.elect_sync(): - for i in cutlass.range_constexpr(my_tiles): - desc_sfb = prims.Tcgen05SmemDesc.build( - smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, - base_offset=0, layout=0, - ) # fmt: skip - desc_b = prims.Tcgen05SmemDesc.build( - smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, - base_offset=0, layout=2, - ) # fmt: skip - sfa_tmem = cutlass.inttoptr( - (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 - ) - sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) - prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) - for kb in cutlass.range_constexpr(K_BLOCKS): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, - cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), - cutlass.inttoptr( - a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 - ), - desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), - i > 0 or kb > 0, - sfa_tmem, sfb_tmem, - ) # fmt: skip - if tx == 64: - _stamp(timers, timed, bx, 3) - if prims.elect_sync(): - prims.tcgen05_commit(acc_done) - else: - if cutlass.const_expr(x_bf16): - # Completed by a thread's arrive: a suspended try_wait wakes late on those. - while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): - pass - else: - while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): - pass - if tx == 64: - _stamp(timers, timed, bx, 2) - for i in cutlass.range_constexpr(my_tiles): - stage = i % ring - while not cute.arch.mbarrier_try_wait( - w_full.subview(stage).data_ptr(), (i // ring) % 2 - ): - pass - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) - sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) - desc_sfa = prims.Tcgen05SmemDesc.build( - smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, - base_offset=0, layout=0, - ) # fmt: skip - desc_sfb = prims.Tcgen05SmemDesc.build( - smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, - base_offset=0, layout=0, - ) # fmt: skip - desc_a = prims.Tcgen05SmemDesc.build( - smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, - base_offset=0, layout=2, - ) # fmt: skip - desc_b = prims.Tcgen05SmemDesc.build( - smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, - base_offset=0, layout=2, - ) # fmt: skip - if prims.elect_sync(): - prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, - multicast=s2t_multicast) # fmt: skip - prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, - multicast=s2t_multicast) # fmt: skip - for kb in cutlass.range_constexpr(K_BLOCKS): - inc = ((MMA_INST_K * 8 // 8) >> 4) * kb - j = i * K_BLOCKS + kb - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, - cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), - desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, - cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), - ) # fmt: skip - prims.tcgen05_commit(mma_done.subview(stage)) - if tx == 64: - _stamp(timers, timed, bx, 3) - if prims.elect_sync(): - prims.tcgen05_commit(acc_done) - elif warp_id >= 4: - lane = tx % 32 - row = (warp_id % 4) * 32 + lane - if cutlass.const_expr(x_bf16): - # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as - # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), - # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, - # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. - qz_tid = tx - cutlass.Int32(128) - prims.griddepcontrol(prims.GridDepAction.WAIT) - din_half_k = k_in // 2 - din_hidden_half = cutlass.Int32(0) - if cutlass.const_expr(din): - # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to - # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 - # threads per token, 16-byte loads, all of a thread's loads in flight at once. - din_hidden_half = cutlass.select_( - rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), - cutlass.Int32(1), - cutlass.Int32(0), - ) - din_tok = qz_tid // cutlass.Int32(16) - din_part = qz_tid % cutlass.Int32(16) - din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) - din_id = din_i.load(idx=din_tok_ld) - din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) - din_hrow = din_tok_ld - if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): - din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) - din_row = cutlass.select_( - din_hidden_half == cutlass.Int32(1), - din_hrow, - cutlass.select_(din_valid, din_id, cutlass.Int32(0)), - ) - # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. - din_copy = ( - ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) - & (bx == cutlass.Int32(split // 2)) - & (din_hidden_half == cutlass.Int32(1)) - & (din_tok < num_tokens) - ) - din_base = din_row * cutlass.Int32(din_half_k // 2) - din_ss = cutlass.Float32(0.0) - din_vecs_per_part = din_half_k // 8 // 16 - for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): - din_n = min(16, din_vecs_per_part - din_round) - din_loaded = [] - for din_j in cutlass.range_constexpr(din_n): - din_widx = din_base + ( - din_part + cutlass.Int32((din_round + din_j) * 16) - ) * cutlass.Int32(4) - din_loaded.append( - _din_load4(din_table, din_hidden, din_hidden_half, din_widx) - ) - if din_copy: - for din_j in cutlass.range_constexpr(din_n): - din_v = din_loaded[din_j] - din_res.store( - (din_v[0], din_v[1], din_v[2], din_v[3]), - idx=din_tok * cutlass.Int32(din_half_k // 2) - + (din_part + cutlass.Int32((din_round + din_j) * 16)) - * cutlass.Int32(4), - alignment=16, - ) - for din_j in cutlass.range_constexpr(din_n): - for din_e in cutlass.range_constexpr(4): - din_word = cutlass.Int32(din_loaded[din_j][din_e]) - din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) - din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) - din_ss = din_ss + din_lo * din_lo + din_hi * din_hi - for din_o in cutlass.range_constexpr(4): - din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) - if din_part == cutlass.Int32(0): - din_eps = din_f.load(idx=din_hidden_half) - din_r = cute.math.rsqrt( - din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True - ) - din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) - din_rs.store( - cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok - ) - prims.barrier_cta_sync(3, thread_count=128) - qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, - alignment=16) # fmt: skip - qz_items = my_tiles * K_BLOCKS * mma_n - qz_rounds = (qz_items + 127) // 128 - # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of - # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. - QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 - qz_loaded = [] - for qz_r in cutlass.range_constexpr(qz_rounds): - qz_item = qz_tid + cutlass.Int32(qz_r * 128) - qz_t = qz_item % cutlass.Int32(mma_n) - qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) - qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) - if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): - qz_i = cutlass.select_( - qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) - ) - qz_k_global = rank + qz_i * cutlass.Int32(split) - if cutlass.const_expr(contig): - qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i - qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) - qz_x_base = ( - qz_t_ld * cutlass.Int32(k_in // 2) - + qz_k_global * cutlass.Int32(CTA_K // 2) - + qz_kb * cutlass.Int32(MMA_INST_K // 2) - ) - qz_wvecs = None - if cutlass.const_expr(din): - # The draft input's raw values (table row of the id / hidden row) and its norm weights. - qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) - qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) - qz_id = din_i.load(idx=qz_t_ld) - qz_hrow = qz_t_ld - if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): - qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) - qz_row = cutlass.select_( - din_hidden_half == cutlass.Int32(1), - qz_hrow, - cutlass.select_( - (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), - qz_id, - cutlass.Int32(0), - ), - ) - qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) - qz_vecs = [ - _din_load4( - din_table, - din_hidden, - din_hidden_half, - qz_src_w + cutlass.Int32(qz_v * 4), - ) - for qz_v in range(4) - ] - qz_wvecs = [ - din_w.load( - idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), - vector_size=4, - alignment=16, - ) - for qz_v in range(4) - ] - else: - qz_vecs = [ - x_words.load( - idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 - ) - for qz_v in range(4) - ] - qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) - if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): - for qz_g in cutlass.range_constexpr(len(qz_loaded)): - qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] - qz_vals = [] - qz_amax = cutlass.Float32(0.0) - if cutlass.const_expr(din): - qz_rs = din_rs.load( - idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) - ) - for qz_v in cutlass.range_constexpr(4): - for qz_e4 in cutlass.range_constexpr(4): - qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) - qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) - qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) - if cutlass.const_expr(din): - # bf16(x * rs * w), the draft-input kernel's rounding. - qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) - qz_lo = _bf16_round( - qz_lo - * qz_rs - * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) - ) - qz_hi = _bf16_round( - qz_hi - * qz_rs - * (qz_wword & cutlass.Int32(-65536)).bitcast( - cutlass.Float32 - ) - ) - qz_vals.append(qz_lo) - qz_vals.append(qz_hi) - qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) - qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( - cutlass.Int32 - ) - qz_e = ( - (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) - ) + cutlass.select_( - (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), - cutlass.Int32(1), - cutlass.Int32(0), - ) - qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) - qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( - cutlass.Float32 - ) - qz_scale = cutlass.select_( - qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) - ) - qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, - qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) - for qz_j in range(8)] # fmt: skip - if qz_item < cutlass.Int32(qz_items): - if qz_t < num_tokens: - qz_row_w = qz_i * cutlass.Int32( - x_tile_bytes // 4 - ) + qz_t * cutlass.Int32(CTA_K // 4) - qz_c0 = qz_kb * cutlass.Int32(2) - qz_b_words.store( - (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), - idx=qz_row_w - + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), - alignment=16, - ) - qz_b_words.store( - (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), - idx=qz_row_w - + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) - * cutlass.Int32(4), - alignment=16, - ) - smem_sfb.store( - cutlass.Uint8(qz_e), - idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) - + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) - + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) - + qz_kb, - ) - qz_loaded = [] - # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - prims.barrier_cta_sync(2, thread_count=128) - if qz_tid == cutlass.Int32(0): - prims.mbarrier_arrive(act_full) - if cutlass.const_expr(st_async): - while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): - pass - else: - while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): - pass - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. - acc_all = [] - ld_col = 0 - for ld_w in (128, 64, 32, 16, 8): - if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): - piece = prims.tcgen05_ld( - "32x32b", - cutlass.inttoptr( - tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 - ), - num=ld_w, - ) - acc_all = acc_all + [piece[c] for c in range(ld_w)] - ld_col += ld_w - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - acc = [] - for t in cutlass.range_constexpr(mma_n): - v = cutlass.Float32(acc_all[t]) - for a in cutlass.range_constexpr(1, nacc): - v = v + cutlass.Float32(acc_all[a * mma_n + t]) - acc.append(v) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - if tx == 128: - _stamp(timers, timed, bx, 4) - n = m_offset + row - if cutlass.const_expr(split == 1): - for t in cutlass.range_constexpr(mma_n): - if cutlass.Int32(t) < num_tokens: - y.store( - cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n - ) - if cutlass.const_expr(ar_push > 0): - prims.griddepcontrol(prims.GridDepAction.WAIT) - # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are - # issued before the first store so they overlap (one L2 round trip, not one per token). - ar_par = [ - _ld_volatile_s32( - ar_flags - + cutlass.Int64(4) - * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) - ) - & cutlass.Int32(1) - for t in range(mma_n) - ] - for t in cutlass.range_constexpr(mma_n): - ar_v = cutlass.Float32(acc[t]) - ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) - if (lane % cutlass.Int32(2)) == cutlass.Int32(0): - if cutlass.Int32(t) < num_tokens: - ar_off = ( - cutlass.Int64(ar_par[t]) - * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) - + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) - * cutlass.Int64(n_out // 2) - + (cutlass.Int64(n) // cutlass.Int64(2)) - ) - _st_global_b32( - ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) - ) - else: - owner = row // cutlass.Int32(rows_per_rank) - local = row % cutlass.Int32(rows_per_rank) - if cutlass.const_expr(rope): - # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. - pairs = rows_per_rank // 2 - owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) - local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( - row % cutlass.Int32(64) - ) % cutlass.Int32(pairs) - if owner != rank: - base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) - if cutlass.const_expr(st_async): - dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) - mb = _mapa_u32(reduce_ready.data_ptr(), owner) - words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] - for g in cutlass.range_constexpr(mma_n // 4): - _st_async_v4( - dst + cutlass.Int32(16 * g), - words[4 * g], - words[4 * g + 1], - words[4 * g + 2], - words[4 * g + 3], - mb, - ) - else: - for t in cutlass.range_constexpr(mma_n): - prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( - cutlass.Float32(acc[t]) - ) - prims.mbarrier_arrive( - prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER - ) - else: - if cutlass.const_expr(st_async): - while not _test_wait_cluster(reduce_ready.data_ptr(), 0): - pass - else: - while not prims.mbarrier_try_wait_parity( - reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER - ): - pass - if lane == 0: - if row // cutlass.Int32(rows_per_rank) == rank: - if (row % cutlass.Int32(rows_per_rank)) == 0: - _stamp(timers, timed, bx, 5) - totals = [] - for g in cutlass.range_constexpr(mma_n // 4): - sums = [cutlass.Float32(0.0)] * 4 - for q in cutlass.range_constexpr(split): - parts = _ld_shared_v4_f32( - mailbox.subview( - (cutlass.Int32(q * rows_per_rank) + local) - * cutlass.Int32(mb_stride) - + cutlass.Int32(4 * g) - ).data_ptr() - ) - sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), - parts[j]) for j in range(4)] # fmt: skip - for j in cutlass.range_constexpr(4): - totals.append(sums[j]) - if cutlass.const_expr(rope): - # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. - mailbox.store( - sums[j], - idx=(rank * cutlass.Int32(rows_per_rank) + local) - * cutlass.Int32(mb_stride) - + cutlass.Int32(4 * g + j), - ) - elif cutlass.Int32(4 * g + j) < num_tokens: - y.store( - cutlass.BFloat16(sums[j]), - idx=cutlass.Int32((4 * g + j) * n_out) + n, - ) - if cutlass.const_expr(ar_push > 0): - # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. - prims.griddepcontrol(prims.GridDepAction.WAIT) - ar_par = [ - _ld_volatile_s32( - ar_flags - + cutlass.Int64(4) - * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) - ) - & cutlass.Int32(1) - for t in range(mma_n) - ] - for t in cutlass.range_constexpr(mma_n): - ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) - if (lane % cutlass.Int32(2)) == cutlass.Int32(0): - if cutlass.Int32(t) < num_tokens: - ar_off = ( - cutlass.Int64(ar_par[t]) - * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) - + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) - * cutlass.Int64(n_out // 2) - + (cutlass.Int64(n) // cutlass.Int64(2)) - ) - _st_global_b32( - ar_ws + ar_off * cutlass.Int64(4), - _bf16x2_word(totals[t], ar_hi), - ) - prims.barrier_cta_sync(1, thread_count=128) - if cutlass.const_expr(rope): - # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over - # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to - # rope_q, K / V into the page of that position. - pairs = rows_per_rank // 2 - rp_e = tx - cutlass.Int32(128) - rp_p = rp_e % cutlass.Int32(pairs) - rp_t = rp_e // cutlass.Int32(pairs) - if rp_t < num_tokens: - rp_max_blocks = rope_dims.load(idx=0) - rp_tpb_log2 = rope_dims.load(idx=1) - rp_hq = rope_dims.load(idx=2) - rp_hkv = rope_dims.load(idx=3) - rp_T = rope_dims.load(idx=4) - rp_d = rank * cutlass.Int32(pairs) + rp_p - rp_a = _bf16_round( - mailbox.load( - idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) - + rp_t - ) - ) - rp_b = _bf16_round( - mailbox.load( - idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) - * cutlass.Int32(mb_stride) - + rp_t - ) - ) - rp_req = rp_t // rp_T - rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) - rp_head = n_tile - if rp_head < rp_hq + rp_hkv: - rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) - rp_s = rope_cs.load( - idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) - + cutlass.Int32(1) - ) - rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) - rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) - rp_a = rp_ra - rp_b = rp_rb - rp_qs = rope_scale.load(idx=0) - rp_w = _e4m3x4( - rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) - ) - rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) - rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) - if rp_head < rp_hq: - rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d - rope_q.store(rp_lo, idx=rp_q) - rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) - else: - rp_kv = cutlass.select_( - rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) - ) - rp_kh = rp_head - rp_hq - rp_kv * rp_hkv - rp_page = rope_bt.load( - idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks - + (rp_tok >> rp_tpb_log2) - ) - # 64-bit byte offset: the pool view can exceed 2^31 bytes. - rp_dst = ( - (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) - << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) - ) + cutlass.Int64( - (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) - * cutlass.Int32(128) - + rp_d - ) - rp_addr = rope_pool.data_ptr(0).toint() + rp_dst - cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) - cutlass.inttoptr( - rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 - ).store(rp_hi) - if tx == 128: - _stamp(timers, timed, bx, 6) - if cutlass.const_expr(ar_push == 2): - # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. - prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) - if warp_id == 4: - prims.tcgen05_dealloc( - cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), - tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), - ) - - -def _bytes_map(t, cols, rows, box_rows, swizzle): - return cuda.create_tensor_map_tiled( - global_address=t.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[cols, rows], - global_strides=[cols // 16], - box_dims=[CTA_K, box_rows], - swizzle=swizzle, - ) - - -def _atom_map(t, atoms): - return cuda.create_tensor_map_tiled( - global_address=t.iterator.toint(), - dtype=cutlass.Uint16, - global_dims=[SF_ATOM_U16, atoms], - global_strides=[SF_ATOM_BYTES // 16], - box_dims=[SF_ATOM_U16, 1], - swizzle=cuda.TensorMapSwizzle.none, - ) - - -@cute.jit -def m3_mx_gemv( - w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous - wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 - x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 - xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 - y: cute.Tensor, # [M * N] bf16 - x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) - timers: cute.Tensor, # [grid * 8] int64, written when `timed` - num_tokens: cutlass.Int32, - n_out: cutlass.Constexpr[int], - k_in: cutlass.Constexpr[int], - split: cutlass.Constexpr[int], - ring: cutlass.Constexpr[int], - trigger_early: cutlass.Constexpr[bool], - contig: cutlass.Constexpr[bool], - timed: cutlass.Constexpr[bool], - nacc: cutlass.Constexpr[int], - a_tmem: cutlass.Constexpr[bool], - st_async: cutlass.Constexpr[bool], - x_bf16: cutlass.Constexpr[bool], - ar_ws: cutlass.Int64, - ar_flags: cutlass.Int64, - ar_rank: cutlass.Int32, - ar_push: cutlass.Constexpr[int], - ar_world: cutlass.Constexpr[int], - ar_maxt: cutlass.Constexpr[int], - mma_n: cutlass.Constexpr[int], - use_pdl: cutlass.Constexpr[bool], - stream: cuda_driver.CUstream, - din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) - din_table: cute.Tensor, - din_hidden: cute.Tensor, - din_w: cute.Tensor, - din_f: cute.Tensor, - din_vocab: cutlass.Int32, - din_g: cute.Tensor, - din_res: cute.Tensor, - din_flags: cutlass.Int32, - rope_cs: cute.Tensor, - rope_i: cute.Tensor, - rope_bt: cute.Tensor, - rope_q: cute.Tensor, - rope_pool: cute.Tensor, - rope_scale: cute.Tensor, - rope_dims: cute.Tensor, - din: cutlass.Constexpr[bool], - rope: cutlass.Constexpr[bool], -) -> None: - k_tiles = num_k_tiles(k_in) - tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) - tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) - tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) - tma_desc_xsf = _atom_map(xsf, k_tiles) - m3_mx_gemv_kernel( - tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, - trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, - ar_maxt, mma_n, - din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, - rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, - ).launch( - grid=((n_out // CTA_M) * split, 1, 1), - block=(THREADS, 1, 1), - cluster=(split, 1, 1), - stream=stream, - use_pdl=use_pdl, - ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py index 3124827082df..421e7b45abfa 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/mx_gemv/op.py @@ -21,12 +21,19 @@ from __future__ import annotations +import sys import threading -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Tuple, cast import torch -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + mxfp8_core as _core, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + mxfp8_host as _host, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl as get_env_enable_pdl MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) MMA_WIDTHS = ( @@ -39,16 +46,11 @@ def _mma_n(m: int) -> int: - return next(w for w in MMA_WIDTHS if m <= w) + return _host.gemv_mma_n(m, mma_widths=MMA_WIDTHS) def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: - return ( - ring * (128 * 128 + 512) - + tiles * (mma_n * 128 + 512) - + split * (128 // split) * (mma_n + 4) * 4 - + 1024 - ) + return _host.gemv_smem_bytes(ring, tiles, split, mma_n) _lock = threading.Lock() @@ -93,17 +95,12 @@ def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: def _dummy_timers(device) -> torch.Tensor: - """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" - t = _DUMMY.get(device) - if t is None: - t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) - return t + "Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on." + return _host.gemv_dummy_timers(device, dummy=_DUMMY) def _arg(t: torch.Tensor, align: int = 16): - from cutlass.cute.runtime import from_dlpack - - return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + return _host.gemv_arg(t, align) # (rule 1b, deterministic) The decode projections at up to STRIDED_MAX_ROWS rows split K over strided k-tiles per @@ -118,38 +115,32 @@ def _arg(t: torch.Tensor, align: int = 16): def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: - cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None - if cfg is None and m is not None and m > 8: - cfg = _CONFIGS_WIDE.get((n, k)) - cfg = cfg or _CONFIGS.get((n, k)) - # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. - if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: - return None - if ( - STRIDED_K_TILES - and m is not None - and m <= STRIDED_MAX_ROWS - and (n, k) in _STRIDED - and len(cfg) == 2 - ): - cfg = cfg + (0,) - return cfg + return _host.gemv_get_config( + n, + k, + m, + configs_by_m=_CONFIGS_BY_M, + configs_wide=_CONFIGS_WIDE, + configs=_CONFIGS, + strided=STRIDED_K_TILES, + strided_max_rows=STRIDED_MAX_ROWS, + strided_shapes=_STRIDED, + ) def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: - """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" - split = cfg[0] - if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: - return True - tiles = k // 128 // split - mma_n = _mma_n(m) - # 17..32 rows at K < 1024 decline (an unsplit K below 1024 is slower than the FlashInfer GEMM there); no TP2 - # projection has K < 1024. - # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; - # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. - if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): - return True - return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + "Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds." + return _host.gemv_declines( + n, + k, + m, + x_bf16, + cfg, + max_tokens=MAX_TOKENS, + smem_limit=_SMEM_LIMIT, + select_mma_n=_mma_n, + smem_bytes=_smem_bytes, + ) def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: @@ -172,137 +163,43 @@ def _resolve( ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the rank are runtime arguments.""" - cfg = cfg or get_config(n, k, m) - if cfg is None or _declines(n, k, m, x_bf16, cfg): - return None - split, ring = cfg[0], cfg[1] - contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG - nacc = int(cfg[3]) if len(cfg) > 3 else _NACC - a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM - st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC - if nacc != 1: - # More than one TMEM accumulator is broken on SM107 (783556 / 783557 / 783579: nondeterministic or garbage - # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. - return None - ring = min(ring, k // 128 // split) - tiles = k // 128 // split - mma_n = _mma_n(m) - while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: - ring -= 1 - if a_tmem: - from .m3_mx_gemv_kernel import tmem_cols - - if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: - a_tmem = False - # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into - # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). - ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 - if ar is not None and split <= 4 and n % 2 == 0: - ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar - if ar_mode == 2: - trigger_early = False - if use_pdl is None: - use_pdl = get_env_enable_pdl() - # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. - key = ( + return _core._resolve( + _CORE_TARGET, + m, n, k, - split, - ring, - trigger_early, - contig, - False, - nacc, - a_tmem, - st_async, x_bf16, - ar_mode, - ar_world, - ar_maxt, - mma_n, + trigger_early, + cfg, use_pdl, + ar, + has_swiglu=False, + has_producer=False, ) - return key, (ar_ws, ar_flags, ar_rank) def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" - import cutlass.cute as cute - - from . import m3_mx_gemv_kernel as kernel - - ( - n, - k, - split, - ring, - trigger_early, - contig, - timed, - nacc, - a_tmem, - st_async, - x_bf16, - ar_mode, - ar_world, - ar_maxt, - mma_n, - use_pdl, - ) = key - ar_ws, ar_flags, ar_rank = runtime_ar - with _lock: - fn = _compiled.get(key) - if fn is None: - fn = _compiled[key] = cute.compile( - kernel.m3_mx_gemv, - *args, - m, - n, - k, - split, - ring, - trigger_early, - contig, - timed, - nacc, - a_tmem, - st_async, - x_bf16, - ar_ws, - ar_flags, - ar_rank, - ar_mode, - ar_world, - ar_maxt, - mma_n, - use_pdl, - stream, - *extra, - False, - False, - ) - return fn + return _core._compile( + _CORE_TARGET, key, args, m, runtime_ar, stream, extra, has_swiglu=False, has_producer=False + ) def _views( x: torch.Tensor, xsf: Optional[torch.Tensor], weight, weight_scale, y, x_bf16: bool ) -> tuple: - k = x.shape[1] - if x_bf16: - x_bytes = x.view(torch.uint8) - xsf = _placeholder_scales(k, x.device) - x_words = x.view(torch.int32) - else: - x_bytes = x.view(torch.uint8) - x_words = _placeholder_words(x.device) - return ( - _arg(weight.view(torch.uint8)), - _arg(weight_scale.view(torch.uint8).reshape(-1)), - _arg(x_bytes), - _arg(xsf.view(torch.uint8).reshape(-1)), - _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), - _arg(x_words.reshape(-1)), - _arg(_dummy_timers(x.device)), + return _host.gemv_views( + x, + xsf, + weight, + weight_scale, + y, + x_bf16, + False, + placeholder_scales=_placeholder_scales, + placeholder_words=_placeholder_words, + arg=_arg, + dummy_timers=_dummy_timers, ) @@ -317,40 +214,20 @@ def _run( use_pdl: Optional[bool], ar: Optional[tuple] = None, ) -> Optional[torch.Tensor]: - global last_pushed, last_uncompiled - last_pushed = False - last_uncompiled = None - import cuda.bindings.driver as cuda_driver - - m, k = x.shape - n = weight.shape[0] - if weight.shape[1] != k: - return None - resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar) - if resolved is None: - return None - key, runtime_ar = resolved - fn = _compiled.get(key) - if fn is None: - if _frozen[0]: - # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller - # reports the miss. - last_uncompiled = key - return None - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "m3_mx_gemv must run once per shape outside CUDA-graph capture first" - ) - y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) - args = _views(x, xsf, weight, weight_scale, y, x_bf16) - extra = _extra_args(x.device) - stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) - if fn is None: - fn = _compile(key, args, m, runtime_ar, stream, extra) - ar_ws, ar_flags, ar_rank = runtime_ar - fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) - last_pushed = key[11] > 0 - return y + return _core._run( + _CORE_TARGET, + x, + xsf, + weight, + weight_scale, + x_bf16, + trigger_early, + cfg, + use_pdl, + ar, + has_swiglu=False, + has_producer=False, + ) def precompile( @@ -364,26 +241,17 @@ def precompile( """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" - import cuda.bindings.driver as cuda_driver - - n, k = weight.shape - push = None if ar is None else (0, 0, 0, *ar) - resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push) - if resolved is None: - return None - key, runtime_ar = resolved - if key in _compiled: - return key - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") - dev = weight.device - x = torch.zeros(m, k, dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, device=dev) - xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) - y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) - args = _views(x, xsf, weight, weight_scale, y, x_bf16) - stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) - _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) - return key + return _core.precompile( + _CORE_TARGET, + weight, + weight_scale, + m, + x_bf16, + use_pdl, + ar, + has_swiglu=False, + has_producer=False, + ) def freeze(frozen: bool = True) -> None: @@ -401,13 +269,8 @@ def is_compiled(key: tuple) -> bool: def _extra_args(device) -> tuple: - """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" - i32 = _placeholder_words(device) - f32 = i32.view(torch.float32) - u8 = i32.view(torch.uint8) - d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) - r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) - return d + r + "Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off)." + return _host.gemv_extra_args(device, placeholder_words=_placeholder_words, arg=_arg) def m3_mx_gemv( @@ -450,47 +313,25 @@ def m3_mx_gemv_bf16( def _placeholder_scales(k: int, device) -> torch.Tensor: - key = ("sf", k, device) - t = _PLACEHOLDERS.get(key) - if t is None: - t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) - return t + return _host.gemv_placeholder_scales(k, device, placeholders=_PLACEHOLDERS) def _placeholder_words(device) -> torch.Tensor: - key = ("w", device) - t = _PLACEHOLDERS.get(key) - if t is None: - t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) - return t + return _host.gemv_placeholder_words(device, placeholders=_PLACEHOLDERS) def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: - """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" - if ( - weight.dim() != 2 - or weight.dtype != torch.float8_e4m3fn - or get_config(*weight.shape) is None - ): - return - n, k = weight.shape - if (n, k, use_pdl) in _ensured: - return - _ensured.add((n, k, use_pdl)) - xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) - xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) - m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) - for width in MMA_WIDTHS[1:]: - m3_mx_gemv( - torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), - xsf, - weight, - weight_scale, - use_pdl=use_pdl, - ) - m3_mx_gemv_bf16( - torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + "Compiles the kernel of this weight's shape outside capture (a one-token dummy call)." + return _host.gemv_ensure_compiled( weight, weight_scale, - use_pdl=use_pdl, + use_pdl, + get_config=get_config, + ensured=_ensured, + mma_widths=MMA_WIDTHS, + mx_gemv=m3_mx_gemv, + mx_gemv_bf16=m3_mx_gemv_bf16, ) + + +_CORE_TARGET = cast(_core.GemvTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py index 94827619dad2..fe319112c436 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/norm/m3_embedding_firstnorm.py @@ -9,6 +9,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.decode_embedding import ( + embedding_firstnorm_reference as _embedding_reference, +) from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm @@ -86,13 +89,7 @@ def reference( use_pdl: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" - valid = (ids >= 0) & (ids < table.shape[0]) - rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( - table.dtype - ) - x = rows.double() - rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) - return (x * rstd * weight.double()).to(torch.bfloat16), rows + return _embedding_reference(ids, table, weight, eps) m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py index e322e1787230..4ce54e9a1d34 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_silu.py @@ -14,7 +14,7 @@ class M3DraftSilu(OpWrapper): """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: - trtllm::minimax_m3_sm107_tp2_silu_and_mul (minimaxM3Sm107Tp2/siluAndMul.cu), bit for bit the Triton + trtllm::minimax_m3_sm107_tp2_silu_and_mul (minimaxM3Sm107/siluAndMul.cu), bit for bit the Triton ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs the Triton op).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py index 708e6c5caa71..3b807ff2b0bb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_step.py @@ -6,19 +6,23 @@ from __future__ import annotations -from typing import Optional +import sys +from typing import Optional, cast import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, + true_fp32_matmul as true_fp32_matmul, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + decode as decode_coded, # noqa: F401 — used by the bound family implementation +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec import ( + draft_head_reference as _draft_head_shared, ) -from ..gemm.m3_target_head import decode as decode_coded from .m3_shard_accept import exchange_tokens, local_argmax, make_keys #: Requests (head rows) per pass, and exchange rows per call half of the workspace. @@ -37,18 +41,12 @@ def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). """ - n, k = weight.shape - assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) - q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() - return q.view(n // 16, k // 16, 256) + return _draft_head_shared.bf16_fragments(_DRAFT_HEAD_TARGET, weight) def from_fragments(wf: torch.Tensor) -> torch.Tensor: """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" - tiles, ks, _ = wf.shape - return ( - wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) - ) + return _draft_head_shared.from_fragments(_DRAFT_HEAD_TARGET, wf) def head_weight( @@ -61,10 +59,9 @@ def head_weight( ) -> torch.Tensor: """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the 12-bit coded stream with its exception lists decoded.""" - if exc_off is None: - return from_fragments(wf) - n = (exc_off.numel() - 1) * 16 - return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + return _draft_head_shared.head_weight( + _DRAFT_HEAD_TARGET, wf, k, exc_off, exc_pos, exc_w, exp_base + ) def shard_logits( @@ -77,10 +74,9 @@ def shard_logits( ) -> torch.Tensor: """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, rounded once to BF16 (the BF16 GEMM's output), as fp32.""" - w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) - with true_fp32_matmul(): - logits = x.float() @ w.float().t() - return logits.to(torch.bfloat16).float() + return _draft_head_shared.shard_logits( + _DRAFT_HEAD_TARGET, x, wf, exc_off, exc_pos, exc_w, exp_base + ) class M3DraftStep(OpWrapper): @@ -430,3 +426,5 @@ def is_valid( m3_draft_step = M3DraftStep() + +_DRAFT_HEAD_TARGET = cast(_draft_head_shared.DraftHeadTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py deleted file mode 100644 index a192e89c5c76..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_force_vec.py +++ /dev/null @@ -1,136 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from ..m3_force_vec import force_vec_kernel, force_vec_req_kernel - - -class M3ForceVec(OpWrapper): - """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write - ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` - (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 - ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). - - Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) - runs it. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), - Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), - ) - - note = """ - Integer arithmetic and one fp32 comparison: `reference` must match exactly. - """ - - def __call__( - self, - pool, - counter, - out, - batch_size, - base_total, - max_total, - frac, - counter_stride, - slot_stride, - pool_mask, - block, - ): - force_vec_kernel[(1,)]( - pool, - counter, - out, - batch_size, - base_total, - max_total, - frac, - counter_stride, - slot_stride, - pool_mask, - BLOCK=block, - ) - return out - - def reference( - self, - pool, - counter, - out, - batch_size, - base_total, - max_total, - frac, - counter_stride, - slot_stride, - pool_mask, - block, - ): - """The counts (a new tensor; ``counter`` and ``out`` are not written).""" - c = int(counter.item()) + 1 - slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) - idx = (c * counter_stride + slots * slot_stride) & pool_mask - r = pool[idx] - return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) - - -m3_force_vec = M3ForceVec() - - -class M3ForceVecReq(OpWrapper): - """For each of the ``batch_size`` requests ``b`` write - ``min(base_total + (pool[(keys[b] * key_stride) & pool_mask] < frac), max_total)`` into ``out`` (int32): - one draw per request at its own key (its KV length at the verify), no shared counter. ``block`` is the launch - width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). - - Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) - runs it. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), - Cell(why="two requests, fraction 0.3", spec=dict(batch=2, block=16, frac=0.3)), - ) - - note = """ - Integer arithmetic and one fp32 comparison: `reference` must match exactly. - """ - - def __call__( - self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block - ): - force_vec_req_kernel[(1,)]( - pool, - keys, - out, - batch_size, - base_total, - max_total, - frac, - key_stride, - pool_mask, - BLOCK=block, - ) - return out - - def reference( - self, pool, keys, out, batch_size, base_total, max_total, frac, key_stride, pool_mask, block - ): - """The counts (a new tensor; ``out`` is not written).""" - k = keys[:batch_size].to(torch.int64) - r = pool[(k * key_stride) & pool_mask] - return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) - - -m3_force_vec_req = M3ForceVecReq() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py index 34fb3387b996..d9f9b9cc989c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_sampled_accept.py @@ -5,14 +5,18 @@ from __future__ import annotations -from typing import Optional +import sys +from typing import Optional, cast import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec import ( + sampling_reference as _sampling_shared, +) -from .m3_shard_accept import strict_accept +from .m3_shard_accept import strict_accept as strict_accept #: Verify rows per call. MAX_ROWS = 48 @@ -25,25 +29,7 @@ def nucleus_probs( ) -> torch.Tensor: """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" - inv = 1.0 / (temperatures.float() + TEMP_EPSILON) - x = logits.float() * inv[:, None] - w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() - z = w.sum(dim=1, keepdim=True) - order = torch.sort(w, dim=1, descending=True) - sw = order.values - # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group - csum = torch.cumsum(sw, dim=1) - sw - first = torch.ones_like(sw, dtype=torch.bool) - first[:, 1:] = sw[:, 1:] != sw[:, :-1] - group = torch.cumsum(first.long(), dim=1) - 1 - starts = torch.zeros_like(csum) - starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) - above = torch.gather(starts, 1, group) - keep_sorted = above <= top_ps.double()[:, None] * z - keep = torch.zeros_like(keep_sorted) - keep.scatter_(1, order.indices, keep_sorted) - p = torch.where(keep, w, torch.zeros_like(w)) - return (p / p.sum(dim=1, keepdim=True)).float() + return _sampling_shared.nucleus_probs(_REFERENCE_TARGET, logits, temperatures, top_ps) class M3SampledAccept(OpWrapper): @@ -239,13 +225,34 @@ def reference( """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" - rows = logits.shape[0] - probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) - - def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: - return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) - - return {"probs": probs, "accept": accept} + return _sampling_shared.sampled_reference( + self, + _REFERENCE_TARGET, + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + top_ks, + min_ps, + ) def is_valid( self, @@ -274,35 +281,36 @@ def is_valid( top_ks: Optional[torch.Tensor] = None, min_ps: Optional[torch.Tensor] = None, ) -> None: - rows = logits.shape[0] - # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. - if top_ks is not None: - k = top_ks[:rows] - assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( - "top-k rows must take the generic path" - ) - if min_ps is not None: - assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" - # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the - # wrong rows. - assert rows == draft.shape[0] * (draft.shape[1] + 1), ( - "logits rows must be nreq * (draft_len + 1)" - ) - # Per-row parameters are read at the row index, unchecked. - for t in (temperatures, top_ps, seeds, offsets): - assert t.numel() >= rows and t.is_contiguous(), ( - "per-row parameters must cover every row" - ) - assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( - "fp32 sampling parameters" - ) - assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( - "int64 Philox seeds / offsets" - ) - # The finished-row counter elects the CTA that runs the acceptance. - assert int(counters[0].item()) == 0, ( - "counters[0] (finished rows) must be zero between calls" + return _sampling_shared.validate_sampling( + self, + _REFERENCE_TARGET, + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + top_ks, + min_ps, ) m3_sampled_accept = M3SampledAccept() + +_REFERENCE_TARGET = cast(_sampling_shared.SamplingTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py index 408759009472..04e13e2acbb0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_shard_accept.py @@ -12,97 +12,18 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from ....ops.spec.acceptance_reference import _LOW32 as _LOW32 +from ....ops.spec.acceptance_reference import exchange_tokens as exchange_tokens +from ....ops.spec.acceptance_reference import key_token as key_token +from ....ops.spec.acceptance_reference import local_argmax as local_argmax +from ....ops.spec.acceptance_reference import make_keys as make_keys +from ....ops.spec.acceptance_reference import ordered_bits as ordered_bits +from ....ops.spec.acceptance_reference import strict_accept as strict_accept +from ....ops.spec.acceptance_reference import unsigned_max as unsigned_max + #: Verify rows per call, and exchange rows per call half of its workspace. SHARD_ROWS = 48 -_LOW32 = 0xFFFFFFFF - - -def ordered_bits(v: torch.Tensor) -> torch.Tensor: - """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" - b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 - b = torch.where(b == 0x80000000, torch.zeros_like(b), b) - return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) - - -def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: - """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" - return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) - - -def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: - """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" - sign = torch.iinfo(torch.int64).min - return (keys ^ sign).max(dim=dim).values ^ sign - - -def key_token(key: torch.Tensor) -> torch.Tensor: - return (_LOW32 - (key & _LOW32)).to(torch.int32) - - -def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: - """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" - best = logits.max(dim=1).values - idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max - return best, idx.to(torch.int64) + vocab_offset - - -def exchange_tokens( - own_keys: torch.Tensor, - ws_uc: torch.Tensor, - counters: torch.Tensor, - rank: int, - world: int, - slot_rows: int, -) -> torch.Tensor: - """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. - - Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A - single-device run stands in for the other ranks by writing their keys there before the call. - """ - m = own_keys.shape[0] - half = int(counters[1].item()) & 1 - slots = ( - ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] - .view(m, world) - .clone() - ) - slots[:, rank] = own_keys - return key_token(unsigned_max(slots, dim=1)) - - -def strict_accept( - tokens: torch.Tensor, - draft: torch.Tensor, - batch_idx: torch.Tensor, - force_total: int, - next_new: Optional[torch.Tensor], - force_vec: Optional[torch.Tensor], -) -> dict[str, torch.Tensor]: - """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). - - accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by - force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and - next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. - """ - nreq, d = draft.shape - k1 = d + 1 - tokens = tokens.to(torch.int32) - accepted = tokens.view(nreq, k1).clone() - match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) - num = 1 + match.cumprod(dim=1).sum(dim=1) - if force_vec is not None: - num = force_vec[:nreq].to(num.dtype) - elif force_total > 0: - num = torch.full_like(num, force_total) - gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 - out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} - if next_new is not None: - nn = next_new.clone() - nn[:, 0] = tokens[gather] - out["next_new"] = nn - return out - class M3ShardAccept(OpWrapper): """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py deleted file mode 100644 index ff07c54de1f0..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_slot_scatter.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> -``m3_spec_tail.slot_scatter``).""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -def padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: - """src rows zero-padded (or cut) to ``width`` columns.""" - out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) - w = min(width, src.shape[1]) - out[:, :w] = src[:, :w] - return out - - -class M3SlotScatter(OpWrapper): - """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, - ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The - destinations are the store's transposed [slots, width] views.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of - # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) - # is zero-padded. - CELLS: tuple[Cell, ...] = ( - Cell( - why="8 rows into shuffled slots, widths equal to the store's", - spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), - ), - Cell( - why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", - spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), - ), - ) - - note = """ - Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. - """ - - def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): - torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) - - def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): - """The four destinations after the call, on copies.""" - s = slots.long() - out = {} - for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): - r = dst.clone() - r[s] = padded_rows(src[: s.numel()], dst.shape[1]) - out[name] = r - lr = dl.clone() - lr[s] = sl[: s.numel()] - out["dl"] = lr - return out - - def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): - assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( - "slots must index the store" - ) - # Two rows into one slot race (which one lands is unordered). - assert slots.unique().numel() == slots.numel(), "slots must be distinct" - assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" - - -m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py index 8925fc8ab5b8..6a766b8dc0ce 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_spec_offsets.py @@ -16,7 +16,7 @@ class M3SpecOffsets(OpWrapper): """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths - per pair) in one plain launch: trtllm::minimax_m3_sm107_tp2_spec_offsets (minimaxM3Sm107Tp2/specOffsets.cu), the + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp2_spec_offsets (minimaxM3Sm107/specOffsets.cu), the same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other layouts (the caller runs the ATen adds).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py index 080c3018e4b6..06e9ecbca36f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_prologue.py @@ -14,7 +14,7 @@ class M3StepPrologue(OpWrapper): """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch - (trtllm::minimax_m3_sm107_tp2_step_prologue, minimaxM3Sm107Tp2/stepPrologue.cu). + (trtllm::minimax_m3_sm107_tp2_step_prologue, minimaxM3Sm107/stepPrologue.cu). ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py index ada666e91599..f473526a6646 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/triton_sparse_decode.py @@ -45,15 +45,57 @@ from typing import Optional import torch -import triton import triton.language as tl +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _NVFP4_TUNED_CAPABILITIES as _NVFP4_TUNED_CAPABILITIES, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _TARGET_GRID as _TARGET_GRID, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + MergeMxfp8 as MergeMxfp8, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _check_nvfp4_inputs as _check_nvfp4_inputs, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_launch_options as _sm100f_nvfp4_launch_options, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_merge_launch_options as _sm100f_nvfp4_merge_launch_options, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_num_topk_chunks as _sm100f_nvfp4_num_topk_chunks, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_query_group_size as _sm100f_nvfp4_query_group_size, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_use_linear_softmax as _sm100f_nvfp4_use_linear_softmax, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + resolve_num_topk_chunks as resolve_num_topk_chunks, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( check_decode_span_shape, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers +from ...kernels.attention.triton_sparse_decode import _dequant_nvfp4_rows as _dequant_nvfp4_rows +from ...kernels.attention.triton_sparse_decode import _e2m1x2_to_f16x2 as _e2m1x2_to_f16x2 +from ...kernels.attention.triton_sparse_decode import ( + _gqa_sparse_decode_kernel as _gqa_sparse_decode_kernel, +) +from ...kernels.attention.triton_sparse_decode import ( + _gqa_sparse_decode_pair_kernel as _gqa_sparse_decode_pair_kernel, +) +from ...kernels.attention.triton_sparse_decode import ( + _merge_topk_attn_out_kernel as _merge_topk_attn_out_kernel, +) +from ...kernels.attention.triton_sparse_decode import _rcp_approx_ftz as _rcp_approx_ftz +from ...kernels.attention.triton_sparse_decode import _store_mxfp8_head as _store_mxfp8_head from . import m3_attn_lc, m3_prequant, m3_report # One sparse block is exactly one KV page. @@ -68,7 +110,6 @@ } # Total CTAs the split-K partitioning aims for before it stops splitting. -_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size # NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale # byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. @@ -93,1151 +134,12 @@ PRODUCER_PATH = "sparse_decode_producer" -class MergeMxfp8: - """MXFP8 copy of the split-K merge output for the o_proj GEMM. - - From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and - swizzled UE8M0 scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips - its own quantization. One buffer of ``MAX_ROWS`` rows (one 128-row scale tile) per width serves every layer: each - o_proj consumes it before the next layer's merge rewrites it, in stream order. - """ - - MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) - - def __init__(self): - self._buffers = {} - - def buffers(self, width: int, device: torch.device): - key = (width, device) - if key not in self._buffers: - self._buffers[key] = ( - torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), - torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), - ) - return self._buffers[key] - - @staticmethod - def handoff(mxfp8, rows: int): - """The hand-off of a ``rows``-row step: its E4M3 rows and the 128-row scale tiles that cover them. The o_proj - GEMM checks the swizzled scales against the step's padded rows, so the step gets the tiles that cover its rows, - never the whole buffer.""" - q8, sf = mxfp8 - width = q8.shape[-1] - return q8[:rows], sf[: (rows + 127) // 128 * 128 * (width // 32)] - - MERGE_MXFP8 = MergeMxfp8() -@triton.jit -def _e2m1x2_to_f16x2(packed): - """Convert one packed E2M1 byte into its low/high FP16 values.""" - return tl.inline_asm_elementwise( - asm="""{ - .reg .b8 src; - .reg .b32 pair; - mov.b32 {src, _, _, _}, $2; - cvt.rn.f16x2.e2m1x2 pair, src; - mov.b32 {$0, $1}, pair; - }""", - constraints="=h,=h,r", - args=(packed,), - dtype=(tl.float16, tl.float16), - is_pure=True, - pack=1, - ) - - -@triton.jit -def _dequant_nvfp4_rows( - packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte - scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements - BLOCK_N: tl.constexpr, - D: tl.constexpr, - SCALE_COLS: tl.constexpr, - SF_VEC: tl.constexpr, -): - """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. - - The caller gathers the block scales in whichever layout its role stores - them, so this handles only the value math and the scale broadcast. The - per-tensor scale is applied by the caller. - """ - # Blackwell converts both E2M1 nibbles with one instruction. Joining the - # low/high results along a new innermost axis and flattening it rebuilds - # the original element order without scalar exponent/sign reconstruction. - value_lo, value_hi = _e2m1x2_to_f16x2(packed) - values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) - - scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) - scales = tl.reshape( - tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) - ) - return values * scales - - -@triton.heuristics( - { - "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), - "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), - } -) -# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary -# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). -@triton.jit( - do_not_specialize=["decode_query_len", "total_q"], - do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], -) -def _gqa_sparse_decode_kernel( - q_ptr, # [total_q, num_heads, head_dim] - k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] - v_ptr, # paged V: same layout as K - kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False - k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] - v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order - k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 - v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 - t_ptr, # topk_idx: [num_kv_heads, total_q, topk] - o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT - lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT - out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused - counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest - block_table_ptr, # [num_reqs, max_blocks] - seq_lens, # [num_reqs] - total_q, - gqa_group_size, - head_dim, - max_topk, - sm_scale, - decode_query_len, - stride_qn, - stride_qh, - stride_qd, - stride_k_blk, - stride_k_h, - stride_k_pos, - stride_k_d, - stride_v_blk, - stride_v_h, - stride_v_pos, - stride_v_d, - stride_ksc_blk, - stride_ksc_h, - stride_vsc_blk, - stride_vsc_h, - stride_th, - stride_tn, - stride_tk, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_bt_b, - num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request - num_pages, # KV pool pages: a page id outside [0, num_pages) is not read - stride_out_n, - stride_out_h, - stride_out_d, - num_kv_heads, - BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_H: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache - WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register - COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math - USE_PDL: tl.constexpr, - KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales - USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence - SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only - DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge - FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place - EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry - IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) -): - sm_scale_log2e = sm_scale * 1.4426950409 - if KV_NVFP4: - # Both per-tensor dequant scales factor out of the inner loop: QK is - # linear in K, and V's scale commutes with the softmax weights. Reading - # them from device memory keeps the launch free of a host sync, so - # CUDA-graph capture stays unconstrained. - sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) - v_global_scale = tl.load(v_gs_ptr) - # Split-K over the topk dimension: pid(0) folds (query token, chunk). - pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) - pid_b = pid_bc % total_q - pid_c = pid_bc // total_q - req_id = pid_b // decode_query_len - q_offset = pid_b - req_id * decode_query_len - pid_h = pid_kh * gqa_group_size - chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_start = pid_c * chunk_size_topk - - if USE_PDL: - tl.extra.cuda.gdc_wait() - if EARLY_TRIGGER: - # Dependents still wait for this grid to complete; triggering here only - # lets them start their input-independent prologue (e.g. weight loads). - tl.extra.cuda.gdc_launch_dependents() - - seq_len = tl.load(seq_lens + req_id) - query_pos = seq_len - decode_query_len + q_offset - # Full-CG padding uses zero-length request rows. Clamp to an empty - # attention range instead of letting padded rows produce negative lengths. - kv_len = tl.maximum(query_pos + 1, 0) - - # Bound the walk by the token's own valid block count rather than a - # sentinel: the selector emits ascending ids and only pads with -1 past - # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. - num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk = tl.minimum(max_topk, num_blocks) - chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) - - off_n = tl.arange(0, BLOCK_SIZE_K) - off_d = tl.arange(0, BLOCK_SIZE_D) - off_h = tl.arange(0, BLOCK_SIZE_H) - if KV_NVFP4: - # Packed byte column and block-scale column. The NVFP4 path runs only at - # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. - off_dp = tl.arange(0, BLOCK_SIZE_D // 2) - off_s = tl.arange(0, SCALE_COLS) - d_mask = off_d < head_dim - h_mask = off_h < gqa_group_size - hd_mask = h_mask[:, None] & d_mask[None, :] - bt_row = block_table_ptr + req_id * stride_bt_b - - m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - if USE_LINEAR_SOFTMAX: - # Keep the linear softmax denominator and reuse the producer's alpha. - # This removes one exp2 plus one log2 from every visited KV page; only - # the final LSE written for split-K merging still needs a logarithm. - l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - else: - lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - q = tl.load( - q_ptr - + pid_b * stride_qn - + (pid_h + off_h[:, None]) * stride_qh - + off_d[None, :] * stride_qd, - mask=hd_mask, - other=0.0, - ) - # Widen before anything reads q.dtype: K, V and p are all cast to it below, - # so leaving q narrow would silently run the whole attention in FP8 and - # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result - # matches a caller that widened q itself before the call. - if WIDEN_Q: - q = q.to(COMPUTE_DTYPE) - kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 - - cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk - for it in tl.range(chunk_start, chunk_end): - if IDENTITY: - # Slot `it` of the identity selection, bounded by this row's own kv extent above. - blk = it - else: - blk = tl.load(cur_idx_ptr) - cur_idx_ptr = cur_idx_ptr + stride_tk - # int64 page offsets: a large cache overflows int32 well before the - # per-page block offsets do. - # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, - # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). - blk_ok = (blk >= 0) & (blk < num_bt_cols) - page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) - blk_ok = blk_ok & (page >= 0) & (page < num_pages) - page = tl.where(blk_ok, page, 0) - pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len - # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block - # would make exp2(-inf - -inf) = NaN). - if blk_ok: - if KV_NVFP4: - # K block scales are token-major linear within the page-head region; - # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. - k_deq = _dequant_nvfp4_rows( - tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[:, None] * stride_k_pos - + off_dp[None, :] * stride_k_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - k_sc_ptr - + page * stride_ksc_blk - + pid_kh * stride_ksc_h - + off_n[:, None] * SCALE_COLS - + off_s[None, :], - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ) - k = tl.trans(k_deq).to(q.dtype) - else: - k = tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[None, :] * stride_k_pos - + off_d[:, None] * stride_k_d, - mask=d_mask[:, None] & pos_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - k = (k * kv_scale).to(q.dtype) - qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) - qk += tl.dot(q, k) * sm_scale_log2e - m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp2(qk - m_ij[:, None]) - l_ij = tl.sum(p, axis=1) - if USE_LINEAR_SOFTMAX: - alpha = tl.exp2(m_i - m_ij) - acc_o = acc_o * alpha[:, None] - else: - acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] - if KV_NVFP4: - # V block scales use vLLM's 4x4 token-quad order rather than K's - # linear one: [token // 4, scale_col, token % 4]. - v = _dequant_nvfp4_rows( - tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_dp[None, :] * stride_v_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - v_sc_ptr - + page * stride_vsc_blk - + pid_kh * stride_vsc_h - + (off_n[:, None] // 4) * (4 * SCALE_COLS) - + 4 * off_s[None, :] - + (off_n[:, None] % 4), - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ).to(q.dtype) - else: - v = tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_d[None, :] * stride_v_d, - mask=pos_mask[:, None] & d_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - v = (v * kv_scale).to(q.dtype) - acc_o += tl.dot(p.to(v.dtype), v) - if USE_LINEAR_SOFTMAX: - l_i = l_i * alpha + l_ij - else: - lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) - m_i = m_ij - - # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. - if USE_LINEAR_SOFTMAX: - scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) - lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) - else: - scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) - lse_out = lse_i - if KV_NVFP4: - # Applying V's scale per partial is exact: the merge is a weighted - # average whose weights sum to one, so a constant survives it. - scale = scale * v_global_scale - acc_o = acc_o * scale[:, None] - o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h - tl.store( - o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, - acc_o.to(o_ptr.dtype.element_ty), - mask=hd_mask, - ) - if not DIRECT_OUTPUT: - lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h - tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) - - if FUSED_MERGE: - # Split-K serial reduction: every chunk publishes its partial, and the - # chunk that arrives last for this (token, KV head) merges all of them - # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. - tl.debug_barrier() - counter = counter_ptr + pid_b * num_kv_heads + pid_kh - arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") - if arrived == NUM_TOPK_CHUNKS - 1: - tl.debug_barrier() - o_rows = ( - o_ptr - + pid_b * stride_o_b - + (pid_h + off_h[:, None]) * stride_o_h - + off_d[None, :] * stride_o_d - ) - l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h - lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - lse_max = tl.maximum(lse_max, lse_c) - # All chunks empty (a CUDA-graph padding row): zero, never NaN. - lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) - denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - denom += tl.exp2(lse_c - lse_max) - denom = tl.where(denom > 0, denom, 1.0) - merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - w_c = tl.exp2(lse_c - lse_max) / denom - o_c = tl.load( - o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" - ).to(tl.float32) - merged += o_c * w_c[:, None] - tl.store( - out_ptr - + pid_b * stride_out_n - + (pid_h + off_h[:, None]) * stride_out_h - + off_d[None, :] * stride_out_d, - merged.to(out_ptr.dtype.element_ty), - mask=hd_mask, - ) - tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") - - # After the stores, never before: this releases either the merge grid or, - # for direct output, the next PDL-enabled consumer of the final tensor. - if USE_PDL and not EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.heuristics( - { - "BLOCK_SIZE_H": lambda args: max( - 16, - triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), - ), - "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), - } -) -@triton.jit( - do_not_specialize=["decode_query_len", "total_q"], - do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], -) -def _gqa_sparse_decode_pair_kernel( - q_ptr, # [total_q, num_heads, head_dim] - k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] - v_ptr, # paged V: same layout as K - kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False - k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] - v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order - k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 - v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 - t_ptr, # topk_idx: [num_kv_heads, total_q, topk] - o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT - lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT - block_table_ptr, # [num_reqs, max_blocks] - seq_lens, # [num_reqs] - total_q, - gqa_group_size, - head_dim, - max_topk, - sm_scale, - decode_query_len, - stride_qn, - stride_qh, - stride_qd, - stride_k_blk, - stride_k_h, - stride_k_pos, - stride_k_d, - stride_v_blk, - stride_v_h, - stride_v_pos, - stride_v_d, - stride_ksc_blk, - stride_ksc_h, - stride_vsc_blk, - stride_vsc_h, - stride_th, - stride_tn, - stride_tk, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_bt_b, - num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request - num_pages, # KV pool pages: a page id outside [0, num_pages) is not read - BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_H: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache - WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register - COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math - USE_PDL: tl.constexpr, - KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales - USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence - QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set - MAX_TOPK: tl.constexpr, - SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only - DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge - IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) -): - sm_scale_log2e = sm_scale * 1.4426950409 - if KV_NVFP4: - # Both per-tensor dequant scales factor out of the inner loop: QK is - # linear in K, and V's scale commutes with the softmax weights. Reading - # them from device memory keeps the launch free of a host sync, so - # CUDA-graph capture stays unconstrained. - sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) - v_global_scale = tl.load(v_gs_ptr) - # Split-K over the topk dimension: pid(0) folds (query group, chunk). - pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) - num_query_groups = total_q // QUERY_GROUP_SIZE - pid_g = pid_bc % num_query_groups - pid_c = pid_bc // num_query_groups - groups_per_req = decode_query_len // QUERY_GROUP_SIZE - req_id = pid_g // groups_per_req - q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE - pid_b = req_id * decode_query_len + q_group_offset - pid_h = pid_kh * gqa_group_size - chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_start = pid_c * chunk_size_topk - - if USE_PDL: - tl.extra.cuda.gdc_wait() - - seq_len = tl.load(seq_lens + req_id) - # Query rows are flattened as [query-in-group, GQA head]. - off_h = tl.arange(0, BLOCK_SIZE_H) - q_lane = off_h // gqa_group_size - h_lane = off_h - q_lane * gqa_group_size - h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE - q_token = pid_b + q_lane - query_pos = seq_len - decode_query_len + q_group_offset + q_lane - # Full-CG padding uses zero-length request rows. Clamp to an empty - # attention range instead of letting padded rows produce negative lengths. - kv_len = tl.maximum(query_pos + 1, 0) - kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) - - # Each query can have a different valid-block count at a page boundary. - # The selector puts those valid entries first and pads only the tail. - kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) - num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk_0 = tl.minimum(max_topk, num_blocks_0) - chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) - if QUERY_GROUP_SIZE == 2: - kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) - num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk_1 = tl.minimum(max_topk, num_blocks_1) - chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) - - off_n = tl.arange(0, BLOCK_SIZE_K) - off_d = tl.arange(0, BLOCK_SIZE_D) - if KV_NVFP4: - # Packed byte column and block-scale column. The NVFP4 path runs only at - # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. - off_dp = tl.arange(0, BLOCK_SIZE_D // 2) - off_s = tl.arange(0, SCALE_COLS) - d_mask = off_d < head_dim - hd_mask = h_mask[:, None] & d_mask[None, :] - bt_row = block_table_ptr + req_id * stride_bt_b - - m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - if USE_LINEAR_SOFTMAX: - # Keep the linear softmax denominator and reuse the producer's alpha. - # This removes one exp2 plus one log2 from every visited KV page; only - # the final LSE written for split-K merging still needs a logarithm. - l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - else: - lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - q = tl.load( - q_ptr - + q_token[:, None] * stride_qn - + (pid_h + h_lane[:, None]) * stride_qh - + off_d[None, :] * stride_qd, - mask=hd_mask, - other=0.0, - ) - # Widen before anything reads q.dtype: K, V and p are all cast to it below, - # so leaving q narrow would silently run the whole attention in FP8 and - # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this - # reproduces the values the caller's standalone widening kernel produced. - if WIDEN_Q: - q = q.to(COMPUTE_DTYPE) - kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 - - idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk - if QUERY_GROUP_SIZE == 2: - idx_ptr_1 = idx_ptr_0 + stride_tn - idx_0 = chunk_start - idx_1 = chunk_start - count_0 = chunk_end_0 - chunk_start - count_1 = chunk_end_1 - chunk_start - chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_offsets = tl.arange(0, chunk_capacity) - sentinel = 2147483647 - if IDENTITY: - blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) - blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) - else: - blocks_0 = tl.load( - idx_ptr_0 + chunk_offsets * stride_tk, - mask=chunk_offsets < count_0, - other=sentinel, - ) - blocks_1 = tl.load( - idx_ptr_1 + chunk_offsets * stride_tk, - mask=chunk_offsets < count_1, - other=sentinel, - ) - chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 - # Equal chunks take the measured shared-page loop. A divergent chunk - # walks the sorted union, bounded by the sum of both row lengths. - num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) - else: - num_candidates = chunk_end_0 - chunk_start - - for cand in tl.range(0, num_candidates): - if QUERY_GROUP_SIZE == 2: - if chunk_equal: - has_block = True - if IDENTITY: - blk = chunk_start + cand - else: - blk = tl.load(idx_ptr_0) - idx_ptr_0 += stride_tk - query_active = h_mask - else: - valid_0 = idx_0 < chunk_end_0 - valid_1 = idx_1 < chunk_end_1 - if IDENTITY: - blk_0 = tl.where(valid_0, idx_0, sentinel) - blk_1 = tl.where(valid_1, idx_1, sentinel) - else: - blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) - blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) - use_0 = valid_0 & (blk_0 <= blk_1) - use_1 = valid_1 & (blk_1 <= blk_0) - has_block = use_0 | use_1 - blk = tl.minimum(blk_0, blk_1) - idx_0 += use_0.to(tl.int32) - idx_1 += use_1.to(tl.int32) - idx_ptr_0 += use_0.to(tl.int64) * stride_tk - idx_ptr_1 += use_1.to(tl.int64) * stride_tk - query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) - else: - has_block = True - if IDENTITY: - blk = chunk_start + cand - else: - blk = tl.load(idx_ptr_0) - idx_ptr_0 += stride_tk - query_active = h_mask - - if has_block: - # int64 page offsets: a large cache overflows int32 well before the - # per-page block offsets do. - # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the - # pool, reads nothing. - blk_ok = (blk >= 0) & (blk < num_bt_cols) - page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) - blk_ok = blk_ok & (page >= 0) & (page < num_pages) - page = tl.where(blk_ok, page, 0) - pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max - q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] - # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block - # would make exp2(-inf - -inf) = NaN). - if blk_ok: - if KV_NVFP4: - # K scales are token-major linear within the page-head region; - # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. - k_deq = _dequant_nvfp4_rows( - tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[:, None] * stride_k_pos - + off_dp[None, :] * stride_k_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - k_sc_ptr - + page * stride_ksc_blk - + pid_kh * stride_ksc_h - + off_n[:, None] * SCALE_COLS - + off_s[None, :], - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ) - k = tl.trans(k_deq).to(q.dtype) - else: - k = tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[None, :] * stride_k_pos - + off_d[:, None] * stride_k_d, - mask=d_mask[:, None] & pos_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - k = (k * kv_scale).to(q.dtype) - active_pos = query_active[:, None] & q_pos_mask - qk = tl.where(active_pos, 0.0, float("-inf")) - qk += tl.dot(q, k) * sm_scale_log2e - m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) - p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) - l_ij = tl.sum(p, axis=1) - alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) - acc_o = acc_o * alpha[:, None] - if KV_NVFP4: - # V scales use vLLM's 4x4 token-quad order rather than K's - # linear one: [token // 4, scale_col, token % 4]. - v = _dequant_nvfp4_rows( - tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_dp[None, :] * stride_v_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - v_sc_ptr - + page * stride_vsc_blk - + pid_kh * stride_vsc_h - + (off_n[:, None] // 4) * (4 * SCALE_COLS) - + 4 * off_s[None, :] - + (off_n[:, None] % 4), - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ).to(q.dtype) - else: - v = tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_d[None, :] * stride_v_d, - mask=pos_mask[:, None] & d_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - v = (v * kv_scale).to(q.dtype) - acc_o += tl.dot(p.to(v.dtype), v) - if USE_LINEAR_SOFTMAX: - l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) - else: - lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) - lse_i = tl.where(query_active, lse_ij, lse_i) - m_i = m_ij - - # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. - if USE_LINEAR_SOFTMAX: - scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) - lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) - else: - scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) - lse_out = lse_i - if KV_NVFP4: - # Applying V's scale per partial is exact: the merge is a weighted - # average whose weights sum to one, so a constant survives it. - scale = scale * v_global_scale - acc_o = acc_o * scale[:, None] - tl.store( - o_ptr - + pid_c * stride_o_c - + q_token[:, None] * stride_o_b - + (pid_h + h_lane[:, None]) * stride_o_h - + off_d[None, :] * stride_o_d, - acc_o.to(o_ptr.dtype.element_ty), - mask=hd_mask, - ) - if not DIRECT_OUTPUT: - tl.store( - lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, - lse_out, - mask=h_mask, - ) - - # After the stores, never before: this releases either the merge grid or, - # for direct output, the next PDL-enabled consumer of the final tensor. - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _rcp_approx_ftz(x): - return tl.inline_asm_elementwise( - "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 - ) - - -@triton.jit -def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): - """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as - torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, - scales in the 128x4-swizzled layout.""" - NUM_SF: tl.constexpr = HEAD_DIM // 32 - xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) - amax = tl.max(tl.abs(xr), axis=1) - sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) - bits = sfv.to(tl.int32, bitcast=True) - e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) - e = tl.minimum(e, 254) - scale = ((254 - e) << 23).to(tl.float32, bitcast=True) - scale = tl.where(amax != 0, scale, 0.0) - q = (xr * scale[:, None]).to(tl.float8e4nv) - off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) - tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) - k_sf = head * NUM_SF + tl.arange(0, NUM_SF) - off_sf = ( - (row // 128) * num_sf_k_tiles * 512 - + (k_sf // 4) * 512 - + (row % 32) * 16 - + ((row % 128) // 32) * 4 - + k_sf % 4 - ) - tl.store(sf_ptr + off_sf, e.to(tl.uint8)) - - -@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) -@triton.jit -def _merge_topk_attn_out_kernel( - o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] - lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] - out_ptr, # merged out: [total_q, num_heads, head_dim] - head_dim, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_out_n, - stride_out_h, - stride_out_d, - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_PDL: tl.constexpr, - EARLY_TRIGGER: tl.constexpr = False, - q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 - sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled - num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) - MXFP8_OUT: tl.constexpr = False, -): - pid_b, pid_h = tl.program_id(0), tl.program_id(1) - - if USE_PDL: - if EARLY_TRIGGER: - # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight - # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits - # for this grid (and therefore the producer) to complete. Because this trigger precedes this grid's own - # wait, it says nothing about what upstream has finished: a dependent may read only step-invariant data - # (weights) before its griddepcontrol.wait, never this step's attention output, q, top-k table or lengths. - tl.extra.cuda.gdc_launch_dependents() - tl.extra.cuda.gdc_wait() - - off_c = tl.arange(0, NUM_TOPK_CHUNKS) - off_d = tl.arange(0, BLOCK_SIZE_D) - o = tl.load( - o_ptr - + pid_b * stride_o_b - + pid_h * stride_o_h - + off_c[:, None] * stride_o_c - + off_d[None, :] * stride_o_d, - mask=off_d[None, :] < head_dim, - other=0.0, - ).to(tl.float32) - # Empty chunks contribute -inf, hence weight 0. - lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) - lse_max = tl.max(lse, axis=0) - # A row whose every chunk is empty (a CUDA-graph padding row) would give - # -inf - -inf here. Zero it instead: the padded output is discarded, but a - # NaN would survive into the residual stream and the all-reduce. - lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) - weights = tl.exp2(lse - lse_max) - denom = tl.sum(weights, axis=0) - weights = weights / tl.where(denom > 0, denom, 1.0) - o_merged = tl.sum(o * weights[:, None], axis=0) - out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d - tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) - if MXFP8_OUT: - _store_mxfp8_head( - o_merged.to(out_ptr.dtype.element_ty), - q8_ptr, - sf_ptr, - pid_b, - pid_h, - num_sf_k_tiles, - BLOCK_SIZE_D, - ) - - if USE_PDL and not EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - -def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: - """Split-K factor over the top-k blocks, as a power of two. - - Depends only on shapes that are fixed for a captured batch size, so the - launch geometry is frozen inside a CUDA graph. - """ - target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) - return 1 << (target.bit_length() - 1) - - # NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. _SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) # Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. -_NVFP4_TUNED_CAPABILITIES = ((10, 3),) - - -def _sm100f_nvfp4_num_topk_chunks( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> Optional[int]: - """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return None - local_batch = total_q // decode_query_len - return { - 2: 8, - 3: 8, - 4: 8, - 5: 16, - 6: 4, - 7: 4, - 8: 2, - 9: 2, - 10: 8, - 11: 4, - 12: 2, - 13: 2, - 14: 2, - }.get(local_batch) - - -def _sm100f_nvfp4_use_linear_softmax( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> bool: - """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads not in (1, 2, 4) - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return False - local_batch = total_q // decode_query_len - # The TP4/DEP32 single-request graph is 8.9% faster with the original - # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral - # or faster with the linear denominator. - return not (num_kv_heads == 1 and local_batch == 1) - - -def _sm100f_nvfp4_query_group_size( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> int: - """Group Eagle queries using the GB300-measured shared-page policy.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return 1 - # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches - # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or - # slower and retain the single-query producer. - return 2 if total_q // decode_query_len in (11, 12, 14) else 1 - - -def _sm100f_nvfp4_launch_options( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> dict[str, int]: - """Select native-NVFP4 launch geometry for the SM100 family. - - The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four - warps/two stages wins through local batch 28, while two warps/one stage - avoids the next residency cliff at local batches 29--32. Leave every - unmeasured architecture, model geometry, and larger batch on Triton's - established defaults. - """ - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - ): - return {} - local_batch = total_q // decode_query_len - if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): - return {"num_warps": 4, "num_stages": 1} - if local_batch == 4: - return {"num_warps": 2, "num_stages": 1} - if 8 <= local_batch <= 28: - return {"num_warps": 4, "num_stages": 2} - if 29 <= local_batch <= 32: - return {"num_warps": 2, "num_stages": 1} - return {} - - -def _sm100f_nvfp4_merge_launch_options( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> dict[str, int]: - """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability in _NVFP4_TUNED_CAPABILITIES - and num_kv_heads == 4 - and gqa_group_size == 16 - and max_topk == 16 - and decode_query_len == 4 - and total_q % decode_query_len == 0 - and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) - ): - return {"num_warps": 1} - return {} - - -def _check_nvfp4_inputs( - k_paged: torch.Tensor, - v_paged: torch.Tensor, - nvfp4_args: tuple[Optional[torch.Tensor], ...], - *, - head_dim: int, - scale_cols: int, -) -> None: - """Validate the packed-NVFP4 cache against what the kernel assumes. - - The kernel indexes block scales by a flat within-page offset and reshapes a - [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows - and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them - rather than let a mismatch read neighboring pages. - """ - k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args - if any(arg is None for arg in nvfp4_args): - raise ValueError( - "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " - "v_block_scale, k_global_scale and v_global_scale." - ) - if head_dim != triton.next_power_of_2(head_dim): - raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") - if v_paged.shape[:3] != k_paged.shape[:3]: - raise ValueError("NVFP4 K and V caches must have matching page, head and token extents.") - for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): - if int(data.shape[3]) * 2 != head_dim: - raise ValueError( - f"NVFP4 {name} must pack two elements per byte: expected " - f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." - ) - for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): - if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): - raise ValueError( - f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " - f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." - ) - if scale.stride(3) != 1 or scale.stride(2) != scale_cols: - raise ValueError( - f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " - f"{tuple(scale.stride())}." - ) - for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): - if gs.dtype != torch.float32 or gs.numel() != 1: - raise ValueError( - f"NVFP4 {name} must be a single-element FP32 tensor; got " - f"{gs.numel()} element(s) of {gs.dtype}." - ) # Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query @@ -1589,7 +491,7 @@ def minimax_m3_sparse_attn_decode( else {} ) if use_mq: - from .triton_sparse_decode_mq import launch_mq_producer + from ...kernels.attention.triton_sparse_decode_mq import launch_mq_producer launch_mq_producer( q, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py index 550316d84740..eb1b89fde8c1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/modeling.py @@ -54,9 +54,9 @@ from __future__ import annotations import atexit -import copy import dataclasses import os +import sys import threading import time from dataclasses import dataclass @@ -72,6 +72,7 @@ Mapping, Optional, Tuple, + cast, ) import torch @@ -99,6 +100,91 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import ( + modeling_support as _modeling_support, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import reporting as _reporting +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _is_vl_config as _is_vl_config, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _moe_layer_ids as _moe_layer_ids, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _sparse_layer_ids as _sparse_layer_ids, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _text_model_config as _text_model_config, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _validate_runtime as _validate_runtime, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.activation.m3_silu_and_mul import ( + m3_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + msa_kv_update_hook as _msa_kv_update_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_dense_decode import ( + m3_dense_decode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_fp8_indexer_qk_norm_rope import ( + m3_fp8_indexer_qk_norm_rope as m3_fp8_indexer_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qk_norm_rope_fp8 import ( + m3_qk_norm_rope_fp8, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qkv_insert import ( + m3_qkv_insert, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_select_blocks import ( + m3_select_blocks, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( + m3_allreduce, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.m3_embedding import ( + m3_embedding, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_bf16_linear import ( + m3_bf16_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + m3_dense_swiglu, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_fp32_linear import ( + m3_fp32_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( + m3_mxfp8_gemm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_linear import ( + m3_mxfp8_linear as m3_mxfp8_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + m3_target_head, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + MxMode as MxMode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + _mx_mode as _mx_mode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + _Mxfp8GemmState as _Mxfp8GemmState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_tc import ( + m3_router_tc, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_draft_input import ( + m3_draft_input, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( + m3_force_vec, +) from tensorrt_llm._torch.attention.backends import AttentionMetadata from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales from tensorrt_llm._torch.attention.backends.interface import ( @@ -121,14 +207,17 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_speculative import ( _ensure_draft_vocab_size, - _get_requested_draft_moe_backend, - _set_draft_kv_cache_quant_algo, _slice_spec_position_ids, ) +from tensorrt_llm._torch.models.modeling_speculative import ( + _get_requested_draft_moe_backend as _get_requested_draft_moe_backend, +) +from tensorrt_llm._torch.models.modeling_speculative import ( + _set_draft_kv_cache_quant_algo as _set_draft_kv_cache_quant_algo, +) from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM from tensorrt_llm._torch.modules.embedding import LMHead from tensorrt_llm._torch.modules.linear import ( - MXFP8LinearMethod, TensorParallelMode, flashinfer_mxfp8_autotune_active, flashinfer_mxfp8_decode_graph_capture_active, @@ -146,7 +235,9 @@ _FORCE_ACCEPT_RNG_COUNTER_STRIDE, _FORCE_ACCEPT_RNG_POOL_SIZE, _FORCE_ACCEPT_RNG_SLOT_STRIDE, - should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.speculative.interface import ( + should_use_separate_draft_kv_cache as should_use_separate_draft_kv_cache, ) from tensorrt_llm._torch.utils import ActType_TrtllmGen, is_torch_compiling from tensorrt_llm._utils import mpi_rank, prefer_pinned @@ -167,33 +258,21 @@ from .ops import m3_shard_accept as m3_shard_accept_kernel from .ops import m3_target_head as m3_target_head_kernel from .ops import m3_tokmap_guard as _tokmap_guard -from .ops import msa_kv_update as _msa_kv_update_kernel -from .ops.activation.m3_silu_and_mul import m3_silu_and_mul -from .ops.attention.m3_dense_decode import m3_dense_decode from .ops.attention.m3_dense_qk_norm_rope_fp8 import m3_dense_qk_norm_rope_fp8 from .ops.attention.m3_draft_attention import m3_draft_attention -from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope from .ops.attention.m3_fullctx_decode import m3_fullctx_decode from .ops.attention.m3_index_score import SPLIT_K as INDEX_SCORE_SPLIT_K from .ops.attention.m3_index_score import m3_index_score from .ops.attention.m3_index_score_t2sc import m3_index_score_t2sc from .ops.attention.m3_msa_attention import m3_msa_attention from .ops.attention.m3_producer import m3_producer -from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 -from .ops.attention.m3_qkv_insert import m3_qkv_insert -from .ops.attention.m3_select_blocks import m3_select_blocks from .ops.attention.m3_select_long import m3_select_long from .ops.attention.m3_select_long_dsl import m3_select_long_dsl from .ops.attention.m3_sparse_decode import m3_sparse_decode -from .ops.comm.m3_allreduce import m3_allreduce from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm from .ops.comm.m3_mc_finalize_allreduce_norm import m3_mc_finalize_allreduce_norm -from .ops.embedding.m3_embedding import m3_embedding -from .ops.gemm.m3_bf16_linear import m3_bf16_linear -from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS from .ops.gemm.m3_draft_lowm import m3_draft_lowm -from .ops.gemm.m3_fp32_linear import m3_fp32_linear from .ops.gemm.m3_mx_gemv import ( DENSE_DOWN, DENSE_GATE_UP, @@ -205,23 +284,16 @@ m3_mx_gemv, ) from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 -from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm -from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear from .ops.gemm.m3_mxfp8_pinned_gemm import m3_mxfp8_pinned_gemm from .ops.gemm.m3_mxfp8_pinned_quantize import m3_mxfp8_pinned_quantize -from .ops.gemm.m3_target_head import MAX_ROWS as TARGET_HEAD_KERNEL_MAX_ROWS -from .ops.gemm.m3_target_head import m3_target_head from .ops.long_select import long_select_blocks from .ops.moe.m3_moe16_engine import M3Moe16State, m3_moe16_engine from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine -from .ops.moe.m3_router_tc import m3_router_tc from .ops.mx_gemv import op as mx_gemv from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm -from .ops.spec.m3_draft_input import m3_draft_input from .ops.spec.m3_draft_silu import m3_draft_silu from .ops.spec.m3_draft_step import MAX_TOKENS as DRAFT_STEP_KERNEL_MAX_TOKENS from .ops.spec.m3_draft_step import m3_draft_step -from .ops.spec.m3_force_vec import m3_force_vec from .ops.spec.m3_force_vec_req import m3_force_vec_req from .ops.spec.m3_shard_accept import m3_shard_accept from .ops.spec.m3_spec_offsets import m3_spec_offsets @@ -286,7 +358,7 @@ class DecodePlan: # 4.79 -> 4.14 at 4 rows, ~ -2..-3.5 us/step. Kill: False (the row kernel, bitwise the stock boundary). # OFF: in-model timing was flat and natural acceptance length was grey (lower bound -1.16% vs the -1% line). DRAFT_MC_AR_CLUSTER = False -# The drafter's o_proj and down projection, on the copied split-K kernel (ops/m3_lowm_splitk), also push their rows +# The drafter's o_proj and down projection, on the shared split-K kernel (kernels/gemm/m3_lowm_splitk), push their rows # into the drafter boundary all-reduce's exchange (the row kernel's slots: the canonical words its own push would store, # at the (token, 768-column slice) call parity), which then runs prepushed 2 instead of pushing x itself; the copy # releases its dependents only after its own grid wait and the pushes (the prepushed >= 2 contract). Bitwise. The TP8 @@ -416,7 +488,7 @@ class DecodePlan: PRODUCER_LATE_TRIGGER_MAX_ROWS = 0 # capacity: fixed: measured trigger crossover (0: off) # The producer reads its per-step inputs (position ids, cache slots, RoPE rows) before its grid wait, leaving only the # packed QKV row after it: this target's copy of the producer (trtllm::minimax_m3_sm107_tp4_producer_pw, -# minimaxM3Sm107Tp4/producerPw4.cu, the TP2 producerPw2 at 16 query heads). Their writers are host-to-device copies +# minimaxM3Sm107/producerPw.cu, specialized at 16 query heads). Their writers are host-to-device copies # before the replay and plain-launched kernels earlier in the step (the input preprocessing, msa_kv_update or its torch # fallback), so they are final before the producer can start; a PDL-launched writer of any of them would have to turn # this off. The call site's packed.contiguous() and position-id copies are plain launches too. Bit-identical outputs. @@ -538,7 +610,7 @@ class DecodePlan: # clock) at 1 row: qkv 8.15 -> 5.90, o 5.36 -> 3.40, gate_up 12.93 -> 9.93, down 7.68 -> 6.16 us. DRAFT_LOWM_MAX_ROWS = 4 # capacity: fixed: the low-M tactic table's certified rows # The drafter MLP's activation (the plain SwiGLU of its in-tree GatedMLP) runs this target's C++ op -# (ops/spec/m3_draft_silu, minimaxM3Sm107Tp4/siluAndMul.cu) instead of the Triton silu_and_mul: bit for bit, one launch +# (ops/spec/m3_draft_silu, minimaxM3Sm107/siluAndMul.cu) instead of the Triton silu_and_mul: bit for bit, one launch # either way (no Triton on the generation decode path). The fused draft loop's MLP (M3Eagle3Draft.mlp) routes the # call; the in-tree draft path keeps the Triton op. Kill: False. DRAFT_SILU_CPP = True @@ -725,7 +797,9 @@ def covers(self, key: Key) -> bool: return key in self.keys def certified(self) -> FrozenSet[Key]: - return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + cells_by_tp = getattr(self.entry, "CELLS_BY_TP", None) + cells = self.entry.CELLS if cells_by_tp is None else cells_by_tp[4] + return frozenset(k for cell in cells for k in self.certifies(cell.spec)) def _rows(lo: int, hi: int) -> FrozenSet[int]: @@ -981,7 +1055,7 @@ def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: # replay (program 0 of its first scatter launch, no extra kernel) and never resets. The per-layer megakernel pieces # tag their per-step device flags with it (contract v1: tag = epoch << 8 | layer). Needs STEP_PACK_IN_GRAPH. STEP_PACK_EPOCH = True -# The prologue is one launch of this target's op (ops/spec/m3_step_prologue, minimaxM3Sm107Tp4/stepPrologue.cu): the +# The prologue is one launch of this target's op (ops/spec/m3_step_prologue, minimaxM3Sm107/stepPrologue.cu): the # host copies, a CTA barrier, the device copies, the header clear and the epoch, where the default prologue runs two # Triton scatter launches and an ATen fill (3 graph kernels -> 1, no Triton at the head of the decode graph). Same bytes # written. Needs STEP_PACK_EARLY_COPY. Kill: False. @@ -1023,7 +1097,7 @@ def _long_select_dsl_widths(spec: Mapping[str, Any]) -> Iterable[Key]: # between the step graphs (M3SpecStepFold.in_forward). STEP_SPEC_FOLD_IN_FORWARD = True # The forward's tail (STEP_SPEC_FOLD_IN_FORWARD) runs as this target's C++ op (ops/m3_fold_tail cpp=, -# minimaxM3Sm107Tp4/foldTail.cu) instead of the Triton kernel: the same integer copies, one launch either way (no +# minimaxM3Sm107/foldTail.cu) instead of the Triton kernel: the same integer copies, one launch either way (no # Triton on the generation decode path). Kill: False. STEP_FOLD_TAIL_CPP = True # With STEP_SPEC_FOLD_IN_FORWARD the spec sampler's host copies of its store (new tokens, lengths, draft tokens) still @@ -1109,7 +1183,7 @@ def moe_path(num_tokens: int) -> Optional[str]: MX_GEMV_PATH = MX_GEMV.path MX_GEMV_BF16_PATH = MX_GEMV_BF16.path # The dense layers' swigluoai activation (where the fold above would apply) as this target's C++ op -# (ops/gemm/m3_dense_swiglu, minimaxM3Sm107Tp8/swigluOai.cu: bit for bit the Triton silu_and_mul), PDL-launched with +# (ops/gemm/m3_dense_swiglu, minimaxM3Sm107/swigluOai.cu: bit for bit the Triton silu_and_mul), PDL-launched with # an early release so the down GEMV stages its weights during it; the down GEMV reads the activation only after its own # grid wait. TP2 chain bench: -0.69 us per dense layer vs the Triton kernel. Kill constant: False # runs the Triton kernel. @@ -1148,24 +1222,20 @@ class FastPathUnavailable(RuntimeError): _lock = threading.Lock() _require = [False] _logged: set = set() +_POLICY_TARGET = cast(_reporting.ReportingTarget, sys.modules[__name__]) def set_require(require: bool) -> None: """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" - _require[0] = bool(require) + return _reporting.set_require(_POLICY_TARGET, require) def require_mode() -> bool: - return _require[0] + return _reporting.require_mode(_POLICY_TARGET) def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: - key = (path, reason) - with _lock: - if key in _logged: - return - _logged.add(key) - level(message) + return _reporting._log_once(_POLICY_TARGET, level, path, reason, message) def long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: @@ -1184,47 +1254,28 @@ def long_select_dsl_keys() -> Tuple[Tuple[bool, bool, int], ...]: def missed(path: str, rows: Optional[int], reason: str) -> None: """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" - at = "" if rows is None else f" at {rows} rows" - if _require[0]: - raise FastPathUnavailable( - f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " - f"{reason}" - ) - _log_once( - logger.warning, - path, - reason, - f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", - ) + return _reporting.missed(_POLICY_TARGET, path, rows, reason) def capture_missed(path: str, rows: Optional[int], what: str) -> None: """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" - missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + return _reporting.capture_missed(_POLICY_TARGET, path, rows, what) def taken(path: str, rows: Optional[int], what: str = "") -> None: """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" - at = "" if rows is None else f" at {rows} rows" - detail = f" ({what})" if what else "" - _log_once( - logger.info, - path, - f"taken: {rows}" + (f": {what}" if what else ""), - f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", - ) + return _reporting.taken(_POLICY_TARGET, path, rows, what) def uncovered(path: str, reason: str) -> None: """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" - _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + return _reporting.uncovered(_POLICY_TARGET, path, reason) def reset_logged() -> None: """Forget which (path, reason) pairs were logged.""" - with _lock: - _logged.clear() + return _reporting.reset_logged(_POLICY_TARGET) # Kernel modules report from inside their calls through ops/m3_report: the same policy. @@ -1282,132 +1333,6 @@ def _print_values() -> None: # -------------------------------------------------------------------------------------------------------------------- -def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: - model_type = getattr(pretrained_config, "model_type", None) - if model_type == "minimax_m3_vl": - return True - if model_type == "minimax_m3": - return False - return any( - "ForConditionalGeneration" in a - for a in (getattr(pretrained_config, "architectures", None) or []) - ) - - -def _text_model_config(model_config: ModelConfig) -> ModelConfig: - """The model config with the language model's ``text_config`` as its pretrained config. - - The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed - sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and - embedding tie carried over. A text config is returned unchanged. - """ - if not _is_vl_config(model_config.pretrained_config): - return model_config - cfg = copy.deepcopy(model_config) - text = getattr(cfg.pretrained_config, "text_config", None) - if text is None: - raise ValueError("MiniMax-M3 VL config has no `text_config`") - if isinstance(text, dict): - out = PretrainedConfig() - for k, v in text.items(): - setattr(out, k, v) - text = out - if ( - getattr(text, "torch_dtype", None) is None - and getattr(cfg.pretrained_config, "torch_dtype", None) is not None - ): - text.torch_dtype = cfg.pretrained_config.torch_dtype - if ( - getattr(text, "tie_word_embeddings", None) is None - and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None - ): - text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings - return dataclasses.replace(cfg, pretrained_config=text) - - -def _validate_runtime(model_config: ModelConfig) -> None: - """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing - sees only its KV cache algorithm; this check at construction refuses every other runtime. - - The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index - projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. - """ - sparse = model_config.sparse_attention_config - if ( - sparse is None - or getattr(sparse, "algorithm", None) != "minimax_m3" - or sparse.implementation != "msa" - ): - raise ValueError( - "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" - ) - if getattr(sparse, "indexer_kv_dtype", None) != "fp8": - raise ValueError( - "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" - ) - mode = model_config.quant_config.quant_mode - if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): - raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") - - -def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: - freq = cfg.sparse_attention_config["sparse_attention_freq"] - return [i for i, f in enumerate(freq) if int(f) != 0] - - -def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: - freq = getattr(cfg, "moe_layer_freq", None) - if freq is None: - return list(range(cfg.num_hidden_layers)) - return [i for i, f in enumerate(freq) if int(f) != 0] - - -class _Mxfp8GemmState(nn.Module): - """The MXFP8 GEMM backend state the engine's warmup drives. - - The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM - tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of - this target dispatches identically, so one method object carries that state for all of them; the forward reads it - at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. - """ - - def __init__(self) -> None: - super().__init__() - self.quant_method = MXFP8LinearMethod() - - -@dataclass(frozen=True) -class MxMode: - """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" - - decode_graph: bool # inside a decode CUDA-graph warmup or capture - autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass - tuned_quantize_gemm: ( - bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) - ) - flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation - native_tuning: ( - bool # the native GEMM through the autotuner (warmup, before its tactics are registered) - ) - - -def _mx_mode(state: _Mxfp8GemmState) -> MxMode: - method = state.quant_method - decode_graph = flashinfer_mxfp8_decode_graph_capture_active() - autotune = flashinfer_mxfp8_autotune_active() - return MxMode( - decode_graph=decode_graph, - autotune=autotune, - tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, - flashinfer=method.backend == "flashinfer" - or ( - method.backend == "auto" - and (autotune or (method._flashinfer_autotuned and decode_graph)) - ), - native_tuning=method.needs_native_autotune, - ) - - # -------------------------------------------------------------------------------------------------------------------- # Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs # -------------------------------------------------------------------------------------------------------------------- @@ -1417,24 +1342,7 @@ def _mxfp8_linear( core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode ) -> torch.Tensor: """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" - _record( - "m3_mxfp8_linear", - tuple(w.shape), - mode.tuned_quantize_gemm, - mode.flashinfer, - mode.native_tuning, - ) - return m3_mxfp8_linear( - x.contiguous(), - w, - s, - core.mx_alpha, - torch.bfloat16, - mode.tuned_quantize_gemm, - mode.autotune, - mode.flashinfer, - mode.native_tuning, - ) + return _modeling_support._mxfp8_linear(_SUPPORT_TARGET, core, x, w, s, mode) def _gemv_variants(n: int, k: int, push: Optional[str], world: int): @@ -1601,13 +1509,18 @@ def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: kernel.""" rows = int(gu.shape[0]) path = DENSE_SWIGLU_CPP_PATH - max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS_BY_TP[4]) if DENSE_SWIGLU_CPP and not is_torch_compiling(): if rows > max_rows: uncovered(path, f"the dense swigluoai op: {rows} rows, past its certified {max_rows}") else: act = m3_dense_swiglu( - gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + gu, + core.swiglu_limit, + core.swiglu_alpha, + core.swiglu_beta, + get_env_enable_pdl(), + tp_size=4, ) if act is not None: taken(path, rows) @@ -1773,8 +1686,7 @@ def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): def _split_chunk(rows: int, num_chunks: int) -> List[int]: """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" - base, extra = divmod(rows, num_chunks) - return [base + 1] * extra + [base] * (num_chunks - extra) + return _modeling_support._split_chunk(_SUPPORT_TARGET, rows, num_chunks) # -------------------------------------------------------------------------------------------------------------------- @@ -1783,7 +1695,7 @@ def _split_chunk(rows: int, num_chunks: int) -> List[int]: def _nvfp4_layer(manager, i: int) -> bool: - return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + return _modeling_support._nvfp4_layer(_SUPPORT_TARGET, manager, i) def _producer_covered(core) -> bool: @@ -1814,31 +1726,7 @@ def _producer_layout(kv_data, kv_scale) -> bool: def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: """Reject invalid HND pools; select the fused producer only for its packed layout.""" - if ( - data.dim() != 5 - or scales.dim() != 5 - or data.shape[0] <= 0 - or data.shape[0] != scales.shape[0] - or tuple(data.shape[1:]) != (2, num_heads, 128, 64) - or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) - or data.dtype not in (torch.int8, torch.uint8) - or scales.dtype != torch.uint8 - ): - raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") - for cache in (data, scales): - if ( - cache.stride(4) != 1 - or cache.stride(3) != cache.shape[4] - or cache.stride(2) < 128 * cache.stride(3) - or cache.stride(1) < num_heads * cache.stride(2) - or cache.stride(0) < 2 * cache.stride(1) - ): - raise ValueError( - "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" - ) - if scales.stride(2) != 128 * 8: - raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") - return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + return _modeling_support._nvfp4_fused_write_layout(_SUPPORT_TARGET, data, scales, num_heads) def _producer_miss(covered: bool, rows: Optional[int], reason: str) -> None: @@ -2082,31 +1970,9 @@ def _fp8_index_qk_norm_rope( """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` are the calling target's ``step_args`` entries.""" - if position_ids is None or idx_qk.dtype != torch.bfloat16: - raise NotImplementedError( - "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." - ) - index_k_cache = msa_idx_k_cache(i) - if index_k_cache.dtype != torch.float8_e4m3fn: - raise ValueError( - "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " - f"got {index_k_cache.dtype}." - ) - rows = int(idx_qk.shape[0]) - _record("m3_fp8_indexer_qk_norm_rope", rows) - return m3_fp8_indexer_qk_norm_rope( - idx_qk.contiguous(), - index_k_cache, - msa_out_cache_loc[:rows], - core.index_heads, - core.index_dim, - core.rotary_dim, - core.rms_eps, - core.w[f"l{i}_idx_q_norm"], - core.w[f"l{i}_idx_k_norm"], - core.rope_theta, - position_ids.reshape(-1).contiguous().to(torch.int32), - ).flatten(1) + return _modeling_support._fp8_index_qk_norm_rope( + _SUPPORT_TARGET, core, i, idx_qk, position_ids, msa_idx_k_cache, msa_out_cache_loc + ) # -------------------------------------------------------------------------------------------------------------------- @@ -3932,6 +3798,7 @@ def _decode_sparse_core(self, i: int, q, idx_q, step, output) -> None: qlen, INDEX_SCORE_SPLIT_K, early, + runner=_index_score_kernel.IndexScoreRunner, ) taken(INDEX_SCORE.path, rows) n_valid_blocks = step["msa_n_valid_blocks"][:num_tokens] @@ -4523,41 +4390,14 @@ def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, captu in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; otherwise it is the in-tree capture. """ - if capture_norm is None: - spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) - return - for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): - if captured_layer_id == layer_id: - assert spec_metadata.hidden_states is not None - num_tokens = spec_metadata.num_capture_tokens - if num_tokens > spec_metadata.hidden_states.shape[0]: - raise RuntimeError( - "EAGLE3 hidden-state capture token count exceeds buffer capacity: " - f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" - ) - if num_tokens > hidden_states.shape[0]: - raise RuntimeError( - "EAGLE3 hidden-state capture token count exceeds available hidden states: " - f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" - ) - hidden_size = spec_metadata.hidden_size - m3_spec_tail.capture_norm( - hidden_states[:num_tokens], - residual[:num_tokens] if residual is not None else None, - capture_norm[0][i], - capture_norm[1], - spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], - ) - break + return _modeling_support._eagle3_capture( + _SUPPORT_TARGET, spec_metadata, layer_id, hidden_states, residual, capture_norm + ) def _sharded_head(lm_head: nn.Module) -> bool: """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" - return ( - getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN - and getattr(lm_head, "padding_size", 0) == 0 - and getattr(lm_head, "tp_size", 1) > 1 - ) + return _modeling_support._sharded_head(_SUPPORT_TARGET, lm_head) class M3Eagle3Worker(Eagle3OneModelWorker): @@ -4751,6 +4591,7 @@ def _forward_impl( max_draft_len=runtime_draft_len, ) + next_new_tokens = None if self._uses_external_shared_target_kv: next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( position_ids=position_ids, @@ -4764,6 +4605,8 @@ def _forward_impl( batch_size=batch_size, ) else: + # Retain the pre-draft tensor before metadata is replaced, as in the parent worker. + pre_draft_seq_lens = attn_metadata._seq_lens_cuda # Save the old attn_metadata and spec_metadata self._prepare_attn_metadata_for_spec_dec(attn_metadata) @@ -4786,7 +4629,7 @@ def _forward_impl( # Get the draft KV cache manager if using separate layouts draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - next_draft_tokens = self._forward_draft_loop( + next_draft_tokens, next_new_tokens = self._forward_draft_loop( inputs, attn_metadata, spec_metadata, @@ -4798,6 +4641,8 @@ def _forward_impl( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ) # restore attn_metadata to support cuda graph self._restore_attn_metadata_from_spec_dec(attn_metadata) @@ -4809,7 +4654,7 @@ def _forward_impl( if self._m3_next_new_tokens is not None: # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). next_new_tokens = self._m3_next_new_tokens - else: + elif next_new_tokens is None: next_new_tokens = self._prepare_next_new_tokens( accepted_tokens, next_draft_tokens, @@ -5502,10 +5347,12 @@ def _forward_draft_loop( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ): """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop.""" if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): - return self._m3_fused_draft_loop( + next_draft_tokens = self._m3_fused_draft_loop( inputs, attn_metadata, spec_metadata, @@ -5514,6 +5361,7 @@ def _forward_draft_loop( batch_size, num_accepted_tokens, ) + return next_draft_tokens, None return super()._forward_draft_loop( inputs, attn_metadata, @@ -5526,6 +5374,8 @@ def _forward_draft_loop( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ) def _m3_fused_draft_loop( @@ -5654,40 +5504,7 @@ def _m3_fused_draft_loop( def _build_speculation(shell, model_config: ModelConfig): """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for the llama3 drafter architecture).""" - spec = model_config.spec_config - if ( - not spec.spec_dec_mode.use_one_engine() - or not spec.spec_dec_mode.is_eagle3_one_model() - or spec.eagle3_model_arch != "llama3" - or getattr(spec, "use_dynamic_tree", False) - ): - raise ValueError( - "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" - ) - draft_config = ModelConfig.from_pretrained( - spec.speculative_model, - trust_remote_code=True, - attn_backend=model_config.attn_backend, - moe_backend=_get_requested_draft_moe_backend(model_config, spec), - mapping=model_config.mapping, - spec_config=model_config.spec_config, - max_num_tokens=model_config.max_num_tokens, - moe_max_num_tokens=model_config.moe_max_num_tokens, - ) - _set_draft_kv_cache_quant_algo(draft_config, model_config) - draft_config.extra_attrs = model_config.extra_attrs - use_separate = should_use_separate_draft_kv_cache(spec) - draft = M3Eagle3Draft( - shell.model, - draft_config, - model_config.pretrained_config.num_hidden_layers, - spec.num_capture_layers, - ) - worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) - worker.set_draft_model(draft) - shell.epilogue.append(draft) - shell.epilogue.append(worker) - return draft_config, draft, worker, use_separate + return _modeling_support._build_speculation(_SUPPORT_TARGET, shell, model_config) # -------------------------------------------------------------------------------------------------------------------- @@ -5934,3 +5751,6 @@ def forward( ): spec_fold.store_in_forward(outputs) return outputs + + +_SUPPORT_TARGET = cast(_modeling_support.ModelingSupport, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py deleted file mode 100644 index 5bcdee3fbb82..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/activation/m3_silu_and_mul.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3SiluAndMul(OpWrapper): - """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with - ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain - SwiGLU ``silu(gate) * up``. - - The MXFP8 dense MLPs (layers 0-2, intermediate 3072 per rank) and the MoE's shared expert (intermediate 768 per - rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 4608 per - rank) runs it plain. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="shared expert, one decode token", spec=dict(m=1, inter=768, oai=True)), - Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=768, oai=True)), - Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=3072, oai=True)), - Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=3072, oai=True)), - Cell( - why="drafter MLP (plain SwiGLU), eight decode rows", - spec=dict(m=8, inter=4608, oai=False), - ), - Cell( - why="drafter MLP (plain SwiGLU), a prefill chunk", - spec=dict(m=1024, inter=4608, oai=False), - ), - ) - - note = """ - The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two - differ only by the fast sigmoid's last bits (the test's band). - """ - - def __call__( - self, - x: torch.Tensor, - limit: Optional[float], - alpha: Optional[float], - beta: Optional[float], - ) -> torch.Tensor: - return torch.ops.trtllm.silu_and_mul( - x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta - ) - - def reference( - self, - x: torch.Tensor, - limit: Optional[float], - alpha: Optional[float], - beta: Optional[float], - ) -> torch.Tensor: - gate, up = x.float().chunk(2, dim=-1) - if limit is not None: - gate = gate.clamp(max=limit) - up = up.clamp(min=-limit, max=limit) - a = 1.0 if alpha is None else alpha - b = 0.0 if beta is None else beta - return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) - - def is_valid( - self, - x: torch.Tensor, - limit: Optional[float], - alpha: Optional[float], - beta: Optional[float], - ) -> None: - assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" - - -m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py deleted file mode 100644 index ba008f493b94..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_dense_decode.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per -request count (``modeling.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( - fp8_e4m3_saturate, -) -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 - minimax_m3_trtllm_gen_dense_decode, -) - -__all__ = ["m3_dense_decode"] - -#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. -HEADS = 16 -KV_HEADS = 1 -HEAD_DIM = 128 -#: The split this target hands trtllm-gen at one request (modeling.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model -#: maximum the in-tree step passes otherwise. -CAPPED = 16384 -MODEL_MAX = 1 << 20 - - -def dense_attention_reference( - q: torch.Tensor, - keys: torch.Tensor, - values: torch.Tensor, - seq_lens: torch.Tensor, - decode_query_len: int, - sm_scale: float, -) -> torch.Tensor: - """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first - ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's - tokens in order). Returns fp32 [total_q, heads, head_dim].""" - batch, kv_heads, max_len, _ = keys.shape - heads = q.shape[1] - group = heads // kv_heads - out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) - pos = torch.arange(max_len, device=q.device) - with true_fp32_matmul(): - for b in range(batch): - full = int(seq_lens[b]) - rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] - limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) - valid = pos[None, :] < limit[:, None] # [qlen, max_len] - for h in range(kv_heads): - qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] - logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale - logits = logits.masked_fill(~valid[:, None, :], float("-inf")) - probs = torch.softmax(logits, dim=-1) - out[ - b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group - ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) - return out - - -def gather_paged_kv( - kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool - (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" - pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") - page = pool.shape[2] - pages = (int(seq_lens.max()) + page - 1) // page - slots = block_table[:, :pages].long() - k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] - v = pool[slots * subpages + 1] - batch, _, kv_heads, _, d = k.shape - - def order(t): - return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) - - return order(k), order(v) - - -class M3DenseDecode(OpWrapper): - """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len - ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. - - Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into - ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: - every CTA still covers its share of each request's actual length, which may exceed it. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", - spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), - ), - Cell( - why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", - spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), - ), - Cell( - why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", - spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), - ), - Cell( - why="two requests keep the model maximum (no cap at two requests)", - spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), - ), - ) - - note = """ - The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped - one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV - band of the in-tree dense decode test. - """ - - def __call__( - self, - q: torch.Tensor, - kv_cache_manager, - layer_idx: int, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - sm_scale: float, - output: torch.Tensor, - decode_query_len: int, - max_seq_len: int, - max_num_requests: int, - ) -> torch.Tensor: - minimax_m3_trtllm_gen_dense_decode( - q, - kv_cache_manager, - layer_idx, - block_table, - seq_lens, - sm_scale=sm_scale, - output=output, - decode_query_len=decode_query_len, - max_seq_len=max_seq_len, - max_num_requests=max_num_requests, - ) - return output - - def reference( - self, - q: torch.Tensor, - kv_cache_manager, - layer_idx: int, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - sm_scale: float, - output: torch.Tensor, - decode_query_len: int, - max_seq_len: int, - max_num_requests: int, - ) -> torch.Tensor: - """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). - - ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" - keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) - if keys.dtype == torch.float8_e4m3fn: - q = fp8_e4m3_saturate(q) - return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) - - def is_valid( - self, - q: torch.Tensor, - kv_cache_manager, - layer_idx: int, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - sm_scale: float, - output: torch.Tensor, - decode_query_len: int, - max_seq_len: int, - max_num_requests: int, - ) -> None: - # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. - assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" - # Each verify row attends a causal prefix of at least one key. - assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" - - -m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py index ed4e3048d780..e7fd3d41dac4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_draft_attention.py @@ -4,17 +4,24 @@ from __future__ import annotations -import torch +import torch as torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.draft_attention import ( + DraftAttentionBase, +) +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs as AttentionForwardArgs, +) from tensorrt_llm._torch.attention.backends.interface import ( - AttentionForwardArgs, - PredefinedAttentionMask, + PredefinedAttentionMask as PredefinedAttentionMask, ) -from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.functional import PositionEmbeddingType as PositionEmbeddingType +from tensorrt_llm.functional import RotaryScalingType as RotaryScalingType -class M3DraftAttention(OpWrapper): +class M3DraftAttention(DraftAttentionBase): """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` @@ -38,17 +45,6 @@ class M3DraftAttention(OpWrapper): kernel's input precision, its P rounding and accumulation order), which the test bounds. """ - def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: - num_tokens = metadata.num_tokens - out = backend.forward( - qkv[:num_tokens, :], - None, - None, - metadata, - forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), - ) - return out[0] if isinstance(out, tuple) else out - def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py deleted file mode 100644 index 48543be75675..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fp8_indexer_qk_norm_rope.py +++ /dev/null @@ -1,114 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q -out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( - fused_qk_norm_rope, -) - - -class M3Fp8IndexerQkNormRope(OpWrapper): - """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and - writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, - then RoPE on its leading ``rotary_dim`` channels (NeoX). - - The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; - the fused layers' producer writes the index cache itself. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="one decode token, one index head per rank", spec=dict(m=1, heads=1)), - Cell(why="a 48-row decode verify", spec=dict(m=48, heads=1)), - Cell(why="a prefill chunk", spec=dict(m=2048, heads=1)), - ) - - note = """ - `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 - before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of - values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. - """ - - def __call__( - self, - qk: torch.Tensor, - index_k_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - num_heads_q: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - base: float, - position_ids: torch.Tensor, - ) -> torch.Tensor: - return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( - qk, - index_k_cache, - out_cache_loc, - num_heads_q, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - position_ids, - ) - - def reference( - self, - qk: torch.Tensor, - index_k_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - num_heads_q: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - base: float, - position_ids: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, - then the E4M3 cast.""" - out = fused_qk_norm_rope( - qk.clone(), - num_heads_q, - 1, - 0, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - True, - position_ids, - 1.0, - 0.0, - 0.0, - 1.0, - True, - True, - False, - 0, - 0, - ) - q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) - return ( - q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), - k.to(torch.float8_e4m3fn), - ) - - -m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py index 2c6c2133b2bb..3ddc0bf2986b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_fullctx_decode.py @@ -8,12 +8,15 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_dense_decode import ( + dense_attention_reference, + gather_paged_kv, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( fp8_e4m3_saturate, ) from .. import m3_fullctx_decode as kernel -from .m3_dense_decode import dense_attention_reference, gather_paged_kv __all__ = ["m3_fullctx_decode"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py index 4fb706182485..099a9c32c6fc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_index_score.py @@ -8,13 +8,9 @@ from __future__ import annotations -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_index_score import ( + M3IndexScoreBase, ) __all__ = ["m3_index_score"] @@ -25,17 +21,7 @@ SPLIT_K = 512 -def compiled_scorer(key: tuple): - """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss - (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" - from ..m3_index_score import IndexScoreRunner - - if key not in IndexScoreRunner.kernel_cache: - IndexScoreRunner._compile(*key) - return IndexScoreRunner.kernel_cache[key] - - -class M3IndexScore(OpWrapper): +class M3IndexScore(M3IndexScoreBase): """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is @@ -82,78 +68,5 @@ class M3IndexScore(OpWrapper): differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. """ - def __call__( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> torch.Tensor: - _, num_heads, head_dim = idx_q.shape - key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) - compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) - return score - - def reference( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> torch.Tensor: - """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and - ``early`` do not change them.""" - total_q, heads, _ = idx_q.shape - qlen = max_decode_query_len - out = torch.full( - (heads, total_q, score.shape[2]), - -float("inf"), - dtype=torch.float32, - device=idx_q.device, - ) - with true_fp32_matmul(): - for b, seq in enumerate(seq_lens.tolist()): - nblocks = (seq + PAGE - 1) // PAGE - if nblocks == 0: - continue - t0 = b * qlen - q = idx_q[t0 : t0 + qlen].float() - k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() - s = torch.einsum("qhd,kd->hqk", q, k) - q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) - k_pos = torch.arange(k.shape[0], device=idx_q.device) - s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) - out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) - return out - - def is_valid( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> None: - # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. - nblocks = (seq_lens + PAGE - 1) // PAGE - assert bool((nblocks <= block_table.shape[1]).all()), ( - "a request runs past its block-table row" - ) - assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" - assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( - "idx_q must hold max_decode_query_len rows per request" - ) - m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py index 0ce3c864d30f..6611af074ea2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_msa_attention.py @@ -8,16 +8,25 @@ import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa import ( + MsaAttentionBase, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa import ( + _selected_blocks as _selected_blocks, +) +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs as AttentionForwardArgs, +) from tensorrt_llm._torch.attention.backends.interface import ( - AttentionForwardArgs, - SparseBackendForwardArgs, + SparseBackendForwardArgs as SparseBackendForwardArgs, ) -from .m3_sparse_decode_lc import dequant_nvfp4_pages +from .m3_sparse_decode_lc import dequant_nvfp4_pages as dequant_nvfp4_pages -class M3MsaAttention(OpWrapper): +class M3MsaAttention(MsaAttentionBase): """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` [T, heads * 128] (bf16, the step's live rows). @@ -62,35 +71,6 @@ class M3MsaAttention(OpWrapper): rounding and accumulation order), which the test bounds. """ - def __call__( - self, - backend, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - idx_q: Optional[torch.Tensor], - metadata, - output: torch.Tensor, - sparse: bool, - kv_scale_orig_quant: Optional[torch.Tensor], - kv_scale_quant_orig: Optional[torch.Tensor], - ) -> torch.Tensor: - if k is not None: - backend.write_layer_caches( - k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant - ) - if sparse: - blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) - args = AttentionForwardArgs( - output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) - ) - else: - args = AttentionForwardArgs(output=output) - args.kv_scale_orig_quant = kv_scale_orig_quant - args.kv_scale_quant_orig = kv_scale_quant_orig - backend.forward(q, None, None, metadata, forward_args=args) - return output - def reference( self, backend, @@ -183,42 +163,4 @@ def reference( return out.reshape(total, heads * dim).to(output.dtype) -def _selected_blocks( - idx_q: torch.Tensor, - idx_keys: torch.Tensor, - allowed: torch.Tensor, - qpos: torch.Tensor, - page: int, - config, -) -> torch.Tensor: - """The blocks each query token of one request attends, [n, blocks] bool. - - A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of - ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last - ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the - backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the - ranking, so none is applied. - """ - n, length = allowed.shape - blocks = (length + page - 1) // page - scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) - scores = scores.masked_fill(~allowed, float("-inf")) - padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) - padded[:, :length] = scores - block_scores = padded.view(n, blocks, page).amax(dim=-1) - ids = torch.arange(blocks, device=scores.device)[None, :] - n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] - block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) - forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( - (ids >= n_valid - config.local_blocks) & (ids < n_valid) - ) - block_scores = block_scores.masked_fill(forced, float("inf")) - order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] - keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid - chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) - rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) - chosen[rows[keep], order[keep]] = True - return chosen - - m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py index 5a3d1b64ac58..f781ebfd0493 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_producer.py @@ -5,10 +5,16 @@ from __future__ import annotations +import sys +from typing import cast + import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + producer_reference as _producer_shared, +) __all__ = ["m3_producer"] @@ -27,29 +33,20 @@ def _norm_rope( ): """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. x [T, 128] fp32; cos / sin [T, 32].""" - y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) - lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] - return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + return _producer_shared._norm_rope(_REFERENCE_TARGET, x, weight, cos, sin, eps) def decode_nvfp4_token( data: torch.Tensor, scales: torch.Tensor, quant_scale: float ) -> torch.Tensor: """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" - table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) - nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() - vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) - s = scales.view(torch.float8_e4m3fn).float() / quant_scale - return vals * s.repeat_interleave(FP4_BLOCK) + return _producer_shared.decode_nvfp4_token(_REFERENCE_TARGET, data, scales, quant_scale) def scale_index(within_page: int, plane: int) -> torch.Tensor: """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token swizzle for V.""" - cols = torch.arange(8) - if plane == 0: - return within_page * 8 + cols - return (within_page // 4) * 32 + cols * 4 + within_page % 4 + return _producer_shared.scale_index(_REFERENCE_TARGET, within_page, plane) class M3Producer(OpWrapper): @@ -108,7 +105,7 @@ def __call__( prewait_step_inputs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if prewait_step_inputs: - # This target's copy (producerPw4): the same outputs bit for bit, with the position ids, slots and RoPE + # The TP4 producer specialization: same outputs bit for bit, with position ids, slots and RoPE # rows read before its grid wait (none of them may be written by a PDL predecessor). return tuple( torch.ops.trtllm.minimax_m3_sm107_tp4_producer_pw( @@ -173,20 +170,27 @@ def reference( """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change them.""" - t = packed.shape[0] - x = packed.float().view(t, heads_q + 4, HEAD_DIM) - rope = rotary_cos_sin[position_ids[:t].long()] - cos, sin = rope[:, 0], rope[:, 1] - q = torch.stack( - [_norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + return _producer_shared.producer_reference_single_kv( + self, + _REFERENCE_TARGET, + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + prewait_step_inputs, ) - return { - "q": q, - "k": _norm_rope(x[:, heads_q], k_weight, cos, sin, eps), - "v": x[:, heads_q + 1], - "index_q": _norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), - "index_k": _norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), - } def is_valid( self, @@ -208,16 +212,29 @@ def is_valid( early: bool = True, prewait_step_inputs: bool = False, ) -> None: - t = packed.shape[0] - # The RoPE row is read at the position with no bound. - pos = position_ids[:t] - assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( - "position_ids must index rotary_cos_sin" + return _producer_shared.validate_producer( + self, + _REFERENCE_TARGET, + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + prewait_step_inputs, ) - # Two tokens writing one slot race. - slots = out_cache_loc[:t] - live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] - assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" m3_producer = M3Producer() + +_REFERENCE_TARGET = cast(_producer_shared.ProducerTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py deleted file mode 100644 index ef79579c0c21..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qkv_insert.py +++ /dev/null @@ -1,134 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into -the paged caches.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3QkvInsert(OpWrapper): - """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the - partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and - index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. - - ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool - [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 - [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp4 - heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), - Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), - ) - - note = """ - The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` - returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. - """ - - def __call__( - self, - packed: torch.Tensor, - kv_cache: torch.Tensor, - kv_scale: Optional[torch.Tensor], - index_k_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - kv_quant_scale: Optional[torch.Tensor], - num_heads: int, - num_kv_heads: int, - num_index_heads: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_norm: torch.Tensor, - k_norm: torch.Tensor, - index_q_norm: torch.Tensor, - index_k_norm: torch.Tensor, - rotary_cos_sin: torch.Tensor, - position_ids: torch.Tensor, - ): - common = ( - num_heads, - num_kv_heads, - num_index_heads, - head_dim, - rotary_dim, - eps, - q_norm, - k_norm, - index_q_norm, - index_k_norm, - rotary_cos_sin, - position_ids, - ) - if kv_scale is not None: - return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( - packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common - ) - return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( - packed, kv_cache, index_k_cache, out_cache_loc, *common - ) - - def reference( - self, - packed: torch.Tensor, - kv_cache: torch.Tensor, - kv_scale: Optional[torch.Tensor], - index_k_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - kv_quant_scale: Optional[torch.Tensor], - num_heads: int, - num_kv_heads: int, - num_index_heads: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_norm: torch.Tensor, - k_norm: torch.Tensor, - index_q_norm: torch.Tensor, - index_k_norm: torch.Tensor, - rotary_cos_sin: torch.Tensor, - position_ids: torch.Tensor, - ): - """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of - Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as - packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" - m = packed.shape[0] - sizes = [ - num_heads * head_dim, - num_kv_heads * head_dim, - num_kv_heads * head_dim, - num_index_heads * head_dim, - head_dim, - ] - q, k, v, iq, ik = packed.float().split(sizes, dim=-1) - cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) - half = rotary_dim // 2 - - def norm_rope(x, w, heads): - x = x.reshape(m, heads, head_dim) - x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) - x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] - c, s = cos.unsqueeze(1), sin.unsqueeze(1) - return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) - - return { - "q": norm_rope(q, q_norm, num_heads), - "index_q": norm_rope(iq, index_q_norm, num_index_heads), - "k": norm_rope(k, k_norm, num_kv_heads), - "v": v.reshape(m, num_kv_heads, head_dim), - "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], - } - - -m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py deleted file mode 100644 index 76841cca732d..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_blocks.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3SelectBlocks(OpWrapper): - """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), - with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the - row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` - backs the result head-major so ``permute(1, 0, 2)`` is contiguous. - - The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="eight decode rows of 17..128 blocks", - spec=dict(total_q=8, n_blocks=128, mixed=True), - ), - Cell( - why="one decode row at the 128-block edge", - spec=dict(total_q=1, n_blocks=128, mixed=False), - ), - Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), - ) - - note = """ - Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. - """ - - def __call__( - self, - max_score_kv: torch.Tensor, - topk: int, - n_valid_blocks: torch.Tensor, - init_blocks: int, - local_blocks: int, - head_major_output: bool, - ) -> torch.Tensor: - nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() - return torch.ops.trtllm.minimax_m3_select_blocks( - max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output - ) - - def reference( - self, - max_score_kv: torch.Tensor, - topk: int, - n_valid_blocks: torch.Tensor, - init_blocks: int, - local_blocks: int, - head_major_output: bool, - ) -> torch.Tensor: - """Plain torch selection; ``head_major_output`` changes the backing only.""" - heads, n_blocks, total_q = max_score_kv.shape - out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) - blocks = torch.arange(n_blocks, device=max_score_kv.device) - for t in range(total_q): - nv = int(n_valid_blocks[t]) - for h in range(heads): - s = max_score_kv[h, :, t].clone().double() - s[blocks >= nv] = float("-inf") - forced = (blocks < min(init_blocks, nv)) | ( - (blocks >= nv - local_blocks) & (blocks < nv) - ) - s[forced] = float("inf") - order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] - ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) - out[t, h, : ids.numel()] = ids - return out - - -m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py index 3d076c8181d6..ce2adce0c3e8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_select_long.py @@ -4,10 +4,19 @@ from __future__ import annotations +import sys +from typing import cast + import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.select_reference import ( + SelectReferenceTarget, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.select_reference import ( + select_reference as _shared_select_reference, +) #: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). MAX_BLOCKS = 9216 @@ -127,32 +136,18 @@ def reference( (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and ``trigger_phase`` and ``fast_tail`` only say where and how the op writes; they do not enter the result. """ - num_kv_heads, num_blocks, total_q = scores.shape - device = scores.device - blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) - raw = n_valid_blocks[:total_q].to(torch.int64) - valid = raw.clamp(0, num_blocks) - local_start = (raw - local_blocks).clamp_min(0) - result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) - for head in range(num_kv_heads): - forced = scores[head].float().t().clone() # [total_q, num_blocks] - forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) - forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) - bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF - key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) - in_row = blocks[None, :] < valid[:, None] - packed = torch.where( - in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max - ) - k = min(TOPK, num_blocks) - picked = torch.topk(packed, k, dim=1, largest=False).values - picked_key = picked >> 16 - keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) - ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) - ids = ids.sort(dim=1).values - ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) - result[:, head, :k] = ids.to(torch.int32) - return result + return _shared_select_reference( + self, + _SELECT_REFERENCE_TARGET, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) def is_valid( self, @@ -184,3 +179,5 @@ def is_valid( m3_select_long = M3SelectLong() + +_SELECT_REFERENCE_TARGET = cast(SelectReferenceTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py index 96171a6a9a29..5313d5d7a672 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode.py @@ -9,6 +9,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_reference import ( + validate_sparse_selection, +) from ..triton_sparse_decode import minimax_m3_sparse_attn_decode from .m3_sparse_decode_lc import PAGE, sparse_attention_reference @@ -183,33 +186,9 @@ def is_valid( num_topk_chunks: Optional[int] = None, identity_blocks: bool = False, ) -> None: - """The inputs this call takes and answers wrongly. - - Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` - table entries and indexes the block table with them unchecked, so they - must be real blocks of the request. ``identity_blocks`` is a promise the - multi-query producer relies on: every token of a request reads the same - blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the - host: a test-only cost. - """ - total_q = q.shape[0] - seq = seq_lens.tolist() - max_topk = topk_idx.shape[-1] - for tok in range(total_q): - req, offset = divmod(tok, decode_query_len) - kv_len = max(seq[req] - decode_query_len + offset + 1, 0) - nblocks = (kv_len + PAGE - 1) // PAGE - ids = topk_idx[:, tok, : min(max_topk, nblocks)] - assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( - f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " - "the block table is indexed with them unchecked" - ) - if identity_blocks: - assert nblocks <= max_topk and bool( - (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() - ), ( - f"token {tok}: identity_blocks promises the table holds every block of the request in order" - ) + validate_sparse_selection( + q, topk_idx, block_table, seq_lens, decode_query_len, identity_blocks, PAGE + ) m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py index 5bafb312b873..77a32cda65e8 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_sparse_decode_lc.py @@ -16,6 +16,13 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from ....ops.attention.sparse_reference import _E2M1 as _E2M1 +from ....ops.attention.sparse_reference import dequant_nvfp4_pages as _dequant_nvfp4_pages +from ....ops.attention.sparse_reference import mxfp8_handoff_reference as _mxfp8_handoff_reference +from ....ops.attention.sparse_reference import ( + sparse_attention_reference as _sparse_attention_reference, +) + #: The TP4 geometry the kernel is compiled for. NUM_HEADS = 16 HEAD_DIM = 128 @@ -27,32 +34,12 @@ MX_BLOCK = 32 WIDTH = NUM_HEADS * HEAD_DIM -_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) - def dequant_nvfp4_pages( packed: torch.Tensor, block_scale: torch.Tensor, pages: torch.Tensor, quad_scales: bool ) -> torch.Tensor: - """Dequantize KV head 0 of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. - - ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low - nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, - HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token - ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at - ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor - global scale is not applied. Every product is exact in float64. - """ - table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) - rows = packed[pages, 0].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] - nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) - cols = HEAD_DIM // SF_VEC - region = block_scale[pages, 0].reshape(rows.shape[0], PAGE * cols) - t = torch.arange(PAGE, device=packed.device)[:, None] - c = torch.arange(cols, device=packed.device)[None, :] - index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c - scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) - scales = scales.view(torch.float8_e4m3fn).to(torch.float64) - return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + """Use the shared reference with this target's geometry.""" + return _dequant_nvfp4_pages(packed, block_scale, pages, quad_scales) def sparse_attention_reference( @@ -69,87 +56,28 @@ def sparse_attention_reference( decode_query_len: int, sm_scale: float, ) -> torch.Tensor: - """Block-sparse decode attention over one KV head of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. - - Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset - ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / - PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at - ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that - attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or - its byte view); softmax and both products are float64. - """ - q64 = ( - q.view(torch.float8_e4m3fn).to(torch.float64) - if q.element_size() == 1 - else q.to(torch.float64) + """Use the shared reference with this target's geometry.""" + return _sparse_attention_reference( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, ) - total_q, num_heads, _ = q64.shape - kgs = float(k_global_scale.reshape(-1)[0]) - vgs = float(v_global_scale.reshape(-1)[0]) - out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) - seq = seq_lens.tolist() - positions = torch.arange(PAGE, device=q.device) - for tok in range(total_q): - req, offset = divmod(tok, decode_query_len) - kv_len = max(seq[req] - decode_query_len + offset + 1, 0) - real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) - blocks = topk[tok, :real_topk] - blocks = blocks[blocks >= 0].to(torch.int64) - if blocks.numel() == 0: - continue - pages = block_table[req, blocks].to(torch.int64) - keys = ( - dequant_nvfp4_pages(k, k_block_scale, pages, quad_scales=False).reshape(-1, HEAD_DIM) - * kgs - ) - values = ( - dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True).reshape(-1, HEAD_DIM) - * vgs - ) - valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) - logits = (q64[tok] @ keys.T) * sm_scale - logits = logits.masked_fill(~valid[None, :], float("-inf")) - values = values.masked_fill(~valid[:, None], 0) - out[tok] = torch.softmax(logits, dim=-1) @ values - return out def mxfp8_handoff_reference( out: torch.Tensor, num_rows_alloc: int ) -> Tuple[torch.Tensor, torch.Tensor]: - """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, - sf [num_rows_alloc * WIDTH // 32] uint8)``. - - Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` - rounded up when its mantissa is non-zero (capped at 254), the elements are - scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands - in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` - for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a - comparison against it is exact. - """ - total_q = out.shape[0] - x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) - amax = x.abs().amax(dim=-1) - sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) - bits = sfv.view(torch.int32) - e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) - e = e.clamp(max=254) - qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) - xq = (x * qs[..., None]).clamp(-448.0, 448.0) - q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) - num_sf_k_tiles = WIDTH // 128 - sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) - tok = torch.arange(total_q, device=out.device)[:, None] - k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] - offset = ( - (tok // 128) * num_sf_k_tiles * 512 - + (k_sf // 4) * 512 - + (tok % 32) * 16 - + ((tok % 128) // 32) * 4 - + k_sf % 4 - ) - sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) - return q8, sf + """Use the shared reference with this target's geometry.""" + return _mxfp8_handoff_reference(out, num_rows_alloc, width=WIDTH) class M3SparseDecodeLc(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py deleted file mode 100644 index 43166b472eca..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/msa_kv_update.py +++ /dev/null @@ -1,200 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" - -from __future__ import annotations - -from typing import Optional, Tuple - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: The MSA pool's page: one sparse block is one 128-token KV page. -TOKENS_PER_BLOCK = 128 - - -class MsaKvUpdate(OpWrapper): - """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. - - ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its - own buffers at the head of every captured speculative decode step. One - single-CTA launch writes, for ``batch`` requests and ``total_q`` query - tokens (all int32): - - * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, - ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); - * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` - attends up to, through ``block_table``; - * ``n_valid_blocks[t]``, when given, the block count top-k selection is - bounded by. - - It returns None; the cells read the four buffers. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="decode without speculation at one request: qo = 1, every buffer written", - spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), - ), - Cell( - why="decode without speculation at the captured batch 8", - spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), - ), - Cell( - why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", - spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), - ), - Cell( - why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", - spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), - ), - Cell( - why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", - spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), - ), - Cell( - why="positions crossing a page boundary: the slot moves to the previous page of the request", - spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), - ), - ) - - note = """ - The kernel covers any batch up to the block size's loop (it strides over - 256 threads), and any page size; only `tokens_per_block = 128`, the MSA - pool's page, is certified. - - The op does not bounds-check the request index a token reads - (`q_batch_row`) or the block-table column it derives: both come from the - metadata's own staging, which keeps them in range by construction, and an - out-of-range value reads whatever memory sits there. `is_valid` states both. - """ - - def __call__( - self, - kv_lens: torch.Tensor, - kv_lens_staged: torch.Tensor, - seq_lens: torch.Tensor, - q_batch_row: torch.Tensor, - q_intra: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - batch: int, - total_q: int, - msa_seq_lens: torch.Tensor, - cu_kv_lens: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - n_valid_blocks: Optional[torch.Tensor], - ) -> None: - torch.ops.trtllm.minimax_m3_msa_kv_update( - kv_lens, - kv_lens_staged, - seq_lens, - q_batch_row, - q_intra, - block_table, - tokens_per_block, - batch, - total_q, - msa_seq_lens, - cu_kv_lens, - out_cache_loc, - n_valid_blocks, - ) - - def reference( - self, - kv_lens: torch.Tensor, - kv_lens_staged: torch.Tensor, - seq_lens: torch.Tensor, - q_batch_row: torch.Tensor, - q_intra: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - batch: int, - total_q: int, - msa_seq_lens: torch.Tensor, - cu_kv_lens: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - n_valid_blocks: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: - """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], - n_valid_blocks[:total_q])`` as the op is supposed to write them. - - The in-tree metadata's own torch re-derivation - (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook - call), statement for statement, returning new tensors rather than - writing the buffers. The buffer arguments only say which outputs exist: - a ``None`` buffer gives a ``None`` result. - """ - page = int(tokens_per_block) - kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) - cu_kv = None - if cu_kv_lens is not None: - cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) - cu_kv[1:] = torch.cumsum(kv_true, 0) - qbr = q_batch_row[:total_q].to(torch.long) - qo_dev = seq_lens[:batch] - pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] - pos_slot = pos.clamp_min(0) - block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( - max=int(block_table.shape[1]) - 1 - ) - slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) - n_valid = None - if n_valid_blocks is not None: - n_valid = torch.div( - (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" - ).to(torch.int32) - return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid - - def is_valid( - self, - kv_lens: torch.Tensor, - kv_lens_staged: torch.Tensor, - seq_lens: torch.Tensor, - q_batch_row: torch.Tensor, - q_intra: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - batch: int, - total_q: int, - msa_seq_lens: torch.Tensor, - cu_kv_lens: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - n_valid_blocks: Optional[torch.Tensor], - ) -> None: - """The inputs this op takes and answers wrongly. - - Every operand but the block table is read as a dense run from its first - element, so a strided view reads the wrong elements. The request index of - each token and the table row it selects are not bounds-checked. Reads - ``q_batch_row`` back to the host: a test-only cost. - """ - for name, t in ( - ("kv_lens", kv_lens), - ("kv_lens_staged", kv_lens_staged), - ("seq_lens", seq_lens), - ("q_batch_row", q_batch_row), - ("q_intra", q_intra), - ("msa_seq_lens", msa_seq_lens), - ("out_cache_loc", out_cache_loc), - ("cu_kv_lens", cu_kv_lens), - ("n_valid_blocks", n_valid_blocks), - ): - assert t is None or t.is_contiguous(), ( - f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" - ) - rows = q_batch_row[:total_q] - assert total_q == 0 or ( - int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) - ), ( - "q_batch_row must index a request below batch (and a block-table row): the op reads the " - "request's lengths and table row unchecked" - ) - - -msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py index 0b999befbadc..f4ac067feb68 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_allreduce_norm.py @@ -4,85 +4,23 @@ from __future__ import annotations -from typing import Optional, Sequence +from typing import Optional import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper -_E4M3_MAX = 448.0 +from ....ops.comm.reference import _E4M3_MAX as _E4M3_MAX +from ....ops.comm.reference import McAllreduceNormReference +from ....ops.comm.reference import boundary as boundary +from ....ops.comm.reference import mxfp8_rows as mxfp8_rows +from ....ops.comm.reference import sf_written as sf_written +from ....ops.comm.reference import swizzle_sf as swizzle_sf -def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. - - One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), - data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the - op leaves them (the comparison masks them out). - """ - m, h = normed.shape - blocks = normed.float().view(m, h // 32, 32) - amax = blocks.abs().amax(-1) - bits = (amax / _E4M3_MAX).view(torch.int32) - byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) - scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) - data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) - data = data.view(m, h).to(torch.float8_e4m3fn) - return data, swizzle_sf(byte.to(torch.uint8)) - - -def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: - """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" - rows, cols = byte.shape - padded_rows = (rows + 127) // 128 * 128 - r = torch.arange(rows, device=byte.device).view(-1, 1) - c = torch.arange(cols, device=byte.device).view(1, -1) - index = ( - (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) - ) - out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) - out[index.reshape(-1)] = byte.reshape(-1) - return out - - -def sf_written(rows: int, cols: int, device) -> torch.Tensor: - """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" - return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() - - -def boundary( - partials: Sequence[torch.Tensor], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - drafter: bool = False, -) -> tuple[torch.Tensor, ...]: - """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). - - ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer - fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; - only the stored residual is rounded to BF16. - - The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are - summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic - and the normalization are done in float64 and rounded to BF16 once. - """ - total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) - for p in partials: - total = total + p.float() - pre = total.to(torch.bfloat16).float() + residual.float() - updated = pre.to(torch.bfloat16) - u = pre.double() if drafter else updated.double() - inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) - normed = (u * inv * norm_w.double()).to(torch.bfloat16) - if not quant: - return normed, updated - return (normed, updated, *mxfp8_rows(normed)) - - -class M3McAllreduceNorm(OpWrapper): +class M3McAllreduceNorm(McAllreduceNormReference): """One TP group's boundary: ``updated = bf16(residual + sum_r bf16(x_r + x2_r))``, ``normed = RMSNorm(updated)``. Every rank pushes its partial through the multicast mapping ``ws_mc_ptr`` into every rank's exchange, then polls @@ -262,29 +200,6 @@ def __call__( cap_eps, ) - def reference( - self, - per_rank_x: Sequence[torch.Tensor], - per_rank_x2: Optional[Sequence[torch.Tensor]], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - drafter: bool = False, - ) -> tuple[torch.Tensor, ...]: - """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). - - The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling - rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. - """ - partials = [] - for r, x in enumerate(per_rank_x): - if per_rank_x2 is None: - partials.append(x) - else: - partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) - return boundary(partials, residual, norm_w, eps, quant, drafter) - def is_valid( self, x: torch.Tensor, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py index 0d457022c38f..58c52cfd8532 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -4,40 +4,20 @@ from __future__ import annotations -from typing import Optional, Sequence +from typing import Optional import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper -from .m3_mc_allreduce_norm import boundary +from ....ops.comm.reference import McFinalizeAllreduceNormReference +from ....ops.comm.reference import boundary as boundary +from ....ops.comm.reference import finalize_partial as finalize_partial -def finalize_partial( - fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] -) -> torch.Tensor: - """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k - order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). - - The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two - FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. - """ - m, topk = perm.shape - acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) - sc = scales.reshape(m, topk).float() - for k in range(topk): - rows = perm[:, k].long() - keep = (rows >= 0).unsqueeze(-1) - gathered = fc2[rows.clamp_min(0)].double() - step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() - acc = torch.where(keep, step, acc) - if shared is not None: - acc = acc.to(torch.bfloat16).float() + shared.float() - return acc.to(torch.bfloat16) - - -class M3McFinalizeAllreduceNorm(OpWrapper): +class M3McFinalizeAllreduceNorm(McFinalizeAllreduceNormReference): """The MoE boundary with the routed experts' finalize folded in. Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, @@ -132,30 +112,6 @@ def __call__( opts, ) - def reference( - self, - per_rank_fc2: Sequence[torch.Tensor], - per_rank_perm: Sequence[torch.Tensor], - per_rank_scales: Sequence[torch.Tensor], - per_rank_shared: Optional[Sequence[torch.Tensor]], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - ) -> tuple[torch.Tensor, ...]: - """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the - collective exception, as in ``m3_mc_allreduce_norm.reference``).""" - partials = [ - finalize_partial( - fc2, - per_rank_perm[r], - per_rank_scales[r], - None if per_rank_shared is None else per_rank_shared[r], - ) - for r, fc2 in enumerate(per_rank_fc2) - ] - return boundary(partials, residual, norm_w, eps, quant) - def is_valid( self, fc2: torch.Tensor, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py deleted file mode 100644 index 78563b4db40d..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py +++ /dev/null @@ -1,128 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. - -128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original -K3 wrapper this target does not promise FlashInfer's bit-identical reduction. -""" - -import cuda.bindings.driver as cuda_driver -import cutlass -import cutlass.cute as cute -from cutlass.experimental import primitives as prims - -NORM_THREADS = 128 -NORM_VEC = 8 -NORM_WARPS = 4 # capacity: fixed: the CTA's warps - - -@cute.kernel -def m3_embed_norm_kernel( - ids: cute.Tensor, # int32 or int64 [N] - table: cute.Tensor, # bf16 [V, H] - weight: cute.Tensor, # bf16 [H] - raw: cute.Tensor, # bf16 [N, H], out: the rows - out: cute.Tensor, # bf16 [N, H], out: the normed rows - vocab: cutlass.Int64, - eps: cutlass.Float32, - tv_layout: cute.Layout, - tiler_mn: cute.Shape, - hidden: cutlass.Constexpr[int], -): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) - prims.griddepcontrol(prims.GridDepAction.WAIT) - - smem = cutlass.utils.SmemAllocator() - reduction_buffer = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 - ) - - tok = cutlass.Int64(ids[bidx]) - valid = (tok >= cutlass.Int64(0)) & (tok < vocab) - row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) - - gX = cute.local_tile(table, tiler_mn, (row, 0)) - gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) - gY = cute.local_tile(out, tiler_mn, (bidx, 0)) - w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) - gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) - - copy_atom_load = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 - ) - copy_atom_store = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 - ) - thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) - thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) - thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) - - tXgX = thr_copy_X.partition_S(gX) - tXrX = cute.make_fragment_like(tXgX) - tWgW = thr_copy_W.partition_S(gW) - tWrW = cute.make_fragment_like(tWgW) - tXrW = thr_copy_X.retile(tWrW) - tXgR = thr_copy_O.partition_D(gR) - tXgO = thr_copy_O.partition_D(gY) - tXrO = cute.make_fragment_like(tXgO) - - # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. - tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) - if valid: - cute.copy(copy_atom_load, tXgX, tXrX) - cute.copy(copy_atom_load, tWgW, tWrW) - cute.copy(copy_atom_store, tXrX, tXgR) - - # flashinfer's RMSNormKernel arithmetic, in its order (see the header). - x = tXrX.load().to(cutlass.Float32) - x_sq = x * x - sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) - for i in cutlass.range_constexpr(5): - sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) - lane = cute.arch.lane_idx() - warp = cute.arch.warp_idx() - if lane == 0: - reduction_buffer[0, warp] = sum_sq - cute.arch.barrier() - total = cutlass.Float32(0.0) - if lane < NORM_WARPS: - total = reduction_buffer[0, lane] - for i in cutlass.range_constexpr(5): - total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) - mean_sq = total / cutlass.Float32(hidden) - rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) - w = tXrW.load().to(cutlass.Float32) - y = x * rstd * (w + cutlass.Float32(0.0)) - tXrO.store(y.to(out.element_type)) - cute.copy(copy_atom_store, tXrO, tXgO) - - -@cute.jit -def m3_embed_norm( - ids: cute.Tensor, - table: cute.Tensor, - weight: cute.Tensor, - raw: cute.Tensor, - out: cute.Tensor, - vocab: cutlass.Int64, - eps: cutlass.Float32, - n_tokens: cutlass.Constexpr[int], - hidden: cutlass.Constexpr[int], - use_pdl: cutlass.Constexpr[bool], - stream: cuda_driver.CUstream, -) -> None: - # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. - blocks = hidden // (NORM_VEC * NORM_THREADS) - tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), - stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip - m3_embed_norm_kernel( - ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden - ).launch( - grid=[n_tokens, 1, 1], - block=[NORM_THREADS, 1, 1], - smem=NORM_WARPS * 4, - stream=stream, - use_pdl=use_pdl, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py index 276f8c32023b..f123d2f97db2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding_firstnorm/op.py @@ -7,15 +7,22 @@ (:class:`DecodeEmbedding` compiles at construction). """ +import sys import threading -from typing import Dict, Tuple +from typing import Dict, Tuple, cast -import cuda.bindings.driver as cuda_driver -import cutlass.cute as cute +import cuda.bindings.driver as cuda_driver # noqa: F401 — used by the bound family implementation +import cutlass.cute as cute # noqa: F401 — used by the bound family implementation import torch -from cutlass.cute.runtime import from_dlpack +from cutlass.cute.runtime import from_dlpack as from_dlpack -from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding import ( + decode_embedding as _embedding_shared, +) + +from ....kernels.embedding_firstnorm.m3_embedding_firstnorm_kernel import ( # noqa: F401 + m3_embed_norm as _m3_embed_norm_jit, +) HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size MAX_ROWS = 8 # capacity: fixed: the kernel's built rows (decode steps) @@ -25,12 +32,7 @@ def arg(tensor): - return from_dlpack( - tensor.detach(), - assumed_align=tensor.element_size() - if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) - else 16, - ) + return _embedding_shared.arg(_EMBED_TARGET, tensor) def _check(ids, table, weight, raw, out) -> None: @@ -57,32 +59,7 @@ def _check(ids, table, weight, raw, out) -> None: def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): - key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) - fn = _compiled.get(key) - if fn is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "m3_embed_norm compiles on its first call per shape: call it once before capture" - ) - with _lock: - fn = _compiled.get(key) - if fn is None: - stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) - fn = _compiled[key] = cute.compile( - _m3_embed_norm_jit, - arg(ids), - arg(table), - arg(weight), - arg(raw), - arg(out), - int(table.shape[0]), - float(eps), - ids.shape[0], - HIDDEN, - bool(use_pdl), - stream, - ) - return fn + return _embedding_shared._compile(_EMBED_TARGET, ids, table, weight, eps, raw, out, use_pdl) def m3_embed_norm( @@ -96,20 +73,9 @@ def m3_embed_norm( ) -> Tuple[torch.Tensor, torch.Tensor]: """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / ``eps`` into ``out``; returns ``(out, raw)``.""" - _check(ids, table, weight, raw, out) - fn = _compile(ids, table, weight, eps, raw, out, use_pdl) - stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) - fn( - arg(ids), - arg(table), - arg(weight), - arg(raw), - arg(out), - int(table.shape[0]), - float(eps), - stream, + return _embedding_shared.m3_embed_norm( + _EMBED_TARGET, ids, table, weight, eps, raw, out, use_pdl ) - return out, raw class DecodeEmbedding: @@ -118,25 +84,12 @@ class DecodeEmbedding: ``eps``. Built (and compiled) before CUDA-graph capture.""" def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): - assert not torch.cuda.is_current_stream_capturing() - assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN - assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) - assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) - self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( - table, - weight, - rows, - id_dtype, - use_pdl, + _embedding_shared.init_decode_embedding( + self, _EMBED_TARGET, table, weight, eps, rows, id_dtype, use_pdl ) - self.eps = float(eps) - self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) - self.output = torch.empty_like(self.raw) - ids = torch.empty(rows, device=table.device, dtype=id_dtype) - _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) def __call__(self, ids): - assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() - return m3_embed_norm( - ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl - ) + return _embedding_shared.call_decode_embedding(self, _EMBED_TARGET, ids) + + +_EMBED_TARGET = cast(_embedding_shared.EmbeddingTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py deleted file mode 100644 index 4fb5af737dc0..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_bf16_linear.py +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" - -from __future__ import annotations - -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3Bf16Linear(OpWrapper): - """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. - - The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), - o_proj (this rank's input columns), gate_up and down projections run it. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), - Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=2304, k=12288)), - Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=2048)), - Cell(why="drafter gate_up, one row", spec=dict(m=1, n=9216, k=6144)), - Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=4608)), - ) - - note = """ - cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. - """ - - def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - return F.linear(x, w) - - def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - return (x.float() @ w.float().t()).to(x.dtype) - - -m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py deleted file mode 100644 index 1fb7a1d759ae..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_dense_swiglu.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The dense MLP's swigluoai activation as one C++ launch with PDL and an early release (DENSE_SWIGLU_CPP).""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -# The dense layers' swigluoai parameters, which the kernel hard-codes (modeling_minimaxm3._build_swiglu_oai_dense_mlp). -LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 - - -class M3DenseSwiglu(OpWrapper): - """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, - rounded once to bf16: trtllm::minimax_m3_sm107_tp4_swiglu_oai (minimaxM3Sm107Tp4/swigluOai.cu), bit for bit the - Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the - grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights - during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even - half width; ``__call__`` returns None otherwise or for other activation parameters (the caller runs the Triton - op).""" - - ARCHS = frozenset({Arch.SM_107}) - - # The TP4 dense MLP shard: gate | up 6144 wide (d = 3072), at the decode rows. - CELLS: tuple[Cell, ...] = ( - Cell(why="the TP4 dense MLP at 1 row", spec=dict(rows=1, d=3072)), - Cell(why="the TP4 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=3072)), - Cell(why="the TP4 dense MLP at 8 rows", spec=dict(rows=8, d=3072)), - ) - - note = """ - Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the - alpha-scaled ex2.approx sigmoid, div.full), checked over every bf16 gate and every bf16 up pattern. - """ - - def __call__( - self, - x: torch.Tensor, - limit: Optional[float] = LIMIT, - alpha: Optional[float] = ALPHA, - beta: Optional[float] = BETA, - use_pdl: bool = True, - ) -> Optional[torch.Tensor]: - if ( - (limit, alpha, beta) != (LIMIT, ALPHA, BETA) - or x.dtype != torch.bfloat16 - or x.dim() != 2 - or not x.is_cuda - or not x.is_contiguous() - or x.shape[1] % 4 != 0 - or x.data_ptr() % 4 != 0 - ): - return None - out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) - torch.ops.trtllm.minimax_m3_sm107_tp4_swiglu_oai(x, out, bool(use_pdl)) - return out - - def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: - return torch.ops.trtllm.silu_and_mul( - x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta - ) - - def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: - assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" - - -m3_dense_swiglu = M3DenseSwiglu() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py index e4cd998d587d..5b640b47014f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_draft_lowm.py @@ -95,9 +95,13 @@ def __call__( run_direct_dense(x, weight.t(), out, use_pdl, DirectTactic(*tactic)) elif own_splitk: - # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` + # The shared split-K kernel (kernels/gemm/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` # it also pushes the rows into a drafter boundary all-reduce (modeling.DRAFT_AR_PUSH). - from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_lowm_splitk import ( + SplitKTactic, + run_splitk_dense, + run_splitk_dense_push, + ) if ar is not None: run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py deleted file mode 100644 index 9702cc0778a4..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_fp32_linear.py +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" - -from __future__ import annotations - -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) - - -class M3Fp32Linear(OpWrapper): - """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the - fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies - (the engine leaves torch's default).""" - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), - Cell( - why="a generation step above the tensor-core router's 48 rows", - spec=dict(m=64, e=128, k=6144), - ), - ) - - note = """ - Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against - the true fp32 reference. - """ - - def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - return F.linear(x.to(torch.float32), w) - - def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - with true_fp32_matmul(): - return x.float() @ w.t() - - -m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py deleted file mode 100644 index 0b422a7167e5..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_gemm.py +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The decode-graph MXFP8 GEMM on an activation its producer already quantized (rows the GEMV does not serve).""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned - -from .m3_mxfp8_linear import mxfp8_reference - - -class M3Mxfp8Gemm(OpWrapper): - """``act @ weight^T`` for an MXFP8 activation (e4m3 [M, K] + 128x4-swizzled UE8M0 scales) and an MXFP8 weight - [N, K] (swizzled scales), bf16 out. The GEMM backend per token bucket is the AutoTuner's (the same runner, op key - and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The decode dispatch hands it the boundary / merge MXFP8 activation on rows the GEMV does not serve: the - # unsplit shared down projection at 17..32 rows. - CELLS: tuple[Cell, ...] = ( - Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=768)), - Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=768)), - ) - - note = """ - Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and - multiplies in fp32; the test's band covers the accumulation order. - """ - - def __call__( - self, - act: torch.Tensor, - act_scale: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - output_dtype: torch.dtype, - tune: bool, - ) -> torch.Tensor: - return mxfp8_gemm_autotuned(act, act_scale, weight, weight_scale, output_dtype, tune=tune) - - def reference( - self, - act: torch.Tensor, - act_scale: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - output_dtype: torch.dtype, - tune: bool, - ) -> torch.Tensor: - """The dequantized product; ``tune`` is ignored.""" - return mxfp8_reference(act, act_scale, weight, weight_scale).to(output_dtype) - - -m3_mxfp8_gemm = M3Mxfp8Gemm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py index ee781b16ff42..542e69f08d53 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_pinned_gemm.py @@ -9,10 +9,11 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_linear import ( + mxfp8_reference, +) from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner -from .m3_mxfp8_linear import mxfp8_reference - # One runner per output dtype (a runner holds no per-call state). _RUNNERS: Dict[Tuple[type, torch.dtype], object] = {} diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py deleted file mode 100644 index f8deb858c843..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_target_head.py +++ /dev/null @@ -1,346 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, -and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). - -The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values -``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 -marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a -per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: -per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 -code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved -it decodes to the BF16 weight bit for bit. -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) - -__all__ = ["m3_target_head"] - -#: Verify rows per call (requests x (draft length + 1)). -MAX_ROWS = 8 -#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). -CTA_WARPS = 16 -#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. -TARGET_SHARD = 50016 -HIDDEN = 6144 -#: Codes per window: 0..14 decode, 15 marks an exception. -WINDOW = 15 -#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. -MAX_BASE = 254 - WINDOW -#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). -ROWS_PER_CHUNK = 4096 - - -def _pad_rows(weight: torch.Tensor) -> torch.Tensor: - n = weight.shape[0] - tiles = (n + 15) // 16 - if tiles * 16 == n: - return weight - return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) - - -def _fragment_view(w: torch.Tensor) -> torch.Tensor: - """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" - n, k = w.shape - # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) - return ( - w.view(n // 16, 2, 8, k // 16, 2, 4, 2) - .permute(0, 3, 2, 5, 4, 1, 6) - .reshape(n // 16, k // 16, 32, 8) - ) - - -def _exponents(weight: torch.Tensor) -> torch.Tensor: - return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF - - -def _window_base(hist: torch.Tensor) -> int: - """The base of the 15-exponent window holding the most weights (the lowest such base).""" - window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) - return int(window.argmax()) - - -def encode(weight: torch.Tensor, base: Optional[int] = None): - """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's - own window. - - Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile - << 16 | k), exc_w fp32 [E, 2] ({w, 0})). - """ - w = _pad_rows(weight) - n, k = w.shape - tiles, ks = n // 16, k // 16 - assert ks % 4 == 0, "K must be a multiple of 64" - bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF - expo = (bits >> 7) & 0xFF - if base is None: - base = _window_base(torch.bincount(expo.flatten(), minlength=256)) - code = expo - base - exc = (code < 0) | (code >= WINDOW) - code = torch.where(exc, torch.full_like(code, WINDOW), code) - sm = ((bits >> 8) & 0x80) | (bits & 0x7F) - sm = torch.where(exc, torch.zeros_like(sm), sm) - # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. - rows, cols = exc.nonzero(as_tuple=True) - order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) - rows, cols = rows[order], cols[order] - exc_w = torch.stack( - [ - w[rows, cols].float(), - torch.zeros((rows.numel(),), device=w.device), - ], - dim=1, - ) - exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) - exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) - exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( - torch.int32 - ) - # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. - smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) - cf = ( - _fragment_view(code.to(torch.uint8).view(n, k)) - .view(tiles, ks // 4, 4, 32, 8) - .to(torch.int32) - ) - # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. - cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) - for p in range(4): - cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) - cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 - # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. - seg01 = ( - smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) - ) # [.., 2, 32, 16] - seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) - packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] - return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() - - -def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): - """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" - n, _ = weight.shape - assert rows_per_chunk % 16 == 0 - hist = torch.zeros(256, dtype=torch.int64, device=weight.device) - for r in range(0, n, rows_per_chunk): - hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) - base = _window_base(hist) - packed, offs, poss, ws = [], [], [], [] - n_exc = 0 - for r in range(0, n, rows_per_chunk): - p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) - packed.append(p) - offs.append(off[:-1] + n_exc) - poss.append(pos) - ws.append(ew) - n_exc += pos.numel() - exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) - return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) - - -def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: - """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" - tiles, ks = (n + 15) // 16, k // 16 - p = packed.view(tiles, ks // 4, 3, 32, 16) - sm = ( - p[:, :, :2] - .permute(0, 1, 3, 2, 4) - .reshape(tiles, ks // 4, 32, 4, 8) - .permute(0, 1, 3, 2, 4) - .to(torch.int32) - ) - cw = ( - p[:, :, 2] - .contiguous() - .view(torch.int32) - .view(tiles, ks // 4, 32, 4) - .permute(0, 1, 3, 2) - .long() - & 0xFFFFFFFF - ) - codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( - torch.int32 - ) - bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) - bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) - frag = bits.reshape(tiles, ks, 32, 8) - # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). - w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) - return w.to(torch.int16).view(torch.bfloat16)[:n] - - -def decode( - packed: torch.Tensor, - base: int, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - n: int, - k: int, - row0: int = 0, - row1: Optional[int] = None, -) -> torch.Tensor: - """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). - ``row0`` must be a multiple of 16.""" - row1 = n if row1 is None else row1 - assert row0 % 16 == 0 - t0, t1 = row0 // 16, (row1 + 15) // 16 - unit_bytes = (k // 64) * 3 * 32 * 16 - dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() - counts = (exc_off[1:] - exc_off[:-1]).long() - tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) - rows = tile_of * 16 + (exc_pos >> 16).long() - cols = (exc_pos & 0xFFFF).long() - sel = (rows >= row0) & (rows < row1) - dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) - return dec - - -def round_trip_ok( - packed: torch.Tensor, - base: int, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - weight: torch.Tensor, - rows_per_chunk: int = ROWS_PER_CHUNK, -) -> bool: - """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent - window below the BF16 infinity exponent.""" - n, k = weight.shape - if not 0 <= base <= MAX_BASE or k % 64: - return False - if ( - packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 - or exc_off.numel() != (n + 15) // 16 + 1 - ): - return False - for r0 in range(0, n, rows_per_chunk): - r1 = min(n, r0 + rows_per_chunk) - dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) - if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): - return False - return True - - -class M3TargetHead(OpWrapper): - """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit - the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in - another summation order). - - ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` - (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are - the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). - """ - - ARCHS = frozenset({Arch.SM_107}) - - # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and - # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with - # a shard whose rows are not a multiple of 16. - CELLS: tuple[Cell, ...] = ( - *( - Cell( - why=f"{m} verify row(s) of the TP4 target head shard", - spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), - ) - for m in range(1, MAX_ROWS + 1) - ), - Cell( - why="4 rows with outliers and exact zeros in many tiles (the exception lists)", - spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), - ), - Cell( - why="2 rows over a shard of 4000 rows (250 tiles), K 1024", - spec=dict(m=2, n=4000, k=1024, exceptions=True), - ), - ) - - note = """ - The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception - adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs - from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); - integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp - order: the result is deterministic. - """ - - def __call__( - self, - w: torch.Tensor, - x: torch.Tensor, - out: torch.Tensor, - part: torch.Tensor, - counter: torch.Tensor, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - exp_base: int, - n: int, - grid: int, - pdl: bool = True, - ) -> torch.Tensor: - torch.ops.trtllm.minimax_m3_target_head( - w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl - ) - return out - - def reference( - self, - w: torch.Tensor, - x: torch.Tensor, - out: torch.Tensor, - part: torch.Tensor, - counter: torch.Tensor, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - exp_base: int, - n: int, - grid: int, - pdl: bool = True, - ) -> torch.Tensor: - """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. - - ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" - weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) - with true_fp32_matmul(): - logits = x.float() @ weight.float().t() - return logits.to(torch.bfloat16) - - def is_valid( - self, - w: torch.Tensor, - x: torch.Tensor, - out: torch.Tensor, - part: torch.Tensor, - counter: torch.Tensor, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - exp_base: int, - n: int, - grid: int, - pdl: bool = True, - ) -> None: - # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count - # stores a partial sum. - assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" - # The exceptions are read per tile in list order with no bound against the coded stream. - assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" - assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" - assert out.data_ptr() != x.data_ptr(), "out must not alias x" - - -m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py index e89f2de8a8fb..40a9b227a6e0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_lowm.py @@ -74,7 +74,7 @@ def draft_linear( weight: torch.Tensor, x: torch.Tensor, bias: Optional[torch.Tensor] ) -> Optional[torch.Tensor]: """x [M, K] @ weight^T through the low-M kernel when the target covers it (``prebuild``'s keys); None for the - plain GEMM. The split-K ('s') tactics run this target's copy of FlashInfer's kernel (ops/m3_lowm_splitk, bit + plain GEMM. The split-K ('s') tactics run the shared FlashInfer kernel (kernels/gemm/m3_lowm_splitk, bit for bit the same output at every certified cell), which also carries the drafter-boundary push epilogue.""" if bias is not None or x.dim() != 2: return None @@ -135,7 +135,7 @@ def prebuild(weights: Iterable[torch.Tensor], keys: Iterable[Tuple[Tuple[int, in if kind == "s": # The push variant a boundary all-reduce may arm (o_proj / down): compiled only, never run here # (its pushes need the all-reduce's exchange and a matching call). - from .m3_lowm_splitk import compile_push + from ...kernels.gemm.m3_lowm_splitk import compile_push compile_push(x, weight.t(), out, use_pdl, tactic) _BUILT.add((n, k, m, use_pdl)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py index 45a64dfe9210..ab636da5aec9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_draft_step.py @@ -27,7 +27,11 @@ import torch -from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + encode_chunked, + round_trip_ok, +) + from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments, max_tile_pieces from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op from . import m3_report diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py index 7936d2497731..29a0b7425682 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_drafter_attn.py @@ -20,27 +20,50 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Mapping, Optional, Tuple +import sys +from typing import TYPE_CHECKING as TYPE_CHECKING +from typing import Mapping as Mapping +from typing import Optional as Optional +from typing import Tuple as Tuple +from typing import cast import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + drafter as _drafter_shared, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + DrafterFullCtxBase, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + drafter_rope_applies as _shared_rope_applies, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + initialize_drafter_state as _shared_initialize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + install as _shared_install, +) from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( - FlashInferTrtllmGenFmha, + FlashInferTrtllmGenFmha as FlashInferTrtllmGenFmha, ) -from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase -from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase as FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams as FmhaParams from tensorrt_llm._torch.attention.backends.fmha.utils import ( - get_attention_chunk_size, - get_bmm1_scale, + get_attention_chunk_size as get_attention_chunk_size, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import get_bmm1_scale as get_bmm1_scale +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionInputType as AttentionInputType, ) -from tensorrt_llm._torch.attention.backends.interface import AttentionInputType from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( - fp8_e4m3_saturate, + fp8_e4m3_saturate as fp8_e4m3_saturate, ) -from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor -from tensorrt_llm.bindings.internal import thop -from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.quantization.mode import QuantMode +from tensorrt_llm._utils import TensorWrapper as TensorWrapper +from tensorrt_llm._utils import convert_to_torch_tensor as convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop as thop +from tensorrt_llm.functional import PositionEmbeddingType as PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode as QuantMode from . import m3_fullctx_decode as fullctx @@ -57,9 +80,15 @@ def fwd_has_no_mrope(fwd) -> bool: from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention -class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): +class M3DrafterFullCtxFmha(DrafterFullCtxBase): """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + FULLCTX = fullctx + GEOMETRY_ERROR = "not the TP4 FP8 drafter geometry" + + def _launch_rope(self, *args) -> None: + torch.ops.trtllm.minimax_m3_sm107tp4_drafter_rope_fp8_kv(*args, ROPE_CARVEOUT) + def __init__( self, attn: "TrtllmAttention", @@ -74,286 +103,25 @@ def __init__( on_rope_taken=None, ) -> None: super().__init__(attn) - self._recipes = recipes - self._max_kv_len = max_kv_len - self._own_reducer = own_reducer - self._fused_rope_mode = fused_rope_mode - self._pdl = pdl - self._on_rope_taken = on_rope_taken - self._on_taken = on_taken - self._on_uncovered = on_uncovered - self._on_missed = on_missed - - def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: - nreq = int(metadata.num_generations) - if nreq <= 0 or int(q.shape[0]) % nreq: - return None - return int(q.shape[0]) // nreq, nreq - - def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: - if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: - return False - if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: - return False - attn = self.attn - if ( - attn.num_heads != fullctx.HEADS - or attn.num_kv_heads != fullctx.KV_HEADS - or attn.head_dim != fullctx.HEAD_DIM - or not QuantMode(attn.quant_mode).has_fp8_kv_cache() - ): - self._on_uncovered("not the TP4 FP8 drafter geometry") - return False - # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's - # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. - if ( - forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant - or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig - ): - self._on_uncovered("per-step KV cache scales") - return False - if int(metadata.max_seq_len) > self._max_kv_len: - self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") - return False - shape = self._shape(q, metadata) - if shape is None or shape not in self._recipes: - self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") - return False - # Compiled before warmup (``install``); the forward never compiles one. - if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: - self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") - return False - return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) - - def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: - attn = params.attn - width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim - rotary_dim = int(attn.rope_params.dim) - return ( - self._fused_rope_mode is not None - and params.qkv_input is not None - and attention_input.dim() == 2 - and attention_input.dtype == torch.bfloat16 - and attention_input.is_contiguous() - and int(attention_input.shape[1]) == width - and not params.is_cross - and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) - and attn.rotary_cos_sin is not None - and attn.rotary_cos_sin.dtype == torch.float32 - and fwd_has_no_mrope(params.fwd) - and 0 < rotary_dim <= attn.head_dim - and rotary_dim % 16 == 0 - and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) - # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo - # runs for several requests only) and reads the workspace's, so it keeps that path. - and not (self._multi_token(params) and int(params.batch_size) == 1) - ) - - @staticmethod - def _multi_token(params: FmhaParams) -> bool: - return ( - params.spec_decoding_generation_lengths is not None - and int(params.attn.predicted_tokens_per_seq) > 1 + _shared_initialize( + self, + _DRAFTER_BINDINGS, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode, + pdl, + on_rope_taken, + has_one_request_multi=False, ) - def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: - """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" - if not self._fused_rope_applies(params, attention_input): - return False - attn = params.attn - meta = params.meta - qlen, nreq = int(params.input_seq_length), int(params.batch_size) - plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) - if plan is None: - return False - # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary - # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the - # block table is the pool's [requests, 2, blocks] K / V sub-page indices. - pool_index, layer_in_pool = ( - int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] - ) - tpb = int(params.tokens_per_block) - sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim - base = ( - int(meta.host_kv_cache_pool_pointers[pool_index, 0]) - + layer_in_pool * int(params.kv_factor) * sub_page_bytes - ) - pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 - kv_pages = convert_to_torch_tensor( - TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) - ) - tables = meta.kv_cache_block_offsets[ - pool_index, params.seq_offset : params.seq_offset + nreq - ] - multi = self._multi_token(params) - if not torch.cuda.is_current_stream_capturing(): - # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess - # path checks; the op itself writes V at the table's V sub-page. - tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb - cols = torch.arange(tables.shape[-1], device=tables.device) - in_use = cols[None, :] < tpb_held[:, None] - offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() - if offsets.numel() != 1 or int(offsets) != 1: - raise RuntimeError( - f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" - ) - q = torch.empty( - int(params.num_tokens), - attn.num_heads, - attn.head_dim, - dtype=torch.float8_e4m3fn, - device=attention_input.device, - ) - - torch.ops.trtllm.minimax_m3_sm107tp4_drafter_rope_fp8_kv( - attention_input, - attn.rotary_cos_sin, - params.sequence_lengths, - params.spec_decoding_position_offsets if multi else None, - kv_pages, - tables, - q, - qlen, - attn.num_heads, - attn.num_kv_heads, - int(attn.rope_params.dim), - tpb, - self._fused_rope_mode, - self._pdl, - ROPE_CARVEOUT, - ) - plan.run( - q, - kv_pages[:-1], - kv_pages[1:], - tables[:, 0], - params.sequence_lengths[:nreq], - get_bmm1_scale(attn), - params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), - ) - if self._on_rope_taken is not None: - self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") - self._on_taken(params.num_tokens) - return True - - def run_generation(self, params: FmhaParams) -> None: - attention_input = params.qkv_input if params.qkv_input is not None else params.query_input - if attention_input is None: - raise RuntimeError("the drafter decode requires QKV or query input.") - attn = params.attn - meta = params.meta - fwd = params.fwd - rope_params = attn.rope_params - if self._run_fused_rope(params, attention_input): - return - # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 - # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). - fp8_context_fmha = False - ( - q_processed, - kv_pool, - block_tables, - _kv_scale_pool, - _bmm1_scale, - _bmm2_scale, - _fmha_workspace, - _cu_seqlens, - _max_q_len, - _max_kv_len, - _window_left, - is_multi_token_gen, - ) = thop.trtllm_gen_generation_preprocess( - attention_input, - params.workspace, - params.sequence_lengths, - params.spec_decoding_generation_lengths, - params.spec_decoding_position_offsets, - meta.kv_cache_block_offsets, - meta.host_kv_cache_pool_pointers, - meta.host_kv_cache_pool_mapping, - fwd.kv_scale_orig_quant, - fwd.kv_scale_quant_orig, - fwd.out_scale, - attn.rotary_inv_freq, - attn.rotary_cos_sin, - fwd.mrope_position_deltas, - attn.local_layer_idx, - params.seq_offset, - attn.num_heads, - attn.num_kv_heads, - attn.head_dim, - params.tokens_per_block, - attn.quant_mode, - params.max_attention_window_size, - params.cyclic_attention_window_size, - params.num_tokens, - params.batch_size, - params.input_seq_length, - params.max_past_kv_length, - rope_params.dim, - rope_params.theta, - int(rope_params.scale_type), - rope_params.scale, - rope_params.max_positions, - attn.position_embedding_type, - get_bmm1_scale(attn), - 1.0, - fp8_context_fmha, - attn.predicted_tokens_per_seq, - get_attention_chunk_size(attn), - self._multi_processor_count, - params.total_num_blocks, - params.kv_factor, - True, - params.is_cross, - ) - qlen, nreq = params.input_seq_length, params.batch_size - plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) - # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the - # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row - # pass is never multi-token. - if plan is None or (is_multi_token_gen and qlen == 1): - raise RuntimeError( - f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" - ) - rows = params.num_tokens * attn.num_heads * attn.head_dim - q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) - q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) - # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 - # of the pool (the two views have the same page count, as the kernel requires). - if not torch.cuda.is_current_stream_capturing(): - # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. - lengths = params.spec_decoding_generation_lengths - if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): - raise RuntimeError( - f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" - ) - # Only the blocks a request holds are laid out; the table's tail entries are unset. - tpb = int(params.tokens_per_block) - held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb - cols = torch.arange(block_tables.shape[-1], device=block_tables.device) - in_use = cols[None, :] < held[:, None] - offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() - if offsets.numel() != 1 or int(offsets) != 1: - raise RuntimeError( - f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" - ) - # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. - # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. - pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 - kv_pages = convert_to_torch_tensor( - TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) - ) - plan.run( - q, - kv_pages[:-1], - kv_pages[1:], - block_tables[:nreq, 0], - params.sequence_lengths[:nreq], - get_bmm1_scale(attn), - params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + return _shared_rope_applies( + self, _DRAFTER_BINDINGS, params, attention_input, has_one_request_multi=False ) - self._on_taken(params.num_tokens) def install( @@ -371,35 +139,20 @@ def install( """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it was installed on.""" - from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention - - count = 0 - device = next(draft_model.parameters()).device - # The backend is not a module: each attention module of the drafter holds it as ``attn``. - for module in draft_model.modules(): - backend = getattr(module, "attn", None) - if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): - continue - if not bool((backend.kv_cache_scaling_factor == 1.0).all()): - on_uncovered("a scaled FP8 KV cache") - continue - fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) - libs = backend._fmha_manager.fmha_libs - if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): - libs.insert( - 0, - M3DrafterFullCtxFmha( - backend, - recipes, - max_kv_len, - own_reducer, - on_taken, - on_uncovered, - on_missed, - fused_rope_mode=fused_rope_mode, - pdl=pdl, - on_rope_taken=on_rope_taken, - ), - ) - count += 1 - return count + return _shared_install( + draft_model=draft_model, + recipes=recipes, + max_kv_len=max_kv_len, + own_reducer=own_reducer, + pdl=pdl, + on_taken=on_taken, + on_uncovered=on_uncovered, + on_missed=on_missed, + fused_rope_mode=fused_rope_mode, + on_rope_taken=on_rope_taken, + prebuild=fullctx.prebuild, + fmha_class=M3DrafterFullCtxFmha, + ) + + +_DRAFTER_BINDINGS = cast(_drafter_shared.DrafterBindings, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py index 79aff5cb80ce..0629283ca734 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fold_tail.py @@ -14,70 +14,22 @@ Integer copies only: the values are the ones the sampler's and the engine's own copies write. With ``cpp`` the same writes run as this target's C++ op (trtllm::minimax_m3_sm107_tp4_fold_tail, -minimaxM3Sm107Tp4/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. +minimaxM3Sm107/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. """ import torch -import triton -import triton.language as tl +import triton as triton +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.spec.m3_fold_tail import ( + _fold_tail_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.fold_tail import ( + fold_tail as _shared_fold_tail, +) -@triton.jit -def _fold_tail_kernel( - ctrl, - slots, - src_new, - sn_s0, - sn_s1, - SN_W: tl.constexpr, - dst_new, - dn_s0, - dn_s1, - DN_W: tl.constexpr, - src_next, - sx_s0, - sx_s1, - SX_W: tl.constexpr, - dst_next, - dx_s0, - dx_s1, - DX_W: tl.constexpr, - src_draft, - sd_s0, - sd_s1, - SD_W: tl.constexpr, - dst_draft, - dd_s0, - dd_s1, - DD_W: tl.constexpr, - src_lens, - dst_lens, - ids, - draft_ids, - pos_off, - kv_off, - T: tl.constexpr, - D: tl.constexpr, - BLOCK: tl.constexpr, -): - i = tl.program_id(0) - n = tl.load(ctrl) - if i < n: - s = tl.load(slots + i).to(tl.int64) - c = tl.arange(0, BLOCK) - new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) - tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) - nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) - tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) - drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) - tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) - ln = tl.load(src_lens + i) - tl.store(dst_lens + s, ln) - # The next step's inputs read the store's padded rows: zero past each source width. - tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) - tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) - tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) - tl.store(kv_off + i, ln - T) + +def _launch_cpp(*args: object) -> None: + torch.ops.trtllm.minimax_m3_sm107_tp4_fold_tail(*args) def fold_tail( @@ -101,72 +53,24 @@ def fold_tail( ) -> bool: """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width].""" - two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) - one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) - if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): - return False - if ( - any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) - or slots.dtype != torch.long - ): - return False - t_step = tokens_per_step - d_step = t_step - 1 - if ( - rows <= 0 - or slots.numel() < rows - or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) - or ids.numel() < rows * t_step - or pos_off.numel() < rows * t_step - or draft_ids.numel() < rows * d_step - or kv_off.numel() < rows - ): - return False - widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) - widths += (src_draft.shape[1], dst_draft.shape[1], t_step) - if cpp and max(widths) <= 32: - torch.ops.trtllm.minimax_m3_sm107_tp4_fold_tail( - ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, - src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, - ) # fmt: skip - return True - block = triton.next_power_of_2(max(widths)) - _fold_tail_kernel[(rows,)]( - ctrl, - slots, - src_new, - src_new.stride(0), - src_new.stride(1), - src_new.shape[1], - dst_new, - dst_new.stride(0), - dst_new.stride(1), - dst_new.shape[1], - src_next, - src_next.stride(0), - src_next.stride(1), - src_next.shape[1], - dst_next, - dst_next.stride(0), - dst_next.stride(1), - dst_next.shape[1], - src_draft, - src_draft.stride(0), - src_draft.stride(1), - src_draft.shape[1], - dst_draft, - dst_draft.stride(0), - dst_draft.stride(1), - dst_draft.shape[1], - src_lens, - dst_lens, - ids, - draft_ids, - pos_off, - kv_off, - t_step, - d_step, - block, - num_warps=1, + return _shared_fold_tail( + ctrl=ctrl, + slots=slots, + rows=rows, + src_new=src_new, + dst_new=dst_new, + src_next=src_next, + dst_next=dst_next, + src_draft=src_draft, + dst_draft=dst_draft, + src_lens=src_lens, + dst_lens=dst_lens, + ids=ids, + draft_ids=draft_ids, + pos_off=pos_off, + kv_off=kv_off, + tokens_per_step=tokens_per_step, + cpp=cpp, + native_launch=_launch_cpp, + triton_kernel=_fold_tail_kernel, ) - return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py deleted file mode 100644 index be04f26f9aff..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_force_vec.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's -``force_num_accepted_tokens``), in one launch. - -Triton: ``force_vec_kernel``: one CTA bumps the shared step counter and draws each request's count from the fixed-seed -pool with the in-tree (counter, slot) hash, exactly as ``_apply_force_accepted_tokens`` does. ``force_vec_req_kernel`` -(``FORCE_AL_PER_REQUEST``): each request draws from the same pool at its own key, its KV length at this verify. -Reached through the ``ops/spec`` entries ``m3_force_vec`` and ``m3_force_vec_req``. - -One compiled variant for every batch and accept count: specializing them (== 1, divisible by 16) would compile -mid-serving the first time a step reaches a new one, past the warm-up. -""" - -import triton -import triton.language as tl - - -@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) -def force_vec_kernel( - pool, - counter, - out, - batch_size, - base_total, - max_total, - frac, - counter_stride, - slot_stride, - pool_mask, - BLOCK: tl.constexpr, -): - c = tl.load(counter) + 1 - slots = tl.arange(0, BLOCK) - tl.store(counter, c) - idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask - live = slots < batch_size - r = tl.load(pool + idx, mask=live, other=1.0) - v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) - tl.store(out + slots, v, mask=live) - - -@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) -def force_vec_req_kernel( - pool, - keys, - out, - batch_size, - base_total, - max_total, - frac, - key_stride, - pool_mask, - BLOCK: tl.constexpr, -): - # One draw per request from its own key (its KV length at this verify): no shared counter, so a request's forced - # counts depend only on its own prompt and progress, not on what else the server ran before or beside it. - slots = tl.arange(0, BLOCK) - live = slots < batch_size - k = tl.load(keys + slots, mask=live, other=0).to(tl.int64) - idx = (k * key_stride) & pool_mask - r = tl.load(pool + idx, mask=live, other=1.0) - v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) - tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py index 3dc0fc4a3888..29b4336f5b8e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_fullctx_decode.py @@ -28,6 +28,18 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + FullCtxGeometry, + FullCtxPlanBase, + check_layout, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + device_index as _index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + kv_page_views as kv_page_views, +) + # The TP4 per-rank geometry: 16 query heads over 1 KV head. HEADS = 16 KV_HEADS = 1 # capacity: fixed: the tp4 KV heads per rank @@ -44,116 +56,10 @@ _lock = threading.Lock() -class FullCtxPlan: +class FullCtxPlan(FullCtxPlanBase): """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" - def __init__( - self, - device: torch.device, - qlen: int, - nreq: int, - max_kv_len: int, - recipe: Mapping, - own_reducer: bool, - pdl: bool, - ): - from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec - from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( - get_max_active_clusters_for_cluster_size, - make_decode_config, - ) - - self._tsdec = tsdec - self.qlen = qlen - self.nreq = nreq - self.pdl = pdl - self.recipe = dict(recipe) - splits = int(self.recipe["splits"]) - fp8 = tsdec._cutlass_dtype("float8_e4m3fn") - # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer - # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents - # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an - # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the - # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own - # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge - # (minimax_m3_sm107tp4_fullctx_reduce) waits as its first instruction. - cfg = make_decode_config( - headdim=HEAD_DIM, - args={"use_pdl": pdl, **self.recipe.get("config", {})}, - seq_len_q=qlen, - seq_len_kv=min(max_kv_len, POLICY_KV_LEN), - batch_size=nreq, - num_heads_q=HEADS, - num_heads_kv=KV_HEADS, - q_dtype=fp8, - k_dtype=fp8, - v_dtype=fp8, - o_dtype=tsdec._cutlass_dtype("float16"), - qkv_layout="pagedKv", - num_tokens_per_page=PAGE, - storage_tokens_per_page=PAGE, - split_kv_mode="gmem_reduction_with_separate_kernel", - splits_kv=splits, - max_splits_kv=splits, - min_loop_iters_per_split=1, - mask_type="causal" if qlen > 1 else "dense", - auto_tuner=True, - split_kv=True, - ) - _check_layout(cfg, qlen, splits) - with torch.cuda.device(device): - max_active = get_max_active_clusters_for_cluster_size(1) - spec = tsdec._decode_launch_spec_from_config( - cfg, - batch_size=nreq, - num_qo_heads=HEADS, - num_kv_heads=KV_HEADS, - head_dim=HEAD_DIM, - seq_len_q=qlen, - max_active_clusters=int(max_active), - ) - compile_spec = tsdec._make_decode_compile_spec( - spec, - device_index=_index(device), - num_qo_heads=HEADS, - num_kv_heads=KV_HEADS, - head_dim=HEAD_DIM, - page_size=PAGE, - max_kv_len=max_kv_len, - seq_len_q=qlen, - q_dtype_key="float8_e4m3fn", - k_dtype_key="float8_e4m3fn", - v_dtype_key="float8_e4m3fn", - output_dtype_key="float16", - use_packed_q=False, - kv_prefix_mode="dynamic", - kv_lengths_mode="dynamic", - ) - self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) - layout = tsdec._make_decode_workspace_layout( - spec.scratch_shapes, - torch.float16, - use_separate_reduction_kernel=True, - use_split_kv=True, - ) - self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) - self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) - self.workspace.split_kv_counter.zero_() - self.workspace.cu_seqlens_q.zero_() - self.workspace.attention_sinks.zero_() - rows = HEADS_PER_KV * qlen - if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): - raise RuntimeError( - f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " - f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" - ) - self.own_reducer = own_reducer - self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) - self.tile_kv = int(cfg.tile_size_kv) - self.insts_kv = int(cfg.num_insts_kv) - self.causal = qlen > 1 - shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) - self.staging = torch.empty(shape, dtype=torch.float16, device=device) + GEOMETRY = FullCtxGeometry(HEADS, KV_HEADS, HEAD_DIM, PAGE, POLICY_KV_LEN, HEADS_PER_KV) def run( self, @@ -240,32 +146,7 @@ def run( def _check_layout(cfg, qlen: int, splits: int) -> None: - """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. - - Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is - ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding - window, attention sinks or packed Q. - """ - problems = [] - if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): - problems.append("not split-KV with the separate GMEM reducer") - if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: - problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") - if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: - problems.append("cluster-SMEM reduction or a persistent scheduler") - if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: - problems.append("packed Q, a sliding window or attention sinks") - if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): - problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") - if qlen > 1: - if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): - problems.append("multi-token Q without the grouped [token, head] tile layout") - elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): - problems.append( - f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" - ) - if problems: - raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + return check_layout(cfg, qlen, splits, heads_per_kv=HEADS_PER_KV) def emits_mxfp8() -> bool: @@ -303,12 +184,3 @@ def prebuild( _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) built += 1 return built - - -def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: - """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" - return pool[0::subpages_per_slot], pool[1::subpages_per_slot] - - -def _index(device: torch.device) -> int: - return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py index fad75d44694a..3ffe8a7af674 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_index_score.py @@ -20,619 +20,25 @@ and called through this target's ``m3_index_score`` entry. """ -import cutlass -import torch -from cuda.bindings.driver import CUstream -from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute -from cutlass._mlir.dialects import llvm -from cutlass.cute.nvgpu import cpasync, warp -from cutlass.cutlass_dsl import dsl_user_op -from cutlass.utils.smem_allocator import SmemAllocator - -from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( - EVICT_FIRST, - fp8x4_to_fp16x4, - mma_sync, - simple_tma_copy, +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.index_score_runtime import ( + IndexScoreRunner as SharedIndexScoreRunner, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + IndexDecodeScoreKernel as IndexDecodeScoreKernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + _fp8_to_f16_mma_fragments as _fp8_to_f16_mma_fragments, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + _l2_prefetch_bulk as _l2_prefetch_bulk, ) -from tensorrt_llm._torch.kernels.blackwell.utils import TRTLLM_ENABLE_PDL, griddepcontrol_wait __all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] -@dsl_user_op -def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: - """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" - llvm.inline_asm( - res=None, - operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], - asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", - constraints="l,r", - has_side_effects=True, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@cute.jit -def _fp8_to_f16_mma_fragments(src: cute.Tensor): - """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" - src_elems = cute.size(src) - src_u32 = cute.recast_tensor(src, Uint32) - src_f16 = cute.make_rmem_tensor(src_elems, Float16) - src_f16_u32 = cute.recast_tensor(src_f16, Uint32) - # Packed conversion; faster and fewer SASS instructions than - # src.load().to(Float16). - for i in cutlass.range_constexpr(src_elems // 4): - converted = fp8x4_to_fp16x4(src_u32[i]) - src_f16_u32[i * 2] = converted[0] - src_f16_u32[i * 2 + 1] = converted[1] - lower = cute.make_rmem_tensor(src_elems // 2, Float16) - upper = cute.make_rmem_tensor(src_elems // 2, Float16) - - # FP8 ldmatrix yields four consecutive values along K. Split each group - # into the lower two and upper two for the two FP16 MMA k-fragments. - for i in cutlass.range_constexpr(src_elems // 2): - lower[i] = src_f16[(i // 2) * 4 + i % 2] - upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] - return lower, upper - - -class IndexDecodeScoreKernel: - """Per-block max index score for one decode step. - - Grid is (batch, split_k): CTA (b, s) walks KV blocks - s, s + split_k, s + 2 * split_k, ... of request b and writes each block's - score directly, so no cross-CTA reduction is needed. CTAs whose split_id - exceeds the request's block count exit immediately. - """ - - BLOCK_K = 128 - BAR_MMA = 1 - num_stages = 2 - - def __init__( - self, - dtype: type[cutlass.Numeric], - num_heads: int, - max_decode_query_len: int, - split_k: int, - head_dim: int = 128, - prefetch_old: bool = False, - early_issue: bool = False, - prefetch_depth: int = 0, - ): - # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that - # hold only tokens older than this step (the predecessor writes only this step's tokens). - self.prefetch_old = prefetch_old - # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency - # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. - # The tiles are read after the wait, so their contents are final; the block table is only written between - # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles - # turn out to lie past the request's true length drains them before exiting. - self.early_issue = early_issue - # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at - # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. - self.prefetch_depth = prefetch_depth - self.dtype = dtype - self.num_heads = num_heads - self.max_decode_query_len = max_decode_query_len - self.split_k = split_k - self.head_dim = head_dim - - def _swizzle_elems(self) -> int: - """Elements spanned by one 128-byte swizzle atom.""" - return 128 * 8 // self.dtype.width - - def _sq_layout(self): - """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" - elems = self._swizzle_elems() - head_dim = self.head_dim - block_q = self.num_heads * self.max_decode_query_len - layout = cute.make_layout( - (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), - stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), - ) - return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) - - def _sk_layout(self): - """Composed SMEM layout for the pipelined K tiles.""" - elems = self._swizzle_elems() - head_dim = self.head_dim - block_k = self.BLOCK_K - layout = cute.make_layout( - (1, block_k, (elems, head_dim // elems), self.num_stages), - stride=(0, elems, (1, block_k * elems), block_k * head_dim), - ) - return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) - - @cute.jit - def __call__( - self, - gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] - gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] - block_table: cute.Tensor, # [bs, max_pages] - score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] - seq_lens: cute.Tensor, # [bs] - stream: CUstream, - ): - num_heads = self.num_heads - head_dim = self.head_dim - MAX_DQL = self.max_decode_query_len - assert num_heads * MAX_DQL <= 32 - - batch = seq_lens.shape[0] - decode_query_len = gQ.shape[0] // batch - grid = (batch, self.split_k, 1) - block = (32 * 5, 1, 1) - - tma_g2s = cpasync.CopyBulkTensorTileG2SOp() - elems = self._swizzle_elems() - - q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( - tma_g2s, - cute.logical_divide(gQ, (None, None, elems)), - self._sq_layout(), - (MAX_DQL, num_heads, head_dim), - ) - k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( - tma_g2s, - cute.logical_divide(gK_cache, (None, None, elems)), - self._sk_layout(), - (1, self.BLOCK_K, head_dim), - ) - - self.kernel( - q_tma_atom, - q_tma_tensor, - k_tma_atom, - k_tma_tensor, - gK_cache, - block_table, - score, - seq_lens, - decode_query_len, - ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) - - @cute.kernel - def kernel( - self, - q_tma_atom: cute.CopyAtom, - q_tma_tensor: cute.Tensor, - k_tma_atom: cute.CopyAtom, - k_tma_tensor: cute.Tensor, - gK_cache: cute.Tensor, - block_table: cute.Tensor, - score: cute.Tensor, - seq_lens: cute.Tensor, - decode_query_len, - ): - batch_id, split_id, _ = cute.arch.block_idx() - _, split_k, _ = cute.arch.grid_dim() - warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane_id = cute.arch.lane_idx() - - NUM_HEADS = self.num_heads - MAX_DQL = self.max_decode_query_len - BLOCK_Q = NUM_HEADS * MAX_DQL - BLOCK_K = self.BLOCK_K - head_dim = self.head_dim - dtype = self.dtype - MMA_N = 8 - num_stages = self.num_stages - Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) - EPI_Q = Q_TILES * MMA_N - - sq_layout = self._sq_layout() - sk_layout = self._sk_layout() - - smem = SmemAllocator() - sK = smem.allocate_tensor( - dtype, - sk_layout.outer, - byte_alignment=128, - swizzle=sk_layout.inner, - )[0, None, None, None] - # sQ aliases the first K stage: Q is consumed into registers before the - # first K tile is needed, so the two never overlap in time. - sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) - # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q - # column mode. The target profile keeps the rank-2 view even for - # degenerate shapes such as MAX_DQL == 1. - q_tma_elems = self._swizzle_elems() - sQ = cute.coalesce( - cute.group_modes(sQ_tma, 0, 2), - target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), - ) - epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) - - tma_full_mbar = smem.allocate_array(Int64, num_stages) - tma_empty_mbar = smem.allocate_array(Int64, num_stages) - - # Barrier init and descriptor prefetch do not depend on the predecessor - # grid, so they are the work PDL can overlap. Every CTA runs them, even - # the ones that own no blocks, since ownership is not known until - # seq_lens is read below. - if warp_id == 0: - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(num_stages): - cute.arch.mbarrier_init(tma_full_mbar + i, 1) - cute.arch.mbarrier_init(tma_empty_mbar + i, 128) - cute.arch.mbarrier_init_fence() - elif warp_id == 1: - cpasync.prefetch_descriptor(q_tma_atom) - cpasync.prefetch_descriptor(k_tma_atom) - cute.arch.sync_threads() - - if cutlass.const_expr(self.prefetch_old): - if warp_id == 4: - with cute.arch.elect_one(): - seqlen0 = seq_lens[batch_id] - num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) - stable = (seqlen0 - decode_query_len) // BLOCK_K - if stable > num_blocks0: - stable = num_blocks0 - if cutlass.const_expr(self.prefetch_depth > 0): - if stable > split_id + self.prefetch_depth * split_k: - stable = split_id + self.prefetch_depth * split_k - for block_id in range(split_id, stable, split_k): - page = block_table[batch_id, block_id] - _l2_prefetch_bulk( - gK_cache[page, None, None].iterator.toint().to(Int64), - Int32(BLOCK_K * head_dim * (dtype.width // 8)), - ) - - # No explicit launch_dependents: the only consumer is the block selector, - # which reduces over every block this grid writes, so triggering it early - # would leave it resident and spinning in its own wait. The implicit - # trigger as the CTAs exit still gives it the launch overlap. - early_q = False - early_page = Int32(0) - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - hint_blocks = cute.ceil_div(seq_lens[batch_id], BLOCK_K) - if hint_blocks > block_table.shape[1]: - hint_blocks = block_table.shape[1] - if split_id < hint_blocks: - early_q = True - early_page = block_table[batch_id, split_id] - - griddepcontrol_wait() - - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - if early_q: - gQ_early = cute.local_tile( - cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), - tiler=(MAX_DQL, NUM_HEADS, head_dim), - coord=(0, 0, 0), - ) - cute.arch.mbarrier_wait(tma_empty_mbar, 1) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) - ) - simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) - cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - tma_full_mbar + 1 % num_stages, BLOCK_K * head_dim * (dtype.width // 8) - ) - simple_tma_copy( - k_tma_atom, - k_tma_tensor[early_page, None, None], - sK[None, None, 1 % num_stages], - tma_full_mbar + 1 % num_stages, - cache_policy=EVICT_FIRST, - ) - - # seq_lens must be read after the wait. It is a graph-stable buffer that - # on_update_kv_lens patches every step, and num_blocks bounds an unmasked - # block_table load and an unmasked score store, so a stale length is an - # out-of-bounds write rather than just a wrong answer. - seqlen = seq_lens[batch_id] - num_blocks = cute.ceil_div(seqlen, BLOCK_K) - - if split_id < num_blocks: - if warp_id == 4: - # TMA warp - tma_stage = 0 - tma_parity = 1 - first_block = split_id - - if not early_q: - gQ_tile = cute.local_tile( - cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), - tiler=(MAX_DQL, NUM_HEADS, head_dim), - coord=(0, 0, 0), - ) - cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) - with cute.arch.elect_one(): - Q_size = BLOCK_Q * head_dim * (dtype.width // 8) - cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) - # TMA bounds-checks rows when the runtime decode_query_len is - # below MAX_DQL; padded Q columns are masked before the stores. - simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) - - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - if early_q: - # The first K tile was issued right after the wait. - first_block = split_id + split_k - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - for block_id in range(first_block, num_blocks, split_k): - page_id = block_table[batch_id, block_id] - gK_tile = k_tma_tensor[page_id, None, None] - k_mbar = tma_full_mbar + tma_stage - - cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) - with cute.arch.elect_one(): - K_size = BLOCK_K * head_dim * (dtype.width // 8) - cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) - simple_tma_copy( - k_tma_atom, - gK_tile, - sK[None, None, tma_stage], - k_mbar, - cache_policy=EVICT_FIRST, - ) - - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - else: - # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T - sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) - q_start = seqlen - decode_query_len - - elems = 128 // dtype.width # 16B - MMA_K = 32 * 8 // dtype.width # 32B - - # Pre-compute ldmatrix addresses. - # sK loads a [16 x 16B] tile: - # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) - # sQ loads an [8 x 32B] tile: - # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) - sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) - sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) - - # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) - # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) - sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] - sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] - - ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) - ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) - - rQ = cute.make_rmem_tensor( - ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype - ) - rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) - rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) - - if warp_id == 0: - cute.arch.mbarrier_wait(tma_full_mbar, 0) - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - for q in cutlass.range_constexpr(Q_TILES): - cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) - cute.arch.mbarrier_arrive(tma_empty_mbar) - - tma_stage = 1 % self.num_stages - tma_parity = 0 - if tma_stage == 0: - tma_parity ^= 1 - - if cutlass.const_expr(dtype is Float8E4M3FN): - rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) - q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) - rQ_f16[None, None, None, 0].store(q_lower.load()) - rQ_f16[None, None, None, 1].store(q_upper.load()) - - for block_id in range(split_id, num_blocks, split_k): - rC.fill(0.0) - - if warp_id == 0: - cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - - for k in cutlass.range_constexpr(head_dim // MMA_K): - cute.copy( - ldsm_atom, - sK_ldsm[None, (None, k, tma_stage)], - rK[None, None, k], - ) - for m in cutlass.range_constexpr(2): - if cutlass.const_expr(dtype is Float8E4M3FN): - rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) - for n in cutlass.range_constexpr(Q_TILES): - rC[None, m, n] = mma_sync( - rK_lower, - rQ_f16[None, k, n, 0], - rC[None, m, n], - ) - rC[None, m, n] = mma_sync( - rK_upper, - rQ_f16[None, k, n, 1], - rC[None, m, n], - ) - else: - for n in cutlass.range_constexpr(Q_TILES): - rC[None, m, n] = mma_sync( - rK[None, m, k], - rQ[(None, k % 2), k // 2, n], - rC[None, m, n], - ) - - cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) - - k_start = block_id * BLOCK_K + warp_id * 32 - - for q in cutlass.range_constexpr(Q_TILES): - for i in cutlass.range_constexpr(4): - for j in cutlass.range_constexpr(2): - col = q * 8 + (lane_id % 4) * 2 + j - q_local_pos = col % MAX_DQL - q_pos = q_start + q_local_pos - k_pos = k_start + i * 8 + lane_id // 4 - rC[q * 8 + i * 2 + j] = ( - rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") - ) - - for q in cutlass.range_constexpr(Q_TILES): - # Thread-local reduction along the BLOCK_K dim. - rScore = cute.make_rmem_tensor(2, Float32) - rScore.fill(float("-inf")) - for i in cutlass.range_constexpr(4): - rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) - rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) - - # Warp reduction among lanes 0, 4, 8, 12, ... - for i in cutlass.range_constexpr(3): - offset = 4 << i - other0 = cute.arch.shuffle_sync_bfly( - rScore[0], offset=offset, mask=-1, mask_and_clamp=31 - ) - other1 = cute.arch.shuffle_sync_bfly( - rScore[1], offset=offset, mask=-1, mask_and_clamp=31 - ) - rScore[0] = cute.arch.fmax(rScore[0], other0) - rScore[1] = cute.arch.fmax(rScore[1], other1) - - if lane_id * 2 < MMA_N: - epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] - epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - - head_id = lane_id // MAX_DQL - q_local_pos = lane_id - head_id * MAX_DQL - valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len - if lane_id < BLOCK_Q and valid_q: - final_score = epi_buffer[lane_id, 0] - for i in cutlass.range_constexpr(1, 4): - final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) - - t = batch_id * decode_query_len + q_local_pos - score[head_id, t, block_id] = final_score - - tma_stage = (tma_stage + 1) % self.num_stages - if tma_stage == 0: - tma_parity ^= 1 - else: - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - if early_q: - # Early tiles past the request's true length: let them land before the CTA exits. - cute.arch.mbarrier_wait(tma_full_mbar, 0) - cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) - - -class IndexScoreRunner: - """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, - split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile - serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block - per CTA.""" +class IndexScoreRunner(SharedIndexScoreRunner): + """TP4 schedule and private compilation cache.""" kernel_cache: dict = {} - SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim - SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K - # BLOCK_Q must fit one warp's worth of epilogue lanes. - MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile - EARLY_PREFETCH_DEPTH = 1 - - @classmethod - def is_supported( - cls, - *, - q_dtype: torch.dtype, - num_heads: int, - head_dim: int, - page_size: int, - max_decode_query_len: int, - ) -> bool: - from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE - from tensorrt_llm._utils import is_sm_100f - - return ( - is_sm_100f() - and q_dtype in TORCH_TO_CUTE_DTYPE - and head_dim == cls.SUPPORTED_HEAD_DIM - and page_size == cls.SUPPORTED_PAGE_SIZE - and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q - and max_decode_query_len >= 1 - ) - - @classmethod - def _compile( - cls, - q_dtype: torch.dtype, - num_heads: int, - max_decode_query_len: int, - head_dim: int, - split_k: int, - early: bool, - ) -> None: - from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE - - key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) - if key in cls.kernel_cache: - return - cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] - page_size = cls.SUPPORTED_PAGE_SIZE - sym_total_tokens = cute.sym_int() - sym_batch = cute.sym_int() - - # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. - def _sym_stride(): - return cute.sym_int64(divisibility=16) - - q_fake = cute.runtime.make_fake_tensor( - cute_dtype, - (sym_total_tokens, num_heads, head_dim), - stride=(_sym_stride(), _sym_stride(), 1), - ) - # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). - k_fake = cute.runtime.make_fake_tensor( - cute_dtype, - (cute.sym_int(), page_size, head_dim), - stride=(_sym_stride(), _sym_stride(), 1), - ) - bt_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) - ) - # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. - score_fake = cute.runtime.make_fake_tensor( - cutlass.Float32, - (num_heads, sym_total_tokens, cute.sym_int()), - stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), - ) - sl_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (sym_batch,), stride_order=(0,) - ) - fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - kernel = IndexDecodeScoreKernel( - cute_dtype, - num_heads, - max_decode_query_len, - split_k, - head_dim, - prefetch_old=True, - early_issue=early, - prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, - ) - cls.kernel_cache[key] = cute.compile( - kernel, - q_fake, - k_fake, - bt_fake, - score_fake, - sl_fake, - fake_stream, - options="--enable-tvm-ffi", - ) + PRESTAGE_OLD = False + EARLY_TRIGGER = False diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py deleted file mode 100644 index a5b626f3fe15..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_lowm_splitk.py +++ /dev/null @@ -1,1255 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Blackwell low-M BF16/FP16 GEMM with an in-kernel cluster split-K reduction. - -This target's copy of FlashInfer's ``flashinfer/gemm/kernels/dense_bf16_gemm_sm100_splitk.py`` (the Eagle3 -drafter's low-M linears, ``ops/m3_draft_lowm``): the same kernel and the same split order at every tactic, -bit for bit FlashInfer's output (spec/test_modeling_v2_m3_lowm_splitk.py). The one change: the owner rank's final -cast and store is one method, ``_store_tile``, the seam for the drafter-boundary push epilogue and, on the A side, the -SwiGLU prologue that follow. - -Each cluster rank accumulates an exact K slice in FP32. Peers publish partials -to rank 0 through DSMEM; rank 0 reduces, casts, and stores once. The public -``A[M, K] @ B[K, N]`` problem is swapped internally, so tile dimensions below -use kernel coordinates: kernel-M carries public N and kernel-N carries public M. -""" - -from __future__ import annotations - -import dataclasses -import functools - -import cuda.bindings.driver as _cuda -import cutlass -import cutlass.cute as cute -import cutlass.utils as utils -import cutlass.utils.blackwell_helpers as sm100_utils -from cutlass import Int32 -from cutlass._mlir.dialects import llvm -from cutlass.cute import experimental as cute_ext -from cutlass.cute.nvgpu import tcgen05 -from cutlass.cute.runtime import from_dlpack -from cutlass.cutlass_dsl import T, dsl_user_op - -#: Per-CTA SMEM capacity reported by CuTeDSL on SM100/SM103. -_SMEM_CAPACITY_BYTES = utils.get_smem_capacity_in_bytes("sm_100") - -#: K extent of one CTA tile. -_CTA_K = 128 # capacity: fixed: a kernel tile constant - -#: Kernel-M tiles; 64 increases CTA count for low-M decode shapes. -_SUPPORTED_MMA_M = (64, 128) - -#: Kernel-N carries public M, which is limited to 32. -_SUPPORTED_MMA_N = (8, 16, 32) - -#: Physical cluster-K sizes; split 1 compiles out the DSMEM path. -_SUPPORTED_SPLIT_K = (1, 2, 3, 4) - -#: Largest public M this low-M policy serves. -_MAX_M = 32 # capacity: fixed: the copied FlashInfer kernel's row bound - -#: Bytes per FP32 partial exchanged through DSMEM. -_FP32_BYTES = 4 - -#: DSMEM mailbox base alignment, in bytes. -_MAILBOX_ALIGN_BYTES = 128 - -#: Size and alignment of one mbarrier, in bytes. -_MBARRIER_BYTES = 8 - -#: Size and alignment of the TMEM base pointer slot. -_TMEM_POINTER_BYTES = 4 - -#: Bytes per BF16/FP16 element. -_AB_ELEMENT_BYTES = 2 - -#: Alignment of the A/B shared-memory buffers. -_AB_BUFFER_ALIGN_BYTES = 1024 # capacity: fixed: an alignment - -#: A/B pipeline stage bounds. One stage is useful only for a single K tile. -_MIN_AB_STAGES = 1 # capacity: fixed: pipeline stages -_DEFAULT_MIN_AB_STAGES = 2 -_MAX_AB_STAGES = 12 # capacity: fixed: pipeline stages - - -@dataclasses.dataclass(frozen=True) -class SplitKTactic: - """One specialization; mma_m carries public N and mma_n carries public M.""" - - mma_m: int - mma_n: int - split_k: int - ab_stages: int - - -def _align_up(value: int, alignment: int) -> int: - return ((value + alignment - 1) // alignment) * alignment - - -def _smem_bytes( - tactic: SplitKTactic, - ab_stages: int, -) -> int: - """Mirror the device allocator's shared-memory layout.""" - cursor = ( - _align_up( - tactic.mma_m * _CTA_K * _AB_ELEMENT_BYTES * ab_stages, - _AB_BUFFER_ALIGN_BYTES, - ) - + tactic.mma_n * _CTA_K * _AB_ELEMENT_BYTES * ab_stages - ) - - cursor = _align_up(cursor, _MBARRIER_BYTES) - cursor += 2 * ab_stages * _MBARRIER_BYTES - cursor += 3 * _MBARRIER_BYTES - cursor = _align_up(cursor, _TMEM_POINTER_BYTES) - cursor += _TMEM_POINTER_BYTES - - if tactic.split_k == 1: - return cursor - - return ( - _align_up( - _align_up(cursor, _MAILBOX_ALIGN_BYTES) - + (tactic.split_k - 1) * tactic.mma_m * tactic.mma_n * _FP32_BYTES, - _MBARRIER_BYTES, - ) - + _MBARRIER_BYTES - ) - - -def _max_ab_stages_for( - tactic: SplitKTactic, - smem_capacity: int, -) -> int: - return next( - ( - stages - for stages in range(_MAX_AB_STAGES, -1, -1) - if _smem_bytes(tactic, stages) <= smem_capacity - ), - 0, - ) - - -def validate_tactic( - tactic: SplitKTactic, - m: int, - n: int, - k: int, - *, - smem_capacity: int = _SMEM_CAPACITY_BYTES, -) -> None: - """Reject a tactic that cannot serve ``(m, n, k)``.""" - if tactic.mma_m not in _SUPPORTED_MMA_M: - raise ValueError(f"unsupported mma_m={tactic.mma_m}") - if tactic.mma_n not in _SUPPORTED_MMA_N: - raise ValueError(f"unsupported mma_n={tactic.mma_n}") - if tactic.split_k not in _SUPPORTED_SPLIT_K: - raise ValueError(f"unsupported split_k={tactic.split_k}") - if not _MIN_AB_STAGES <= tactic.ab_stages <= _MAX_AB_STAGES: - raise ValueError( - f"ab_stages must be in [{_MIN_AB_STAGES}, {_MAX_AB_STAGES}], got {tactic.ab_stages}" - ) - if not 1 <= m <= _MAX_M: - raise ValueError(f"this low-M policy requires 1 <= M <= {_MAX_M}, got {m}") - if n <= 0: - raise ValueError(f"N must be positive, got {n}") - if k <= 0 or k % _CTA_K or (k // _CTA_K) % tactic.split_k: - raise ValueError( - f"K={k} with CTA_K={_CTA_K} does not divide evenly across split_k={tactic.split_k}" - ) - smem_bytes = _smem_bytes(tactic, tactic.ab_stages) - if smem_bytes > smem_capacity: - raise ValueError( - f"tactic {tactic} needs {smem_bytes} B of shared memory but only " - f"{smem_capacity} B are available; max ab_stages is " - f"{_max_ab_stages_for(tactic, smem_capacity)}" - ) - - -def autotune_tactics( - m: int, - n: int, - k: int, - *, - smem_capacity: int = _SMEM_CAPACITY_BYTES, -) -> list[SplitKTactic]: - """Return valid tactics in the shape-specific stage window.""" - tactics: list[SplitKTactic] = [] - for mma_m in _SUPPORTED_MMA_M: - for mma_n in _SUPPORTED_MMA_N: - for split_k in _SUPPORTED_SPLIT_K: - min_stages = _MIN_AB_STAGES if k == _CTA_K else _DEFAULT_MIN_AB_STAGES - base = SplitKTactic(mma_m, mma_n, split_k, min_stages) - try: - validate_tactic(base, m, n, k, smem_capacity=smem_capacity) - except ValueError: - continue - max_stages = _max_ab_stages_for(base, smem_capacity) - # Short K favors shallow pipelines; long K stays near the cap. - tactics.extend( - dataclasses.replace(base, ab_stages=ab_stages) - for ab_stages in ( - range(min_stages, min(max_stages, 6) + 1) - if k <= 4 * _CTA_K - else range( - min(max(5, max_stages - 2), max_stages), - max_stages + 1, - ) - ) - ) - return tactics - - -def default_tactic(m: int, n: int, k: int) -> SplitKTactic: - """Choose the default occupancy-oriented tactic.""" - if n <= 512: - mma_m = 64 - mma_n = 16 if n == 512 and m > 24 else 8 - requested_split = 4 - else: - mma_n = 8 if m <= 8 else 16 if m <= 16 else 32 - if n <= 3072: - mma_m = 128 if m <= 16 else 64 - requested_split = 4 if m <= 16 else 2 - elif n < 8192: - mma_m = 64 - requested_split = 2 - else: - mma_m = 128 if k <= 1024 and m <= 24 else 64 - requested_split = 1 - - if k <= 4 * _CTA_K: - requested_split = 1 - split_k = next( - split_k - for split_k in reversed(_SUPPORTED_SPLIT_K) - if split_k <= requested_split and (k // _CTA_K) % split_k == 0 - ) - tactic = SplitKTactic(mma_m, mma_n, split_k, _MIN_AB_STAGES) - max_stages = _max_ab_stages_for(tactic, _SMEM_CAPACITY_BYTES) - tactic = dataclasses.replace( - tactic, - ab_stages=( - _MIN_AB_STAGES - if k == _CTA_K - else _DEFAULT_MIN_AB_STAGES - if k <= 2 * _CTA_K and m > 8 - else min(max_stages, 6) - if k <= 4 * _CTA_K - else max_stages - ), - ) - validate_tactic(tactic, m, n, k) - return tactic - - -__all__ = [ - "SplitKTactic", - "autotune_tactics", - "default_tactic", - "run_splitk_dense", -] - - -@dsl_user_op -def _map_shared_rank( - smem_ptr: "cute.Pointer", - peer_cta_rank_in_cluster: "Int32", - *, - loc=None, - ip=None, -) -> "Int32": - """Map an SMEM pointer into a peer CTA's address space.""" - return Int32( - llvm.inline_asm( - T.i32(), - [ - smem_ptr.toint(loc=loc, ip=ip).ir_value(), - peer_cta_rank_in_cluster.ir_value(), - ], - "mapa.shared::cluster.u32 $0, $1, $2;", - "=r,r,r", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - ) - - -@dsl_user_op -def _store_shared_remote_v4( - value0, - value1, - value2, - value3, - smem_ptr: "cute.Pointer", - mbar_ptr: "cute.Pointer", - peer_cta_rank_in_cluster: "Int32", - *, - loc=None, - ip=None, -) -> None: - """Publish four FP32 partials into a peer's SMEM, crediting 16 bytes.""" - llvm.inline_asm( - None, - [ - _map_shared_rank(smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), - value0.bitcast(Int32).ir_value(loc=loc, ip=ip), - value1.bitcast(Int32).ir_value(loc=loc, ip=ip), - value2.bitcast(Int32).ir_value(loc=loc, ip=ip), - value3.bitcast(Int32).ir_value(loc=loc, ip=ip), - _map_shared_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value(), - ], - "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " - "[$0], {$1, $2, $3, $4}, [$5];", - "r,r,r,r,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - - -#: Rank that gathers partials and stores the output. -OWNER_RANK = 0 # capacity: fixed: a split index, not a size - - -# The drafter-boundary push (``ar_push``, modeling.DRAFT_AR_PUSH): the exchange words of -# ops/m3_mc_allreduce, as m3_mx_gemv's ``ar`` push stores them. -@dsl_user_op -def _ld_volatile_s32(addr, *, loc=None, ip=None): - return Int32( - llvm.inline_asm( - T.i32(), - [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], - "ld.volatile.global.s32 $0, [$1];", - "=r,l", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - ) - - -@dsl_user_op -def _st_global_b32(addr, value, *, loc=None, ip=None): - llvm.inline_asm( - None, - [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), Int32(value).ir_value(loc=loc, ip=ip)], - "st.global.b32 [$0], $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - - -@dsl_user_op -def _bf16x2_word(lo, hi, *, loc=None, ip=None): - """bf16(lo) | bf16(hi) << 16 (RN), -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" - return Int32( - llvm.inline_asm( - T.i32(), - [ - cutlass.Float32(hi).ir_value(loc=loc, ip=ip), - cutlass.Float32(lo).ir_value(loc=loc, ip=ip), - ], - "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" - "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" - "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" - "mov.b32 $0, w;\n\t}", - "=r,f,f", - has_side_effects=False, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - ) - - -class SplitKDenseGemmKernel: - """Standalone BF16/FP16 GEMM with a cluster-local split-K reduction.""" - - def __init__( - self, - *, - tactic: SplitKTactic, - use_pdl: bool, - has_bias: bool, - ar_push: bool = False, - ) -> None: - self.ar_push = ar_push - self.acc_dtype = cutlass.Float32 - self.cta_m = tactic.mma_m - self.cta_n = tactic.mma_n - self.cta_k = _CTA_K - self.num_ab_stage = tactic.ab_stages - self.split_k = tactic.split_k - self.use_pdl = use_pdl - self.has_bias = has_bias - - self.threads_per_cta = 256 - self.epilog_threads = 128 - self.mma_tiler_mn = (tactic.mma_m, tactic.mma_n) - self.cta_group = tcgen05.CtaGroup.ONE - self.tma_op = cute_ext.OperationTypeEnum.SM90_TMA_LOAD - self.cluster_shape = (1, tactic.split_k, 1) - - values_per_thread = (tactic.mma_m * tactic.mma_n) // self.epilog_threads - if values_per_thread % 4: - raise ValueError( - f"CTA tile ({tactic.mma_m}, {tactic.mma_n}) gives " - f"{values_per_thread} " - "values per epilogue thread; remote stores require a multiple of 4" - ) - self.mailbox_elements = (tactic.split_k - 1) * self.epilog_threads * values_per_thread - self.expected_transaction_bytes = self.mailbox_elements * _FP32_BYTES - - @cute.experimental.jit - def __call__( - self, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - bias: cute.Tensor, - stream: _cuda.CUstream, - ar_ws: cutlass.Int64 = cutlass.Int64(0), - ar_flags: cutlass.Int64 = cutlass.Int64(0), - ar_rank: Int32 = Int32(0), - ar_world: Int32 = Int32(1), - ar_maxt: Int32 = Int32(8), - ar_tokens: Int32 = Int32(0), - ar_n: Int32 = Int32(0), - ): - # Grid-y packs output-N tile and cluster rank. - self.kernel( - a, b, c, bias, ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_tokens, ar_n - ).launch( - grid=( - cute.ceil_div(c.layout.shape[0], self.cta_m), - cute.ceil_div(c.layout.shape[1], self.cta_n) * self.split_k, - c.layout.shape[2], - ), - block=(self.threads_per_cta, 1, 1), - cluster=self.cluster_shape, - smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")), - stream=stream, - use_pdl=self.use_pdl, - ) - - @cute.experimental.kernel - def kernel( - self, - mA: cute.Tensor, # (Gemm_M, Gemm_K, Gemm_L), K-major - mB: cute.Tensor, # (Gemm_N, Gemm_K, Gemm_L), K-major - mC: cute.Tensor, # (Gemm_M, Gemm_N, Gemm_L), M-major - mBias: cute.Tensor, # Broadcast bias; dead when has_bias=False - ar_ws: cutlass.Int64, # ar_push: the boundary all-reduce's multicast exchange words (int32) - ar_flags: cutlass.Int64, # ar_push: its per-(token, cluster CTA) call flags (parity = flag & 1) - ar_rank: Int32, - ar_world: Int32, - ar_maxt: Int32, # the bucket's tokens - ar_tokens: Int32, # rows pushed (public M) - ar_n: Int32, # public N (the hidden size): the exchange row width - ): - """Allocate storage and dispatch the specialized warps.""" - stages = self.num_ab_stage - - ab_dtype = mA.element_type - tiled_mma = sm100_utils.make_trivial_tiled_mma( - ab_dtype, - ab_dtype, - utils.LayoutEnum.from_tensor(mA).mma_major_mode(), - utils.LayoutEnum.from_tensor(mB).mma_major_mode(), - self.acc_dtype, - self.cta_group, - self.mma_tiler_mn, - ) - - mnk_tiler = (self.mma_tiler_mn[0], self.mma_tiler_mn[1], self.cta_k) - block_idx = cute.arch.block_idx() - bidx = block_idx[0] - split_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) - n_idx = block_idx[1] // self.split_k - l_idx = block_idx[2] - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - - sA = cute_ext.allocate( - ab_dtype, - cute.AddressSpace.smem, - sm100_utils.make_smem_layout_a(tiled_mma, mnk_tiler, ab_dtype, stages), - alignment=_AB_BUFFER_ALIGN_BYTES, - ) - sB = cute_ext.allocate( - ab_dtype, - cute.AddressSpace.smem, - sm100_utils.make_smem_layout_b(tiled_mma, mnk_tiler, ab_dtype, stages), - alignment=_AB_BUFFER_ALIGN_BYTES, - ) - - acc_layout = cute_ext.make_tmem_layout_acc(tiled_mma, self.mma_tiler_mn, acc_stage=1) - c_tiler_mn = (self.cta_m, self.cta_n) - - bar_full = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(stages), - alignment=_MBARRIER_BYTES, - ).iterator - bar_empty = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(stages), - alignment=_MBARRIER_BYTES, - ).iterator - bar_tma_epilog = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - bar_mma_epilog = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - bar_tmem_alloc = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - tmem_base_ptr = cute_ext.allocate( - cutlass.Int32, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_TMEM_POINTER_BYTES, - ).iterator - # ar_push: the owner's BF16 tile (kernel-M = output column, kernel-N = token), M-major. - sD = cute_ext.allocate( - mC.element_type, - cute.AddressSpace.smem, - cute.make_layout((self.cta_m, self.cta_n), stride=(1, self.cta_m)), - alignment=16, - ) - - if cutlass.const_expr(self.split_k > 1): - mailbox = cute_ext.allocate( - cutlass.Float32, - cute.AddressSpace.smem, - cute.make_layout(self.mailbox_elements), - alignment=_MAILBOX_ALIGN_BYTES, - ) - bar_reduce = cute_ext.allocate( - cutlass.Int64, - cute.AddressSpace.smem, - cute.make_layout(1), - alignment=_MBARRIER_BYTES, - ).iterator - else: - # Dummy operands for the compile-time-elided reduction. - mailbox = sA - bar_reduce = bar_mma_epilog - - if warp_idx == 0: - with cute.arch.elect_one(): - for i in range(stages): - cute.arch.mbarrier_init(bar_full + i, 2) - cute.arch.mbarrier_init(bar_empty + i, 1) - cute.arch.mbarrier_init(bar_tma_epilog, 32) - cute.arch.mbarrier_init(bar_mma_epilog, 1) - cute.arch.mbarrier_init(bar_tmem_alloc, 160) - - if cutlass.const_expr(self.split_k > 1): - # Owner arrival plus peer transaction-byte credits. - cute.arch.mbarrier_init(bar_reduce, 1) - - cute.arch.mbarrier_init_fence() - if cutlass.const_expr(self.split_k > 1): - # Publish peer barriers before cross-CTA stores. - cute.arch.cluster_arrive_relaxed() - else: - cute.arch.barrier() - - # Host validation guarantees an equal, tail-free K partition. - k_tile_count = cute.size(mA, mode=[1]) // self.cta_k // self.split_k - k_tile_start = split_rank * k_tile_count - - if cutlass.const_expr(self.split_k > 1): - cute.arch.cluster_wait() - - # Warp 3 is idle; warps 4-7 run the epilogue. - if warp_idx == 0: - self.dma_warp( - bar_full, - bar_empty, - bar_tma_epilog, - cute.local_tile(mA, (self.cta_m, self.cta_k), (bidx, None, l_idx)), - sA, - cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A"), - k_tile_start, - k_tile_count, - True, - ) - elif warp_idx == 1: - self.dma_warp( - bar_full, - bar_empty, - bar_tma_epilog, - cute.local_tile(mB, (self.cta_n, self.cta_k), (n_idx, None, l_idx)), - sB, - cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B"), - k_tile_start, - k_tile_count, - False, - ) - elif warp_idx == 2: - self.mma_warp( - bar_full, - bar_empty, - bar_mma_epilog, - bar_tmem_alloc, - tiled_mma, - sA, - sB, - tmem_base_ptr, - acc_layout, - self.cta_k // cute.size(tiled_mma.shape_mnk, mode=[2]), - k_tile_count, - ) - elif warp_idx >= 4: - self.epilog_warp( - bar_tma_epilog, - bar_mma_epilog, - bar_tmem_alloc, - tmem_base_ptr, - acc_layout, - cute.local_tile(mC, c_tiler_mn, (bidx, n_idx, l_idx)), - cute.local_tile(mBias, c_tiler_mn, (bidx, n_idx, l_idx)), - cute.arch.thread_idx()[0] - 128, - mC.element_type, - utils.LayoutEnum.from_tensor(mC), - mailbox, - bar_reduce, - split_rank, - sD, - bidx, - n_idx * self.cta_n, - ar_ws, - ar_flags, - ar_rank, - ar_world, - ar_maxt, - ar_tokens, - ar_n, - ) - - @cute.experimental.jit - def dma_warp( - self, - bar_full, - bar_empty, - bar_tma_epilog, - g_tile: cute.Tensor, - s_tile: cute.Tensor, - cta_v_map: cute.Layout, - k_tile_start: cutlass.Int32, - k_tile_count: cutlass.Int32, - is_a: cutlass.Constexpr, - ): - stages = self.num_ab_stage - if cutlass.const_expr(not is_a and self.use_pdl): - cute.arch.griddepcontrol_wait() - - empty_phase = cutlass.Int32(1) - for k_tile in cutlass.range(k_tile_count, unroll=1): - stage = k_tile % stages - cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - bar_full + stage, - cute.size_in_bytes( - s_tile.element_type, - cute.slice_(s_tile.layout, (None, None, None, 0)), - ), - ) - cute_ext.tma_load( - g_tile[None, None, k_tile_start + k_tile], - s_tile[None, None, None, stage], - (bar_full + stage).value, - cta_v_map=cta_v_map, - tma_operation_type=self.tma_op, - update_expect_tx=False, - ) - if stage == stages - 1: - empty_phase = empty_phase ^ 1 - - if cutlass.const_expr(is_a and self.use_pdl and not self.ar_push): - # ar_push: the dependent is a prepushed all-reduce, which skips its grid wait: it is released only after - # this grid's own grid wait and its pushes, in the owner's _store_tile (the prepushed >= 2 contract). - cute.arch.griddepcontrol_launch_dependents() - if cutlass.const_expr(not is_a and self.has_bias): - cute.arch.mbarrier_arrive(bar_tma_epilog) - self._drain_producer(bar_empty, empty_phase, k_tile_count) - - @cute.experimental.jit - def _drain_producer( - self, - bar_empty, - empty_phase: cutlass.Int32, - k_tile_count: cutlass.Int32, - ): - stages = self.num_ab_stage - for tail in cutlass.range(stages, unroll=1): - stage = (tail + k_tile_count) % stages - cute.arch.mbarrier_wait(bar_empty + stage, empty_phase) - if stage == stages - 1: - empty_phase = empty_phase ^ 1 - - @cute.experimental.jit - def mma_warp( - self, - bar_full, - bar_empty, - bar_mma_epilog, - bar_tmem_alloc, - tiled_mma: cute.TiledMma, - sA: cute.Tensor, - sB: cute.Tensor, - tmem_base_ptr, - acc_layout: cutlass.Constexpr, - mma_inst_tile_k: cutlass.Constexpr, - k_tile_count: cutlass.Int32, - ): - num_tmem_cols = 256 - 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) - - tmem_ptr = cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr) - accumulator = cute.make_tensor(tmem_ptr, acc_layout)[None, None, None, 0] - mma_atom = cute.make_mma_atom(tiled_mma.op) - full_phase = cutlass.Int32(0) - for k_tile in cutlass.range(k_tile_count, unroll=1): - stage = k_tile % self.num_ab_stage - cute.arch.mbarrier_wait(bar_full + stage, full_phase) - for k_block in range(mma_inst_tile_k): - if k_block == 0: - mma_atom.set(tcgen05.Field.ACCUMULATE, k_tile != 0) - else: - mma_atom.set(tcgen05.Field.ACCUMULATE, True) - cute_ext.dot( - mma_atom, - cute.append_ones(sA[None, None, k_block, stage], up_to_rank=3), - cute.append_ones(sB[None, None, k_block, stage], up_to_rank=3), - accumulator, - ) - with cute.arch.elect_one(): - tcgen05.commit(bar_empty + stage, None, self.cta_group) - if stage == self.num_ab_stage - 1: - full_phase = full_phase ^ 1 - - with cute.arch.elect_one(): - tcgen05.commit(bar_mma_epilog, None, self.cta_group) - cute.arch.mbarrier_arrive(bar_tmem_alloc) - cute.arch.mbarrier_wait(bar_tmem_alloc, 1) - cute.arch.dealloc_tmem(tmem_ptr, num_tmem_cols, is_two_cta=False) - - @cute.experimental.jit - def epilog_warp( - self, - bar_tma_epilog, - bar_mma_epilog, - bar_tmem_alloc, - tmem_base_ptr, - acc_layout: cutlass.Constexpr, - gD_tile: cute.Tensor, - gBias_tile: cute.Tensor, - epi_tid: cutlass.Int32, - c_dtype: cutlass.Constexpr, - d_layout: cutlass.Constexpr, - mailbox, - bar_reduce, - split_rank: cutlass.Int32, - sD: cute.Tensor, - bidx: cutlass.Int32, - row_offset: cutlass.Int32, - ar_ws: cutlass.Int64, - ar_flags: cutlass.Int64, - ar_rank: cutlass.Int32, - ar_world: cutlass.Int32, - ar_maxt: cutlass.Int32, - ar_tokens: cutlass.Int32, - ar_n: cutlass.Int32, - ): - # Wait until MMA publishes the TMEM base pointer. - cute.arch.mbarrier_arrive(bar_tmem_alloc) - cute.arch.mbarrier_wait(bar_tmem_alloc, 0) - - acc_view = cute.make_tensor( - cute.arch.retrieve_tmem_ptr(self.acc_dtype, 16, tmem_base_ptr), - acc_layout, - )[((None, None), 0, 0, 0)] - - epi_tile = (self.cta_m, self.cta_n) - tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy( - sm100_utils.get_tmem_load_op( - (self.cta_m, self.cta_n, self.cta_k), - d_layout, - c_dtype, - self.acc_dtype, - epi_tile, - False, - ), - acc_view, - ) - gD_epi = cute.flat_divide(gD_tile, epi_tile) - - # Match each epilogue thread's TMEM partition in RMEM. - rmem_layout = cute_ext.make_t2r_rmem_layout(tiled_copy_t2r, gD_epi, epi_tid) - rAcc = cute_ext.allocate( - self.acc_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - rD = cute_ext.allocate( - c_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - thr_t2r = tiled_copy_t2r.get_slice(epi_tid) - - if cutlass.const_expr(self.has_bias): - bias_dtype = gBias_tile.element_type - rBias = cute_ext.allocate( - bias_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - rBiasAcc = cute_ext.allocate( - self.acc_dtype, - cute.AddressSpace.rmem, - rmem_layout, - alignment=32, - ) - if split_rank == OWNER_RANK: - cute.arch.mbarrier_wait(bar_tma_epilog, 0) - cute_ext.partition_and_copy( - cute.make_tiled_copy_D( - cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), bias_dtype), - tiled_copy_t2r, - ).get_slice(epi_tid), - cute.flat_divide(gBias_tile, epi_tile)[None, None, 0, 0], - rBias, - ) - rBiasAcc.store(rBias.load().to(self.acc_dtype)) - - cute.arch.mbarrier_wait(bar_mma_epilog, 0) - cute_ext.partition_and_copy(thr_t2r, acc_view, rAcc) - # Make tcgen05.ld visible before TMEM release and RMEM use. - cute.arch.fence_view_async_tmem_load() - cute.arch.mbarrier_arrive(bar_tmem_alloc) - - # Peers publish FP32 partials; only rank 0 reduces and stores. - if cutlass.const_expr(self.split_k > 1): - assert cute.size(rmem_layout) == self.mailbox_elements // ( - (self.split_k - 1) * self.epilog_threads - ) - values_per_thread = cutlass.const_expr(cute.size(rmem_layout)) - values_per_peer = cutlass.const_expr(self.epilog_threads * values_per_thread) - if split_rank != OWNER_RANK: - for value_idx in cutlass.range_constexpr(0, values_per_thread, 4): - _store_shared_remote_v4( - rAcc[value_idx], - rAcc[value_idx + 1], - rAcc[value_idx + 2], - rAcc[value_idx + 3], - mailbox.iterator - + (split_rank - Int32(1)) * values_per_peer - + epi_tid * values_per_thread - + value_idx, - bar_reduce, - Int32(OWNER_RANK), - ) - else: - if epi_tid == 0: - cute.arch.mbarrier_arrive_and_expect_tx( - bar_reduce, self.expected_transaction_bytes - ) - cute.arch.mbarrier_wait(bar_reduce, 0) - for peer in cutlass.range_constexpr(self.split_k - 1): - for value_idx in cutlass.range_constexpr(values_per_thread): - rAcc[value_idx] = ( - rAcc[value_idx] - + mailbox[ - peer * values_per_peer + epi_tid * values_per_thread + value_idx - ] - ) - - if split_rank == OWNER_RANK: - if cutlass.const_expr(self.has_bias): - rAcc.store(rAcc.load() + rBiasAcc.load()) - self._store_tile(thr_t2r, rAcc, rD, gD_epi, c_dtype) - if cutlass.const_expr(self.ar_push): - self._push_tile( - thr_t2r, - rD, - sD, - epi_tid, - bidx, - row_offset, - ar_ws, - ar_flags, - ar_rank, - ar_world, - ar_maxt, - ar_tokens, - ar_n, - ) - - # The reduction mbarrier covers remote stores; no cluster barrier needed. - - @cute.experimental.jit - def _store_tile(self, thr_t2r, rAcc, rD, gD_epi, c_dtype: cutlass.Constexpr): - """The owner rank's reduced FP32 tile: one cast to ``c_dtype`` and the store to global memory.""" - rD.store(rAcc.load().to(c_dtype)) - # Preserve TMEM coordinates; the copy predicates output tails. - cute_ext.partition_and_copy(thr_t2r, rD, gD_epi[None, None, 0, 0]) - - @cute.experimental.jit - def _push_tile( - self, - thr_t2r, - rD, - sD, - epi_tid: cutlass.Int32, - bidx: cutlass.Int32, - row_offset: cutlass.Int32, - ar_ws: cutlass.Int64, - ar_flags: cutlass.Int64, - ar_rank: cutlass.Int32, - ar_world: cutlass.Int32, - ar_maxt: cutlass.Int32, - ar_tokens: cutlass.Int32, - ar_n: cutlass.Int32, - ): - """ar_push: the owner's tile rows into the boundary all-reduce's exchange, as its own push of the output - would store them (ops/m3_mc_allreduce: word (t, column pair) at the call parity of the (t, column / (N / 8)) - cluster CTA, -0.0 halves as +0.0), then the release of the dependents (the prepushed all-reduce).""" - cute_ext.partition_and_copy(thr_t2r, rD, sD) - cute.arch.barrier(barrier_id=1, number_of_threads=self.epilog_threads) - if cutlass.const_expr(self.use_pdl): - cute.arch.griddepcontrol_wait() - pairs = self.cta_m // 2 - col0 = bidx * self.cta_m - half = cutlass.Int64(ar_maxt) * cutlass.Int64(ar_world) * cutlass.Int64(ar_n // 2) - for w0 in cutlass.range_constexpr(0, pairs * self.cta_n, self.epilog_threads): - word = epi_tid + w0 - local_t = word // pairs - t = row_offset + local_t - p = word % pairs - if t < ar_tokens: - n = col0 + 2 * p - par = _ld_volatile_s32( - ar_flags + cutlass.Int64(4) * cutlass.Int64(t * 8 + n // (ar_n // 8)) - ) & Int32(1) - off = ( - cutlass.Int64(par) * half - + cutlass.Int64(t * ar_world + ar_rank) * cutlass.Int64(ar_n // 2) - + cutlass.Int64(n // 2) - ) - _st_global_b32( - ar_ws + off * cutlass.Int64(4), - _bf16x2_word( - sD[2 * p, local_t].to(cutlass.Float32), - sD[2 * p + 1, local_t].to(cutlass.Float32), - ), - ) - if cutlass.const_expr(self.use_pdl): - cute.arch.griddepcontrol_launch_dependents() - - -import torch as _torch # noqa: E402 (as in FlashInfer) - -_SUPPORTED_TORCH_DTYPES = (_torch.bfloat16, _torch.float16) - - -@cute.experimental.jit -def _bmm_no_bias( - gemm_op: cutlass.Constexpr, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - stream: _cuda.CUstream, -): - c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) - gemm_op( - cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), - cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), - c, - cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), - stream, - ) - - -@cute.experimental.jit -def _bmm_push( - gemm_op: cutlass.Constexpr, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - stream: _cuda.CUstream, - ar_ws: cutlass.Int64, - ar_flags: cutlass.Int64, - ar_rank: Int32, - ar_world: Int32, - ar_maxt: Int32, - ar_tokens: Int32, - ar_n: Int32, -): - c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])) - gemm_op( - cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), - cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), - c, - cute.make_tensor(c.iterator, cute.select(c.layout, mode=[0, 1, 2])), - stream, - ar_ws, - ar_flags, - ar_rank, - ar_world, - ar_maxt, - ar_tokens, - ar_n, - ) - - -@cute.experimental.jit -def _bmm_bias( - gemm_op: cutlass.Constexpr, - a: cute.Tensor, - b: cute.Tensor, - c: cute.Tensor, - bias: cute.Tensor, - stream: _cuda.CUstream, -): - gemm_op( - cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0])), - cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0])), - cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0])), - cute.make_tensor(bias.iterator, cute.select(bias.layout, mode=[1, 2, 0])), - stream, - ) - - -def _from_dlpack_dynamic(tensor, leading_dim: int, assumed_align: int = 32): - return from_dlpack(tensor, assumed_align=assumed_align).mark_layout_dynamic( - leading_dim=leading_dim - ) - - -def _detect_leading_dim(tensor: _torch.Tensor) -> int: - # Ignore synthetic batch stride, including 1x1. - for dim, stride in enumerate(tensor.stride()[1:], start=1): - if stride == 1: - return dim - raise ValueError("tensor has no stride-1 dimension") - - -def _make_layout_tensor( - shape: tuple[int, ...], dtype: _torch.dtype, leading_dim: int -) -> _torch.Tensor: - permutation = [dim for dim in range(len(shape)) if dim != leading_dim] + [leading_dim] - return _torch.empty( - tuple(shape[dim] for dim in permutation), dtype=dtype, device="cuda" - ).permute([permutation.index(dim) for dim in range(len(shape))]) - - -def _make_compile_repr_tensors( - dtype: _torch.dtype, - has_bias: bool, - a_leading: int, - b_leading: int, - c_leading: int, -): - m, n, k, batch = 64, 8, _CTA_K, 1 - tensors = tuple( - _from_dlpack_dynamic(_make_layout_tensor(shape, dtype, leading_dim), leading_dim) - for shape, leading_dim in ( - ((batch, n, k), a_leading), - ((batch, k, m), b_leading), - ((batch, n, m), c_leading), - ) - ) - if not has_bias: - return (*tensors, None) - return ( - *tensors, - _from_dlpack_dynamic( - _torch.empty((n,), dtype=dtype, device="cuda").as_strided( - size=(batch, n, m), stride=(0, 1, 0) - ), - 1, - 2, - ), - ) - - -def _to_cute_swap(a, b, out, bias): - a_swap = b.unsqueeze(0).transpose(-2, -1) - b_swap = a.unsqueeze(0).transpose(-2, -1) - c_swap = out.unsqueeze(0).transpose(-2, -1) - leading_dims = tuple(_detect_leading_dim(tensor) for tensor in (a_swap, b_swap, c_swap)) - cute_tensors = tuple( - _from_dlpack_dynamic(tensor, leading_dim) - for tensor, leading_dim in zip((a_swap, b_swap, c_swap), leading_dims, strict=True) - ) - if bias is None: - return (*cute_tensors, None, leading_dims) - return ( - *cute_tensors, - _from_dlpack_dynamic( - bias.as_strided(size=(1, c_swap.shape[1], c_swap.shape[2]), stride=(0, 1, 0)), - 1, - 2, - ), - leading_dims, - ) - - -@functools.cache -def _get_compiled_splitk_kernel( - dtype, - tactic: SplitKTactic, - use_pdl: bool, - has_bias: bool, - leading_dims: tuple[int, int, int], -): - if dtype not in _SUPPORTED_TORCH_DTYPES: - raise ValueError(f"split-K dense GEMM supports {_SUPPORTED_TORCH_DTYPES}; got {dtype}") - - kernel = SplitKDenseGemmKernel( - tactic=tactic, - use_pdl=use_pdl, - has_bias=has_bias, - ) - compile_tensors = _make_compile_repr_tensors(dtype, has_bias, *leading_dims) - stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) - if has_bias: - compiled = cute_ext.compile(_bmm_bias, kernel, *compile_tensors, stream) - else: - compiled = cute_ext.compile(_bmm_no_bias, kernel, *compile_tensors[:3], stream) - return compiled - - -def _validate_runtime_tensors(a, b, bias, out) -> tuple[int, int, int]: - tensors = (a, b, out) + ((bias,) if bias is not None else ()) - if any(not isinstance(tensor, _torch.Tensor) for tensor in tensors): - raise ValueError("a, b, out, and bias must be torch tensors") - if a.ndim != 2 or b.ndim != 2 or out.ndim != 2: - raise ValueError("split-K dense GEMM accepts only 2D 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, b, out, and bias must share BF16 or FP16 dtype") - if any(not (tensor.is_contiguous() or tensor.t().is_contiguous()) for tensor in (a, b, out)): - raise ValueError("a, b, and out must be dense row-major or column-major matrices") - if any(tensor.data_ptr() % 32 for tensor in (a, b, out)): - raise ValueError("a, b, and out must be 32-byte aligned") - - m, k = a.shape - if b.shape[0] != k: - raise ValueError(f"incompatible shapes: a is {tuple(a.shape)}, b is {tuple(b.shape)}") - n = b.shape[1] - if out.shape != (m, n): - raise ValueError(f"out must have shape {(m, n)}, got {tuple(out.shape)}") - if bias is not None and (bias.ndim != 1 or bias.shape[0] != n or not bias.is_contiguous()): - raise ValueError( - f"bias must be contiguous with shape {(n,)}, " - f"got shape {tuple(bias.shape)} and stride {bias.stride()}" - ) - - return m, n, k - - -def run_splitk_dense( - a, - b, - bias, - out, - pdl: bool, - tactic: SplitKTactic, -): - """Run ``A[M,K] @ B[K,N]`` with the ``mm_bf16`` layouts.""" - validate_tactic(tactic, *_validate_runtime_tensors(a, b, bias, out)) - has_bias = bias is not None - cute_tensors = _to_cute_swap(a, b, out, bias) - compiled = _get_compiled_splitk_kernel( - dtype=a.dtype, - tactic=tactic, - use_pdl=pdl, - has_bias=has_bias, - leading_dims=cute_tensors[4], - ) - stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) - if has_bias: - compiled(*cute_tensors[:4], stream) - else: - compiled(*cute_tensors[:3], stream) - return out - - -@functools.cache -def _get_compiled_push( - dtype, tactic: SplitKTactic, use_pdl: bool, leading_dims: tuple[int, int, int] -): - kernel = SplitKDenseGemmKernel(tactic=tactic, use_pdl=use_pdl, has_bias=False, ar_push=True) - compile_tensors = _make_compile_repr_tensors(dtype, False, *leading_dims) - stream = _cuda.CUstream(_torch.cuda.current_stream().cuda_stream) - z64, z32 = cutlass.Int64(0), Int32(0) - return cute_ext.compile( - _bmm_push, - kernel, - *compile_tensors[:3], - stream, - z64, - z64, - z32, - Int32(1), - Int32(8), - z32, - Int32(6144), - ) - - -def compile_push(a, b, out, pdl: bool, tactic) -> None: - """Compile (only) the push variant of ``run_splitk_dense_push`` for these operands' shapes and layouts.""" - cute_tensors = _to_cute_swap(a, b, out, None) - _get_compiled_push(a.dtype, SplitKTactic(*tactic), pdl, cute_tensors[4]) - - -def run_splitk_dense_push(a, b, out, pdl: bool, tactic: SplitKTactic, ar: tuple): - """``run_splitk_dense`` without bias that also pushes ``out``'s rows into a drafter boundary all-reduce's exchange; - ``ar`` = (multicast words address, flags address, rank, world, bucket tokens) of the 8-token bucket.""" - validate_tactic(tactic, *_validate_runtime_tensors(a, b, None, out)) - cute_tensors = _to_cute_swap(a, b, out, None) - compiled = _get_compiled_push(a.dtype, tactic, pdl, cute_tensors[4]) - stream = _cuda.CUstream(_torch.cuda.current_stream(a.device).cuda_stream) - ws, flags, rank, world, maxt = ar - compiled( - *cute_tensors[:3], - stream, - cutlass.Int64(ws), - cutlass.Int64(flags), - Int32(rank), - Int32(world), - Int32(maxt), - Int32(out.shape[0]), - Int32(out.shape[1]), - ) - return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py index 142bcb9a9133..a35e1f59c78d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_mc_allreduce.py @@ -28,11 +28,12 @@ epilogue-option choice and the MXFP8 publication. """ -from dataclasses import dataclass -from typing import Any - import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.mc_workspace import ( + McAllReduceWorkspace as McAllReduceWorkspace, +) + from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op from ..ops.comm.m3_mc_finalize_allreduce_norm import ( m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, @@ -84,74 +85,6 @@ def _ar_opts(m: int, bucket: int, finalize: bool = False, oneshot_cluster: bool return opts | _OPT_CLUSTER_BARRIER -@dataclass(eq=False) -class McAllReduceWorkspace: - """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call - counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" - - uc: torch.Tensor - mc: torch.Tensor - flags: dict - rank: int - world_size: int - hidden: int - handle: Any - comm: Any - max_tokens: int = 8 - - @classmethod - def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": - from tensorrt_llm._torch.distributed.ops import ( - _make_mnnvl_mcast_buffer, - _mnnvl_workspace_all_succeeded, - ) - from tensorrt_llm._utils import mpi_comm - - session = mpi_comm() - session_group = session.Get_group() - group = session_group.Incl(mapping.tp_group) - session_group.Free() - comm = session.Create_group(group) - group.Free() - words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves - ok = not torch.cuda.is_current_stream_capturing() - if not _mnnvl_workspace_all_succeeded(comm, ok): - comm.Free() - raise RuntimeError( - "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" - ) - handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) - uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) - mc = handle.get_mc_buffer((words,), torch.int32, 0) - uc.fill_(EMPTY_WORD) - torch.cuda.synchronize() - if not _mnnvl_workspace_all_succeeded(comm, True): - raise RuntimeError("McAllReduceWorkspace.create failed on another rank") - return cls( - uc=uc, - mc=mc, - flags={}, - rank=mapping.tp_rank, - world_size=mapping.tp_size, - hidden=hidden, - handle=handle, - comm=comm, - max_tokens=max_tokens, - ) - - def flags_for(self, cluster: int): - """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. - get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) - that None never happens and cannot differ between ranks.""" - flags = self.flags.get(cluster) - if flags is None: - if torch.cuda.is_current_stream_capturing(): - return None - flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) - self.flags[cluster] = flags - return flags - - class McBuckets: """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py deleted file mode 100644 index 81122a06c18c..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/m3_moe_kernel.py +++ /dev/null @@ -1,3914 +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. -"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP4 rank (routed and shared intermediate 768): -``m3_moe`` (group-phased weight stream). - -One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), -FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It -writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. - -- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then - makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. -- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a - stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 - slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group - p - LAG's intermediate between CTAs. -- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. - FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction - atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). -- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data - bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). -- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the - call's tag; warp 5 copies the scales to TMEM. -- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its - top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange - (AR_PRE: call parities and descriptor loaded during the last group's FC2). - -Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", -"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. -Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before -griddepcontrol.wait; every producer-output read and every global write after it. -""" - -from __future__ import annotations - -import cuda.bindings.driver as cuda_driver -import cutlass -import cutlass.cute as cute -import cutlass.experimental.cuda as cuda -from cutlass._mlir.dialects import llvm as _llvm -from cutlass._mlir.extras import types as _T -from cutlass.cutlass_dsl import dsl_user_op -from cutlass.experimental import primitives as prims - -_CFG = globals().get("M3_CONFIG") or {} - - -def _cfg(key: str, default): - return type(default)(_CFG.get(key, default)) - - -# ============================================================================= -# Problem shape (trace-time constants). -# ============================================================================= -H = 6144 # capacity: fixed: the checkpoint's hidden size -# routed intermediate per rank: 1536 at TP2 (this target: FC2 K = 6 full 256-K tiles), 768 at TP4, 384 at TP8 (FC2 K -# = 1.5 256-K tiles: the last tile's A box is half out of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 -# MMAs / scale atoms are issued; no padding, no extra bytes) -I_TP = _cfg("i_tp", 768) -TWO_I = 2 * I_TP -assert I_TP % 128 == 0 and 128 <= I_TP <= 1536, I_TP -E_LOCAL = 128 -TOP_K = 4 # capacity: fixed: the checkpoint's top-k -ROUTE_SCALE = 2.0 -SUM_EPS = 1e-20 -M_MAX = _cfg("m_max", 1) -NCTA = _cfg("num_ctas", 152) -FUSED_ROUTER = bool(_cfg("fused_router", 0)) -NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes -NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. -TIMING = bool(_cfg("timing", 0)) -HANG_DBG = bool(_cfg("hang_dbg", 0)) -XVDUMP = bool( - _cfg("xvdump", 0) -) # debug only (needs a plan without split units): FC1 sums into `partials` -YDUMP = bool( - _cfg("ydump", 0) -) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) -LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG -G_CAP = TOP_K * M_MAX # distinct experts per call -NG_CAP = M_MAX # groups of 4 slots -XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. -THREADS = 256 + 32 * (XW - 1) -GEMV_THREADS = 256 # capacity: fixed: the CTA size -# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more -# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). -# The math is the low-M MXFP8 GEMM's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales -# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns -# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; -# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before -# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). -SHARED = bool(_cfg("shared", 0)) -# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next -# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed -# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. -LATE_TRIGGER = bool(_cfg("late_trigger", 1)) -# TAIL_TRIGGER (with LATE_TRIGGER, no all-reduce push): the CTAs that own more FC2 blocks than the rest (the first -# H / 32 mod NCTA, which end last) release the dependents right after the grid dependency, so the dependents launch -# when the other CTAs tear down, onto the SMs they leave, instead of after the last of the heavy CTAs. -TAIL_TRIGGER = bool(_cfg("tail_trigger", 1)) -NHEAVY = (H // 32) % NCTA if NCTA < H // 32 else 0 -# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's -# FC2 runs (not two dependent global loads per push on the kernel's tail). -AR_PRE = bool(_cfg("ar_pre", 1)) -# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the -# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). -XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) -NSH = 24 if SHARED else 0 -# the shared intermediate per rank: 1536 (TP2), 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate -# positions: its gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP2: 1 x (64 + -# 64), full K in one accumulator, the order of the stock tcgen05 MXFP8 GEMM; TP4: 2 x (32 + 32), the order of -# m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. -SH_I = _cfg("sh_i", 768) -SH_ROWS = SH_I // 24 -# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) -SH_SPLIT = 128 // (2 * SH_ROWS) -SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N -SH_KB = H // 32 # capacity: fixed: a kernel tile constant (gate_up K blocks) (gate_up K blocks) -SH_DKB = SH_I // 32 # down K blocks -SH_DKT = SH_I // 128 # down k-tiles -SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice -SH_STAGES_GU = SH_GT -SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles -SHX_BAR = 6 -# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the -# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction -# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. -SH_PRETILE = bool( - _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) -) # (in-chain: a win at M <= 2, not 4..8) -if SHARED: - assert SH_I in (384, 768, 1536), SH_I - assert SH_PRETILE or SH_SPLIT == 2, ( - "the un-pretiled shared stream is the 32-row (TP4) layout only" - ) -SH_STAGE_BYTES = 16384 + 512 -MMA_M = 128 -KT_BYTES = 128 -INC_K = 2 -NUM_KB = 4 -# FC1 -N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) -QS = N1 // 4 -R1 = 32 -C1 = 4 -U1 = TWO_I // R1 # 48 units per slot -K1_TILES = H // 256 -S1 = K1_TILES // C1 # 6 stages per unit # capacity: fixed: pipeline stages -NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes -NUM_BYTES_BOX1 = R1 * KT_BYTES -SF1_KA = H // 64 -SF1_RB = TWO_I // 128 -NUM_BYTES_SFA_RAW = 16 * 512 -NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 -NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes -NUM_WORDS_SFA = NUM_BYTES_SFA // 4 -NUM_BYTES_B1 = N1 * KT_BYTES -NUM_SFB1_ATOMS = S1 * NUM_KB -# FC2 -R2 = 32 -NB2 = H // R2 # 192 -KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) -# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) -KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 - - -def _nkb2(t2: int) -> int: - """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" - return NUM_KB if t2 < KT2 - 1 else KB2_LAST - - -SF2_KA = I_TP // 64 -SF2_RB = H // 128 -N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t -NUM_BYTES_B2 = N2 * KT_BYTES -NB2_CAP = _cfg("nb2_cap", 2) -UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row -H_ROW = UNITS_ROW * 16 -# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy -# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few -# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. -HB_REPL = _cfg("hb_repl", 1) -assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) -HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 -PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size -EPI_THREADS = 128 -EPI_BAR = 1 -ROUTE_BAR = 2 -GEMV_BAR = 3 -B2_BAR = 5 -TD_STRIDE = 64 # tdbg columns per CTA -TOPK_BAR = 4 -B1_AFTER_TOPK = bool( - _cfg("b1_after_topk", 0) -) # (experiment) the B1 build waits for the routing warps' top-k -TOPK2_BAR = 7 -EARLY0 = bool( - _cfg("early0", 1) -) # warp 4 streams group 0 (token 0's experts) right after its own top-k -EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) -# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token -RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 -NRW0 = ( - min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 -) # warps 4.. with a token (no idle pollers in the barrier) -NRW = NRW0 + RW_EXTRA -LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) -GEMV_SMEM = bool( - _cfg("gemv_smem", 1) -) # (XSTAGE) the router GEMV reads the staged hidden states from smem -LG_STRIDE = M_MAX * E_LOCAL + 32 -POLL_NS = _cfg("poll_ns", 150) -# Hang guard: every cross-CTA spin (the kernel needs all its CTAs resident at once) traps after this long, so a lost -# co-residency or a missing producer surfaces as a device exception instead of a silent hang. The timer is read only -# inside a spin iteration; the arithmetic is unchanged. -HANG_NS = _cfg("hang_ms", 2000) * 1000000 -B2_POLL_NS = _cfg("b2_poll_ns", 200) -# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight -# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled -# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one -# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains -# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. -B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) -IMG_B2 = KT2 * NUM_BYTES_B2 -IMG_SF = KT2 * NUM_KB * 512 -IMG_BYTES = IMG_B2 + IMG_SF -CNT_PER_CALL = 4 * (TWO_I // R1) -B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off -B2_DLY = _cfg( - "b2_dly", 0 -) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first -W_WARP = 4 # capacity: fixed: a warp role index -X_WARP = 5 # capacity: fixed: a warp role index -S_WARP = 6 # capacity: fixed: a warp role index -MMA_WARP = 7 # capacity: fixed: a warp role index -# E0_M4: at M = 4 the MMA warp routes -# token 3 (NRW0 = 4), so the early group-0 stages go to FC2 B builder warp 8 instead (idle until the routing is done; -# RW_EXTRA = 0 at M = 4). Issue order only: exact. -E0_M4 = ( - bool(_cfg("e0_m4", 1)) and NRW0 == 4 and RW_EXTRA == 0 and THREADS // 32 > 8 -) # M == 4 only: at M > 4 warp 8 routes (RW_EXTRA) -E0_WARP = ( - 8 if E0_M4 else MMA_WARP -) # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) -E0_BAR = 8 -E1_BAR = 9 -E0_OFFLOAD = EARLY0 and (NRW0 <= 3 or E0_M4) and bool(_cfg("e0_off", 1)) -# B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing -# warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the -# two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. -if B1_AFTER_TOPK and XSTAGE and EARLY0 and not E0_OFFLOAD: - raise ValueError( - "m3_moe: b1_after_topk=1 with xstage, early0 and no early-stage offload deadlocks (warp 4 waits on x_bar[1] " - "before TOPK2_BAR, the B1 build waits on TOPK2_BAR before x_bar[1]); set b1_after_topk=0 or early0=0" - ) -EVICT_FIRST = 0x12F0000000000000 -_LOG2E = 1.4426950408889634 -LD_POW2 = ( - 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 -) # tcgen05.ld .num covering M columns - -_FIXED_SMEM = ( - S1 * NUM_BYTES_B1 - + 2 * KT2 * NUM_BYTES_B2 - + 2 * KT2 * NUM_KB * 512 - + NB2_CAP * M_MAX * TOP_K * R2 * 4 - + 2 * M_MAX * C1 * R1 * 4 - + (M_MAX * (H // 16) if FUSED_ROUTER else 0) - + 6 * 1024 -) -# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no -# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB -SFA_CP = bool(_cfg("sfa_cp", 0)) -_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW -_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA -_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM -ACC1_COL = 0 -ACC2_COL = 2 * N1 -SFB1_COL = ACC2_COL + 2 * N2 -SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 -SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 -STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) -TMEM_COLS = 512 -assert 1 <= M_MAX <= 8, M_MAX -assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) -# TP2 at 8 tokens: the 98 KB of hidden states fit no 5-stage ring; the router GEMV then reads them from global (the -# kernel's non-XSTAGE path, same math) -if XSTAGE and M_MAX * H * 2 > STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA): - XSTAGE = False -assert not XSTAGE or M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( - "x staging must fit the ring" -) -# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) -_OFF_SFAR = STAGES * NUM_BYTES_A -_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE -_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA -_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 -_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 -BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 -# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales -# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] -# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange -_SB_A = 0 -_SB_RAW = None -# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 -_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 -_SH_FIX = ( - SH_GT * SH_NB * 128 - + SH_DKT * 8 * 128 - + M_MAX * SH_KB - + M_MAX * SH_DKB - + _SH_XCH_BYTES - + 2 * SH_ROWS * 8 * 4 - + 1024 -) -SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 -_SB_RAW = SH_NST * NUM_BYTES_A -_SB_SPL = _SB_RAW + SH_NST * 2048 -_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 -_SB_BH = _SB_BX + SH_GT * SH_NB * 128 -_SB_XC = _SB_BH + SH_DKT * 8 * 128 -_SB_HC = _SB_XC + M_MAX * SH_KB -_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 -_SB_GU = _SB_XCH + _SH_XCH_BYTES -_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 -# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales -SH_ACC_GU = 0 -SH_ACC_DN = SH_NB -SH_SFA = SH_ACC_DN + 16 -SH_SFB_GU = SH_SFA + 4 * 8 -SH_SFB_DN = SH_SFB_GU + 4 * SH_GT -# pre-tiled stages beyond the ring that warp 0 prefetches into L2, per lane (TP2: 72 stages) -SH_PF = max(2, (SH_STAGES - SH_NST + 31) // 32) if SHARED else 2 -# TP2 (SH_GT 48 gate_up k-tiles, SH_DKT 12) needs more than 256 columns: allocate the next power of two -SH_TMEM = 256 if SH_SFB_DN + 4 * SH_DKT <= 256 else 512 -if SHARED: - assert FUSED_ROUTER, "shared expert: fused router only" - assert SH_STAGES - SH_NST <= 32 * SH_PF, SH_STAGES - assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 4 * SH_DKT <= SH_TMEM, ( - SH_NST, - _SB_END, - BIG, - ) - - -# ============================================================================= -# Inline-PTX helpers. -# ============================================================================= -def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): - return _llvm.inline_asm( - ret, - [a.ir_value(loc=loc, ip=ip) for a in args], - ptx, - cons, - has_side_effects=side, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -# ---- all-reduce push (the MoE boundary's multicast exchange, set A): -# arp = int64 [enable, mc, flags, rank, world, maxt] -@dsl_user_op -def _ld_volatile_s32(addr, *, loc=None, ip=None): - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", - has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): - """st.global.b32 when pred != 0.""" - _llvm.inline_asm( - None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), - cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", - has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) # fmt: skip - - -@dsl_user_op -def _bf16x2_word(lo, hi, *, loc=None, ip=None): - """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" - "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" - "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" - "mov.b32 $0, w;\n\t}", "=r,f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -def _ar_push(arp, v, low, t, cl): - """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on - lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call - parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" - hi = cute.arch.shuffle_sync_bfly(v, 8) - pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) - cs = cutlass.select_(low, cl, cutlass.Int32(0)) - world = arp.load(idx=4) - parity = cutlass.Int64( - _ld_volatile_s32( - arp.load(idx=2) - + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) - ) - & cutlass.Int32(1) - ) - off = ( - parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) - + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) - + cutlass.Int64(cs // cutlass.Int32(2)) - ) - _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) - - -def _ar_parity(arp, t, cl): - """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the - next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" - return _ld_volatile_s32( - arp.load(idx=2) - + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) - ) - - -def _ar_push_pre(ard, parity32, v, low, t, cl): - """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" - hi = cute.arch.shuffle_sync_bfly(v, 8) - pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) - cs = cutlass.select_(low, cl, cutlass.Int32(0)) - mc, rk, world, maxt = ard - parity = cutlass.Int64(parity32 & cutlass.Int32(1)) - off = ( - parity * maxt * world * cutlass.Int64(H // 2) - + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) - + cutlass.Int64(cs // cutlass.Int32(2)) - ) - _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) - - -@dsl_user_op -def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): - """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" - _llvm.inline_asm( - None, - [], - f"bar.arrive {bar_id}, {threads};", - "", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def _st_relaxed_u32(addr, val, *, loc=None, ip=None): - _asm( - None, - [cutlass.Int64(addr), cutlass.Int32(val)], - "st.relaxed.gpu.global.u32 [$0], $1;", - "l,r", - True, - loc, - ip, - ) - - -@dsl_user_op -def _ld_relaxed_u32(addr, *, loc=None, ip=None): - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Int64(addr)], - "ld.relaxed.gpu.global.u32 $0, [$1];", - "=r,l", - True, - loc, - ip, - ) - ) - - -@dsl_user_op -def _st_relaxed_u64(addr, val, *, loc=None, ip=None): - _asm( - None, - [cutlass.Int64(addr), cutlass.Int64(val)], - "st.relaxed.gpu.global.u64 [$0], $1;", - "l,l", - True, - loc, - ip, - ) - - -@dsl_user_op -def _ld_relaxed_u64(addr, *, loc=None, ip=None): - return cutlass.Int64( - _asm( - _T.i64(), - [cutlass.Int64(addr)], - "ld.relaxed.gpu.global.u64 $0, [$1];", - "=l,l", - True, - loc, - ip, - ) - ) - - -@dsl_user_op -def _nanosleep(ns: int, *, loc=None, ip=None): - _llvm.inline_asm( - None, - [], - f"nanosleep.u32 {ns};", - "", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def _gtime(*, loc=None, ip=None): - return cutlass.Int64( - _llvm.inline_asm( - _T.i64(), - [], - "mov.u64 $0, %globaltimer;", - "=l", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _trap_after(elapsed, *, loc=None, ip=None): - """trap when elapsed (ns, int64) exceeds HANG_NS.""" - _llvm.inline_asm( - None, - [cutlass.Int64(elapsed).ir_value(loc=loc, ip=ip)], - f"{{ .reg .pred p; setp.gt.s64 p, $0, {HANG_NS}; @p trap; }}", - "l", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -def _hang_tick(t0): - """One spin iteration of a cross-CTA wait: the wait's start (latched on the first call; t0 = 0 before it), after - trapping once the wait has lasted longer than HANG_NS.""" - now = _gtime() - t1 = cutlass.select_(t0 == cutlass.Int64(0), now, t0) - _trap_after(now - t1) - return t1 - - -@dsl_user_op -def _smid(*, loc=None, ip=None): - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), - [], - "mov.u32 $0, %smid;", - "=r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) - - -@cute.jit -def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): - if cutlass.const_expr(TIMING): - tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) - - -@dsl_user_op -def _rcp(a, *, loc=None, ip=None): - """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" - return cutlass.Float32( - _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) - ) - - -@dsl_user_op -def _div_rn(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "div.rn.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _fma_rn(a, b, c, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], - "fma.rn.f32 $0, $1, $2, $3;", - "=f,f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _add_rn(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "add.rn.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _sub_rn(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "sub.rn.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _mul_rn(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "mul.rn.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _e4m3_bits(x, *, loc=None, ip=None): - """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(x)], - "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", - "=r,f", - False, - loc, - ip, - ) - ) & cutlass.Int32(0xFF) - - -@dsl_user_op -def _e4m3_to_f32(b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Int32(b)], - "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" - "cvt.f32.f16 $0, l;\n}", - "=f,r", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _e2m1x2(hi, lo, *, loc=None, ip=None): - """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(hi), cutlass.Float32(lo)], - "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", - "=r,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _ex2(x, *, loc=None, ip=None): - return cutlass.Float32( - _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) - ) - - -@dsl_user_op -def _bf16_round(x, *, loc=None, ip=None): - """x rounded to bf16 (nearest even), as f32.""" - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(x)], - "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", - "=f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _redux_max_u32(v, *, loc=None, ip=None): - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Int32(v)], - "redux.sync.max.u32 $0, $1, 0xffffffff;", - "=r,r", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _bf16_bits(x, *, loc=None, ip=None): - """bf16 (nearest even) bits of x in the low 16 bits.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(x)], - "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", - "=r,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _ue8m0_rp(x, *, loc=None, ip=None): - """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(x)], - "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", - "=r,f", - False, - loc, - ip, - ) - ) & cutlass.Int32(0xFF) - - -@dsl_user_op -def _e4m3x2(hi, lo, *, loc=None, ip=None): - """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(hi), cutlass.Float32(lo)], - "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", - "=r,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _min_xorsign_abs(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "min.xorsign.abs.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _div_full(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "div.full.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _ex2_full(x, *, loc=None, ip=None): - """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" - return cutlass.Float32( - _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) - ) - - -@dsl_user_op -def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): - """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" - _llvm.inline_asm( - None, - [ - cutlass.Int64(addr).ir_value(loc=loc, ip=ip), - cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), - ], - "cp.async.bulk.prefetch.L2.global [$0], $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -def _mx_osc(amax): - """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" - code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) - cf = cutlass.Int32( - cutlass.select_( - code == cutlass.Int32(255), - cutlass.Int32(0x7FFFFFFF), - cutlass.select_( - code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) - ), - ) - ).bitcast(cutlass.Float32) - return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) - - -def _mx_words(v, osc): - """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" - words = [] - for t in range(len(v) // 4): - lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) - hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) - words.append(lo | (hi << cutlass.Int32(16))) - return words - - -def _twiddle(score): - """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" - bits = ( - score.bitcast(cutlass.Int32) - if hasattr(score, "bitcast") - else cutlass.Float32(score).bitcast(cutlass.Int32) - ) - return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) - - -@cute.jit -def _wait(bar: cutlass.Array, parity): - while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): - pass - cute.arch.sync_warp() - - -def _tptr(base_row, col, dtype=cutlass.Int32): - return cutlass.inttoptr((base_row << 16) | col, 6, dtype) - - -@dsl_user_op -def _match_any(value, *, loc=None, ip=None): - """The mask of the warp's lanes holding the same 32-bit value.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Int32(value)], - "match.any.sync.b32 $0, $1, 0xffffffff;", - "=r,r", - False, - loc, - ip, - ) - ) - - -@cute.jit -def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): - """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" - if cutlass.const_expr(HANG_DBG): - n = cutlass.Int32(0) - while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( - n < cutlass.Int32(1 << 22) - ): - n = n + cutlass.Int32(1) - if n >= cutlass.Int32(1 << 22): - bx, _, _ = cute.arch.block_idx() - w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tdbg.store( - cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w - ) - else: - while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): - pass - cute.arch.sync_warp() - - -@cute.jit -def _adv(stage, phase): - """The next ring position's (stage, phase).""" - stage = stage + cutlass.Int32(1) - wrap = stage == cutlass.Int32(STAGES) - return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( - wrap, phase ^ cutlass.Int32(1), phase - ) - - -@cute.jit -def _seg(i0, i1, k): - """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: - the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" - uf = i0 // cutlass.Int32(S1) - ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - unit = uf + k - cutlass.Int32(1) - unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) - unit = cutlass.select_(ul == uf, uf, unit) - j0 = i0 - unit * cutlass.Int32(S1) - j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) - j1 = i1 - unit * cutlass.Int32(S1) - j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) - return unit, j0, j1 - - -@cute.jit -def _nsegs(i0, i1): - n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) - return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) - - -@cute.jit -def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): - """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile - 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" - q = lane // cutlass.Int32(QS) - tl = lane - q * cutlass.Int32(QS) - sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) - sfv = [] - for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): - j = a // NUM_KB - kb = a % NUM_KB - wsf = src.load( - idx=cutlass.select_( - sf_ok, - tl * cutlass.Int32(H // 64) - + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) - + cutlass.Int32(kb), - cutlass.Int32(0), - ) - ) - sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) - for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): - prims.tcgen05_st( - "32x32b", - _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), - sfv[a], - ) - prims.tcgen05_wait(prims.Tcgen05Wait.STORE) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - prims.mbarrier_arrive(b1_ready) - - -@cute.jit -def _b1_addr(t, kb16): - """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" - kt = kb16 >> cutlass.Int32(4) - j = kt >> cutlass.Int32(2) - n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t - cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) - return ( - j * cutlass.Int32(NUM_BYTES_B1) - + n * cutlass.Int32(KT_BYTES) - + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) - + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) - ) // cutlass.Int32(4) - - -@cute.jit -def _build_b1( - tidx, - warp, - lane, - base_row, - base_col, - x32: cutlass.Array, - xsf32: cutlass.Array, - sB1: cutlass.Array, - b1_ready: cutlass.Array, -): - """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" - sB1_32 = cutlass.Array( - sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 - ) - for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): - b = tidx + cutlass.Int32(r * EPI_THREADS) - if b < cutlass.Int32(M_MAX * (H // 16)): - t = b // cutlass.Int32(H // 16) - kb = b - t * cutlass.Int32(H // 16) - v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) - sB1_32.store( - (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 - ) - _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) - - -@cute.jit -def _build_b1_fused( - tidx, - warp, - lane, - base_row, - base_col, - xb32: cutlass.Array, - sfsc: cutlass.Array, - sB1: cutlass.Array, - s_xsf: cutlass.Array, - b1_ready: cutlass.Array, -): - """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" - sB1_32 = cutlass.Array( - sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 - ) - s_xsf8 = cutlass.Array( - s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 - ) - sc = cutlass.Float32(sfsc.load(idx=0)) - rsc = _rcp(sc) - rsix = _rcp(cutlass.Float32(6.0)) - NBLK = M_MAX * (H // 16) // EPI_THREADS - QB = ( - 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) - ) # blocks whose loads are in flight together - # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) - for rg in cutlass.range(NBLK // QB, unroll=1): - wall = [] - for rr in cutlass.range_constexpr(QB): - bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) - for h in cutlass.range_constexpr(2): - wall.append( - xb32.load( - idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), - vector_size=4, - alignment=16, - ) - ) - for rr in cutlass.range_constexpr(QB): - b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) - t = b // cutlass.Int32(H // 16) - kb = b - t * cutlass.Int32(H // 16) - w = [] - for h in cutlass.range_constexpr(2): - v4 = wall[rr * 2 + h] - for z in cutlass.range_constexpr(4): - w.append(cutlass.Int32(v4[z])) - xs = [] - for z in cutlass.range_constexpr(8): - xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) - xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) - vmax = cute.math.abs(xs[0]) - for z in cutlass.range_constexpr(1, 16): - vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) - sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) - osc = cutlass.select_( - vmax != cutlass.Float32(0.0), - _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), - cutlass.Float32(0.0), - ) - by = [ - _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) - for z in range(8) - ] - lo = ( - by[0] - | (by[1] << cutlass.Int32(8)) - | (by[2] << cutlass.Int32(16)) - | (by[3] << cutlass.Int32(24)) - ) - hi = ( - by[4] - | (by[5] << cutlass.Int32(8)) - | (by[6] << cutlass.Int32(16)) - | (by[7] << cutlass.Int32(24)) - ) - sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) - s_xsf8.store(cutlass.Int8(sfb), idx=b) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) - _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) - - -@cute.jit -def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): - """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, - gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 - and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" - neg_max = cutlass.Float32(-3.40282346638528859812e38) - gc = cutlass.select_(gate_raw > lim, lim, gate_raw) - gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) - uc = cutlass.select_(xv > lim, lim, xv) - uc = cutlass.select_(xv < -lim, -lim, uc) - c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) - sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) - gs = _mul_rn(gc, sgm) - h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) - h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) - amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) - for off in (16, 8, 4, 2, 1): - amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) - sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) - sff = _e4m3_to_f32(sfb) - osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) - return _mul_rn(_mul_rn(h, osc), c), sfb - - -@cute.jit -def _plan_range(s_plan: cutlass.Array, p): - w = s_plan.load(idx=p) - return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) - - -@cute.jit -def _sh_issue( - lane, - sid, - g, - s_big: cutlass.Array, - sh_full: cutlass.Array, - tma_gu, - tma_dn, - gusf: cutlass.Array, - dnsf: cutlass.Array, -): - """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both - slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its - raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile - (g - SH_GT) % 6.""" - st = g % cutlass.Int32(SH_NST) - bar = sh_full.subview(st) - if cutlass.const_expr(SH_PRETILE): - if lane == 0: - src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) - prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) - prims.cp_async_bulk_shared_cluster_global( - s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 - ) - prims.cp_async_bulk_shared_cluster_global( - s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), - gusf.subview(src + cutlass.Int32(16384)), - bar, - 512, - ) - cute.arch.sync_warp() - else: - if g < cutlass.Int32(SH_GT): - if lane == 0: - prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) - cute.arch.sync_warp() - if lane < 4: - q = lane >> cutlass.Int32(1) - u = lane & cutlass.Int32(1) - row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) - kt = q * cutlass.Int32(SH_GT) + g - prims.cp_async_bulk_tensor_shared_cta_global( - s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), - tma_gu, - (kt * cutlass.Int32(128), row0), - bar, - l2_cache_hint=EVICT_FIRST, - ) - prims.cp_async_bulk_shared_cluster_global( - s_big.subview( - cutlass.Int32(_SB_RAW) - + st * cutlass.Int32(2048) - + lane * cutlass.Int32(512) - ), - gusf.subview( - ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) - * cutlass.Int32(512) - ), - bar, - 512, - ) - else: - gd = g - cutlass.Int32(SH_GT) - jt = gd // cutlass.Int32(SH_DKT) - kt = gd - jt * cutlass.Int32(SH_DKT) - row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) - if lane == 0: - prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) - prims.cp_async_bulk_tensor_shared_cta_global( - s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), - tma_dn, - (kt * cutlass.Int32(128), row0), - bar, - l2_cache_hint=EVICT_FIRST, - ) - prims.cp_async_bulk_shared_cluster_global( - s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), - dnsf.subview( - ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) - * cutlass.Int32(512) - ), - bar, - 512, - ) - cute.arch.sync_warp() - - -@dsl_user_op -def _red_release_add(addr, val, *, loc=None, ip=None): - _asm( - None, - [cutlass.Int64(addr), cutlass.Int32(val)], - "red.release.gpu.global.add.u32 [$0], $1;", - "l,r", - True, - loc, - ip, - ) - - -@dsl_user_op -def _ld_acquire_u32(addr, *, loc=None, ip=None): - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Int64(addr)], - "ld.acquire.gpu.global.u32 $0, [$1];", - "=r,l", - True, - loc, - ip, - ) - ) - - -@dsl_user_op -def _fence_proxy_async_global(*, loc=None, ip=None): - _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) - - -@cute.jit -def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): - """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - sd0 = prims.Tcgen05SmemDesc.build( - sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - for kb in cutlass.range_constexpr(NUM_KB): - prims.tcgen05_cp( - s2t_shape, - _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), - sd0 + (kb * 512 >> 4), - multicast=s2t_mc, - ) - - -@cute.jit -def _sf1_cp( - lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array -): - """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) - word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their - arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" - src0 = ( - (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) - + lane * cutlass.Int32(4) - + (c & cutlass.Int32(3)) - ) - dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) - for q in cutlass.range_constexpr(C1): - for kb in cutlass.range_constexpr(NUM_KB): - prims.cp_async_shared_global( - sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), - s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), - 4, - "ca", - ) - prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) - cute.arch.sync_warp() - - -@cute.jit -def _sf2_cp( - lane, - stage, - g, - blk, - t2: cutlass.Constexpr, - n_here, - s_el: cutlass.Array, - sSFA: cutlass.Array, - s2w: cutlass.Array, - ab_full: cutlass.Array, -): - """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot - jj's expert, word [r][blk % 4].""" - dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) - for jj in cutlass.range_constexpr(4): - if cutlass.Int32(jj) < n_here: - e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) - src0 = ( - ( - (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) - + cutlass.Int32(4 * t2) - ) - * cutlass.Int32(128) - + lane * cutlass.Int32(4) - + (blk & cutlass.Int32(3)) - ) - # (a partial last k-tile has no atoms past I_TP) - for kb in cutlass.range_constexpr(_nkb2(t2)): - prims.cp_async_shared_global( - sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), - s2w.subview(src0 + cutlass.Int32(kb * 128)), - 4, - "ca", - ) - prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) - cute.arch.sync_warp() - - -@cute.jit -def _shared_cta( - tidx, - warp, - lane, - sid, - s_big: cutlass.Array, - sh_bar: cutlass.Array, - epochs: cutlass.Array, - tdbg: cutlass.Array, - bx, - xb32: cutlass.Array, - tma_gu, - tma_dn, - gusf: cutlass.Array, - dnsf: cutlass.Array, - hx: cutlass.Array, - dnx: cutlass.Array, - arp: cutlass.Array, - tmem_ptr_i32: cutlass.Array, - gu8: cutlass.Array, - dn8: cutlass.Array, -): - """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one - 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles - of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" - big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) - bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) - big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) - sh_full = sh_bar # [0, SH_NST): stage full - sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) - b_gu_rdy = sh_bar.subview(16) - b_dn_rdy = sh_bar.subview(17) - acc_gu = sh_bar.subview(18) - acc_dn = sh_bar.subview(19) - if warp == 0: - prims.prefetch_tensormap(tma_gu) - prims.prefetch_tensormap(tma_dn) - if lane < SH_NST: - prims.mbarrier_init(sh_full.subview(lane), 1) - prims.mbarrier_init( - sh_empty.subview(lane), 33 - ) # the MMAs' commit + each splicer lane's reads done - if lane == 0: - prims.mbarrier_init(b_gu_rdy, 1) - prims.mbarrier_init(b_dn_rdy, 1) - prims.mbarrier_init(acc_gu, 1) - prims.mbarrier_init(acc_dn, 1) - if warp == 1: - prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) - prims.tcgen05_relinquish_alloc_permit() - prims.fence_mbarrier_init() - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - prims.barrier_cta_sync(0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - tmem_raw = tmem_ptr_i32.load() - base_col = tmem_raw & 0xFFFF - base_row = tmem_raw >> 16 - if warp == 0: - for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency - _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) - # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) - if cutlass.const_expr(SH_PRETILE): - # this CTA's stages beyond the ring, 4 per lane - for r in cutlass.range_constexpr(SH_PF): - pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(SH_PF) + cutlass.Int32(r) - if pg < cutlass.Int32(SH_STAGES): - _l2_prefetch( - gusf.subview( - (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) - ) - .data_ptr() - .toint(), - SH_STAGE_BYTES, - ) - else: - for r in cutlass.range_constexpr(2): - prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane - _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) - if lane == 0: - _l2_prefetch( - dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I - ) - cute.arch.griddepcontrol_wait() - if cutlass.const_expr(not LATE_TRIGGER): - if arp.load(idx=0) == cutlass.Int64(0): - cute.arch.griddepcontrol_launch_dependents() - if tidx == 0: - _stamp(tdbg, bx, 1) - ep = epochs.load(idx=bx) - tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) - NQ = THREADS - 64 # activation quant threads (warps 2 ..) - if warp == 0: - for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): - _wait( - sh_empty.subview(g % cutlass.Int32(SH_NST)), - ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), - ) - _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) - elif warp == 1: - # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down - idesc_gu = prims.Tcgen05MxInstrDesc.build( - a_dtype=cutlass.Float8E4M3FN, - b_dtype=cutlass.Float8E4M3FN, - scale_format=1, - n_dim=SH_NB, - m_dim=128, - ) - idesc_dn = prims.Tcgen05MxInstrDesc.build( - a_dtype=cutlass.Float8E4M3FN, - b_dtype=cutlass.Float8E4M3FN, - scale_format=1, - n_dim=8, - m_dim=128, - ) - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - desc_a0 = prims.Tcgen05SmemDesc.build( - s_big.subview(0), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - desc_bx = prims.Tcgen05SmemDesc.build( - s_big.subview(_SB_BX), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - desc_bh = prims.Tcgen05SmemDesc.build( - s_big.subview(_SB_BH), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - tcol = (base_row << 16) | base_col - sq = sid & cutlass.Int32(3) - _wait(b_gu_rdy, 0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - for g in cutlass.range(SH_STAGES, unroll=1): - st = g % cutlass.Int32(SH_NST) - if g == cutlass.Int32(SH_GT): - _wait(b_dn_rdy, 0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) - raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) - if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): - # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 - w4 = [ - big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) - for qq in range(4) - ] - big32.store( - ( - cutlass.Int32(w4[0]), - cutlass.Int32(w4[1]), - cutlass.Int32(w4[2]), - cutlass.Int32(w4[3]), - ), - idx=cutlass.Int32(_SB_SPL // 4) - + st * cutlass.Int32(128) - + lane * cutlass.Int32(4), - alignment=16, - ) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - cute.arch.sync_warp() - prims.mbarrier_arrive( - sh_empty.subview(st) - ) # this lane's generic reads of the slot are done - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - if prims.elect_sync(): - if g < cutlass.Int32(SH_GT): - dsf = prims.Tcgen05SmemDesc.build( - s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - prims.tcgen05_cp( - s2t_shape, - _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), - dsf, - multicast=s2t_mc, - ) - da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) - db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) - for kb in cutlass.range_constexpr(4): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF8F6F4, - prims.CTAGroup.CTA_1, - cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), - da + 2 * kb, - db + 2 * kb, - idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), - (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), - cutlass.inttoptr( - tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 - ), - cutlass.inttoptr( - tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 - ), - ) - prims.tcgen05_commit(sh_empty.subview(st)) - if g == cutlass.Int32(SH_GT - 1): - prims.tcgen05_commit(acc_gu) - else: - gd = g - cutlass.Int32(SH_GT) - jt = gd // cutlass.Int32(SH_DKT) - kt = gd - jt * cutlass.Int32(SH_DKT) - dsf = prims.Tcgen05SmemDesc.build( - s_big.subview( - cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) - + st * cutlass.Int32(512 if SH_PRETILE else 2048) - ), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - prims.tcgen05_cp( - s2t_shape, - _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), - dsf, - multicast=s2t_mc, - ) - da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) - db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) - for kb in cutlass.range_constexpr(4): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF8F6F4, - prims.CTAGroup.CTA_1, - cutlass.inttoptr( - tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 - ), - da + 2 * kb, - db + 2 * kb, - idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), - (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), - cutlass.inttoptr( - tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 - ), - cutlass.inttoptr( - tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 - ), - ) - prims.tcgen05_commit(sh_empty.subview(st)) - if g == cutlass.Int32(SH_STAGES - 1): - prims.tcgen05_commit(acc_dn) - cute.arch.sync_warp() - if warp >= 2: - # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = - # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes - qt = tidx - cutlass.Int32(64) - for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): - q_ = qt + cutlass.Int32(r * NQ) - if q_ < cutlass.Int32(M_MAX * SH_KB): - t = q_ // cutlass.Int32(SH_KB) - b = q_ - t * cutlass.Int32(SH_KB) - wv = [ - xb32.load( - idx=t * cutlass.Int32(H // 2) - + b * cutlass.Int32(16) - + cutlass.Int32(4 * z), - vector_size=4, - alignment=16, - ) - for z in range(4) - ] - v = [] - for z in cutlass.range_constexpr(4): - for zz in cutlass.range_constexpr(4): - w = cutlass.Int32(wv[z][zz]) - v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) - v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) - amax = cutlass.Float32(0.0) - for z in cutlass.range_constexpr(32): - amax = cute.arch.fmax(amax, cute.math.abs(v[z])) - code, osc = _mx_osc(amax) - words = _mx_words(v, osc) - kt = b >> cutlass.Int32(2) - sl = kt // cutlass.Int32(SH_GT) - it = kt - sl * cutlass.Int32(SH_GT) - n = sl * cutlass.Int32(8) + t - c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) - rw = ( - cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) - ) >> cutlass.Int32(2) - big32.store( - (words[0], words[1], words[2], words[3]), - idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), - alignment=16, - ) - big32.store( - (words[4], words[5], words[6], words[7]), - idx=rw - + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), - alignment=16, - ) - big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) - if warp >= 4: - if warp < 8: - qd = warp - cutlass.Int32(4) - # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i - n = lane - t = n & cutlass.Int32(7) - ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) - for i in cutlass.range_constexpr(SH_GT): - kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) - wsf = big32.load( - idx=( - cutlass.Int32(_SB_XC) - + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) - + kt * cutlass.Int32(4) - ) - >> cutlass.Int32(2) - ) - prims.tcgen05_st( - "32x32b", - _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), - cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), - ) - prims.tcgen05_wait(prims.Tcgen05Wait.STORE) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - et = tidx - cutlass.Int32(128) - if et == 0: - prims.mbarrier_arrive(b_gu_rdy) - _stamp(tdbg, bx, 54) - # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t - _wait(acc_gu, 0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - if cutlass.const_expr(SH_SPLIT == 1): - # TP2: one K slice; quadrant qd = rows 32 qd .. (gate positions 0..63, then up 0..63), no partials: - # the owner sum is (0 + acc), then bf16, at _SB_GU word (u SH_ROWS + p) 8 + t = (32 qd + l) 8 + t - av = prims.tcgen05_ld( - "32x32b", - _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), - num=8, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - for tt in cutlass.range_constexpr(8): - tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) - bigf.store( - _bf16_round(tot), - idx=cutlass.Int32(_SB_GU // 4) - + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - elif cutlass.const_expr(SH_SPLIT == 2): - av = prims.tcgen05_ld( - "32x32b", - _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), - num=16, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - if qd >= cutlass.Int32(2): - for tt in cutlass.range_constexpr(8): - bigf.store( - cutlass.Float32(av[8 + tt]), - idx=cutlass.Int32(_SB_XCH // 4) - + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) - * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - if qd < cutlass.Int32(2): - # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 - for tt in cutlass.range_constexpr(8): - tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) - tot = _add_rn( - tot, - bigf.load( - idx=cutlass.Int32(_SB_XCH // 4) - + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) - + cutlass.Int32(tt) - ), - ) - bigf.store( - _bf16_round(tot), - idx=cutlass.Int32(_SB_GU // 4) - + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - else: - # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are - # accumulator columns 8 qd + t - av = prims.tcgen05_ld( - "32x32b", - _tptr( - base_row + qd * cutlass.Int32(32), - base_col + SH_ACC_GU + qd * cutlass.Int32(8), - cutlass.Float32, - ), - num=8, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - if qd >= cutlass.Int32(1): - for tt in cutlass.range_constexpr(8): - bigf.store( - cutlass.Float32(av[tt]), - idx=cutlass.Int32(_SB_XCH // 4) - + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) - * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - if qd == cutlass.Int32(0): - # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 - for tt in cutlass.range_constexpr(8): - tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) - for s_ in cutlass.range_constexpr(1, SH_SPLIT): - tot = _add_rn( - tot, - bigf.load( - idx=cutlass.Int32(_SB_XCH // 4) - + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) - + cutlass.Int32(tt) - ), - ) - bigf.store( - _bf16_round(tot), - idx=cutlass.Int32(_SB_GU // 4) - + lane * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag - for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): - pq = et + cutlass.Int32(r * 128) - if pq < cutlass.Int32(SH_ROWS * M_MAX): - p = pq // cutlass.Int32(M_MAX) - tt = pq - p * cutlass.Int32(M_MAX) - ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) - up = bigf.load( - idx=cutlass.Int32(_SB_GU // 4) - + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) - + tt - ) - q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) - q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) - q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) - q_e2 = _ex2_full( - _mul_rn( - _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), - cutlass.Float32(1.4426950216293334961), - ) - ) - q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) - q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) - _st_relaxed_u32( - hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) - .data_ptr() - .toint(), - (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), - ) - if et == 0: - _stamp(tdbg, bx, 50) - # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes - for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): - q_ = et + cutlass.Int32(rr * 128) - if q_ < cutlass.Int32(M_MAX * SH_DKB): - t = q_ // cutlass.Int32(SH_DKB) - kb = q_ - t * cutlass.Int32(SH_DKB) - base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) - hv = [cutlass.Int32(0) for _ in range(32)] - pend = cutlass.Int32(1) - t0h = cutlass.Int64(0) - while pend != cutlass.Int32(0): - t0h = _hang_tick(t0h) - for q4 in cutlass.range_constexpr(8): - v4 = prims.load_ext( - hx.subview(base + cutlass.Int32(4 * q4)), - dtype=cutlass.Int32, - count=4, - order="relaxed", - scope="gpu", - ) - for z in cutlass.range_constexpr(4): - hv[4 * q4 + z] = cutlass.Int32(v4[z]) - okw = cutlass.Boolean(True) - for z in cutlass.range_constexpr(32): - okw = okw & ( - ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 - ) - pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) - vals = [ - (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) - ] - amax = cutlass.Float32(0.0) - for z in cutlass.range_constexpr(32): - amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) - code, osc = _mx_osc(amax) - words = _mx_words(vals, osc) - kt = kb >> cutlass.Int32(2) - c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) - rw = ( - cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) - ) >> cutlass.Int32(2) - big32.store( - (words[0], words[1], words[2], words[3]), - idx=rw + ((c0 ^ t) << cutlass.Int32(2)), - alignment=16, - ) - big32.store( - (words[4], words[5], words[6], words[7]), - idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), - alignment=16, - ) - big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - okh = lane < cutlass.Int32(M_MAX) - for kt in cutlass.range_constexpr(SH_DKT): - wsf = big32.load( - idx=( - cutlass.Int32(_SB_HC) - + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) - + cutlass.Int32(4 * kt) - ) - >> cutlass.Int32(2) - ) - prims.tcgen05_st( - "32x32b", - _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), - cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), - ) - prims.tcgen05_wait(prims.Tcgen05Wait.STORE) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - if et == 0: - prims.mbarrier_arrive(b_dn_rdy) - _stamp(tdbg, bx, 51) - # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l - _wait(acc_dn, 0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - for jt in cutlass.range_constexpr(2): - dv = prims.tcgen05_ld( - "32x32b", - _tptr( - base_row + qd * cutlass.Int32(32), - base_col + SH_ACC_DN + jt * 8, - cutlass.Float32, - ), - num=8, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - rowo = ( - sid * cutlass.Int32(256) - + cutlass.Int32(128 * jt) - + qd * cutlass.Int32(32) - + lane - ) - for tt in cutlass.range_constexpr(M_MAX): - _st_relaxed_u32( - dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), - (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), - ) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - if et == 0: - _stamp(tdbg, bx, 52) - cute.arch.sync_warp() - prims.barrier_cta_sync(0) - if warp == 1: - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) - if tidx == 0: - epochs.store(ep + cutlass.Int32(1), idx=bx) - - -@cute.kernel -def m3_moe_kernel( - tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], - tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], - tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines - tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines - tma_gu_desc: cutlass.GridConstant[ - cuda.TensorMap - ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows - tma_dn_desc: cutlass.GridConstant[ - cuda.TensorMap - ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows - logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) - bias: cutlass.Array, # f32 [128] - x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) - xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) - escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha - plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 - partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums - hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records - epochs: cutlass.Array, # int32 [NCTA] - out: cutlass.Array, # bf16 [M, 6144] - route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) - tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) - xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] - rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] - sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale - lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits - guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] - gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales - dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] - dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales - hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written - dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written - arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off - s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale - s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale - img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) - b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) -): - tidx, _, _ = cute.arch.thread_idx() - bx, _, _ = cute.arch.block_idx() - warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane = tidx % 32 - if tidx == 0: - _stamp(tdbg, bx, 0) - if cutlass.const_expr(TIMING): - if tidx == 0: - tdbg.store( - cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) - ) - - s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) - sA = cutlass.Array( - s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 - ) - sSFAraw = cutlass.Array( - s_big.data_ptr(_OFF_SFAR), - shape=(max(1, _RAW_STAGE // 4) * STAGES,), - dtype=cutlass.Int32, - alignment=1024, - ) - sSFA = cutlass.Array( - s_big.data_ptr(_OFF_SFA), - shape=(NUM_WORDS_SFA * STAGES,), - dtype=cutlass.Int32, - alignment=1024, - ) - sB1 = cutlass.Array( - s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 - ) - sB2 = cutlass.Array( - s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 - ) - sSFB2 = cutlass.Array( - s_big.data_ptr(_OFF_SFB2), - shape=(2 * KT2 * NUM_KB * 128,), - dtype=cutlass.Int32, - alignment=1024, - ) - ys = cutlass.Array( - cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_part = cutlass.Array( - cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) - s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) - s_xsf = cutlass.Array( - cutlass.Int32, - M_MAX * (H // 64) if FUSED_ROUTER else 4, - space=cutlass.AddressSpace.smem, - alignment=16, - ) - ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) - ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) - sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) - acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice - # while warp 5 still works on an earlier group, and its parity wait would then miss them) - fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) - b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) - x_bar = cutlass.Array( - cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 - ) # x landed, x consumed - img_bar = cutlass.Array( - cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 - ) # (B2IMG) image landed - tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) - s_el = cutlass.Array( - cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 - ) # slots, then G - s_kpos = cutlass.Array( - cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_w = cutlass.Array( - cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_ids = cutlass.Array( - cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) - # this CTA's plan rows for every G (static data: loaded before the grid dependency) - s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) - s_plan_all = cutlass.Array( - cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 - ) - s_esc = cutlass.Array( - cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 - ) # slot s at 8 s - - is_routed = cutlass.Boolean(True) - if cutlass.const_expr(SHARED): - sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) - is_routed = bx < cutlass.Int32(NCTA) - if bx >= cutlass.Int32(NCTA): - _shared_cta( - tidx, - warp, - lane, - bx - cutlass.Int32(NCTA), - s_big, - sh_full, - epochs, - tdbg, - bx, - xb32, - tma_gu_desc.get_ptr(), - tma_dn_desc.get_ptr(), - gusf, - dnsf, - hx, - dnx, - arp, - tmem_ptr_i32, - guw, - dnw, - ) - if is_routed: - # ---- before the grid dependency - if warp == 0: - if tidx < STAGES: - prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) - prims.mbarrier_init(ab_empty.subview(tidx), 1) - prims.mbarrier_init(sfa_ready.subview(tidx), 1) - if tidx < 2: - prims.mbarrier_init(acc_full.subview(tidx), 1) - prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) - prims.mbarrier_init(acc2_full.subview(tidx), 1) - prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) - prims.mbarrier_init(b2_ready.subview(tidx), 1) - prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) - if cutlass.const_expr(B2IMG): - prims.mbarrier_init(img_bar.subview(tidx), 1) - if tidx < 4: - prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) - if tidx == 0: - prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) - if cutlass.const_expr(XSTAGE): - prims.mbarrier_init(x_bar.subview(0), 1) - prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) - if warp == MMA_WARP: - prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) - prims.tcgen05_relinquish_alloc_permit() - if warp == W_WARP: - prims.prefetch_tensormap(tma_a1_desc.get_ptr()) - prims.prefetch_tensormap(tma_a2_desc.get_ptr()) - prims.prefetch_tensormap(tma_s1_desc.get_ptr()) - prims.prefetch_tensormap(tma_s2_desc.get_ptr()) - for i in cutlass.range_constexpr(E_LOCAL // 32): - s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) - for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): - q = lane + cutlass.Int32(32 * i) - if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): - gq = q // cutlass.Int32(NG_CAP + 1) - s_plan_all.store( - plan.load( - idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) - + q - - gq * cutlass.Int32(NG_CAP + 1) - ), - idx=q, - ) - rwv = [] - if cutlass.const_expr(FUSED_ROUTER): - tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) - for rg in cutlass.range_constexpr(NGEMV): - eg = bx + cutlass.Int32(rg * NCTA) - rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) - rwv.append( - [ - rw.load( - idx=rbx * cutlass.Int32(H) - + tg * cutlass.Int32(4) - + cutlass.Int32(1024 * j), - vector_size=4, - alignment=16, - ) - for j in range(H // 1024) - ] - ) - if cutlass.const_expr(TIMING): - if tidx == 0: - s_el.store(cutlass.Int32(-7), idx=G_CAP) - prims.fence_mbarrier_init() - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - prims.barrier_cta_sync(0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - tmem_raw = tmem_ptr_i32.load() - base_col = tmem_raw & 0xFFFF - base_row = tmem_raw >> 16 - - cute.arch.griddepcontrol_wait() - if cutlass.const_expr(XSTAGE): - if warp == 0: - if lane == 0: - prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) - cute.arch.sync_warp() - if lane < M_MAX: - prims.cp_async_bulk_shared_cluster_global( - s_big.subview(lane * cutlass.Int32(H * 2)), - xb32.subview(lane * cutlass.Int32(H // 2)), - x_bar.subview(0), - H * 2, - ) - ar_on = arp.load(idx=0) != cutlass.Int64(0) - ar_off = arp.load(idx=0) == cutlass.Int64(0) - if cutlass.const_expr(not LATE_TRIGGER): - if ar_off: - cute.arch.griddepcontrol_launch_dependents() - if cutlass.const_expr(LATE_TRIGGER and TAIL_TRIGGER and NHEAVY > 0): - if ar_off & (bx < cutlass.Int32(NHEAVY)): - cute.arch.griddepcontrol_launch_dependents() - if tidx == 0: - _stamp(tdbg, bx, 1) - - ep = epochs.load(idx=bx) - tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) - tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) - ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) - if cutlass.const_expr(FUSED_ROUTER): - # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly - for rg in cutlass.range_constexpr(NGEMV): - eg = bx + cutlass.Int32(rg * NCTA) - if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): - accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] - xsrc = xb32 - if cutlass.const_expr(GEMV_SMEM and XSTAGE): - _wait(x_bar.subview(0), 0) - xsrc = cutlass.Array( - s_big.data_ptr(0), - shape=(M_MAX * H // 2,), - dtype=cutlass.Int32, - alignment=16, - ) - for j in cutlass.range_constexpr(H // 1024): - for t in cutlass.range_constexpr(M_MAX): - xw = xsrc.load( - idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), - vector_size=2, - alignment=8, - ) - x0 = cutlass.Int32(xw[0]) - x1 = cutlass.Int32(xw[1]) - a = accs[t] - a = _fma_rn( - (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), - cutlass.Float32(rwv[rg][j][0]), - a, - ) - a = _fma_rn( - (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), - cutlass.Float32(rwv[rg][j][1]), - a, - ) - a = _fma_rn( - (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), - cutlass.Float32(rwv[rg][j][2]), - a, - ) - a = _fma_rn( - (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), - cutlass.Float32(rwv[rg][j][3]), - a, - ) - accs[t] = a - for t in cutlass.range_constexpr(M_MAX): - v = accs[t] - for off in (16, 8, 4, 2, 1): - v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) - if lane == 0: - s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) - if tidx < M_MAX: - sacc = cutlass.Float32(0.0) - for w in cutlass.range_constexpr(8): - sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) - s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) - if tidx < LG_REPL: - for t in cutlass.range_constexpr(M_MAX): - lw = ltag | ( - cutlass.Int64( - cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( - cutlass.Int32 - ) - ) - & cutlass.Int64(0xFFFFFFFF) - ) - _st_relaxed_u64( - lgbuf.subview( - tidx * cutlass.Int32(LG_STRIDE) - + cutlass.Int32(t * E_LOCAL) - + eg - ) - .data_ptr() - .toint(), - lw, - ) - if cutlass.const_expr(NGEMV > 1): - cute.arch.sync_warp() - cute.arch.barrier( - barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS - ) # s_red reused next - if tidx == 0: - _stamp(tdbg, bx, 21) - - # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile - if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( - (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) - ): - rw_i = cutlass.select_( - warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) - ) - # ILP: poll every token's logits first, then run the independent top-k chains back to back - lgv_all = [] - t_all = [] - for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): - t = rw_i + cutlass.Int32(NRW * tr) - if cutlass.const_expr(M_MAX % NRW != 0): - t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) - t_all.append(t) - lgv = [] - if cutlass.const_expr(FUSED_ROUTER): - want = ltag >> cutlass.Int64(32) - rbase = ( - (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) - + t * cutlass.Int32(E_LOCAL) - + lane - ) - done = cutlass.Int32(0) - res = [cutlass.Int64(0) for _ in range(4)] - t0h = cutlass.Int64(0) - while done == cutlass.Int32(0): - t0h = _hang_tick(t0h) - ok = cutlass.Boolean(True) - for z in cutlass.range_constexpr(4): - res[z] = _ld_relaxed_u64( - lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() - ) - ok = ok & ((res[z] >> cutlass.Int64(32)) == want) - if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( - -1 - ): - done = cutlass.Int32(1) - else: - _nanosleep(POLL_NS) - for z in cutlass.range_constexpr(4): - lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) - if (lane == 0) & (warp == cutlass.Int32(W_WARP)): - _stamp(tdbg, bx, 22) - else: - for i in cutlass.range_constexpr(4): - lgv.append( - cutlass.Float32( - logits.load( - idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane - ) - ) - ) - lgv_all.append(lgv) - if (lane == 0) & (warp == cutlass.Int32(W_WARP)): - _stamp(tdbg, bx, 61) - res_all = [] - for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): - t = t_all[tr] - lgv = lgv_all[tr] - own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) - # top-4 (keys sigmoid + bias, ties to the lower id) - keys = [] - for i in cutlass.range_constexpr(4): - sg = _fma_rn( - cutlass.Float32(0.5), - cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), - cutlass.Float32(0.5), - ) - keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) - hi = [_twiddle(keys[i]) for i in range(4)] - taken = [cutlass.Boolean(False) for _ in range(4)] - ids_t = [] - kt_t = [] - for k in cutlass.range_constexpr(TOP_K): - bh = cutlass.Int32(0) - bl = cutlass.Int32(0) - for i in cutlass.range_constexpr(4): - lo_i = cutlass.Int32(65535 - i * 32) - lane - hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) - lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) - better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( - (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) - ) - bh = cutlass.select_(better, hv, bh) - bl = cutlass.select_(better, lv, bl) - mh = _redux_max_u32(bh) - ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) - e_k = cutlass.Int32(65535) - ml - for i in cutlass.range_constexpr(4): - taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) - ids_t.append(e_k) - kv = cutlass.Float32(0.0) - for i in cutlass.range_constexpr(4): - kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) - kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) - sig = cutlass.Float32(0.0) - e_l = cutlass.Int32(0) - for k in cutlass.range_constexpr(TOP_K): - sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) - sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) - e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) - total = sig - for off in (16, 8, 4, 2, 1): - total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) - wv = _bf16_round( - _div_rn( - _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), - _add_rn(total, cutlass.Float32(SUM_EPS)), - ) - ) - res_all.append((t, own, wv, e_l)) - if (lane == 0) & (warp == cutlass.Int32(W_WARP)): - _stamp(tdbg, bx, 62) - for tr in cutlass.range_constexpr( - (M_MAX + NRW - 1) // NRW - ): # stores after every chain (one branch-free block) - t, own, wv, e_l = res_all[tr] - if (lane < TOP_K) & own: - s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) - s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) - if bx == 0: - route_dbg.store( - wv.bitcast(cutlass.Int32), - idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, - ) - cute.arch.sync_warp() - if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): - if warp == cutlass.Int32(W_WARP): - # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same - # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now - if cutlass.const_expr(XSTAGE): - _wait(x_bar.subview(1), 0) - if lane == 0: - _stamp(tdbg, bx, 56) - n0 = cutlass.Int32(0) - w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) - ei0 = w0 & cutlass.Int32(0xFFFF) - ei1 = w0 >> cutlass.Int32(16) - for k in range(_nsegs(ei0, ei1)): - ug, j0, j1 = _seg(ei0, ei1, k) - c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - e = s_ids.load(idx=ug // cutlass.Int32(U1)) - sf_line = ( - (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) - * cutlass.Int32(SF1_KA) - + j0 * cutlass.Int32(C1 * NUM_KB) - ) * cutlass.Int32(4) - kx = j0 * cutlass.Int32(C1 * KT_BYTES) - row = c * cutlass.Int32(R1) - for j in cutlass.range(j0, j1, unroll=1): - if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx( - ab_full.subview(n0), - NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, - ) - for q in cutlass.range_constexpr(C1): - prims.cp_async_bulk_tensor_shared_cta_global( - sA.subview( - n0 * cutlass.Int32(NUM_BYTES_A) - + cutlass.Int32(q * NUM_BYTES_BOX1) - ), - tma_a1_desc.get_ptr(), - (kx + cutlass.Int32(q * KT_BYTES), row, e), - ab_full.subview(n0), - l2_cache_hint=EVICT_FIRST, - ) - if cutlass.const_expr(not SFA_CP): - prims.cp_async_bulk_tensor_shared_cta_global( - sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), - tma_s1_desc.get_ptr(), - (cutlass.Int32(0), sf_line), - ab_full.subview(n0), - ) - if cutlass.const_expr(SFA_CP): - _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) - n0 = n0 + cutlass.Int32(1) - kx = kx + cutlass.Int32(C1 * KT_BYTES) - sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) - if lane == 0: - s_misc.store(n0, idx=0) - cute.arch.sync_warp() - if cutlass.const_expr(E0_OFFLOAD): - if warp == cutlass.Int32(W_WARP): - _bar_arrive( - E0_BAR, 64 - ) # token 0's experts are in s_ids: warp 7 issues the early stages - if cutlass.const_expr(B1_AFTER_TOPK): - _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) - cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) - if cutlass.const_expr(E0_OFFLOAD): - if warp == cutlass.Int32(E0_WARP): - cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) - # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same - # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now - if cutlass.const_expr(XSTAGE): - _wait(x_bar.subview(1), 0) - if lane == 0: - _stamp(tdbg, bx, 56) - n0 = cutlass.Int32(0) - w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) - ei0 = w0 & cutlass.Int32(0xFFFF) - ei1 = w0 >> cutlass.Int32(16) - for k in range(_nsegs(ei0, ei1)): - ug, j0, j1 = _seg(ei0, ei1, k) - c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - e = s_ids.load(idx=ug // cutlass.Int32(U1)) - sf_line = ( - (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) - * cutlass.Int32(SF1_KA) - + j0 * cutlass.Int32(C1 * NUM_KB) - ) * cutlass.Int32(4) - kx = j0 * cutlass.Int32(C1 * KT_BYTES) - row = c * cutlass.Int32(R1) - for j in cutlass.range(j0, j1, unroll=1): - if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx( - ab_full.subview(n0), - NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, - ) - for q in cutlass.range_constexpr(C1): - prims.cp_async_bulk_tensor_shared_cta_global( - sA.subview( - n0 * cutlass.Int32(NUM_BYTES_A) - + cutlass.Int32(q * NUM_BYTES_BOX1) - ), - tma_a1_desc.get_ptr(), - (kx + cutlass.Int32(q * KT_BYTES), row, e), - ab_full.subview(n0), - l2_cache_hint=EVICT_FIRST, - ) - if cutlass.const_expr(not SFA_CP): - prims.cp_async_bulk_tensor_shared_cta_global( - sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), - tma_s1_desc.get_ptr(), - (cutlass.Int32(0), sf_line), - ab_full.subview(n0), - ) - if cutlass.const_expr(SFA_CP): - _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) - n0 = n0 + cutlass.Int32(1) - kx = kx + cutlass.Int32(C1 * KT_BYTES) - sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) - if lane == 0: - s_misc.store(n0, idx=0) - cute.arch.sync_warp() - _bar_arrive(E1_BAR, 64) - if warp == W_WARP: - if lane == 0: - _stamp(tdbg, bx, 20) - # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) - valid = lane < cutlass.Int32(TOP_K * M_MAX) - my_e = cutlass.select_( - valid, - s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), - cutlass.Int32(-1) - lane, - ) - same = _match_any(my_e) - leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) - first = valid & (lane == leader) - bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) - rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) - slot_l = cute.arch.shuffle_sync(rank, leader) - for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): - if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): - s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) - cute.arch.sync_warp() - if first: - s_el.store(my_e, idx=rank) - if valid: - s_kpos.store( - lane % cutlass.Int32(TOP_K), - idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), - ) - if bx == 0: - route_dbg.store( - my_e, - idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) - + lane % cutlass.Int32(TOP_K), - ) - gcnt = cute.arch.popc(bal) - if lane == 0: - s_el.store(gcnt, idx=G_CAP) - if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) - s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) - cute.arch.sync_warp() - if cutlass.const_expr(B2IMG): - # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) - if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): - nh = gcnt - lane * cutlass.Int32(4) - nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) - nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) - pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) - if pad > cutlass.Int32(0): - _red_release_add( - b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad - ) - cute.arch.sync_warp() - if lane == 0: - _stamp(tdbg, bx, 24) - cute.arch.sync_warp() - _bar_arrive(ROUTE_BAR, THREADS) - if warp < 4: - if cutlass.const_expr(B1_AFTER_TOPK): - cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) - if cutlass.const_expr(FUSED_ROUTER): - if tidx == 0: - _stamp(tdbg, bx, 57) - if cutlass.const_expr(XSTAGE): - _wait(x_bar.subview(0), 0) - if tidx == 0: - _stamp(tdbg, bx, 58) - xs32 = cutlass.Array( - s_big.data_ptr(0), - shape=(M_MAX * H // 2,), - dtype=cutlass.Int32, - alignment=16, - ) - _build_b1_fused( - tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready - ) - prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights - else: - _build_b1_fused( - tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready - ) - if tidx == 0: - _stamp(tdbg, bx, 23) - else: - _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) - if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) - if tidx == 0: - _stamp(tdbg, bx, 2) - if cutlass.const_expr(TIMING): - tdbg.store( - cutlass.Int64(s_el.load(idx=G_CAP)), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), - ) - - g_n = s_el.load(idx=G_CAP) - ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) - pr2 = s_plan.load(idx=NG_CAP) - blk0 = pr2 & cutlass.Int32(0xFFFF) - nb = pr2 >> cutlass.Int32(16) - - # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG - if warp == W_WARP: - if cutlass.const_expr(XSTAGE): - _wait(x_bar.subview(1), 0) # the staged hidden states are quantized - stage = cutlass.Int32(0) - phase = cutlass.Int32(1) - skip = cutlass.Int32(0) - if cutlass.const_expr(E0_OFFLOAD): - cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) - if cutlass.const_expr(EARLY0): - skip = s_misc.load(idx=0) # group-0 stages already in the ring - stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) - phase = cutlass.select_( - skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) - ) - for p in range(ng + cutlass.Int32(LAG)): - if p < ng: - i0, i1 = _plan_range(s_plan, p) - for k in range(_nsegs(i0, i1)): - ug, j0, j1 = _seg(i0, i1, k) - slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) - c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - e = s_el.load(idx=slot) - sf_line = ( - (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) - * cutlass.Int32(SF1_KA) - + j0 * cutlass.Int32(C1 * NUM_KB) - ) * cutlass.Int32(4) - kx = j0 * cutlass.Int32(C1 * KT_BYTES) - row = c * cutlass.Int32(R1) - for j in cutlass.range(j0, j1, unroll=1): - if skip > cutlass.Int32(0): - skip = skip - cutlass.Int32(1) - else: - _waitd(ab_empty.subview(stage), phase, tdbg, 1) - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx( - ab_full.subview(stage), - NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, - ) - for q in cutlass.range_constexpr(C1): - prims.cp_async_bulk_tensor_shared_cta_global( - sA.subview( - stage * cutlass.Int32(NUM_BYTES_A) - + cutlass.Int32(q * NUM_BYTES_BOX1) - ), - tma_a1_desc.get_ptr(), - (kx + cutlass.Int32(q * KT_BYTES), row, e), - ab_full.subview(stage), - l2_cache_hint=EVICT_FIRST, - ) - if cutlass.const_expr(not SFA_CP): - prims.cp_async_bulk_tensor_shared_cta_global( - sSFAraw.subview( - stage * cutlass.Int32(NUM_WORDS_SFA_RAW) - ), - tma_s1_desc.get_ptr(), - (cutlass.Int32(0), sf_line), - ab_full.subview(stage), - ) - if cutlass.const_expr(SFA_CP): - _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) - stage, phase = _adv(stage, phase) - kx = kx + cutlass.Int32(C1 * KT_BYTES) - sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) - if p >= cutlass.Int32(LAG): - g = p - cutlass.Int32(LAG) - n_here = g_n - g * cutlass.Int32(4) - n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) - for bb in range(nb): - blk = blk0 + bb - for t2 in cutlass.range_constexpr(KT2): - _waitd(ab_empty.subview(stage), phase, tdbg, 2) - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx( - ab_full.subview(stage), - n_here - * cutlass.Int32( - NUM_BYTES_BOX1 - if SFA_CP - else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 - ), - ) - for jj in cutlass.range_constexpr(4): - if cutlass.Int32(jj) < n_here: - e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) - prims.cp_async_bulk_tensor_shared_cta_global( - sA.subview( - stage * cutlass.Int32(NUM_BYTES_A) - + cutlass.Int32(jj * NUM_BYTES_BOX1) - ), - tma_a2_desc.get_ptr(), - ( - cutlass.Int32(t2 * KT_BYTES), - blk * cutlass.Int32(R2), - e, - ), - ab_full.subview(stage), - l2_cache_hint=EVICT_FIRST, - ) - if cutlass.const_expr(not SFA_CP): - prims.cp_async_bulk_tensor_shared_cta_global( - sSFAraw.subview( - stage * cutlass.Int32(NUM_WORDS_SFA_RAW) - + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) - ), - tma_s2_desc.get_ptr(), - ( - cutlass.Int32(0), - ( - ( - e * cutlass.Int32(SF2_RB) - + (blk >> cutlass.Int32(2)) - ) - * cutlass.Int32(SF2_KA) - + cutlass.Int32(4 * t2) - ) - * cutlass.Int32(4), - ), - ab_full.subview(stage), - ) - if cutlass.const_expr(SFA_CP): - _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) - stage, phase = _adv(stage, phase) - if lane == 0: - _stamp(tdbg, bx, 3) - - # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] - if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - stage = cutlass.Int32(0) - phase = cutlass.Int32(0) - for p in range(ng + cutlass.Int32(LAG)): - i0 = cutlass.Int32(0) - i1 = cutlass.Int32(0) - for part in cutlass.range_constexpr(2): - cnt = cutlass.Int32(0) - m1 = cutlass.Int32(0) - nk = cutlass.Int32(0) - if cutlass.const_expr(part == 0): - if p < ng: - i0, i1 = _plan_range(s_plan, p) - nk = _nsegs(i0, i1) - else: - if p >= cutlass.Int32(LAG): - nk = nb - for k in range(nk): - if cutlass.const_expr(part == 0): - ug, j0, j1 = _seg(i0, i1, k) - m1 = ( - ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - ) & cutlass.Int32(3) - cnt = j1 - j0 - else: - m1 = (blk0 + k) & cutlass.Int32(3) - cnt = cutlass.Int32(KT2) - for _j in cutlass.range(cnt, unroll=1): - _waitd(ab_full.subview(stage), phase, tdbg, 3) - raw0 = ( - stage * cutlass.Int32(NUM_WORDS_SFA_RAW) - + lane * cutlass.Int32(4) - + m1 - ) - dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) - for kb in cutlass.range_constexpr(NUM_KB): - w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) - w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) - w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) - w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) - sSFA.store( - ( - cutlass.Int32(w0), - cutlass.Int32(w1), - cutlass.Int32(w2), - cutlass.Int32(w3), - ), - idx=dst0 + cutlass.Int32(kb * 128), - alignment=16, - ) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - cute.arch.sync_warp() - if prims.elect_sync(): - sd0 = prims.Tcgen05SmemDesc.build( - sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - for kb in cutlass.range_constexpr(NUM_KB): - prims.tcgen05_cp( - s2t_shape, - _tptr( - base_row, - base_col - + SFA1_COL - + stage * cutlass.Int32(16) - + kb * 4, - ), - sd0 + (kb * 512 >> 4), - multicast=s2t_mc, - ) - prims.tcgen05_commit(sfa_ready.subview(stage)) - stage, phase = _adv(stage, phase) - - if cutlass.const_expr(B2IMG): - # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group - if warp == cutlass.Int32(S_WARP): - for p in range(ng): - i0, i1 = _plan_range(s_plan, p) - n_done = cutlass.Int32(0) - for k in range(_nsegs(i0, i1)): - ug, j0, j1 = _seg(i0, i1, k) - n_done = n_done + cutlass.select_( - (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), - cutlass.Int32(1), - cutlass.Int32(0), - ) - _wait( - fc1_done.subview(p & cutlass.Int32(3)), - (p >> cutlass.Int32(2)) & cutlass.Int32(1), - ) - if cutlass.const_expr(B2_DLY > 0): - t0d = _gtime() - dl = cutlass.Int64( - ( - (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) - & cutlass.Int32(63) - ) - * cutlass.Int32(B2_DLY // 64) - ) - while _gtime() - t0d < dl: - pass - if (lane == 0) & (n_done > cutlass.Int32(0)): - _fence_proxy_async_global() - _red_release_add( - b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done - ) - cute.arch.sync_warp() - # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete - if warp == cutlass.Int32(X_WARP): - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) - for g in range(ng): - buf = g & cutlass.Int32(1) - if g >= cutlass.Int32(2): - _waitd( - b2_free.subview(buf), - ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), - tdbg, - 101, - ) - if lane == 0: - _stamp(tdbg, bx, 29) - caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() - if lane == 0: - cv = _ld_acquire_u32(caddr) - t0h = cutlass.Int64(0) - while (cv - cnt_target) < cutlass.Int32(0): - t0h = _hang_tick(t0h) - _nanosleep(B2IMG_NS) - cv = _ld_acquire_u32(caddr) - _fence_proxy_async_global() - prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) - prims.cp_async_bulk_shared_cluster_global( - sB2.subview(buf * cutlass.Int32(IMG_B2)), - img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), - img_bar.subview(buf), - IMG_B2, - ) - prims.cp_async_bulk_shared_cluster_global( - sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), - img32.subview( - g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) - ), - img_bar.subview(buf), - IMG_SF, - ) - cute.arch.sync_warp() - _waitd( - img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 - ) - if prims.elect_sync(): - sd0 = prims.Tcgen05SmemDesc.build( - sSFB2.subview(0), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - for a in cutlass.range_constexpr(KT2 * NUM_KB): - prims.tcgen05_cp( - s2t_shape, - _tptr( - base_row, - base_col - + SFB2_COL - + buf * cutlass.Int32(KT2 * NUM_KB * 4) - + a * 4, - ), - sd0 - + cutlass.Int64( - (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) - * cutlass.Int32(512 >> 4) - ), - multicast=s2t_mc, - ) - prims.tcgen05_commit(b2_ready.subview(buf)) - cute.arch.sync_warp() - if lane == 0: - _stamp(tdbg, bx, 25) - # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on - # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then - # tcgen05.cp to TMEM. - if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): - xi = cutlass.select_( - warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) - ) - if cutlass.const_expr(G_CAP > TOP_K): - if (bx == 0) & (xi == cutlass.Int32(0)): - # rows of slots this call does not use: tag 0, never taken for a later call's - for i in range( - g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, - cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), - 32, - ): - for rr in cutlass.range_constexpr(HB_REPL): - hbuf32.store( - ( - cutlass.Int32(0), - cutlass.Int32(0), - cutlass.Int32(0), - cutlass.Int32(0), - ), - idx=i * cutlass.Int32(4) - + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), - alignment=16, - ) - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - sSFB2_8 = cutlass.Array( - sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 - ) - sB2_32 = cutlass.Array( - sB2.data_ptr(0), - shape=(2 * KT2 * NUM_BYTES_B2 // 4,), - dtype=cutlass.Int32, - alignment=16, - ) - nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group - hb_off = cutlass.Int32(0) - if cutlass.const_expr(HB_REPL > 1): - # this CTA's copy of the records - hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) - NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) - for g in range(ng): - buf = g & cutlass.Int32(1) - if g >= cutlass.Int32(2): - _waitd( - b2_free.subview(buf), - ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), - tdbg, - 101, - ) - _waitd( - fc1_done.subview(g & cutlass.Int32(3)), - (g >> cutlass.Int32(2)) & cutlass.Int32(1), - tdbg, - 4, - ) # this CTA's own share of group g is out - if cutlass.const_expr(TIMING): - if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): - tdbg.store( - _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g - ) - n_here = g_n - g * cutlass.Int32(4) - n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) - nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) - rec0 = g * cutlass.Int32( - 4 * M_MAX * UNITS_ROW - ) # rows (4 g + j) M + t are contiguous - done = cutlass.Int32(0) - rspin = cutlass.Int32(0) - if (lane == 0) & (xi == cutlass.Int32(0)): - _stamp(tdbg, bx, 29) - npoll = cutlass.Int32(0) - got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] - t0h = cutlass.Int64(0) - while done == cutlass.Int32(0): - t0h = _hang_tick(t0h) - ok = cutlass.Boolean(True) - for z in cutlass.range_constexpr(NZ): - qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi - if qi < nrec: - v4 = prims.load_ext( - hbuf32.subview( - (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) - ), - dtype=cutlass.Int32, - count=4, - order="relaxed", - scope="gpu", - ) - for wq in cutlass.range_constexpr(3): - got[z][wq] = cutlass.Int32(v4[wq]) - ok = ok & ( - ( - ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) - & cutlass.Int32(-16777216) - ) - == cutlass.Int32(0) - ) - npoll = npoll + cutlass.Int32(1) - if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): - done = cutlass.Int32(1) - else: - _nanosleep(B2_POLL_NS) - if cutlass.const_expr(HANG_DBG): - rspin = rspin + cutlass.Int32(1) - if rspin > cutlass.Int32(1 << 14): - done = cutlass.Int32(1) - tdbg.store( - cutlass.Int64(200), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, - ) - for z in cutlass.range_constexpr(NZ): - qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi - if qi < nrec: - w0, w1, w2 = got[z][0], got[z][1], got[z][2] - rr = qi // cutlass.Int32( - UNITS_ROW - ) # row within the group: j M + t = B row n - c = qi - rr * cutlass.Int32(UNITS_ROW) - t2 = c >> cutlass.Int32(4) - cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) - lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) - hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( - (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) - ) - sB2_32.store( - (lo, hi), - idx=( - (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) - + rr * cutlass.Int32(KT_BYTES) - + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) - + (c & cutlass.Int32(1)) * cutlass.Int32(8) - ) - // cutlass.Int32(4), - alignment=8, - ) - # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 - sSFB2_8.store( - cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), - idx=( - ( - (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) - + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) - ) - * cutlass.Int32(128) - + rr * cutlass.Int32(4) - ) - * cutlass.Int32(4) - + (c & cutlass.Int32(3)), - ) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - if cutlass.const_expr(XW > 1): - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) - cute.arch.sync_warp() - if (xi == cutlass.Int32(0)) & prims.elect_sync(): - sd0 = prims.Tcgen05SmemDesc.build( - sSFB2.subview(0), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - for a in cutlass.range_constexpr(KT2 * NUM_KB): - prims.tcgen05_cp( - s2t_shape, - _tptr( - base_row, - base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, - ), - sd0 - + cutlass.Int64( - (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) - * cutlass.Int32(512 >> 4) - ), - multicast=s2t_mc, - ) - prims.tcgen05_commit(b2_ready.subview(buf)) - if (lane == 0) & (xi == cutlass.Int32(0)): - _stamp(tdbg, bx, 25) - if cutlass.const_expr(TIMING): - if g < cutlass.Int32(8): - tdbg.store( - _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g - ) - if cutlass.const_expr(TIMING): - tdbg.store( - cutlass.Int64(npoll), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), - ) - - # ---- MMA (warp 7) - if warp == MMA_WARP: - idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( - a_dtype=cutlass.Float4E2M1FN, - b_dtype=cutlass.Float4E2M1FN, - scale_format=0, - n_dim=N1, - m_dim=MMA_M, - sparsity_version=0, - ) - idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( - a_dtype=cutlass.Float4E2M1FN, - b_dtype=cutlass.Float4E2M1FN, - scale_format=0, - n_dim=N2, - m_dim=MMA_M, - sparsity_version=0, - ) - desc_a0 = prims.Tcgen05SmemDesc.build( - sA.subview(0), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - desc_b10 = prims.Tcgen05SmemDesc.build( - sB1.subview(0), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - desc_b20 = prims.Tcgen05SmemDesc.build( - sB2.subview(0), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - tcol = (base_row << 16) | base_col - _waitd(b1_ready, 0, tdbg, 5) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - stage = cutlass.Int32(0) - phase = cutlass.Int32(0) - k1 = cutlass.Int32(0) - k2 = cutlass.Int32(0) - acc_b2 = cutlass.Int64(0) - acc_sf = cutlass.Int64(0) - for p in range(ng + cutlass.Int32(LAG)): - if p < ng: - i0, i1 = _plan_range(s_plan, p) - for k in range(_nsegs(i0, i1)): - ug, j0, j1 = _seg(i0, i1, k) - ab = k1 & cutlass.Int32(1) - _waitd( - acc_empty.subview(ab), - ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), - tdbg, - 102, - ) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - acc_ptr = cutlass.inttoptr( - tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 - ) - for j in cutlass.range(j0, j1, unroll=1): - tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) - if cutlass.const_expr(SFA_CP): - _waitd(ab_full.subview(stage), phase, tdbg, 6) - prims.fence_proxy( - "async_shared", space=prims.SharedSpace.shared_cta - ) - else: - _waitd(sfa_ready.subview(stage), phase, tdbg, 6) - if cutlass.const_expr(TIMING): - acc_sf = acc_sf + (_gtime() - tw0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) - db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) - sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) - sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) - if prims.elect_sync(): - if cutlass.const_expr(SFA_CP): - _sfa_to_tmem(sSFA, stage, base_row, base_col) - for kb in cutlass.range_constexpr(NUM_KB): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF4NVF4, - prims.CTAGroup.CTA_1, - acc_ptr, - da + INC_K * kb, - db + INC_K * kb, - idesc1, - (j > j0) | cutlass.Boolean(kb > 0), - cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), - cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), - scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, - ) - prims.tcgen05_commit(ab_empty.subview(stage)) - stage, phase = _adv(stage, phase) - if prims.elect_sync(): - prims.tcgen05_commit(acc_full.subview(ab)) - k1 = k1 + cutlass.Int32(1) - if p >= cutlass.Int32(LAG): - g = p - cutlass.Int32(LAG) - buf = g & cutlass.Int32(1) - tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) - _waitd( - b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 - ) - if cutlass.const_expr(TIMING): - acc_b2 = acc_b2 + (_gtime() - tw2) - if (lane == 0) & (g == ng - cutlass.Int32(1)): - _stamp(tdbg, bx, 4) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - for bb in range(nb): - a2 = k2 & cutlass.Int32(1) - _waitd( - acc2_empty.subview(a2), - ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), - tdbg, - 104, - ) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - acc2_ptr = cutlass.inttoptr( - tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 - ) - for t2 in cutlass.range_constexpr(KT2): - tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) - if cutlass.const_expr(SFA_CP): - _waitd(ab_full.subview(stage), phase, tdbg, 7) - prims.fence_proxy( - "async_shared", space=prims.SharedSpace.shared_cta - ) - else: - _waitd(sfa_ready.subview(stage), phase, tdbg, 7) - if cutlass.const_expr(TIMING): - acc_sf = acc_sf + (_gtime() - tw1) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) - db = desc_b20 + cutlass.Int64( - (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) - * cutlass.Int32(NUM_BYTES_B2 >> 4) - ) - sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) - sfb_c = ( - tcol - + SFB2_COL - + buf * cutlass.Int32(KT2 * NUM_KB * 4) - + cutlass.Int32(t2 * NUM_KB * 4) - ) - if prims.elect_sync(): - if cutlass.const_expr(SFA_CP): - _sfa_to_tmem(sSFA, stage, base_row, base_col) - for kb in cutlass.range_constexpr(_nkb2(t2)): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF4NVF4, - prims.CTAGroup.CTA_1, - acc2_ptr, - da + INC_K * kb, - db + INC_K * kb, - idesc2, - cutlass.Boolean(t2 > 0 or kb > 0), - cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), - cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), - scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, - ) - prims.tcgen05_commit(ab_empty.subview(stage)) - stage, phase = _adv(stage, phase) - if prims.elect_sync(): - prims.tcgen05_commit(acc2_full.subview(a2)) - k2 = k2 + cutlass.Int32(1) - if cutlass.const_expr(TIMING): - if lane == 0: - tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) - tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) - - # ---- epilogue warps (0-3) - if warp < 4: - row_w = base_row + warp * 32 - for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): - qi = tidx + cutlass.Int32(r * EPI_THREADS) - sl_e = qi // cutlass.Int32(6) - if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): - fld = qi - sl_e * cutlass.Int32(6) - s_esc.store( - cutlass.Float32( - escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) - ), - idx=sl_e * cutlass.Int32(8) + fld, - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) - k1 = cutlass.Int32(0) - k2 = cutlass.Int32(0) - SH_ADD = SHARED - wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] - pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] - ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) - for p in range(ng + cutlass.Int32(LAG)): - if cutlass.const_expr(AR_PRE): - if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: - # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads - # at 8-byte offsets: misaligned) - ard = tuple( - cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) - for i_ in (1, 3, 4, 5) - ) - for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): - f_pair = warp + cutlass.Int32(bt * 4) - f_bb = f_pair // cutlass.Int32(M_MAX) - f_t = f_pair - f_bb * cutlass.Int32(M_MAX) - f_cl = cutlass.select_( - f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) - ) - pars[bt] = _ar_parity(arp, f_t, f_cl) - if cutlass.const_expr(SH_ADD): - if p == ng + cutlass.Int32(LAG - 1): - # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) - for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): - f_pair = warp + cutlass.Int32(bt * 4) - f_bb = f_pair // cutlass.Int32(M_MAX) - f_t = f_pair - f_bb * cutlass.Int32(M_MAX) - f_ch = cutlass.select_( - f_bb < nb, - f_t * cutlass.Int32(H) - + (blk0 + f_bb) * cutlass.Int32(R2) - + cutlass.Int32(4) * (lane & cutlass.Int32(7)) - + (lane >> cutlass.Int32(3)), - cutlass.Int32(0), - ) - wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) - if p < ng: - i0, i1 = _plan_range(s_plan, p) - for k in range(_nsegs(i0, i1)): - ug, j0, j1 = _seg(i0, i1, k) - ab = k1 & cutlass.Int32(1) - slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) - cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - unit = slot * cutlass.Int32(U1) + cblk - _waitd( - acc_full.subview(ab), - (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), - tdbg, - 105, - ) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - acc = prims.tcgen05_ld( - "32x32b", - _tptr( - row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 - ), - num=N1, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - prims.mbarrier_arrive(acc_empty.subview(ab)) - pb = ab * cutlass.Int32(M_MAX * C1 * R1) - for t in cutlass.range_constexpr(M_MAX): - part = cutlass.Float32(0.0) - for w in cutlass.range_constexpr(4): - part = cutlass.select_( - warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part - ) - s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) - for tr in cutlass.range_constexpr((M_MAX + 3) // 4): - t = warp + cutlass.Int32(4 * tr) - if t < cutlass.Int32(M_MAX): - xv = _add_rn( - _add_rn( - _add_rn( - s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), - s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), - ), - s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), - ), - s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), - ) - finish = cutlass.Int32(1) - if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): - paddr = ( - partials.subview( - (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) - + lane - ) - .data_ptr() - .toint() - ) - if j0 == cutlass.Int32(0): - _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) - finish = cutlass.Int32(0) - else: - pbits = _ld_relaxed_u32(paddr) - spins = cutlass.Int32(0) - t0h = cutlass.Int64(0) - while pbits == cutlass.Int32(PARTIAL_EMPTY): - t0h = _hang_tick(t0h) - pbits = _ld_relaxed_u32(paddr) - spins = spins + cutlass.Int32(1) - _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) - xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) - cute.arch.sync_warp() - if cutlass.const_expr( - XVDUMP - ): # debug: FC1 sums into the (unused) partial slots - partials.store( - xv.bitcast(cutlass.Int32), - idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) - + lane, - ) - if finish == cutlass.Int32(1): - gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) - is_up = ( - (lane >> cutlass.Int32(3)) & cutlass.Int32(1) - ) == cutlass.Int32(0) - es = slot * cutlass.Int32(8) - qv, sfb = _swiglu_quant( - xv, - gate_raw, - is_up, - s_esc.load(idx=es), - s_esc.load(idx=es + cutlass.Int32(1)), - s_esc.load(idx=es + cutlass.Int32(2)), - s_esc.load(idx=es + cutlass.Int32(3)), - s_esc.load(idx=es + cutlass.Int32(4)), - ) - byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) - bt = [ - cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) - for i in range(8) - ] - rw0 = ( - bt[0] - | (bt[1] << cutlass.Int32(8)) - | (bt[2] << cutlass.Int32(16)) - | tagv - ) - rw1 = ( - bt[3] - | (bt[4] << cutlass.Int32(8)) - | (bt[5] << cutlass.Int32(16)) - | tagv - ) - rw2 = ( - bt[6] - | (bt[7] << cutlass.Int32(8)) - | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) - | tagv - ) - if cutlass.const_expr(B2IMG): - if lane == 0: - # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 - # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), - # row n at byte 16 n, byte c % 4 - n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t - t2i = cblk >> cutlass.Int32(4) - cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) - gbase = p * cutlass.Int32(IMG_BYTES) - lo = ( - bt[0] - | (bt[1] << cutlass.Int32(8)) - | (bt[2] << cutlass.Int32(16)) - | (bt[3] << cutlass.Int32(24)) - ) - hi = ( - bt[4] - | (bt[5] << cutlass.Int32(8)) - | (bt[6] << cutlass.Int32(16)) - | (bt[7] << cutlass.Int32(24)) - ) - img32.store( - (lo, hi), - idx=( - gbase - + t2i * cutlass.Int32(NUM_BYTES_B2) - + n * cutlass.Int32(KT_BYTES) - + (cci ^ (n & cutlass.Int32(7))) - * cutlass.Int32(16) - + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) - ) - // cutlass.Int32(4), - alignment=8, - ) - img8 = cutlass.Array( - img32.data_ptr(0), - shape=(NG_CAP * IMG_BYTES,), - dtype=cutlass.Int8, - alignment=16, - ) - img8.store( - cutlass.Int8(sfb & cutlass.Int32(0xFF)), - idx=gbase - + cutlass.Int32(IMG_B2) - + ( - ( - t2i * cutlass.Int32(NUM_KB) - + ( - (cblk & cutlass.Int32(15)) - >> cutlass.Int32(2) - ) - ) - * cutlass.Int32(512) - ) - + n * cutlass.Int32(16) - + (cblk & cutlass.Int32(3)), - ) - else: - if cutlass.const_expr(HB_REPL == 1): - if lane == 0: - hbuf32.store( - (rw0, rw1, rw2, tagv), - idx=( - (slot * cutlass.Int32(M_MAX) + t) - * cutlass.Int32(UNITS_ROW) - + cblk - ) - * cutlass.Int32(4), - alignment=16, - ) - else: - if lane < HB_REPL: # copy `lane` - hbuf32.store( - (rw0, rw1, rw2, tagv), - idx=( - (slot * cutlass.Int32(M_MAX) + t) - * cutlass.Int32(UNITS_ROW) - + cblk - ) - * cutlass.Int32(4) - + lane * cutlass.Int32(HB_STRIDE), - alignment=16, - ) - k1 = k1 + cutlass.Int32(1) - prims.mbarrier_arrive( - fc1_done.subview(p & cutlass.Int32(3)) - ) # warp 5 may poll group p - if tidx == 0: - _stamp(tdbg, bx, 10) - if cutlass.const_expr(TIMING): - if p < cutlass.Int32(8): - tdbg.store( - _gtime(), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, - ) - if p >= cutlass.Int32(LAG): - g = p - cutlass.Int32(LAG) - for bb in range(nb): - a2 = k2 & cutlass.Int32(1) - _waitd( - acc2_full.subview(a2), - (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), - tdbg, - 106, - ) - if cutlass.const_expr(TIMING): - if (tidx == 0) & (g == ng - cutlass.Int32(1)): - tdbg.store( - _gtime(), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, - ) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - yv = prims.tcgen05_ld( - "32x32b", - _tptr( - row_w, - base_col - + ACC2_COL - + a2 * cutlass.Int32(N2) - + warp * cutlass.Int32(M_MAX), - cutlass.Float32, - ), - num=LD_POW2, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - prims.mbarrier_arrive(acc2_empty.subview(a2)) - sl = g * cutlass.Int32(4) + warp - if sl < g_n: - a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) - for t in cutlass.range_constexpr(M_MAX): - kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) - if kp >= cutlass.Int32(0): - ys.store( - _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), - idx=( - (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) - * cutlass.Int32(TOP_K) - + kp - ) - * cutlass.Int32(R2) - + lane, - ) - k2 = k2 + cutlass.Int32(1) - prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) - if tidx == 0: - _stamp(tdbg, bx, 6) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) - if tidx == 0: - _stamp(tdbg, bx, 19) - # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order - for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): - pair = warp + cutlass.Int32(bt * 4) - bb = pair // cutlass.Int32(M_MAX) - t = pair - bb * cutlass.Int32(M_MAX) - if bb < nb: - acc_o = cutlass.Float32(0.0) - for k in cutlass.range_constexpr(TOP_K): - acc_o = _fma_rn( - s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), - ys.load( - idx=( - (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) - + cutlass.Int32(k) - ) - * cutlass.Int32(R2) - + lane - ), - acc_o, - ) - if cutlass.const_expr( - YDUMP - ): # debug: y of (token, k) into the partials workspace - partials.store( - cutlass.Float32( - ys.load( - idx=( - (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) - + cutlass.Int32(k) - ) - * cutlass.Int32(R2) - + lane - ) - ).bitcast(cutlass.Int32), - idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) - + (blk0 + bb) * cutlass.Int32(R2) - + cutlass.Int32(4) * (lane & cutlass.Int32(7)) - + (lane >> cutlass.Int32(3)), - ) - ch = ( - t * cutlass.Int32(H) - + (blk0 + bb) * cutlass.Int32(R2) - + cutlass.Int32(4) * (lane & cutlass.Int32(7)) - + (lane >> cutlass.Int32(3)) - ) - if cutlass.const_expr(SH_ADD): - # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) - wy = wys[bt] - t0h = cutlass.Int64(0) - while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: - t0h = _hang_tick(t0h) - wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) - acc_o = _add_rn( - (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) - ) - out.store(cutlass.BFloat16(acc_o), idx=ch) - if ar_on: - cl = ( - (blk0 + bb) * cutlass.Int32(R2) - + cutlass.Int32(4) * (lane & cutlass.Int32(7)) - + (lane >> cutlass.Int32(3)) - ) - if cutlass.const_expr(AR_PRE): - _ar_push_pre( - ard, - pars[bt], - _bf16_round(acc_o), - ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), - t, - cl, - ) - else: - _ar_push( - arp, - _bf16_round(acc_o), - ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), - t, - cl, - ) - - # ---- teardown - if tidx == 0: - _stamp(tdbg, bx, 7) - epochs.store(ep + cutlass.Int32(1), idx=bx) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - prims.barrier_cta_sync(0) - if ar_on | cutlass.Boolean(LATE_TRIGGER): - # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA - cute.arch.griddepcontrol_launch_dependents() - if warp == MMA_WARP: - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - prims.tcgen05_dealloc( - cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS - ) - - -@cute.jit -def m3_moe( - a1_tensor: cute.Tensor, - sfa1_tensor: cute.Tensor, - a2_tensor: cute.Tensor, - sfa2_tensor: cute.Tensor, - logits: cute.Tensor, - bias: cute.Tensor, - x32: cute.Tensor, - xsf32: cute.Tensor, - escal: cute.Tensor, - plan: cute.Tensor, - partials: cute.Tensor, - hbuf32: cute.Tensor, - epochs: cute.Tensor, - out: cute.Tensor, - route_dbg: cute.Tensor, - tdbg: cute.Tensor, - xb32: cute.Tensor, - rw: cute.Tensor, - sfsc: cute.Tensor, - lgbuf: cute.Tensor, - guw: cute.Tensor, - gusf: cute.Tensor, - dnw: cute.Tensor, - dnsf: cute.Tensor, - hx: cute.Tensor, - dnx: cute.Tensor, - arp: cute.Tensor, - s1w: cute.Tensor, - s2w: cute.Tensor, - img32: cute.Tensor, - b2cnt: cute.Tensor, - stream: cuda_driver.CUstream, -): - tma_a1_desc = cuda.create_tensor_map_tiled( - global_address=a1_tensor.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[H // 2, TWO_I, E_LOCAL], - global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], - box_dims=(KT_BYTES, R1, 1), - swizzle=cuda.TensorMapSwizzle.s128b, - ) - tma_a2_desc = cuda.create_tensor_map_tiled( - global_address=a2_tensor.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[I_TP // 2, H, E_LOCAL], - global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], - box_dims=(KT_BYTES, R2, 1), - swizzle=cuda.TensorMapSwizzle.s128b, - ) - tma_s1_desc = cuda.create_tensor_map_tiled( - global_address=sfa1_tensor.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], - global_strides=[128 // 16], - box_dims=(128, NUM_BYTES_SFA_RAW // 128), - swizzle=cuda.TensorMapSwizzle.none, - ) - tma_s2_desc = cuda.create_tensor_map_tiled( - global_address=sfa2_tensor.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], - global_strides=[128 // 16], - box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), - swizzle=cuda.TensorMapSwizzle.none, - ) - tma_gu_desc = tma_a1_desc - tma_dn_desc = tma_a2_desc - if cutlass.const_expr(SHARED and not SH_PRETILE): - tma_gu_desc = cuda.create_tensor_map_tiled( - global_address=guw.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[H, 2 * SH_I], - global_strides=[H // 16], - box_dims=(KT_BYTES, 32), - swizzle=cuda.TensorMapSwizzle.s128b, - ) - tma_dn_desc = cuda.create_tensor_map_tiled( - global_address=dnw.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[SH_I, H], - global_strides=[SH_I // 16], - box_dims=(KT_BYTES, 128), - swizzle=cuda.TensorMapSwizzle.s128b, - ) - m3_moe_kernel( - tma_a1_desc, - tma_a2_desc, - tma_s1_desc, - tma_s2_desc, - tma_gu_desc, - tma_dn_desc, - logits, - bias, - x32, - xsf32, - escal, - plan, - partials, - hbuf32, - epochs, - out, - route_dbg, - tdbg, - xb32, - rw, - sfsc, - lgbuf, - guw, - gusf, - dnw, - dnsf, - hx, - dnx, - arp, - s1w, - s2w, - img32, - b2cnt, - ).launch( - grid=(NCTA + NSH, 1, 1), - block=(THREADS, 1, 1), - cluster=(1, 1, 1), - stream=stream, - use_pdl=True, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py index 281b56869436..fedde40b3af7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op.py @@ -53,6 +53,33 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.phase_plan import ( + MoePlanGeometry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.phase_plan import ( + make_phase_plan as _shared_make_phase_plan, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + M3MoeLayer as M3MoeLayer, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + _view as _view, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + ar_push_desc as ar_push_desc, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + fine_ranges, + pretile_shared, + shared_views, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + is_supported as is_supported, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.state import ( + M3MoeStateBase, + MoeStateBindings, +) from tensorrt_llm.logger import logger HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size @@ -130,7 +157,9 @@ def flat_applies(num_tokens: int) -> bool: M3_MOE_LAG = 2 -_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") +_ENGINE_PATH = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../../kernels/moe/m3_moe_kernel.py") +) def _kernel_module(config: dict, path: str): @@ -145,128 +174,13 @@ def _kernel_module(config: dict, path: str): ) spec = importlib.util.spec_from_file_location(name, path) mod = importlib.util.module_from_spec(spec) - mod.M3_CONFIG = dict(config) + mod.M3_CONFIG = dict(config, tp_size=4) sys.modules[name] = mod spec.loader.exec_module(mod) _modules[key] = mod return mod -def _view(t: torch.Tensor, align: int): - from cutlass.cute.runtime import from_dlpack - - return from_dlpack(t.detach(), assumed_align=align) - - -def is_supported(device: Optional[torch.device] = None) -> tuple: - if not torch.cuda.is_available(): - return False, "no CUDA" - major, _ = torch.cuda.get_device_capability(device) - if major != 10: - return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" - try: - import cutlass # noqa: F401 - except ImportError: - return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" - return True, "" - - -def ar_push_desc( - mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device -) -> torch.Tensor: - """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A - of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). - With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old - behaviour.""" - return torch.tensor( - [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device - ) - - -def _shared_views(shared_weights, who): - """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" - if shared_weights is None: - raise ValueError(f"{who} (shared): shared_weights are required") - gu_w, gu_sf, dn_w, dn_sf = shared_weights - want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} - for name, (t, shape) in want.items(): - if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): - raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") - if ( - gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 - or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 - ): - raise ValueError( - f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" - ) - return ( - gu_w.view(torch.uint8).view(-1), - gu_sf.view(torch.uint8).reshape(-1), - dn_w.view(torch.uint8).view(-1), - dn_sf.view(torch.uint8).reshape(-1), - ) - - -def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: - """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB - shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and - its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 - = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row - R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then - down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per - layer.""" - dev = gu8.device - H_, I_ = HIDDEN, SHARED_INTER - P = I_ // n_cta - split = 128 // (2 * P) - gt = H_ // 128 // split - dkt = I_ // 128 - gu = gu8.view(2 * I_, H_) - dn = dn8.view(H_, I_) - gkb, dkb = H_ // 32, I_ // 32 - - def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast - return sf[ - ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) - ] - - r = torch.arange(128, device=dev) - q = r >> 5 - ri = r & 31 - octant = r // P - chunk = torch.arange(8, device=dev) - swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) - out = torch.empty( - n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev - ) # capacity: fixed: a workspace of the kernel's fixed geometry - for s in range(n_cta): - # gate_up: MMA row R -> weight row, slice - nrow = (octant & 1) * I_ + P * s + r % P - sl = octant >> 1 - for g in range(gt): - kt = sl * gt + g # per row - tile = gu[ - nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) - ] # [128, 128] - img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) - out[s, g, :16384] = img.reshape(-1) - kb = torch.arange(4, device=dev) - codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] - atom = torch.empty(512, dtype=torch.uint8, device=dev) - atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) - out[s, g, 16384:] = atom - for gd in range(2 * dkt): - jt, kt = gd // dkt, gd % dkt - row0 = 256 * s + 128 * jt - tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] - img = tile.reshape(128, 8, 16)[r[:, None], swz] - out[s, gt + gd, :16384] = img.reshape(-1) - out[s, gt + gd, 16384:] = dnsf8[ - ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 - ] - return out.view(-1) - - # --------------------------------------------------------------------------------------------------------------------- # m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. # --------------------------------------------------------------------------------------------------------------------- @@ -283,26 +197,6 @@ def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: return hit[1] -def _fine_ranges(weights, total: int): - """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item - counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary - strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit - runs at least to that unit's end. Returns the item counts.""" - n = len(weights) - items = [0] * n - s = 0 - acc = 0.0 - for c in range(n): - acc += weights[c] - e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) - if s % _S1 and e < total: # started inside a unit: run to (at least) its end - e = max(e, min(total, (s // _S1 + 1) * _S1)) - items[c] = e - s - s = e - assert s == total - return items - - def make_phase_plan( num_ctas: int, m_max: int, @@ -319,87 +213,27 @@ def make_phase_plan( group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" - gcap = TOP_K * m_max - nb2 = HIDDEN // 32 - plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) - nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] - b0s = [sum(nbs[:c]) for c in range(num_ctas)] - # (heavy_fc2_only) CTAs with more FC2 blocks than the minimum stream no FC1: they run each group's FC2 as soon as - # its hand-off lands, so their trailing FC2 after the last hand-off matches the others' two lagged groups - min_nb = min(nbs) - no_fc1 = [heavy_fc2_only and nbs[c] > min_nb for c in range(num_ctas)] - BIG_W = 1 << 40 - blk_bytes = 32 * _ROW2_BYTES # one block, one slot - for g in range(TOP_K, gcap + 1): - ng = (g + 3) // 4 - for c in range(num_ctas): - plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) - for p in range(ng): - n_p = min(4, g - 4 * p) - total = n_p * U1 * _S1 - fc2 = [0] * num_ctas - for c in range(num_ctas): - if p >= lag: - fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes - if p == ng - 1 and not ( - early0 and p == 0 and m_max >= 2 - ): # early0: group 0's ranges the same for every G - for q in range(max(0, ng - lag), ng): - fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) - for c in range(num_ctas): - if no_fc1[c]: - fc2[c] += BIG_W - items = [0] * num_ctas - if total >= _S1 * num_ctas: - - def counts_for(limit): - return [ - 0 if no_fc1[c] else max(_S1, (limit - f) // _ITEM_BYTES) - for c, f in enumerate(fc2) - ] - - lo, hi = 0, max(f for c, f in enumerate(fc2) if not no_fc1[c]) + total * _ITEM_BYTES - while lo < hi: - mid = (lo + hi) // 2 - if sum(counts_for(mid)) >= total: - hi = mid - else: - lo = mid + 1 - items = counts_for(lo) - extra = sum(items) - total - while extra > 0: - c = min( - (c for c in range(num_ctas) if items[c] > _S1), - key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), - ) - items[c] -= 1 - extra -= 1 - elif fine: - # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 - lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) - for _ in range(100): - mid = 0.5 * (lo_f + hi_f) - if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: - hi_f = mid - else: - lo_f = mid - wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] - sw = sum(wts) - items = _fine_ranges([w * total / sw for w in wts], total) - else: - chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) - for c in chosen: - items[c] = _S1 - i0 = 0 - for c in range(num_ctas): - if items[c]: - plan[g, c, p] = i0 | ((i0 + items[c]) << 16) - i0 += items[c] - assert i0 == total - return plan, max(nbs) - - -class M3MoeState: + return _shared_make_phase_plan( + num_ctas=num_ctas, + m_max=m_max, + lag=lag, + tail_w=tail_w, + early0=early0, + fine=fine, + heavy_fc2_only=heavy_fc2_only, + geometry=MoePlanGeometry( + top_k=TOP_K, + hidden=HIDDEN, + u1=U1, + stages=_S1, + row2_bytes=_ROW2_BYTES, + item_bytes=_ITEM_BYTES, + ), + fine_ranges_fn=_fine_ranges, + ) + + +class M3MoeState(M3MoeStateBase): """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated @@ -408,6 +242,15 @@ class M3MoeState: fused_router = True + BINDINGS = MoeStateBindings( + inter=INTER, + shared_inter=SHARED_INTER, + engine_path=_ENGINE_PATH, + compiled=_compiled, + pretiled=_pretiled, + cache_module_config=True, + ) + def __init__( self, device: torch.device, @@ -556,227 +399,14 @@ def __init__( self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) self._fn = None - def layer( - self, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - e_score_correction_bias, - router_weight=None, - fc31_input_scale=None, - shared_weights=None, - ) -> "M3MoeLayer": - """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). - ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [3072, 6144], its swizzled UE8M0 - scales, down e4m3 [6144, 768], its swizzled UE8M0 scales).""" - layer = M3MoeLayer( - self, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - e_score_correction_bias, - router_weight, - fc31_input_scale, - ) - if self.shared: - layer.shared_weights = _shared_views(shared_weights, "m3_moe") - layer.shared_mxfp8 = tuple(shared_weights) - if self.mod.SH_PRETILE: - # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy - gu8, _, dn8, _ = layer.shared_weights - layer.shared_weights = ( - gu8, - _pretiled(*layer.shared_weights), - dn8, - self._dummy.view(torch.uint8), - ) - return layer - - def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: - m = self.m - if ( - hidden is None - or hidden.dtype != torch.bfloat16 - or tuple(hidden.shape) != (m, HIDDEN) - or not hidden.is_contiguous() - ): - raise ValueError( - f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" - ) - fz = ( - hidden.view(-1).view(torch.int32), - layer.router_weight.view(-1), - layer.fc31_input_scale.view(-1), - self.lgbuf, - ) - lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy - if out is None: - out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) - sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 - args = ( - layer.w3_w1_weight, - layer.w3_w1_weight_scale.view(-1), - layer.w2_weight, - layer.w2_weight_scale.view(-1), - lgv, - layer.bias, - xq, - xs, - layer.escal, - self.plan.view(-1), - self.partials, - self.hbuf32, - self.epochs, - out.view(-1), - self.route_dbg.view(-1), - self.tdbg.view(-1), - *fz, - *sh, - self.hx, - self.dnx, - self._ar_off if arp is None else arp, - layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), - layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), - self.img32, - self.cnt, - ) - if self._fn is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("m3_moe compiles on its first call: call it once before capture") - import cutlass.cute as cute - - aligns = [ - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 4, - 16, - 16, - 16, - 16, - 16, - 4, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - ] - sig = [_view(t, a) for t, a in zip(args, aligns)] - # keyed on the config the kernel module was built with (not self.config), so a module built with any other - # trace-time config can never reuse another build's compiled kernel - key = (_ENGINE_PATH,) + tuple(sorted(self.mod.M3_CONFIG.items())) - fn = _compiled.get(key) - if fn is None: - fn = _compiled[key] = cute.compile( - self.mod.m3_moe, - *sig, - cute.runtime.make_fake_stream(), - options="--enable-tvm-ffi", - ) - self._fn = fn - self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) - return out - - -class M3MoeLayer: - """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in - place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading - weights).""" - - shared_weights = None # the kernel's view of the shared expert (a shared state) - shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given - def __init__( - self, - state, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - bias, - router_weight=None, - fc31_input_scale=None, - ): - e = NUM_EXPERTS - expected = { - "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), - "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), - "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), - "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), - } - for name, (t, shape) in expected.items(): - if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): - raise ValueError( - f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" - ) - self.state = state - self.w3_w1_weight = w3_w1_weight.view(torch.uint8) - self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) - self.w2_weight = w2_weight.view(torch.uint8) - self.w2_weight_scale = w2_weight_scale.view(torch.uint8) - self.escal = torch.stack( - [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] - ).float() - self.escal = self.escal.reshape(6, e).contiguous() - self.bias = bias.float().contiguous() - if ( - router_weight is None - or fc31_input_scale is None - or router_weight.dtype != torch.float32 - or tuple(router_weight.shape) != (e, HIDDEN) - or not router_weight.is_contiguous() - or fc31_input_scale.dtype != torch.float32 - or fc31_input_scale.numel() != 1 - ): - raise ValueError( - "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" - ) - self.router_weight = router_weight - self.fc31_input_scale = fc31_input_scale +def _shared_views(shared_weights, who): + return shared_views(shared_weights, who, shared_inter=SHARED_INTER) - def __call__( - self, - hidden: torch.Tensor, - out: Optional[torch.Tensor] = None, - arp: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output - into an all-reduce exchange.""" - return self.state.run(self, hidden, out, arp=arp) + +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24): + return pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta, shared_inter=SHARED_INTER) + + +def _fine_ranges(weights, total: int): + return fine_ranges(weights, total, stages=_S1) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py index 24a73860a7f0..3cb0d6e97d5b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_moe/op16.py @@ -40,6 +40,7 @@ from .op import ( HIDDEN, + INTER, SHARED_CTAS, SHARED_INTER, U1, @@ -222,6 +223,7 @@ def layer( e_score_correction_bias, router_weight, fc31_input_scale, + inter=INTER, ) layer.shared_weights = _shared_views(shared_weights, "m3_moe16") layer.shared_mxfp8 = tuple(shared_weights) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py index 48a7827f59b3..a89f96fe0b10 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_prequant.py @@ -1,111 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Hand-off of an MXFP8-quantized activation from the kernel that produced it to the next MXFP8 GEMM. -A producer that already holds a BF16 row in registers (the multicast all-reduce + RMSNorm, the split-K attention -merge) also writes its MXFP8 form (E4M3 rows + UE8M0 scales in the 128x4-swizzled layout of -``trtllm::mxfp8_quantize(x, True)``) and publishes it here keyed by the BF16 tensor. The MXFP8 linear whose input -is that tensor takes it and skips its own quantize kernel. Every take clears the slot, so a hand-off is only ever -consumed by the first MXFP8 linear that runs after the producer. - -Module state shared across files (one instance each per process): - -* ``SLOT`` (PrequantSlot): published by ``m3_mc_allreduce`` (``quant=True``), the sparse-decode merge and the fused - decode attention (MXFP8 output); cleared by the sparse decode before it runs; taken by the MXFP8 linear dispatch - and by the one-launch QKV + attention + o_proj path (which also peeks ``pending``). -* ``AR_PUSH`` (ArPushSlot): armed by the fast forward around the attention call; the o_proj GEMV takes the target - and marks ``pushed``; the fast forward checks ``take_pushed`` before the attention boundary all-reduce. -* ``MOE_PUSH`` (MoePushSlot): armed by the fast forward before the MoE; the routed-expert engine takes ``engine`` - and the shared expert's down projection takes ``shared``, each marks its pushed output; the fast forward - checks ``take_pushed`` before the MoE boundary all-reduce. -""" - - -class PrequantSlot: - def __init__(self): - self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) - - def publish(self, x, xq, xsf): - self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) - - def clear(self): - self.pending = None - - def take(self, x): - """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" - pending, self.pending = self.pending, None - if ( - pending is None - or x.dim() != 2 - or x.data_ptr() != pending[0] - or x.shape[0] != pending[1] - or x.shape[1] != pending[2] - ): - return None - return pending[3], pending[4] +"""Target-local activation hand-off state.""" +from ...ops.gemm.prequant import ArPushSlot, MoePushSlot, PrequantSlot SLOT = PrequantSlot() - - -class ArPushSlot: - """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. - - The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output - it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" - - def __init__(self): - self.target = None - self.pushed = None # data_ptr of the pushed output - - def arm(self, target): - self.target = target - self.pushed = None - - def take_target(self): - target, self.target = self.target, None - return target - - def take_pushed(self, x): - pushed, self.pushed = self.pushed, None - return pushed is not None and x.data_ptr() == pushed - - AR_PUSH = ArPushSlot() - - -class MoePushSlot: - """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, - ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv - ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" - - def __init__(self): - self.engine = None - self.shared = None - self.engine_pushed = None # data_ptr of the pushed routed output - self.shared_pushed = None # data_ptr of the pushed shared output - - def arm(self, engine, shared): - self.engine, self.shared = engine, shared - self.engine_pushed = self.shared_pushed = None - - def disarm(self): - self.engine = self.shared = None - - def take_engine(self): - engine, self.engine = self.engine, None - return engine - - def take_shared(self): - shared, self.shared = self.shared, None - return shared - - def take_pushed(self, shared, routed): - """(routed pushed, shared pushed) for these outputs; clears the marks.""" - e = routed is not None and self.engine_pushed == routed.data_ptr() - s = shared is not None and self.shared_pushed == shared.data_ptr() - self.engine_pushed = self.shared_pushed = None - return e, s - - MOE_PUSH = MoePushSlot() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py index 49bcef462239..07010dceb26b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_select_long_dsl.py @@ -1,299 +1,48 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). +"""Target-specific select policy and independent compilation cache.""" -Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp4_select_long``, -one 1024-thread CTA per row), from a CTA of 384 threads: +import sys +from typing import Dict, Tuple, cast -* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and - mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's - valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned - key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped - block and masks the result, so a thread's loads issue back to back. -* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of - the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. -* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by - that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate - against the whole row instead. -* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. - -The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait -(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are -compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and -loads them before the executor's warmup, so a decode step never compiles one. -""" - -from typing import Dict, Tuple - -import cutlass -import cutlass.cute as cute +import cutlass as cutlass +import cutlass.cute as cute # noqa: F401 — used by the bound family implementation import torch -from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor -from cutlass.utils.smem_allocator import SmemAllocator +from cutlass.cute.runtime import make_fake_compact_tensor as make_fake_compact_tensor +from cutlass.cute.runtime import make_fake_stream as make_fake_stream +from cutlass.cute.runtime import make_fake_tensor as make_fake_tensor -from tensorrt_llm._torch.kernels.blackwell.utils import ( - griddepcontrol_launch_dependents, - griddepcontrol_wait, +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + long_select_host as _select_shared, ) -SELECT_THREADS = 384 -SELECT_WARPS = SELECT_THREADS // 32 -# 1M tokens + draft tokens at 128 tokens per block. -MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row -MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging -TOPK = 16 # capacity: fixed: MSA's top-k -NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) -# The key of -inf (the reference's 0xff800000): never emitted. -NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size -INIT_SCORE = 1.0e30 -LOCAL_SCORE = 1.0e29 -# (keys per thread, rows up to this many valid blocks) -ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) - -# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare -# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group -# minima [16], per-warp candidate counts [SELECT_WARPS]. -_CAND = 0 -_PICK_KEY = _CAND + 2 * MAX_CAND -_PICK_BLOCK = _PICK_KEY + TOPK -_WARP_MIN = _PICK_BLOCK + TOPK -_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 -_WARP_CNT = _GROUP_MIN + 16 -SMEM_WORDS = _WARP_CNT + SELECT_WARPS -SMEM_BYTES = 4 * SMEM_WORDS - - -def _key(score): - """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" - bits = cutlass.Float32(score).bitcast(cutlass.Int32) - return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) - - -def _imin(a, b): - return cutlass.select_(b < a, b, a) - - -def _imax(a, b): - return cutlass.select_(b > a, b, a) - - -def _forced_key(scores, head, block, query, init_blocks, local_start): - raw = scores[head, block, query] - s = cutlass.select_( - block >= local_start, - cutlass.Float32(LOCAL_SCORE), - cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), - ) - return _key(s) - - -@cute.jit -def _select_row( - row, - scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides - n_valid: cute.Tensor, # int32 [>= total_q] - out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] - init_blocks, - local_blocks, - num_kv_heads, - smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory - bar_id: cutlass.Constexpr, - loaded_trigger: cutlass.Constexpr, - raw_prewait, - n_valid_prewait: cutlass.Constexpr, -): - """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier - ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the - row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by - the caller before its grid wait; otherwise it is read here.""" - tid, _, _ = cute.arch.thread_idx() - lane = tid % cutlass.Int32(32) - warp = tid // cutlass.Int32(32) - query = row // num_kv_heads - head = row % num_kv_heads - if cutlass.const_expr(n_valid_prewait): - raw = raw_prewait - else: - raw = n_valid[query] - valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) - local_start = _imax(raw - local_blocks, cutlass.Int32(0)) - cand = cute.make_tensor( - cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) - ) - - if tid < TOPK: - smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) - smem[_PICK_BLOCK + tid] = cutlass.Int32(0) - - if valid <= TOPK: - # Every valid block is picked (thread t holds block t). - if tid < valid: - smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) - smem[_PICK_BLOCK + tid] = tid - if cutlass.const_expr(loaded_trigger): - griddepcontrol_launch_dependents() - for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): - items = ITEM_CLASSES[ci][0] - lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] - hi = ITEM_CLASSES[ci][1] - if (valid > lo) & (valid <= hi): - key = cute.make_rmem_tensor(items, cutlass.Int32) - local_min = cutlass.Int32(NO_KEY) - for s in cutlass.range_constexpr(items): - # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. - block = tid + cutlass.Int32(s * SELECT_THREADS) - k = _forced_key( - scores, head, _imin(block, valid - 1), query, init_blocks, local_start - ) - key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) - local_min = _imin(local_min, key[s]) - local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) - if lane < 16: - smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - if cutlass.const_expr(loaded_trigger): - griddepcontrol_launch_dependents() - if tid < 16: - g = cutlass.Int32(NO_KEY) - for w in cutlass.range_constexpr(SELECT_WARPS): - g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) - smem[_GROUP_MIN + tid] = g - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - bound = cutlass.Int32( - -2147483648 - ) # the smallest key: signed keys, unlike the reference's unsigned ones - for g in cutlass.range_constexpr(16): - bound = _imax(bound, smem[_GROUP_MIN + g]) - - # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. - cnt = cutlass.Int32(0) - for s in cutlass.range_constexpr(items): - cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) - incl = cnt - for d in cutlass.range_constexpr(5): - up = cute.arch.shuffle_sync( - incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) - ) - incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) - if lane == 31: - smem[_WARP_CNT + warp] = incl - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - before = cutlass.Int32(0) - total = cutlass.Int32(0) - for w in cutlass.range_constexpr(SELECT_WARPS): - c = smem[_WARP_CNT + w] - before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) - total += c - slot = before + incl - cnt - for s in cutlass.range_constexpr(items): - if key[s] <= bound: - if slot < MAX_CAND: - cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( - tid + cutlass.Int32(s * SELECT_THREADS) - ) - slot += 1 - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - - if total <= MAX_CAND: - for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): - mine = cand[i] - rank = cutlass.Int32(0) - for j in cutlass.range(total, unroll=8): - rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) - if rank < TOPK: - smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) - smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) - else: - # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. - for s in cutlass.range_constexpr(items): - if key[s] <= bound: - block = tid + cutlass.Int32(s * SELECT_THREADS) - rank = cutlass.Int32(0) - for j in cutlass.range(valid, unroll=1): - kj = _forced_key(scores, head, j, query, init_blocks, local_start) - less = (kj < key[s]) | ((kj == key[s]) & (j < block)) - rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) - if rank < TOPK: - smem[_PICK_KEY + rank] = key[s] - smem[_PICK_BLOCK + rank] = block - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - - if warp == 0: - # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. - v = cutlass.Int32(NO_KEY) - if lane < TOPK: - pk = smem[_PICK_KEY + lane] - if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): - v = smem[_PICK_BLOCK + lane] - for size_log in cutlass.range_constexpr(1, 6): - for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): - other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) - ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) - lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) - v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) - if lane < TOPK: - out[query, head, lane] = cutlass.select_( - v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v - ) - - -class _SelectLongDslKernel: - """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count - before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because - ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's - first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it - (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" +from ...kernels.attention.m3_select_long_dsl import _CAND as _CAND +from ...kernels.attention.m3_select_long_dsl import _GROUP_MIN as _GROUP_MIN +from ...kernels.attention.m3_select_long_dsl import _PICK_BLOCK as _PICK_BLOCK +from ...kernels.attention.m3_select_long_dsl import _PICK_KEY as _PICK_KEY +from ...kernels.attention.m3_select_long_dsl import _WARP_CNT as _WARP_CNT +from ...kernels.attention.m3_select_long_dsl import _WARP_MIN as _WARP_MIN +from ...kernels.attention.m3_select_long_dsl import INIT_SCORE as INIT_SCORE +from ...kernels.attention.m3_select_long_dsl import ITEM_CLASSES as ITEM_CLASSES +from ...kernels.attention.m3_select_long_dsl import LOCAL_SCORE as LOCAL_SCORE +from ...kernels.attention.m3_select_long_dsl import MAX_BLOCKS as MAX_BLOCKS +from ...kernels.attention.m3_select_long_dsl import MAX_CAND as MAX_CAND +from ...kernels.attention.m3_select_long_dsl import NEG_INF_KEY as NEG_INF_KEY +from ...kernels.attention.m3_select_long_dsl import NO_KEY as NO_KEY +from ...kernels.attention.m3_select_long_dsl import SELECT_THREADS as SELECT_THREADS +from ...kernels.attention.m3_select_long_dsl import SELECT_WARPS as SELECT_WARPS +from ...kernels.attention.m3_select_long_dsl import SMEM_BYTES as SMEM_BYTES +from ...kernels.attention.m3_select_long_dsl import SMEM_WORDS as SMEM_WORDS +from ...kernels.attention.m3_select_long_dsl import TOPK as TOPK +from ...kernels.attention.m3_select_long_dsl import _forced_key as _forced_key +from ...kernels.attention.m3_select_long_dsl import _imax as _imax +from ...kernels.attention.m3_select_long_dsl import _imin as _imin +from ...kernels.attention.m3_select_long_dsl import _key as _key +from ...kernels.attention.m3_select_long_dsl import _select_row as _select_row +from ...kernels.attention.m3_select_long_dsl import _SelectLongDslKernel as _SelectLongDslKernel - def __init__( - self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool - ): - self.use_pdl = use_pdl - self.early_trigger = early_trigger - self.loaded_trigger = (not early_trigger) and trigger_phase == 1 - self.n_valid_prewait = n_valid_prewait - - @cute.kernel - def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): - row, _, _ = cute.arch.block_idx() - smem = SmemAllocator() - words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) - raw_prewait = cutlass.Int32(0) - if cutlass.const_expr(self.n_valid_prewait): - raw_prewait = n_valid[row // num_kv_heads] - griddepcontrol_wait() - if cutlass.const_expr(self.early_trigger): - griddepcontrol_launch_dependents() - _select_row( - row, - scores, - n_valid, - out, - init_blocks, - local_blocks, - num_kv_heads, - words, - bar_id=1, - loaded_trigger=self.loaded_trigger, - raw_prewait=raw_prewait, - n_valid_prewait=self.n_valid_prewait, - ) - - @cute.jit - def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): - num_kv_heads = scores.shape[0] - self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( - grid=(scores.shape[2] * num_kv_heads, 1, 1), - block=(SELECT_THREADS, 1, 1), - stream=stream, - use_pdl=self.use_pdl, - ) - - -# The select reads its row's valid-block count before its grid wait instead of after it, which takes one dependent -# global load off the post-wait path; the selection is unchanged bit for bit. Exact because the count -# (MsaMetadata.msa_n_valid_blocks) is written only before the step's first layer, behind a plain launch -# (attention/test_modeling_v2_m3_n_valid_contract.py). N_VALID_PREWAIT = True + #: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} @@ -312,28 +61,7 @@ def kernel_key( def _compile(key: Tuple[bool, bool, int, bool]): - if key not in kernel_cache: - heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() - scores = make_fake_tensor( - cutlass.Float32, - (heads, blocks, rows), - stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), - ) - n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) - out = make_fake_tensor( - cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) - ) - kernel_cache[key] = cute.compile( - _SelectLongDslKernel(*key), - scores, - n_valid, - out, - cutlass.Int32(0), - cutlass.Int32(0), - make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) - return kernel_cache[key] + return _select_shared._compile(_SELECT_TARGET, key) def select_long_dsl( @@ -347,21 +75,22 @@ def select_long_dsl( trigger_phase: int, ) -> None: """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" - _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( - scores, n_valid_blocks, out, init_blocks, local_blocks + return _select_shared.select_long_dsl( + _SELECT_TARGET, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, ) def prebuild(keys, device) -> list: """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" - scores = torch.zeros( - 1, 2048, 4, dtype=torch.float32, device=device - ) # capacity: fixed: a dummy prebuild row, not a size - n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) - out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) - built = [] - for key in keys: - key = kernel_key(*key) - _compile(key)(scores, n_valid, out, 0, 1) - built.append(key) - return built + return _select_shared.prebuild(_SELECT_TARGET, keys, device) + + +_SELECT_TARGET = cast(_select_shared.LongSelectTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py index 2b7d05ad920f..7a2204394472 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_shard_accept.py @@ -21,6 +21,10 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.accept_workspace import ( + create_accept_workspace, +) + from ..ops.spec.m3_shard_accept import m3_shard_accept as _shard_accept_op #: Verify rows per call, and exchange rows per call half of the workspace. @@ -56,49 +60,7 @@ def _num_sms() -> int: def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: - """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the - local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call - outside CUDA-graph capture, in the same order on every rank.""" - words = 2 * slot_rows * max(mapping.tp_size, 1) - if mapping.tp_size > 1: - from tensorrt_llm._torch.distributed.ops import ( - _make_mnnvl_mcast_buffer, - _mnnvl_workspace_all_succeeded, - ) - from tensorrt_llm._utils import mpi_comm - - session = mpi_comm() - session_group = session.Get_group() - group = session_group.Incl(mapping.tp_group) - session_group.Free() - comm = session.Create_group(group) - group.Free() - if not _mnnvl_workspace_all_succeeded(comm, True): - comm.Free() - raise RuntimeError("exchange workspace creation failed on another rank") - handle = _make_mnnvl_mcast_buffer( - comm, max(words * 8, 4096), mapping, mapping.is_multi_node() - ) - uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) - mc = handle.get_mc_buffer((words,), torch.int64, 0) - uc.zero_() - torch.cuda.synchronize() - if not _mnnvl_workspace_all_succeeded(comm, True): - raise RuntimeError("exchange workspace creation failed on another rank") - mc_ptr, rank = mc.data_ptr(), mapping.tp_rank - else: - uc = torch.zeros(words, dtype=torch.int64, device=device) - mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None - return AcceptWorkspace( - uc=uc, - mc_ptr=mc_ptr, - keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), - counters=torch.zeros(2, dtype=torch.int32, device=uc.device), - rank=rank, - world_size=max(mapping.tp_size, 1), - handle=handle, - comm=comm, - ) + return create_accept_workspace(mapping, device, slot_rows, AcceptWorkspace) def _workspace(mapping, device) -> Optional[AcceptWorkspace]: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py index 0d4dc9cd1b64..0bfe1c46c58c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_spec_tail.py @@ -15,11 +15,21 @@ # The kernels are compiled into TensorRT-LLM (cpp/tensorrt_llm/kernels/minimaxM3/specTail.cu) and reached through the # ops/spec entries. -from ..ops.spec.m3_capture_norm import m3_capture_norm as _capture_norm_op -from ..ops.spec.m3_chunk_norm import m3_chunk_norm as _chunk_norm_op -from ..ops.spec.m3_draft_input import m3_draft_input as _draft_input_op -from ..ops.spec.m3_slot_scatter import m3_slot_scatter as _slot_scatter_op -from ..ops.spec.m3_step_gather import m3_step_gather as _step_gather_op +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( + m3_capture_norm as _capture_norm_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_chunk_norm import ( + m3_chunk_norm as _chunk_norm_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_draft_input import ( + m3_draft_input as _draft_input_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_slot_scatter import ( + m3_slot_scatter as _slot_scatter_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_step_gather import ( + m3_step_gather as _step_gather_op, +) def draft_input( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py index f5a1b229c777..03cae9c855c9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_target_head.py @@ -12,13 +12,30 @@ from __future__ import annotations +import sys from dataclasses import dataclass -from typing import Optional +from typing import Optional, cast import torch -from ..ops.gemm.m3_target_head import CTA_WARPS, MAX_ROWS, encode_chunked, round_trip_ok -from ..ops.gemm.m3_target_head import m3_target_head as _target_head_op +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + target_head_state as _head_shared, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + CTA_WARPS as CTA_WARPS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + MAX_ROWS as MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + encode_chunked as encode_chunked, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + m3_target_head as _target_head_op, # noqa: F401 — used by the bound family implementation +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + round_trip_ok as round_trip_ok, +) @dataclass(eq=False) @@ -43,71 +60,24 @@ class TargetHead: def _num_sms() -> int: - global _NUM_SMS - if _NUM_SMS is None: - _NUM_SMS = torch.cuda.get_device_properties( - torch.cuda.current_device() - ).multi_processor_count - return _NUM_SMS + return _head_shared._num_sms(_HEAD_TARGET) def head_supported(lm_head) -> bool: """A bias-free 2-D BF16 head of whole 16-row tiles whose K the coding covers (a multiple of 64).""" - weight = getattr(lm_head, "weight", None) - if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: - return False - return ( - getattr(lm_head, "bias", None) is None - and weight.shape[0] % 16 == 0 - and weight.shape[1] % 64 == 0 - ) + return _head_shared.head_supported(_HEAD_TARGET, lm_head) def prepare_head(lm_head, grid: Optional[int] = None) -> TargetHead: """Code ``lm_head``'s local shard and check it decodes to the BF16 weight bit for bit (fails closed: a shard that does not is returned with ``coded=False`` and nothing kept). Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_target_head``).""" - weight = lm_head.weight - with torch.no_grad(): - w = weight.detach() - packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) - if not round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): - return TargetHead(weight=weight, coded=False) - if grid is None: - grid = _num_sms() - n = weight.shape[0] - return TargetHead( - weight=weight, - coded=True, - w=packed, - base=base, - exc_off=exc_off, - exc_pos=exc_pos, - exc_w=exc_w.view(-1), - part=torch.zeros(grid * CTA_WARPS * 2 * 4 * 32, dtype=torch.float32, device=weight.device), - counter=torch.zeros((n + 15) // 16, dtype=torch.int32, device=weight.device), - grid=grid, - ) + return _head_shared.prepare_head(_HEAD_TARGET, lm_head, grid) def logits(head: TargetHead, x: torch.Tensor, use_pdl: bool = True) -> torch.Tensor: """This rank's BF16 logits shard [rows, n] of the rows x [rows <= MAX_ROWS, K] (BF16) from the coded shard.""" - assert head.coded and 1 <= x.shape[0] <= MAX_ROWS, (head.coded, x.shape) - if x.stride(1) != 1 or x.stride(0) % 8 or x.data_ptr() % 16: - x = x.contiguous() - n = head.weight.shape[0] - out = torch.empty((x.shape[0], n), dtype=torch.bfloat16, device=x.device) - return _target_head_op( - head.w, - x, - out, - head.part, - head.counter, - head.exc_off, - head.exc_pos, - head.exc_w, - head.base, - n, - head.grid, - use_pdl, - ) + return _head_shared.logits(_HEAD_TARGET, head, x, use_pdl) + + +_HEAD_TARGET = cast(_head_shared.HeadTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py index db0077117fe3..5d99aab162d7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/m3_tokmap_guard.py @@ -17,61 +17,20 @@ from __future__ import annotations -from typing import Callable, Dict, Optional, Tuple +import sys +from typing import Callable, Dict, Optional, Tuple, cast import torch -from tensorrt_llm.bindings.internal import thop -from tensorrt_llm.logger import logger +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + token_map as _shared, +) +from tensorrt_llm.bindings.internal import thop as thop +from tensorrt_llm.logger import logger as logger #: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). -_ARG_NAMES = ( - "qkv_input", - "workspace", - "sequence_lengths", - "spec_decoding_generation_lengths", - "spec_decoding_position_offsets", - "kv_cache_block_offsets", - "host_kv_cache_pool_pointers", - "host_kv_cache_pool_mapping", - "kv_scale_orig_quant", - "kv_scale_quant_orig", - "attention_output_orig_quant", - "rotary_inv_freq", - "rotary_cos_sin", - "mrope_position_deltas", - "layer_idx", - "seq_offset", - "num_heads", - "num_kv_heads", - "head_size", - "tokens_per_block", - "kv_cache_quant_mode", - "max_attention_window_size", - "cyclic_attention_window_size", - "num_tokens", - "batch_beam", - "input_seq_length", - "max_past_kv_length", - "rotary_embedding_dim", - "rotary_embedding_base", - "rotary_embedding_scale_type", - "rotary_embedding_scale", - "rotary_embedding_max_positions", - "position_embedding_type", - "bmm1_scale", - "bmm2_scale", - "fp8_context_fmha", - "predicted_tokens_per_seq", - "attention_chunk_size", - "multi_processor_count", - "total_num_blocks", - "kv_factor", - "need_build_kv_cache_metadata", - "cross_attention", - "skip_fmha_workspace", -) -_INT2_BYTES = 8 # capacity: fixed: an element size +_ARG_NAMES = _shared._ARG_NAMES +_INT2_BYTES = _shared._INT2_BYTES _original: Optional[Callable] = None _max_rows = 0 @@ -91,70 +50,20 @@ def _no_missed(rows: int, reason: str) -> None: _on_missed: Callable[[int, str], None] = _no_missed -def _indexed(device: torch.device) -> torch.device: - """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" - if device.type == "cuda" and device.index is None: - return torch.device("cuda", torch.cuda.current_device()) - return device +_indexed = _shared._indexed def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" - device = _indexed(device) - m = _maps.get(device) - if m is None or m.shape[0] < rows: - reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" - _on_missed(rows, reason) - raise RuntimeError(f"m3_tokmap_guard: {reason}") - return m[:rows] + return _shared._one_request_map(_TARGET, device, rows) def _map_offset(args: dict, rows: int) -> int: - x = args["qkv_input"] - skip = bool(args.get("skip_fmha_workspace", False)) - key = ( - x.dtype, - rows, - int(args["num_heads"]), - int(args["head_size"]), - int(args["rotary_embedding_dim"]), - int(args["num_kv_heads"]), - skip, - ) - off = _offsets.get(key) - if off is None: - layout = thop.get_trtllm_gen_generation_workspace_layout( - x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip - ) - off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) - if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: - raise RuntimeError( - f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " - f"(offset {off}, {size} bytes): the layout changed under this guard" - ) - _offsets[key] = off - return off + return _shared._map_offset(_TARGET, args, rows) def _guarded(*args, **kwargs): - a = dict(zip(_ARG_NAMES, args)) - a.update(kwargs) - rows = int(a["num_tokens"]) - if ( - a.get("spec_decoding_generation_lengths") is not None - and int(a["predicted_tokens_per_seq"]) > 1 - and int(a["batch_beam"]) == 1 - and rows > 1 - and not bool(a.get("cross_attention", False)) - ): - ws = a["workspace"] - src = _one_request_map(ws.device, rows) - off = _map_offset(a, rows) - # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. - dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] - dst.view(torch.int32).view(rows, 2).copy_(src) - _on_taken(rows) - return _original(*args, **kwargs) + return _shared._guarded(_TARGET, *args, **kwargs) def install( @@ -166,20 +75,7 @@ def install( """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this call installed the wrapper.""" - global _original, _max_rows, _on_taken, _on_missed - device = _indexed(device) - if device not in _maps or _maps[device].shape[0] < max_rows: - t = torch.arange(max_rows, dtype=torch.int32, device=device) - _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() - _max_rows = max(_max_rows, max_rows) - _on_taken = on_taken - _on_missed = on_missed - if _original is not None: - return False - _original = thop.trtllm_gen_generation_preprocess - thop.trtllm_gen_generation_preprocess = _guarded - logger.info( - "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " - f"(up to {max_rows} rows)" - ) - return True + return _shared.install(_TARGET, device, max_rows, on_taken, on_missed) + + +_TARGET = cast(_shared.TokenMapTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py index 7fe1bcb50976..17c92b77fe7a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_moe_engine.py @@ -10,21 +10,34 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm.quantization.utils.fp4_utils import ( - get_reorder_rows_for_gated_act_gemm_row_indices, - get_shuffle_matrix_a_row_indices, -) -from ..m3_moe.op import ( - HIDDEN, - INTER, - SHARED_INTER, - TOP_K, - M3MoeLayer, - M3MoeState, - ar_push_desc, - is_supported, -) +from ....ops.moe.reference import _E2M1 as _E2M1 +from ....ops.moe.reference import _F32_MAX as _F32_MAX +from ....ops.moe.reference import _LOG2E as _LOG2E +from ....ops.moe.reference import _RCP6 as _RCP6 +from ....ops.moe.reference import ROUTED_SCALING as ROUTED_SCALING +from ....ops.moe.reference import SHARED_ALPHA as SHARED_ALPHA +from ....ops.moe.reference import SHARED_LIMIT as SHARED_LIMIT +from ....ops.moe.reference import _bf16 as _bf16 +from ....ops.moe.reference import _e2m1_decode as _e2m1_decode +from ....ops.moe.reference import _e2m1_round as _e2m1_round +from ....ops.moe.reference import _e4m3_round as _e4m3_round +from ....ops.moe.reference import _f32 as _f32 +from ....ops.moe.reference import _fc1_epilogue as _fc1_epilogue +from ....ops.moe.reference import _gemm as _gemm +from ....ops.moe.reference import _nvfp4_act as _nvfp4_act +from ....ops.moe.reference import _nvfp4_expert as _nvfp4_expert +from ....ops.moe.reference import _row_perm as _row_perm +from ....ops.moe.reference import _shared_expert as _shared_expert_reference +from ....ops.moe.reference import _ue8m0_up as _ue8m0_up +from ....ops.moe.reference import moe_reference as _moe_reference +from ....ops.moe.reference import moe_route as moe_route +from ....ops.moe.reference import mxfp8_dequant as mxfp8_dequant +from ....ops.moe.reference import mxfp8_quant_dequant as mxfp8_quant_dequant +from ....ops.moe.reference import unswizzle_sf as unswizzle_sf +from ..m3_moe.op import HIDDEN as HIDDEN +from ..m3_moe.op import INTER, SHARED_INTER, M3MoeLayer, M3MoeState, ar_push_desc, is_supported +from ..m3_moe.op import TOP_K as TOP_K __all__ = [ "M3MoeLayer", @@ -39,204 +52,16 @@ "unswizzle_sf", ] -# The model's constants the kernels are built with: routed_scaling_factor, the shared expert's SwiGLU-OAI -# (swiglu_alpha, swiglu_limit, the (up + 1) offset). -ROUTED_SCALING = 2.0 -SHARED_ALPHA = 1.702 -SHARED_LIMIT = 7.0 -_LOG2E = 1.4426950216293334961 -_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum -_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) -_F32_MAX = 3.40282346638528859812e38 - - -def _f32(x: float) -> torch.Tensor: - return torch.tensor(x, dtype=torch.float32) - - -def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: - """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" - lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) - lo = lut[(packed & 0xF).long()] - hi = lut[(packed >> 4).long()] - return torch.stack((lo, hi), dim=-1).flatten(-2) - - -def _e2m1_round(x: torch.Tensor) -> torch.Tensor: - """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" - a = x.abs() - q = torch.full_like(a, 6.0) - for hi, value, inclusive in ( - (5.0, 4.0, True), - (3.5, 3.0, False), - (2.5, 2.0, True), - (1.75, 1.5, False), - (1.25, 1.0, True), - (0.75, 0.5, False), - (0.25, 0.0, True), - ): - q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) - return torch.where(x < 0, -q, q) - - -def _e4m3_round(x: torch.Tensor) -> torch.Tensor: - """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" - return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() - - -def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: - """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" - bits = x.contiguous().view(torch.int32) - code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) - return code.clamp(0, 254) - - -def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: - """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" - flat = sf.reshape(-1).view(torch.uint8) - r = torch.arange(rows, device=flat.device).view(-1, 1) - b = torch.arange(blocks, device=flat.device).view(1, -1) - padded_blocks = (blocks + 3) // 4 * 4 - idx = ( - (r // 128) * (128 * padded_blocks) - + (b // 4) * 512 - + (r % 32) * 16 - + ((r % 128) // 32) * 4 - + (b % 4) - ) - return flat[idx] - - -def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: - """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" - rows, k = q.shape - codes = unswizzle_sf(sf, rows, k // 32).float() - return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( - 32, dim=1 - ) - - -def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: - """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to - fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" - rows, k = x.shape - blocks = x.float().view(rows, k // 32, 32) - code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) - scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) - return (_e4m3_round(blocks / scale) * scale).view(rows, k) - - -def _bf16(x: torch.Tensor) -> torch.Tensor: - return x.to(torch.bfloat16).float() - - -def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" - return (a.double() @ b.double().t()).float() - - -def _row_perm(rows: int, gated: bool) -> torch.Tensor: - """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle - of epilogue tile 128).""" - probe = torch.empty(rows, 1) - perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) - if gated: - perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] - return perm - - -def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: - """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row - order (the global scale is not applied: it is the per-expert alpha).""" - rows, half = w.shape - codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() - vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) - return vals[perm_inv] - - -def moe_route(x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor): - """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" - logits = _gemm(x, router_weight) - sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() - keys = sig + bias - ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] - sk = keys.gather(1, ids) - bias[ids] - total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) - w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) - return ids, _bf16(w) - - -def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: - """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" - rows, k = x.shape - blocks = x.view(rows, k // 16, 16) - vmax = blocks.abs().amax(dim=-1, keepdim=True) - sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) - # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element - # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). - osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) - return (_e2m1_round(blocks * osc) * sff).view(rows, k) - - -def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: - """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 - input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with - act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" - a, c, alpha, beta, lim = (esc[i] for i in range(5)) - gc = torch.where(gate > lim, lim, gate) - gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) - uc = torch.minimum(torch.maximum(up, -lim), lim) - c1 = a * (alpha * _f32(_LOG2E).to(up.device)) - sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) - h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) - blocks = h.view(-1, 16) - amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) - sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) - osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) - return (_e2m1_round((blocks * osc) * c) * sff).view(-1) - - -def _shared_expert(x: torch.Tensor, shared_mxfp8) -> torch.Tensor: - """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the - model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" - gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 - gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) - gate, up = gu[:, :SHARED_INTER], gu[:, SHARED_INTER:] - a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) - ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 - sg = 1.0 / ( - torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 - ) - h = _bf16(ub * (sg * a)) - return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) +def _shared_expert( + x: torch.Tensor, + shared_mxfp8: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], +) -> torch.Tensor: + return _shared_expert_reference(x, shared_mxfp8, shared_inter=SHARED_INTER) def moe_reference(hidden: torch.Tensor, layer: M3MoeLayer, shared: bool) -> torch.Tensor: - """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" - x = hidden.float() - m = x.shape[0] - ids, weights = moe_route(x, layer.router_weight, layer.bias) - xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) - inv31 = torch.argsort(_row_perm(2 * INTER, gated=True)).to(x.device) - inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) - y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) - for e in torch.unique(ids).tolist(): - w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) - w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) - esc = layer.escal[:, e] - for t, k in (ids == e).nonzero().tolist(): - acc1 = _gemm(xr[t : t + 1], w31)[0] - # The buffer holds [w3 (up); w1 (gate)]. - h = _fc1_epilogue(acc1[:INTER], acc1[INTER:], esc) - y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) - acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) - for k in range(TOP_K): - acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() - routed = acc.float() - if not shared: - return routed.to(torch.bfloat16) - return (_shared_expert(x, layer.shared_mxfp8) + _bf16(routed)).to(torch.bfloat16) + return _moe_reference(hidden, layer, shared, inter=INTER, shared_inter=SHARED_INTER) class M3MoeEngine(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py deleted file mode 100644 index 6e0a44f38b30..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/moe/m3_router_tc.py +++ /dev/null @@ -1,119 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""TF32 tensor-core router of the unfused decode MoE (9..16 generation rows), optionally with the NVFP4 activation.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) - -_E4M3_MAX = 448.0 -_E2M1_MAX = 6.0 -#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. -_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) -_TIE_UP = (False, True, False, True, False, True, False) - - -def nvfp4_quantize( - x: torch.Tensor, sf_scale: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], - the pre-rounding scaled values [M, K]). - - Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. - """ - gs = sf_scale.float().reshape(()) - m, k = x.shape - blocks = x.float().view(m, k // 16, 16) - vecmax = blocks.abs().amax(-1) - sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) - out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) - scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) - mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) - tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) - a = scaled.abs().unsqueeze(-1) - code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) - negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) - code = code | (negative.to(torch.uint8) << 3) - packed = code[:, 0::2] | (code[:, 1::2] << 4) - return packed, sf.view(torch.uint8).reshape(-1), scaled - - -def tf32_rna(w: torch.Tensor) -> torch.Tensor: - """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" - bits = w.contiguous().view(torch.int32) - rounded = (bits + 0x1000) & ~0x1FFF - finite = (bits & 0x7F800000) != 0x7F800000 - return torch.where(finite, rounded, bits).view(torch.float32) - - -class M3RouterTc(OpWrapper): - """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. - - ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' - partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result - does not depend on it. With ``sf_scale`` also the NVFP4 activation (``fp4_quantize``'s bytes): returns - ``[logits]`` or ``[logits, x_fp4, x_sf]``. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # The target routes the unfused MoE's pure generation steps of 9..16 rows here (ROUTER_TC_GEN), logits only. - CELLS: tuple[Cell, ...] = ( - Cell( - why="9 rows: a 3-request verify past the MoE engine", spec=dict(m=9, k=6144, fp4=False) - ), - Cell(why="10 rows", spec=dict(m=10, k=6144, fp4=False)), - Cell(why="11 rows", spec=dict(m=11, k=6144, fp4=False)), - Cell(why="12 rows: 3 requests x 4 tokens", spec=dict(m=12, k=6144, fp4=False)), - Cell(why="13 rows: a partial second token tile", spec=dict(m=13, k=6144, fp4=False)), - Cell(why="14 rows", spec=dict(m=14, k=6144, fp4=False)), - Cell(why="15 rows", spec=dict(m=15, k=6144, fp4=False)), - Cell( - why="16 rows: 4 requests x 4 tokens, two full token tiles", - spec=dict(m=16, k=6144, fp4=False), - ), - ) - - note = """ - The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both - sum the same 16 slice partials in the same order. Only the one the test device picks is driven. - """ - - def __call__( - self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool - ) -> list[torch.Tensor]: - return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) - - def reference( - self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool - ) -> tuple[torch.Tensor, ...]: - """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" - with true_fp32_matmul(): - logits = x.float() @ tf32_rna(w).t() - if sf_scale is None: - return (logits,) - packed, sf, _ = nvfp4_quantize(x, sf_scale) - return logits, packed, sf - - def is_valid( - self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool - ) -> None: - assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" - if sf_scale is not None: - value = float(sf_scale.reshape(())) - assert value > 0.0 and value != float("inf"), ( - f"sf_scale must be positive and finite, got {value}" - ) - - -m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py deleted file mode 100644 index 8c2a67356d3f..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/msa_kv_update.py +++ /dev/null @@ -1,89 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""One kernel for MSA's device-side KV-length correction (MsaAttentionMetadata.on_update_kv_lens). - -The torch formulation is ~23 one-CTA kernels (min, cumsum, gathers, div / clamp / remainder, the block-table lookup, -the valid-block counts) and runs at the head of every captured decode step, twice. This kernel computes the same -int32 results in one launch: - - kv_true[b] = min(kv_lens[b], staged[b]) -> seq_lens_out[b], cu_kv[b + 1] = prefix sum (cu_kv[0] = 0) - pos[t] = kv_true[r] - qo[r] + intra[t] (r = q_batch_row[t]) - out_loc[t] = table[r, min(max(pos, 0) // page, ncols - 1)] * page + max(pos, 0) % page - n_valid[t] = (max(pos + 1, 0) + page - 1) // page - -``msa_kv_update`` is the target's ``MsaKvUpdate`` (``tensorrt_llm._torch.pyexecutor.runtime_hooks``): the MSA -metadata calls it with its own buffers and falls back to the torch formulation when it returns False. The kernel is -``torch.ops.trtllm.minimax_m3_msa_kv_update``, called through this target's ``msa_kv_update`` entry. -""" - -from typing import Optional - -import torch - -from ..ops.attention.msa_kv_update import msa_kv_update as _kv_update_op - - -def msa_kv_update( - kv_lens: torch.Tensor, - kv_lens_staged: torch.Tensor, - seq_lens: torch.Tensor, - q_batch_row: torch.Tensor, - q_intra: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - batch: int, - total_q: int, - msa_seq_lens: torch.Tensor, - cu_kv_lens: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - n_valid_blocks: Optional[torch.Tensor], -) -> bool: - """The fused correction (see module doc); False, with nothing written, when an operand is outside what the - kernel covers. Arguments in ``MsaKvUpdate`` order; ``seq_lens`` are the per-request query lengths (``qo``).""" - ts = [ - kv_lens, - kv_lens_staged, - seq_lens, - q_batch_row, - q_intra, - block_table, - msa_seq_lens, - out_cache_loc, - ] - ts += [t for t in (cu_kv_lens, n_valid_blocks) if t is not None] - if ( - any(t.dtype != torch.int32 or not t.is_cuda for t in ts) - or block_table.dim() != 2 - or block_table.stride(1) != 1 - ): - return False - if any(not t.is_contiguous() for t in ts if t is not block_table): - return False - if ( - kv_lens.numel() < batch - or kv_lens_staged.numel() < batch - or seq_lens.numel() < batch - or msa_seq_lens.numel() < batch - or q_batch_row.numel() < total_q - or q_intra.numel() < total_q - or out_cache_loc.numel() < total_q - or (cu_kv_lens is not None and cu_kv_lens.numel() < batch + 1) - or (n_valid_blocks is not None and n_valid_blocks.numel() < total_q) - ): - return False - _kv_update_op( - kv_lens, - kv_lens_staged, - seq_lens, - q_batch_row, - q_intra, - block_table, - tokens_per_block, - batch, - total_q, - msa_seq_lens, - cu_kv_lens, - out_cache_loc, - n_valid_blocks, - ) - return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py deleted file mode 100644 index 932971278c49..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/m3_mx_gemv_kernel.py +++ /dev/null @@ -1,1299 +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. -"""MXFP8 decode GEMV on tcgen05: ``y[M, N] = x[M, K] @ W[N, K]^T`` for M <= 8, block-scaled E4M3 x E4M3. - -The weight is the MMA's A operand (128 output features per CTA), the activation is B (8 token columns), -fp32 accumulation in TMEM. Both operands come in their MXFP8 storage: E4M3 values (K contiguous) and UE8M0 -scales in the 128x4-swizzled layout, whose 512-byte atom per (128 rows, 128 K) is exactly the source of a -``tcgen05.cp`` 32x128b.warpx4 scale copy, so both scale tensors are bulk-copied unchanged. For an activation of at -most 32 rows the atom of k-tile ``k`` sits at byte ``512 k``. - -A cluster of SPLIT CTAs shares one 128-row weight tile; rank r streams the interleaved k-tiles r, r + SPLIT, ... -through a RING-stage ring. Before the grid dependency each CTA issues its first RING weight tiles (and their scale -atoms) and prefetches the rest into L2; after it, the activation slice of the CTA (values and scale atoms) lands in -one buffer. Rank r owns rows [r * 128 / SPLIT, (r + 1) * 128 / SPLIT) of the tile: the other ranks store their fp32 -partials of those rows into its shared memory and arrive on its barrier (release, cluster scope); the owner adds the -SPLIT partials in rank order and stores bf16. - -Warps: 0 weight TMA, 1 activation TMA, 2 TMEM allocation + scale copies + MMA, 3 idle, 4-7 epilogue (warp 4 + i -reads TMEM lanes 32 i .. 32 i + 31). -""" - -from __future__ import annotations - -import cuda.bindings.driver as cuda_driver -import cutlass -import cutlass.cute as cute -import cutlass.experimental.cuda as cuda -from cutlass import dsl_user_op -from cutlass.experimental import primitives as prims - -CTA_M = 128 # weight rows per CTA = the MMA's M # capacity: fixed: a kernel tile constant -MMA_N = 8 # token columns -CTA_K = 128 # one k-tile: 128 E4M3 bytes per row = one 128-byte swizzle row # capacity: fixed: a kernel tile constant -MMA_INST_K = 32 -K_BLOCKS = ( - CTA_K // MMA_INST_K -) # 4 MMAs per k-tile, one per scale byte of the atom # capacity: fixed: a kernel tile constant -THREADS = 256 -SF_ATOM_BYTES = 512 -SF_ATOM_U16 = SF_ATOM_BYTES // 2 -SF_COLS = 4 # TMEM columns of one scale atom -W_TILE_BYTES = CTA_M * CTA_K -X_TILE_BYTES = MMA_N * CTA_K -EVICT_FIRST = 0x12F0000000000000 -A_COLS_PER_MMA = MMA_INST_K // 4 # TMEM columns of one 128 x K32 E4M3 granule of A - -a_dtype = cutlass.Float8E4M3FN -b_dtype = cutlass.Float8E4M3FN - - -def num_k_tiles(k_in: int) -> int: - return k_in // CTA_K - - -def acc_cols(nacc: int, mma_n: int = MMA_N) -> int: - # nacc independent accumulators of MMA_N columns: MMA j of the CTA accumulates into j % nacc, so the M128 x N8 - # MMAs do not form one serial dependency chain; the epilogue adds them. - return max(16, nacc * mma_n) - - -def a_col0(nacc: int, my_tiles: int, mma_n: int = MMA_N) -> int: - # a_tmem: [acc | SFA per k-tile | one SFB slot (reused in issue order) | A: 32 columns (128 x 128 E4M3) per k-tile] - return (acc_cols(nacc, mma_n) + (my_tiles + 1) * SF_COLS + 31) // 32 * 32 - - -def tmem_cols( - ring: int, nacc: int, a_tmem: bool = False, my_tiles: int = 0, mma_n: int = MMA_N -) -> int: - need = acc_cols(nacc, mma_n) + 2 * ring * SF_COLS - if a_tmem: - need = a_col0(nacc, my_tiles, mma_n) + my_tiles * K_BLOCKS * A_COLS_PER_MMA - cols = 32 - while cols < need: - cols *= 2 - return cols - - -@dsl_user_op -def read_globaltimer(*, loc=None, ip=None): - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int64( - _llvm.inline_asm( - _T.i64(), [], "mov.u64 $0, %globaltimer;", "=l", has_side_effects=True, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _mapa_u32(smem_ptr, peer, *, loc=None, ip=None): - """The shared::cluster address of this CTA's shared-memory location in cluster CTA ``peer``.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [smem_ptr.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(peer).ir_value(loc=loc, ip=ip)], - "mapa.shared::cluster.u32 $0, $1, $2;", "=r,r,r", has_side_effects=False, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _st_async_v4(dst, w0, w1, w2, w3, mbar, *, loc=None, ip=None): - """st.async of four int32 words to a shared::cluster address, completing ``mbar`` there by 16 bytes.""" - from cutlass._mlir.dialects import llvm as _llvm - - _llvm.inline_asm( - None, - [cutlass.Int32(dst).ir_value(loc=loc, ip=ip), cutlass.Int32(w0).ir_value(loc=loc, ip=ip), - cutlass.Int32(w1).ir_value(loc=loc, ip=ip), cutlass.Int32(w2).ir_value(loc=loc, ip=ip), - cutlass.Int32(w3).ir_value(loc=loc, ip=ip), cutlass.Int32(mbar).ir_value(loc=loc, ip=ip)], - "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 [$0], {$1, $2, $3, $4}, [$5];", "r,r,r,r,r,r", - has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) # fmt: skip - - -@dsl_user_op -def _ld_shared_v4_f32(smem_ptr, *, loc=None, ip=None): - """ld.shared.v4.f32 of four fp32 at a 16-byte-aligned shared-memory pointer.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - out = _llvm.inline_asm( - _llvm.StructType.get_literal([_T.f32()] * 4), [smem_ptr.toint(loc=loc, ip=ip).ir_value()], - "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];", "=f,=f,=f,=f,r", has_side_effects=True, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) # fmt: skip - return [ - cutlass.Float32(_llvm.extractvalue(_T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4) - ] - - -@dsl_user_op -def _test_wait_cluster(mbar, parity, *, loc=None, ip=None): - """mbarrier.test_wait.parity.acquire.cluster on a barrier of this CTA completed by other CTAs' st.async.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - done = cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [mbar.toint(loc=loc, ip=ip).ir_value(), cutlass.Int32(parity).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .pred p;\n\tmbarrier.test_wait.parity.acquire.cluster.shared::cta.b64 p, [$1], $2;\n\t" - "selp.u32 $0, 1, 0, p;\n\t}", "=r,r,r,~{memory}", has_side_effects=True, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - return done != cutlass.Int32(0) - - -@dsl_user_op -def _absmax(acc, v, *, loc=None, ip=None): - """max(acc, |v|).""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Float32( - _llvm.inline_asm( - _T.f32(), [cutlass.Float32(acc).ir_value(loc=loc, ip=ip), cutlass.Float32(v).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .f32 a;\n\tabs.f32 a, $2;\n\tmax.f32 $0, $1, a;\n\t}", "=f,f,f", has_side_effects=False, - is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _rcp_approx_ftz(x, *, loc=None, ip=None): - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Float32( - _llvm.inline_asm( - _T.f32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _e4m3x4(v0, v1, v2, v3, *, loc=None, ip=None): - """Four f32 to E4M3 (RN, satfinite), packed little-endian into one int32 word (v0 in the low byte).""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (v0, v1, v2, v3)], - "{\n\t.reg .b16 lo, hi;\n\tcvt.rn.satfinite.e4m3x2.f32 lo, $2, $1;\n\t" - "cvt.rn.satfinite.e4m3x2.f32 hi, $4, $3;\n\tmov.b32 $0, {lo, hi};\n\t}", "=r,f,f,f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _ld_volatile_s32(addr, *, loc=None, ip=None): - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", - has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _st_global_b32(addr, value, *, loc=None, ip=None): - from cutlass._mlir.dialects import llvm as _llvm - - _llvm.inline_asm( - None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip)], - "st.global.b32 [$0], $1;", "l,r", has_side_effects=True, is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) # fmt: skip - - -@dsl_user_op -def _bf16x2_word(lo, hi, *, loc=None, ip=None): - """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" - "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" - "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" - "mov.b32 $0, w;\n\t}", "=r,f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _bf16_bits(x, *, loc=None, ip=None): - """bf16(x) (RN) as the fp32 bit pattern of the rounded value.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(x).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .b16 h, z;\n\tcvt.rn.bf16.f32 h, $1;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -def _bf16_round(x): - return _bf16_bits(x).bitcast(cutlass.Float32) - - -# The swigluoai activation of MiniMax-M3's dense MLP, gate * sigmoid(1.702 * gate) * (up + 1) with gate clamped to <= 7 -# and up to [-7, 7], in exactly the instruction sequence the Triton silu_and_mul kernel compiles to (modules/swiglu.py -# on SM107): the gate clamp on the packed bf16 pair, the up clamp as min.xorsign.abs, two rounded multiplies into -# ex2.approx, div.full for the reciprocal, two rounded multiplies, round to bf16. No multiply feeds an add, so the -# explicit .rn forms compile to what the Triton kernel's contractable forms do. -SWIGLU_OAI_ALPHA = 1.702 -SWIGLU_OAI_LIMIT = 7.0 - - -@dsl_user_op -def _swiglu_oai_bf16x2(gate_word, up_word, *, loc=None, ip=None): - """The activation of one packed bf16 pair of gate and up values: its two bf16 results as one int32 word (the low - half from the low halves).""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), - [cutlass.Int32(w).ir_value(loc=loc, ip=ip) for w in (gate_word, up_word)], - "{\n\t.reg .b32 g, glim, gl, gh, ul, uh, one, lim, t0, t1, e0, e1, s0, s1, r0, r1;\n\t.reg .b16 hl, hh;\n\t" - "mov.b32 glim, 0x40E040E0;\n\tmin.bf16x2 g, $1, glim;\n\t" - "mov.b32 {hl, hh}, g;\n\tcvt.f32.bf16 gl, hl;\n\tcvt.f32.bf16 gh, hh;\n\t" - "mov.b32 {hl, hh}, $2;\n\tcvt.f32.bf16 ul, hl;\n\tcvt.f32.bf16 uh, hh;\n\t" - "mov.b32 lim, 0f40E00000;\n\tmin.xorsign.abs.f32 ul, ul, lim;\n\tmin.xorsign.abs.f32 uh, uh, lim;\n\t" - "mov.b32 one, 0f3F800000;\n\tadd.rn.f32 ul, ul, one;\n\tadd.rn.f32 uh, uh, one;\n\t" - "mul.rn.f32 t0, gl, 0fBFD9DB23;\n\tmul.rn.f32 t1, gh, 0fBFD9DB23;\n\t" - "mul.rn.f32 t0, t0, 0f3FB8AA3B;\n\tmul.rn.f32 t1, t1, 0f3FB8AA3B;\n\t" - "ex2.approx.f32 e0, t0;\n\tex2.approx.f32 e1, t1;\n\t" - "add.rn.f32 e0, e0, one;\n\tadd.rn.f32 e1, e1, one;\n\t" - "div.full.f32 s0, one, e0;\n\tdiv.full.f32 s1, one, e1;\n\t" - "mul.rn.f32 r0, s0, gl;\n\tmul.rn.f32 r1, s1, gh;\n\t" - "mul.rn.f32 r0, ul, r0;\n\tmul.rn.f32 r1, uh, r1;\n\t" - "cvt.rn.bf16x2.f32 $0, r1, r0;\n\t}", "=r,r,r", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _rope_bits(c, a, s, b, *, loc=None, ip=None): - """bf16(c * a + s * b) with separately rounded products (no FMA), as fp32 bits: the stock RoPE arithmetic.""" - from cutlass._mlir.dialects import llvm as _llvm - from cutlass._mlir.extras import types as _T - - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(v).ir_value(loc=loc, ip=ip) for v in (c, a, s, b)], - "{\n\t.reg .f32 p, q, r;\n\t.reg .b16 h, z;\n\tmul.rn.f32 p, $1, $2;\n\tmul.rn.f32 q, $3, $4;\n\t" - "add.rn.f32 r, p, q;\n\tcvt.rn.bf16.f32 h, r;\n\tmov.b16 z, 0;\n\tmov.b32 $0, {z, h};\n\t}", "=r,f,f,f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -def _din_load4(table, hidden, hidden_half, word_idx): - """Four int32 words at word_idx of the hidden rows (hidden_half == 1) or the embedding table: one CTA-uniform - base-address select instead of a branch.""" - addr = cutlass.select_( - hidden_half == cutlass.Int32(1), hidden.data_ptr(0).toint(), table.data_ptr(0).toint() - ) - ptr = cutlass.inttoptr( - addr + cutlass.Int64(word_idx) * cutlass.Int64(4), mem_space=1, dtype=cutlass.Int32 - ) - return ptr.load(count=4, alignment=16) - - -def _stamp(timers, timed: bool, bx, slot: int): - # Diagnostic: %globaltimer of milestone `slot` of CTA bx (timers [grid, 8] int64). - if timed: - timers.store(read_globaltimer(), idx=bx * cutlass.Int32(8) + cutlass.Int32(slot)) - - -def _k_tile(rank, i: int, split: int, my_tiles: int, contig: bool): - # contig: rank r streams k-tiles [r * my_tiles, (r + 1) * my_tiles); else the interleaved r, r + split, ... - if contig: - return rank * cutlass.Int32(my_tiles) + cutlass.Int32(i) - return rank + cutlass.Int32(i * split) - - -@cute.kernel -def m3_mx_gemv_kernel( - tma_desc_w: cutlass.GridConstant[cuda.TensorMap], # W [N, K] E4M3 as bytes, box 128 x 128 - tma_desc_wsf: cutlass.GridConstant[ - cuda.TensorMap - ], # W scale atoms [N / 128 * K / 128, 256] u16 - tma_desc_x: cutlass.GridConstant[cuda.TensorMap], # x [M, K] E4M3 as bytes, box 128 x 8 - tma_desc_xsf: cutlass.GridConstant[cuda.TensorMap], # x scale atoms [K / 128, 256] u16 - y: cutlass.Array, # [M * N] bf16, token-major - x_words: cutlass.Array, # x_bf16: the bf16 activation [M, K] as int32 words - timers: cutlass.Array, # [grid, 8] int64 milestones when `timed` - num_tokens: cutlass.Int32, - n_out: cutlass.Constexpr[int], - k_in: cutlass.Constexpr[int], - split: cutlass.Constexpr[int], - ring: cutlass.Constexpr[int], - trigger_early: cutlass.Constexpr[bool], - contig: cutlass.Constexpr[bool], - timed: cutlass.Constexpr[bool], - nacc: cutlass.Constexpr[int], - a_tmem: cutlass.Constexpr[bool], - st_async: cutlass.Constexpr[bool], - x_bf16: cutlass.Constexpr[bool], - ar_ws: cutlass.Int64, # ar_push: the all-reduce workspace's multicast mapping (int32 words) - ar_flags: cutlass.Int64, # ar_push: its call flags (parity of the next call = flags[0] & 1) - ar_rank: cutlass.Int32, - ar_push: cutlass.Constexpr[ - int - ], # 1: push the output rows into the next all-reduce; 2: and trigger after it - ar_world: cutlass.Constexpr[int], - ar_maxt: cutlass.Constexpr[int], - mma_n: cutlass.Constexpr[int], # token columns of the MMA (8, 16 or 48): M <= mma_n - din_i: cutlass.Array, # din: int32 token ids [M] - din_table: cutlass.Array, # din: embedding [V, K / 2] bf16 as int32 words - din_hidden: cutlass.Array, # din: hidden rows [M, K / 2] bf16 as int32 words - din_w: cutlass.Array, # din: [embedding norm | hidden norm] weights [K] bf16 as int32 words - din_f: cutlass.Array, # din: fp32 [eps_e, eps_h] - din_vocab: cutlass.Int32, - din_g: cutlass.Array, # din, flag 1: int64 hidden row of each token, as int32 words - din_res: cutlass.Array, # din, flag 2: the gathered hidden rows out [M, K / 2] bf16 as int32 words - din_flags: cutlass.Int32, # din: 1 = gather the hidden rows, 2 = write them to din_res - rope_cs: cutlass.Array, # rope: fp32 (cos, sin) [positions][64][2] - rope_i: cutlass.Array, # rope: int32 cache lengths [M / T] - rope_bt: cutlass.Array, # rope: int32 K / V page tables [M / T][2][max_blocks] - rope_q: cutlass.Array, # rope: E4M3 Q out [M][hq][128] as uint8 - rope_pool: cutlass.Array, # rope: paged E4M3 cache [pages][hkv][tokens per page][128] as uint8 - rope_scale: cutlass.Array, # rope: fp32 [1] E4M3 quantization scale - rope_dims: cutlass.Array, # rope: int32 [max_blocks, tokens-per-page log2, hq, hkv, T] - din: cutlass.Constexpr[bool], - rope: cutlass.Constexpr[bool], - x_swiglu: cutlass.Constexpr[ - bool - ], # x_bf16: x_words is [M, 2K] gate | up; the activation is swigluoai(gate, up) -): - k_tiles = num_k_tiles(k_in) - x_tile_bytes = mma_n * CTA_K - my_tiles = k_tiles // split - rows_per_rank = CTA_M // split - tx, _, _ = cute.arch.thread_idx() - bx, _, _ = cute.arch.block_idx() - warp_id = cute.arch.warp_idx() - rank = cutlass.Int32(0) - if cutlass.const_expr(split > 1): - rank = cute.arch.block_idx_in_cluster() - tma_ptr_w = tma_desc_w.get_ptr() - tma_ptr_wsf = tma_desc_wsf.get_ptr() - tma_ptr_x = tma_desc_x.get_ptr() - tma_ptr_xsf = tma_desc_xsf.get_ptr() - n_tile = bx // cutlass.Int32(split) - m_offset = n_tile * cutlass.Int32(CTA_M) - - smem_a = cutlass.Array( - cutlass.Uint8, ring * W_TILE_BYTES, space=cutlass.AddressSpace.smem, alignment=1024 - ) - smem_b = cutlass.Array( - cutlass.Uint8, my_tiles * x_tile_bytes, space=cutlass.AddressSpace.smem, alignment=1024 - ) - smem_sfa = cutlass.Array( - cutlass.Uint8, ring * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 - ) - smem_sfb = cutlass.Array( - cutlass.Uint8, my_tiles * SF_ATOM_BYTES, space=cutlass.AddressSpace.smem, alignment=128 - ) - w_full = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) - mma_done = cutlass.Array(cutlass.Int64, ring, space=cutlass.AddressSpace.smem, alignment=8) - act_full = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) - acc_done = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) - reduce_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) - tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) - # Partials of the rows this rank owns: [source rank][row][token] fp32 (its own slot stays unused). Rows are padded - # to mma_n + 4 words: the owner's 16-byte row reads (lane = row) then hit all 32 banks once per 8 lanes; at a - # stride of mma_n they were 8- to 16-way conflicted (4.3 us of the 48-token QKV tail). - mb_stride = mma_n + 4 - mailbox = cutlass.Array( - cutlass.Float32, - split * rows_per_rank * mb_stride, - space=cutlass.AddressSpace.smem, - alignment=16, - ) - # din: per-token RMS factor of the half of x this CTA's k-tiles lie in. - din_rs = None - if cutlass.const_expr(din): - din_rs = cutlass.Array(cutlass.Float32, 8, space=cutlass.AddressSpace.smem, alignment=16) - - if warp_id == 0: - prims.prefetch_tensormap(tma_ptr_w) - prims.prefetch_tensormap(tma_ptr_wsf) - prims.prefetch_tensormap(tma_ptr_x) - prims.prefetch_tensormap(tma_ptr_xsf) - if prims.elect_sync(): - for s in cutlass.range_constexpr(ring): - prims.mbarrier_init(w_full.subview(s), 1) - prims.mbarrier_init(mma_done.subview(s), 1) - prims.mbarrier_init(act_full, 1) - prims.mbarrier_init(acc_done, 1) - if cutlass.const_expr(split > 1 and st_async): - # Completed by the peers' st.async bytes; this CTA's own arrival carries the expected count. - prims.mbarrier_init(reduce_ready, 1) - prims.mbarrier_arrive_expect_tx( - reduce_ready, (split - 1) * rows_per_rank * mma_n * 4 - ) - elif cutlass.const_expr(split > 1): - prims.mbarrier_init(reduce_ready, (split - 1) * rows_per_rank) - if warp_id == 2: - prims.tcgen05_alloc(tmem_ptr_i32, tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n)) - prims.tcgen05_relinquish_alloc_permit() - prims.fence_mbarrier_init() - if cutlass.const_expr(split > 1): - prims.barrier_cluster_arrive_relaxed() - prims.barrier_cluster_wait() - prims.barrier_cta_sync(0) - if tx == 0: - _stamp(timers, timed, bx, 0) - - if warp_id == 0: - if prims.elect_sync(): - for i in cutlass.range_constexpr(ring, my_tiles): - k_global = _k_tile(rank, i, split, my_tiles, contig) - prims.cp_async_bulk_tensor_prefetch( - tma_ptr_w, [k_global * cutlass.Int32(CTA_K), m_offset], [] - ) - if cutlass.const_expr((n_out // CTA_M) * split >= 80): - # Wide grids launch in part only after the predecessor (e.g. the fused attention) frees its SMs, which - # loses their pre-wait weight prefetch: the CTAs of the first half also stage their partner's - # (bx + grid / 2) weights into L2. - if bx < cutlass.Int32((n_out // CTA_M) * split // 2): - pb = bx + cutlass.Int32((n_out // CTA_M) * split // 2) - p_rank = pb % cutlass.Int32(split) - p_m_offset = (pb // cutlass.Int32(split)) * cutlass.Int32(CTA_M) - for i in cutlass.range_constexpr(my_tiles): - p_k = _k_tile(p_rank, i, split, my_tiles, contig) - prims.cp_async_bulk_tensor_prefetch( - tma_ptr_w, [p_k * cutlass.Int32(CTA_K), p_m_offset], [] - ) - for i in cutlass.range_constexpr(my_tiles): - stage = i % ring - if cutlass.const_expr(i >= ring): - while not cute.arch.mbarrier_try_wait( - mma_done.subview(stage).data_ptr(), (i // ring - 1) % 2 - ): - pass - k_global = _k_tile(rank, i, split, my_tiles, contig) - prims.mbarrier_arrive_expect_tx(w_full.subview(stage), W_TILE_BYTES + SF_ATOM_BYTES) - prims.cp_async_bulk_tensor_shared_cta_global( - smem_a.subview(stage * W_TILE_BYTES), - tma_ptr_w, - (k_global * cutlass.Int32(CTA_K), m_offset), - w_full.subview(stage), - l2_cache_hint=EVICT_FIRST, - ) - prims.cp_async_bulk_tensor_shared_cta_global( - smem_sfa.subview(stage * SF_ATOM_BYTES), - tma_ptr_wsf, - (cutlass.Int32(0), n_tile * cutlass.Int32(k_tiles) + k_global), - w_full.subview(stage), - l2_cache_hint=EVICT_FIRST, - ) - elif warp_id == 1: - prims.griddepcontrol(prims.GridDepAction.WAIT) - if cutlass.const_expr(trigger_early): - # Dependents launch once this grid's own dependency is satisfied (not at entry): consumers that read - # step metadata or caches before their own grid wait rely on every grid ahead of them having passed its - # wait, so the chain of releases never runs ahead of a producer still in flight. - prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) - if tx == 32: - _stamp(timers, timed, bx, 1) - if cutlass.const_expr(not x_bf16): - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx(act_full, my_tiles * (x_tile_bytes + SF_ATOM_BYTES)) - for i in cutlass.range_constexpr(my_tiles): - k_global = _k_tile(rank, i, split, my_tiles, contig) - prims.cp_async_bulk_tensor_shared_cta_global( - smem_b.subview(i * x_tile_bytes), - tma_ptr_x, - (k_global * cutlass.Int32(CTA_K), cutlass.Int32(0)), - act_full, - ) - prims.cp_async_bulk_tensor_shared_cta_global( - smem_sfb.subview(i * SF_ATOM_BYTES), - tma_ptr_xsf, - (cutlass.Int32(0), k_global), - act_full, - ) - elif warp_id == 2: - tmem_raw = tmem_ptr_i32.load() - base_col = tmem_raw & 0xFFFF - base_row = tmem_raw >> 16 - sfa_col0 = base_col + acc_cols(nacc, mma_n) - sfb_col0 = sfa_col0 + ring * SF_COLS - idesc = prims.Tcgen05MxInstrDesc.build( - a_dtype=a_dtype, b_dtype=b_dtype, scale_format=1, n_dim=mma_n, m_dim=CTA_M - ) - s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 - if cutlass.const_expr(a_tmem): - # A staged in TMEM: before the grid wait every weight k-tile and its scale atom go on into TMEM as they - # land (the copies of a stage refilled later are committed to its mma_done); once the activation lands, - # its scale atoms and the MMAs, A read from TMEM. - base_row = tmem_raw >> 16 - base_col = tmem_raw & 0xFFFF - sfa_col0 = base_col + acc_cols(nacc, mma_n) - sfb_col0 = sfa_col0 + my_tiles * SF_COLS - a_base = tmem_raw + cutlass.Int32(a_col0(nacc, my_tiles, mma_n)) - s2t_shape, s2t_multicast = prims.S2TCopyMode.S2T_32x128b_WARPX4 - for i in cutlass.range_constexpr(my_tiles): - stage = i % ring - while not cute.arch.mbarrier_try_wait( - w_full.subview(stage).data_ptr(), (i // ring) % 2 - ): - pass - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - desc_a = prims.Tcgen05SmemDesc.build( - smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, - base_offset=0, layout=2, - ) # fmt: skip - desc_sfa = prims.Tcgen05SmemDesc.build( - smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, - base_offset=0, layout=0, - ) # fmt: skip - if prims.elect_sync(): - for kb in cutlass.range_constexpr(K_BLOCKS): - prims.tcgen05_cp( - prims.Tcgen05CpShape.SHAPE_128X256B, - cutlass.inttoptr( - a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), - 6, - cutlass.Int32, - ), - desc_a + ((MMA_INST_K >> 4) * kb), - group=prims.CTAGroup.CTA_1, - ) - prims.tcgen05_cp( - s2t_shape, - cutlass.inttoptr((base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32), - desc_sfa, multicast=s2t_multicast, - ) # fmt: skip - if cutlass.const_expr(i + ring < my_tiles): - prims.tcgen05_commit(mma_done.subview(stage)) - if cutlass.const_expr(x_bf16): - # Completed by a thread's arrive: a suspended try_wait wakes late on those. - while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): - pass - else: - while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): - pass - if tx == 64: - _stamp(timers, timed, bx, 2) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - if prims.elect_sync(): - for i in cutlass.range_constexpr(my_tiles): - desc_sfb = prims.Tcgen05SmemDesc.build( - smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, - base_offset=0, layout=0, - ) # fmt: skip - desc_b = prims.Tcgen05SmemDesc.build( - smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, - base_offset=0, layout=2, - ) # fmt: skip - sfa_tmem = cutlass.inttoptr( - (base_row << 16) | (sfa_col0 + i * SF_COLS), 6, cutlass.Int32 - ) - sfb_tmem = cutlass.inttoptr((base_row << 16) | sfb_col0, 6, cutlass.Int32) - prims.tcgen05_cp(s2t_shape, sfb_tmem, desc_sfb, multicast=s2t_multicast) - for kb in cutlass.range_constexpr(K_BLOCKS): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, - cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), - cutlass.inttoptr( - a_base + cutlass.Int32((i * K_BLOCKS + kb) * A_COLS_PER_MMA), 6, cutlass.Int32 - ), - desc_b + ((MMA_INST_K >> 4) * kb), idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), - i > 0 or kb > 0, - sfa_tmem, sfb_tmem, - ) # fmt: skip - if tx == 64: - _stamp(timers, timed, bx, 3) - if prims.elect_sync(): - prims.tcgen05_commit(acc_done) - else: - if cutlass.const_expr(x_bf16): - # Completed by a thread's arrive: a suspended try_wait wakes late on those. - while not cute.arch.mbarrier_test_wait(act_full.data_ptr(), 0): - pass - else: - while not cute.arch.mbarrier_try_wait(act_full.data_ptr(), 0): - pass - if tx == 64: - _stamp(timers, timed, bx, 2) - for i in cutlass.range_constexpr(my_tiles): - stage = i % ring - while not cute.arch.mbarrier_try_wait( - w_full.subview(stage).data_ptr(), (i // ring) % 2 - ): - pass - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - sfa_tmem = (base_row << 16) | (sfa_col0 + stage * SF_COLS) - sfb_tmem = (base_row << 16) | (sfb_col0 + stage * SF_COLS) - desc_sfa = prims.Tcgen05SmemDesc.build( - smem_sfa.subview(stage * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, - base_offset=0, layout=0, - ) # fmt: skip - desc_sfb = prims.Tcgen05SmemDesc.build( - smem_sfb.subview(i * SF_ATOM_BYTES), leading_byte_offset=16, stride_byte_offset=128, - base_offset=0, layout=0, - ) # fmt: skip - desc_a = prims.Tcgen05SmemDesc.build( - smem_a.subview(stage * W_TILE_BYTES), leading_byte_offset=16, stride_byte_offset=1024, - base_offset=0, layout=2, - ) # fmt: skip - desc_b = prims.Tcgen05SmemDesc.build( - smem_b.subview(i * x_tile_bytes), leading_byte_offset=16, stride_byte_offset=1024, - base_offset=0, layout=2, - ) # fmt: skip - if prims.elect_sync(): - prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), desc_sfa, - multicast=s2t_multicast) # fmt: skip - prims.tcgen05_cp(s2t_shape, cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), desc_sfb, - multicast=s2t_multicast) # fmt: skip - for kb in cutlass.range_constexpr(K_BLOCKS): - inc = ((MMA_INST_K * 8 // 8) >> 4) * kb - j = i * K_BLOCKS + kb - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF8F6F4, prims.CTAGroup.CTA_1, - cutlass.inttoptr(tmem_raw + cutlass.Int32((j % nacc) * mma_n), 6, cutlass.Float32), - desc_a + inc, desc_b + inc, idesc.set_sf_ids(a_sf_id=kb, b_sf_id=kb), j >= nacc, - cutlass.inttoptr(sfa_tmem, 6, cutlass.Int32), cutlass.inttoptr(sfb_tmem, 6, cutlass.Int32), - ) # fmt: skip - prims.tcgen05_commit(mma_done.subview(stage)) - if tx == 64: - _stamp(timers, timed, bx, 3) - if prims.elect_sync(): - prims.tcgen05_commit(acc_done) - elif warp_id >= 4: - lane = tx % 32 - row = (warp_id % 4) * 32 + lane - if cutlass.const_expr(x_bf16): - # The bf16 activation of this CTA's k-tiles to MXFP8 in shared memory, as - # torch.ops.trtllm.mxfp8_quantize(x, True) does (UE8M0 = round-up(qz_amax * rcp(448)), - # E4M3 RN-satfinite): one 32-element block (token qz_t, k-tile qz_i, block qz_kb) per thread and round, - # the E4M3 bytes into the 128-byte-swizzled B tile, the qz_scale byte into the tile's qz_scale atom. - qz_tid = tx - cutlass.Int32(128) - prims.griddepcontrol(prims.GridDepAction.WAIT) - din_half_k = k_in // 2 - din_hidden_half = cutlass.Int32(0) - if cutlass.const_expr(din): - # Draft input x = [rmsnorm(table[ids]) * w_e, rmsnorm(hidden) * w_h] (an out-of-range id embeds to - # zero): this CTA's k-tiles lie in one half, so it needs that half's statistic of every token. 16 - # threads per token, 16-byte loads, all of a thread's loads in flight at once. - din_hidden_half = cutlass.select_( - rank * cutlass.Int32(my_tiles * CTA_K) >= cutlass.Int32(din_half_k), - cutlass.Int32(1), - cutlass.Int32(0), - ) - din_tok = qz_tid // cutlass.Int32(16) - din_part = qz_tid % cutlass.Int32(16) - din_tok_ld = cutlass.select_(din_tok < num_tokens, din_tok, cutlass.Int32(0)) - din_id = din_i.load(idx=din_tok_ld) - din_valid = (din_id >= cutlass.Int32(0)) & (din_id < din_vocab) - din_hrow = din_tok_ld - if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): - din_hrow = din_g.load(idx=din_tok_ld * cutlass.Int32(2)) - din_row = cutlass.select_( - din_hidden_half == cutlass.Int32(1), - din_hrow, - cutlass.select_(din_valid, din_id, cutlass.Int32(0)), - ) - # The first hidden-half CTA of tile 0 reads whole hidden rows here: it also writes them out. - din_copy = ( - ((din_flags & cutlass.Int32(2)) != cutlass.Int32(0)) - & (bx == cutlass.Int32(split // 2)) - & (din_hidden_half == cutlass.Int32(1)) - & (din_tok < num_tokens) - ) - din_base = din_row * cutlass.Int32(din_half_k // 2) - din_ss = cutlass.Float32(0.0) - din_vecs_per_part = din_half_k // 8 // 16 - for din_round in cutlass.range_constexpr(0, din_vecs_per_part, 16): - din_n = min(16, din_vecs_per_part - din_round) - din_loaded = [] - for din_j in cutlass.range_constexpr(din_n): - din_widx = din_base + ( - din_part + cutlass.Int32((din_round + din_j) * 16) - ) * cutlass.Int32(4) - din_loaded.append( - _din_load4(din_table, din_hidden, din_hidden_half, din_widx) - ) - if din_copy: - for din_j in cutlass.range_constexpr(din_n): - din_v = din_loaded[din_j] - din_res.store( - (din_v[0], din_v[1], din_v[2], din_v[3]), - idx=din_tok * cutlass.Int32(din_half_k // 2) - + (din_part + cutlass.Int32((din_round + din_j) * 16)) - * cutlass.Int32(4), - alignment=16, - ) - for din_j in cutlass.range_constexpr(din_n): - for din_e in cutlass.range_constexpr(4): - din_word = cutlass.Int32(din_loaded[din_j][din_e]) - din_lo = (din_word << cutlass.Int32(16)).bitcast(cutlass.Float32) - din_hi = (din_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) - din_ss = din_ss + din_lo * din_lo + din_hi * din_hi - for din_o in cutlass.range_constexpr(4): - din_ss = din_ss + cute.arch.shuffle_sync_bfly(din_ss, offset=1 << din_o) - if din_part == cutlass.Int32(0): - din_eps = din_f.load(idx=din_hidden_half) - din_r = cute.math.rsqrt( - din_ss / cutlass.Float32(din_half_k) + din_eps, fastmath=True - ) - din_keep = din_valid | (din_hidden_half == cutlass.Int32(1)) - din_rs.store( - cutlass.select_(din_keep, din_r, cutlass.Float32(0.0)), idx=din_tok - ) - prims.barrier_cta_sync(3, thread_count=128) - qz_b_words = cutlass.Array(smem_b.data_ptr(0), shape=(my_tiles * x_tile_bytes // 4,), dtype=cutlass.Int32, - alignment=16) # fmt: skip - qz_items = my_tiles * K_BLOCKS * mma_n - qz_rounds = (qz_items + 127) // 128 - # A group's loads first (clamped to a valid row), then its math: one L2 round trip per group of - # QZ_GROUP rounds instead of one per round, without holding every round's data in registers. - QZ_GROUP = qz_rounds if qz_rounds <= 8 else 4 - qz_loaded = [] - for qz_r in cutlass.range_constexpr(qz_rounds): - qz_item = qz_tid + cutlass.Int32(qz_r * 128) - qz_t = qz_item % cutlass.Int32(mma_n) - qz_kb = (qz_item // cutlass.Int32(mma_n)) % cutlass.Int32(K_BLOCKS) - qz_i = qz_item // cutlass.Int32(mma_n * K_BLOCKS) - if cutlass.const_expr(qz_r * 128 + 127 >= qz_items): - qz_i = cutlass.select_( - qz_item < cutlass.Int32(qz_items), qz_i, cutlass.Int32(0) - ) - qz_k_global = rank + qz_i * cutlass.Int32(split) - if cutlass.const_expr(contig): - qz_k_global = rank * cutlass.Int32(my_tiles) + qz_i - qz_t_ld = cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) - qz_x_base = ( - qz_t_ld * cutlass.Int32(k_in if x_swiglu else k_in // 2) - + qz_k_global * cutlass.Int32(CTA_K // 2) - + qz_kb * cutlass.Int32(MMA_INST_K // 2) - ) - qz_wvecs = None - if cutlass.const_expr(din): - # The draft input's raw values (table row of the id / hidden row) and its norm weights. - qz_col = qz_k_global * cutlass.Int32(CTA_K) + qz_kb * cutlass.Int32(MMA_INST_K) - qz_hcol = qz_col - din_hidden_half * cutlass.Int32(din_half_k) - qz_id = din_i.load(idx=qz_t_ld) - qz_hrow = qz_t_ld - if (din_flags & cutlass.Int32(1)) != cutlass.Int32(0): - qz_hrow = din_g.load(idx=qz_t_ld * cutlass.Int32(2)) - qz_row = cutlass.select_( - din_hidden_half == cutlass.Int32(1), - qz_hrow, - cutlass.select_( - (qz_id >= cutlass.Int32(0)) & (qz_id < din_vocab), - qz_id, - cutlass.Int32(0), - ), - ) - qz_src_w = qz_row * cutlass.Int32(din_half_k // 2) + qz_hcol // cutlass.Int32(2) - qz_vecs = [ - _din_load4( - din_table, - din_hidden, - din_hidden_half, - qz_src_w + cutlass.Int32(qz_v * 4), - ) - for qz_v in range(4) - ] - qz_wvecs = [ - din_w.load( - idx=qz_col // cutlass.Int32(2) + cutlass.Int32(qz_v * 4), - vector_size=4, - alignment=16, - ) - for qz_v in range(4) - ] - else: - qz_vecs = [ - x_words.load( - idx=qz_x_base + cutlass.Int32(qz_v * 4), vector_size=4, alignment=16 - ) - for qz_v in range(4) - ] - if cutlass.const_expr(x_swiglu): - # The up half of the same columns, K values after the gate in the row. - qz_wvecs = [ - x_words.load( - idx=qz_x_base + cutlass.Int32(k_in // 2 + qz_v * 4), - vector_size=4, - alignment=16, - ) - for qz_v in range(4) - ] - qz_loaded.append((qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs)) - if cutlass.const_expr((qz_r + 1) % QZ_GROUP == 0 or qz_r + 1 == qz_rounds): - for qz_g in cutlass.range_constexpr(len(qz_loaded)): - qz_item, qz_t, qz_kb, qz_i, qz_vecs, qz_wvecs = qz_loaded[qz_g] - qz_vals = [] - qz_amax = cutlass.Float32(0.0) - if cutlass.const_expr(din): - qz_rs = din_rs.load( - idx=cutlass.select_(qz_t < num_tokens, qz_t, cutlass.Int32(0)) - ) - for qz_v in cutlass.range_constexpr(4): - for qz_e4 in cutlass.range_constexpr(4): - qz_word = cutlass.Int32(qz_vecs[qz_v][qz_e4]) - if cutlass.const_expr(x_swiglu): - qz_word = _swiglu_oai_bf16x2( - qz_word, cutlass.Int32(qz_wvecs[qz_v][qz_e4]) - ) - qz_lo = (qz_word << cutlass.Int32(16)).bitcast(cutlass.Float32) - qz_hi = (qz_word & cutlass.Int32(-65536)).bitcast(cutlass.Float32) - if cutlass.const_expr(din): - # bf16(x * rs * w), the draft-input kernel's rounding. - qz_wword = cutlass.Int32(qz_wvecs[qz_v][qz_e4]) - qz_lo = _bf16_round( - qz_lo - * qz_rs - * (qz_wword << cutlass.Int32(16)).bitcast(cutlass.Float32) - ) - qz_hi = _bf16_round( - qz_hi - * qz_rs - * (qz_wword & cutlass.Int32(-65536)).bitcast( - cutlass.Float32 - ) - ) - qz_vals.append(qz_lo) - qz_vals.append(qz_hi) - qz_amax = _absmax(_absmax(qz_amax, qz_lo), qz_hi) - qz_bits = (qz_amax * _rcp_approx_ftz(cutlass.Float32(448.0))).bitcast( - cutlass.Int32 - ) - qz_e = ( - (qz_bits >> cutlass.Int32(23)) & cutlass.Int32(0xFF) - ) + cutlass.select_( - (qz_bits & cutlass.Int32(0x7FFFFF)) != cutlass.Int32(0), - cutlass.Int32(1), - cutlass.Int32(0), - ) - qz_e = cutlass.select_(qz_e > cutlass.Int32(254), cutlass.Int32(254), qz_e) - qz_scale = ((cutlass.Int32(254) - qz_e) << cutlass.Int32(23)).bitcast( - cutlass.Float32 - ) - qz_scale = cutlass.select_( - qz_amax != cutlass.Float32(0.0), qz_scale, cutlass.Float32(0.0) - ) - qz_q = [_e4m3x4(qz_vals[4 * qz_j] * qz_scale, qz_vals[4 * qz_j + 1] * qz_scale, - qz_vals[4 * qz_j + 2] * qz_scale, qz_vals[4 * qz_j + 3] * qz_scale) - for qz_j in range(8)] # fmt: skip - if qz_item < cutlass.Int32(qz_items): - if qz_t < num_tokens: - qz_row_w = qz_i * cutlass.Int32( - x_tile_bytes // 4 - ) + qz_t * cutlass.Int32(CTA_K // 4) - qz_c0 = qz_kb * cutlass.Int32(2) - qz_b_words.store( - (qz_q[0], qz_q[1], qz_q[2], qz_q[3]), - idx=qz_row_w - + (qz_c0 ^ (qz_t & cutlass.Int32(7))) * cutlass.Int32(4), - alignment=16, - ) - qz_b_words.store( - (qz_q[4], qz_q[5], qz_q[6], qz_q[7]), - idx=qz_row_w - + ((qz_c0 + cutlass.Int32(1)) ^ (qz_t & cutlass.Int32(7))) - * cutlass.Int32(4), - alignment=16, - ) - smem_sfb.store( - cutlass.Uint8(qz_e), - idx=qz_i * cutlass.Int32(SF_ATOM_BYTES) - + (qz_t % cutlass.Int32(32)) * cutlass.Int32(16) - + (qz_t // cutlass.Int32(32)) * cutlass.Int32(4) - + qz_kb, - ) - qz_loaded = [] - # Generic-proxy stores made visible to the MMA's async-proxy reads, then one arrival for all. - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - prims.barrier_cta_sync(2, thread_count=128) - if qz_tid == cutlass.Int32(0): - prims.mbarrier_arrive(act_full) - if cutlass.const_expr(st_async): - while not cute.arch.mbarrier_test_wait(acc_done.data_ptr(), 0): - pass - else: - while not cute.arch.mbarrier_try_wait(acc_done.data_ptr(), 0): - pass - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - # tcgen05.ld widths are powers of two: the accumulator columns in descending power-of-two pieces. - acc_all = [] - ld_col = 0 - for ld_w in (128, 64, 32, 16, 8): - if cutlass.const_expr((mma_n * nacc - ld_col) >= ld_w): - piece = prims.tcgen05_ld( - "32x32b", - cutlass.inttoptr( - tmem_ptr_i32.load() + cutlass.Int32(ld_col), 6, cutlass.Float32 - ), - num=ld_w, - ) - acc_all = acc_all + [piece[c] for c in range(ld_w)] - ld_col += ld_w - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - acc = [] - for t in cutlass.range_constexpr(mma_n): - v = cutlass.Float32(acc_all[t]) - for a in cutlass.range_constexpr(1, nacc): - v = v + cutlass.Float32(acc_all[a * mma_n + t]) - acc.append(v) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - if tx == 128: - _stamp(timers, timed, bx, 4) - n = m_offset + row - if cutlass.const_expr(split == 1): - for t in cutlass.range_constexpr(mma_n): - if cutlass.Int32(t) < num_tokens: - y.store( - cutlass.BFloat16(cutlass.Float32(acc[t])), idx=cutlass.Int32(t * n_out) + n - ) - if cutlass.const_expr(ar_push > 0): - prims.griddepcontrol(prims.GridDepAction.WAIT) - # The parity of each (token, column slice) slot: the all-reduce's CTA t * 8 + slice flag. All loads are - # issued before the first store so they overlap (one L2 round trip, not one per token). - ar_par = [ - _ld_volatile_s32( - ar_flags - + cutlass.Int64(4) - * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) - ) - & cutlass.Int32(1) - for t in range(mma_n) - ] - for t in cutlass.range_constexpr(mma_n): - ar_v = cutlass.Float32(acc[t]) - ar_hi = cute.arch.shuffle_sync_bfly(ar_v, offset=1) - if (lane % cutlass.Int32(2)) == cutlass.Int32(0): - if cutlass.Int32(t) < num_tokens: - ar_off = ( - cutlass.Int64(ar_par[t]) - * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) - + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) - * cutlass.Int64(n_out // 2) - + (cutlass.Int64(n) // cutlass.Int64(2)) - ) - _st_global_b32( - ar_ws + ar_off * cutlass.Int64(4), _bf16x2_word(ar_v, ar_hi) - ) - else: - owner = row // cutlass.Int32(rows_per_rank) - local = row % cutlass.Int32(rows_per_rank) - if cutlass.const_expr(rope): - # Rank r owns head dims [r * P, (r + 1) * P) and their rotation partners 64 + ..., P = rows / 2. - pairs = rows_per_rank // 2 - owner = (row % cutlass.Int32(64)) // cutlass.Int32(pairs) - local = (row // cutlass.Int32(64)) * cutlass.Int32(pairs) + ( - row % cutlass.Int32(64) - ) % cutlass.Int32(pairs) - if owner != rank: - base = (rank * cutlass.Int32(rows_per_rank) + local) * cutlass.Int32(mb_stride) - if cutlass.const_expr(st_async): - dst = _mapa_u32(mailbox.subview(base).data_ptr(), owner) - mb = _mapa_u32(reduce_ready.data_ptr(), owner) - words = [cutlass.Float32(acc[t]).bitcast(cutlass.Int32) for t in range(mma_n)] - for g in cutlass.range_constexpr(mma_n // 4): - _st_async_v4( - dst + cutlass.Int32(16 * g), - words[4 * g], - words[4 * g + 1], - words[4 * g + 2], - words[4 * g + 3], - mb, - ) - else: - for t in cutlass.range_constexpr(mma_n): - prims.mapa(mailbox.subview(base + cutlass.Int32(t)), owner).store( - cutlass.Float32(acc[t]) - ) - prims.mbarrier_arrive( - prims.mapa(reduce_ready, owner), scope=prims.MemScope.CLUSTER - ) - else: - if cutlass.const_expr(st_async): - while not _test_wait_cluster(reduce_ready.data_ptr(), 0): - pass - else: - while not prims.mbarrier_try_wait_parity( - reduce_ready, 0, scope=prims.MBarrierScope.CLUSTER - ): - pass - if lane == 0: - if row // cutlass.Int32(rows_per_rank) == rank: - if (row % cutlass.Int32(rows_per_rank)) == 0: - _stamp(timers, timed, bx, 5) - totals = [] - for g in cutlass.range_constexpr(mma_n // 4): - sums = [cutlass.Float32(0.0)] * 4 - for q in cutlass.range_constexpr(split): - parts = _ld_shared_v4_f32( - mailbox.subview( - (cutlass.Int32(q * rows_per_rank) + local) - * cutlass.Int32(mb_stride) - + cutlass.Int32(4 * g) - ).data_ptr() - ) - sums = [sums[j] + cutlass.select_(rank == cutlass.Int32(q), cutlass.Float32(acc[4 * g + j]), - parts[j]) for j in range(4)] # fmt: skip - for j in cutlass.range_constexpr(4): - totals.append(sums[j]) - if cutlass.const_expr(rope): - # The own-rank mailbox slot is unused: it carries the totals to the RoPE stage. - mailbox.store( - sums[j], - idx=(rank * cutlass.Int32(rows_per_rank) + local) - * cutlass.Int32(mb_stride) - + cutlass.Int32(4 * g + j), - ) - elif cutlass.Int32(4 * g + j) < num_tokens: - y.store( - cutlass.BFloat16(sums[j]), - idx=cutlass.Int32((4 * g + j) * n_out) + n, - ) - if cutlass.const_expr(ar_push > 0): - # Every lane of an owner warp is an owner (rows_per_rank >= 32), so lane pairs hold row pairs. - prims.griddepcontrol(prims.GridDepAction.WAIT) - ar_par = [ - _ld_volatile_s32( - ar_flags - + cutlass.Int64(4) - * (cutlass.Int64(t * 8) + cutlass.Int64(n // cutlass.Int32(n_out // 8))) - ) - & cutlass.Int32(1) - for t in range(mma_n) - ] - for t in cutlass.range_constexpr(mma_n): - ar_hi = cute.arch.shuffle_sync_bfly(totals[t], offset=1) - if (lane % cutlass.Int32(2)) == cutlass.Int32(0): - if cutlass.Int32(t) < num_tokens: - ar_off = ( - cutlass.Int64(ar_par[t]) - * cutlass.Int64(ar_maxt * ar_world * (n_out // 2)) - + (cutlass.Int64(t * ar_world) + cutlass.Int64(ar_rank)) - * cutlass.Int64(n_out // 2) - + (cutlass.Int64(n) // cutlass.Int64(2)) - ) - _st_global_b32( - ar_ws + ar_off * cutlass.Int64(4), - _bf16x2_word(totals[t], ar_hi), - ) - prims.barrier_cta_sync(1, thread_count=128) - if cutlass.const_expr(rope): - # Token t of request t / T sits at cache position kv_len - T + t % T. Q / K heads: GPT-NeoX RoPE over - # the 128-wide head, bf16(cos * x -/+ sin * pair) on the bf16 projection; then E4M3 of x * scale, Q to - # rope_q, K / V into the page of that position. - pairs = rows_per_rank // 2 - rp_e = tx - cutlass.Int32(128) - rp_p = rp_e % cutlass.Int32(pairs) - rp_t = rp_e // cutlass.Int32(pairs) - if rp_t < num_tokens: - rp_max_blocks = rope_dims.load(idx=0) - rp_tpb_log2 = rope_dims.load(idx=1) - rp_hq = rope_dims.load(idx=2) - rp_hkv = rope_dims.load(idx=3) - rp_T = rope_dims.load(idx=4) - rp_d = rank * cutlass.Int32(pairs) + rp_p - rp_a = _bf16_round( - mailbox.load( - idx=(rank * cutlass.Int32(rows_per_rank) + rp_p) * cutlass.Int32(mb_stride) - + rp_t - ) - ) - rp_b = _bf16_round( - mailbox.load( - idx=(rank * cutlass.Int32(rows_per_rank) + cutlass.Int32(pairs) + rp_p) - * cutlass.Int32(mb_stride) - + rp_t - ) - ) - rp_req = rp_t // rp_T - rp_tok = rope_i.load(idx=rp_req) - rp_T + (rp_t - rp_req * rp_T) - rp_head = n_tile - if rp_head < rp_hq + rp_hkv: - rp_c = rope_cs.load(idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2)) - rp_s = rope_cs.load( - idx=(rp_tok * cutlass.Int32(64) + rp_d) * cutlass.Int32(2) - + cutlass.Int32(1) - ) - rp_ra = _rope_bits(rp_c, rp_a, -rp_s, rp_b).bitcast(cutlass.Float32) - rp_rb = _rope_bits(rp_c, rp_b, rp_s, rp_a).bitcast(cutlass.Float32) - rp_a = rp_ra - rp_b = rp_rb - rp_qs = rope_scale.load(idx=0) - rp_w = _e4m3x4( - rp_a * rp_qs, rp_b * rp_qs, cutlass.Float32(0.0), cutlass.Float32(0.0) - ) - rp_lo = cutlass.Uint8(rp_w & cutlass.Int32(0xFF)) - rp_hi = cutlass.Uint8((rp_w >> cutlass.Int32(8)) & cutlass.Int32(0xFF)) - if rp_head < rp_hq: - rp_q = (rp_t * rp_hq + rp_head) * cutlass.Int32(128) + rp_d - rope_q.store(rp_lo, idx=rp_q) - rope_q.store(rp_hi, idx=rp_q + cutlass.Int32(64)) - else: - rp_kv = cutlass.select_( - rp_head < rp_hq + rp_hkv, cutlass.Int32(0), cutlass.Int32(1) - ) - rp_kh = rp_head - rp_hq - rp_kv * rp_hkv - rp_page = rope_bt.load( - idx=(rp_req * cutlass.Int32(2) + rp_kv) * rp_max_blocks - + (rp_tok >> rp_tpb_log2) - ) - # 64-bit byte offset: the pool view can exceed 2^31 bytes. - rp_dst = ( - (cutlass.Int64(rp_page) * cutlass.Int64(rp_hkv) + cutlass.Int64(rp_kh)) - << cutlass.Int64(rp_tpb_log2 + cutlass.Int32(7)) - ) + cutlass.Int64( - (rp_tok & ((cutlass.Int32(1) << rp_tpb_log2) - cutlass.Int32(1))) - * cutlass.Int32(128) - + rp_d - ) - rp_addr = rope_pool.data_ptr(0).toint() + rp_dst - cutlass.inttoptr(rp_addr, mem_space=1, dtype=cutlass.Uint8).store(rp_lo) - cutlass.inttoptr( - rp_addr + cutlass.Int64(64), mem_space=1, dtype=cutlass.Uint8 - ).store(rp_hi) - if tx == 128: - _stamp(timers, timed, bx, 6) - if cutlass.const_expr(ar_push == 2): - # The all-reduce polls the pushed words; it launches after this CTA's rows are out, not at grid end. - prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) - if warp_id == 4: - prims.tcgen05_dealloc( - cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Int32), - tmem_cols(ring, nacc, a_tmem, my_tiles, mma_n), - ) - - -def _bytes_map(t, cols, rows, box_rows, swizzle): - return cuda.create_tensor_map_tiled( - global_address=t.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[cols, rows], - global_strides=[cols // 16], - box_dims=[CTA_K, box_rows], - swizzle=swizzle, - ) - - -def _atom_map(t, atoms): - return cuda.create_tensor_map_tiled( - global_address=t.iterator.toint(), - dtype=cutlass.Uint16, - global_dims=[SF_ATOM_U16, atoms], - global_strides=[SF_ATOM_BYTES // 16], - box_dims=[SF_ATOM_U16, 1], - swizzle=cuda.TensorMapSwizzle.none, - ) - - -@cute.jit -def m3_mx_gemv( - w: cute.Tensor, # [N, K] E4M3 as uint8, K contiguous - wsf: cute.Tensor, # weight scales, 128x4-swizzled, as uint8 - x: cute.Tensor, # [M, K] E4M3 as uint8, M <= 8 - xsf: cute.Tensor, # activation scales, 128x4-swizzled (atom k at byte 512 k), as uint8 - y: cute.Tensor, # [M * N] bf16 - x_words: cute.Tensor, # x_bf16: bf16 x [M, K] as int32 words (x / xsf are then unused placeholders) - timers: cute.Tensor, # [grid * 8] int64, written when `timed` - num_tokens: cutlass.Int32, - n_out: cutlass.Constexpr[int], - k_in: cutlass.Constexpr[int], - split: cutlass.Constexpr[int], - ring: cutlass.Constexpr[int], - trigger_early: cutlass.Constexpr[bool], - contig: cutlass.Constexpr[bool], - timed: cutlass.Constexpr[bool], - nacc: cutlass.Constexpr[int], - a_tmem: cutlass.Constexpr[bool], - st_async: cutlass.Constexpr[bool], - x_bf16: cutlass.Constexpr[bool], - ar_ws: cutlass.Int64, - ar_flags: cutlass.Int64, - ar_rank: cutlass.Int32, - ar_push: cutlass.Constexpr[int], - ar_world: cutlass.Constexpr[int], - ar_maxt: cutlass.Constexpr[int], - mma_n: cutlass.Constexpr[int], - use_pdl: cutlass.Constexpr[bool], - stream: cuda_driver.CUstream, - din_i: cute.Tensor, # din / rope inputs: see m3_mx_gemv_kernel (placeholders when off) - din_table: cute.Tensor, - din_hidden: cute.Tensor, - din_w: cute.Tensor, - din_f: cute.Tensor, - din_vocab: cutlass.Int32, - din_g: cute.Tensor, - din_res: cute.Tensor, - din_flags: cutlass.Int32, - rope_cs: cute.Tensor, - rope_i: cute.Tensor, - rope_bt: cute.Tensor, - rope_q: cute.Tensor, - rope_pool: cute.Tensor, - rope_scale: cute.Tensor, - rope_dims: cute.Tensor, - din: cutlass.Constexpr[bool], - rope: cutlass.Constexpr[bool], - x_swiglu: cutlass.Constexpr[bool] = False, -) -> None: - # x_swiglu's up half rides the slot din uses for its norm weights, and it activates the bf16 x. - assert not (din and x_swiglu) and (x_bf16 or not x_swiglu), ( - "x_swiglu needs x_bf16 and excludes din" - ) - k_tiles = num_k_tiles(k_in) - tma_desc_w = _bytes_map(w, k_in, n_out, CTA_M, cuda.TensorMapSwizzle.s128b) - tma_desc_wsf = _atom_map(wsf, (n_out // CTA_M) * k_tiles) - tma_desc_x = _bytes_map(x, k_in, num_tokens, mma_n, cuda.TensorMapSwizzle.s128b) - tma_desc_xsf = _atom_map(xsf, k_tiles) - m3_mx_gemv_kernel( - tma_desc_w, tma_desc_wsf, tma_desc_x, tma_desc_xsf, y, x_words, timers, num_tokens, n_out, k_in, split, ring, - trigger_early, contig, timed, nacc, a_tmem, st_async, x_bf16, ar_ws, ar_flags, ar_rank, ar_push, ar_world, - ar_maxt, mma_n, - din_i, din_table, din_hidden, din_w, din_f, din_vocab, din_g, din_res, din_flags, rope_cs, - rope_i, rope_bt, rope_q, rope_pool, rope_scale, rope_dims, din, rope, x_swiglu, - ).launch( - grid=((n_out // CTA_M) * split, 1, 1), - block=(THREADS, 1, 1), - cluster=(split, 1, 1), - stream=stream, - use_pdl=use_pdl, - ) # fmt: skip diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py index 3e6b07326658..727a01ccbdc2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/mx_gemv/op.py @@ -21,12 +21,19 @@ from __future__ import annotations +import sys import threading -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Tuple, cast import torch -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + mxfp8_core as _core, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + mxfp8_host as _host, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl as get_env_enable_pdl MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) MMA_WIDTHS = ( @@ -39,16 +46,11 @@ def _mma_n(m: int) -> int: - return next(w for w in MMA_WIDTHS if m <= w) + return _host.gemv_mma_n(m, mma_widths=MMA_WIDTHS) def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: - return ( - ring * (128 * 128 + 512) - + tiles * (mma_n * 128 + 512) - + split * (128 // split) * (mma_n + 4) * 4 - + 1024 - ) + return _host.gemv_smem_bytes(ring, tiles, split, mma_n) _lock = threading.Lock() @@ -93,45 +95,41 @@ def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: def _dummy_timers(device) -> torch.Tensor: - """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" - t = _DUMMY.get(device) - if t is None: - t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) - return t + "Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on." + return _host.gemv_dummy_timers(device, dummy=_DUMMY) def _arg(t: torch.Tensor, align: int = 16): - from cutlass.cute.runtime import from_dlpack - - return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + return _host.gemv_arg(t, align) def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: - cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None - if cfg is None and m is not None and m > 8: - cfg = _CONFIGS_WIDE.get((n, k)) - cfg = cfg or _CONFIGS.get((n, k)) - # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. - if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: - return None - return cfg + return _host.gemv_get_config( + n, + k, + m, + configs_by_m=_CONFIGS_BY_M, + configs_wide=_CONFIGS_WIDE, + configs=_CONFIGS, + strided=False, + strided_max_rows=0, + strided_shapes=(), + ) def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: - """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" - split = cfg[0] - if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: - return True - tiles = k // 128 // split - mma_n = _mma_n(m) - # GB300 exposed time vs the FlashInfer prequantized GEMM (behind a PDL predecessor), at 24 / 48 tokens: QKV 3.76 / - # 4.63 vs 7.46 / 11.44 us, o_proj 2.74 / 3.33 vs 4.46 / 5.84, shared gate_up 4.57 / 5.71 vs 7.16 / 11.25. Behind - # only for K = 768 (unsplit shared down) at 17-32 tokens: 3.1 vs 2.9. - # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; - # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. - if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): - return True - return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + "Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds." + return _host.gemv_declines( + n, + k, + m, + x_bf16, + cfg, + max_tokens=MAX_TOKENS, + smem_limit=_SMEM_LIMIT, + select_mma_n=_mma_n, + smem_bytes=_smem_bytes, + ) def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: @@ -155,119 +153,27 @@ def _resolve( ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the rank are runtime arguments.""" - cfg = cfg or get_config(n, k, m) - if cfg is None or _declines(n, k, m, x_bf16, cfg): - return None - split, ring = cfg[0], cfg[1] - contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG - nacc = int(cfg[3]) if len(cfg) > 3 else _NACC - a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM - st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC - if nacc != 1: - # More than one TMEM accumulator is broken on SM107 (nondeterministic or garbage - # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. - return None - ring = min(ring, k // 128 // split) - tiles = k // 128 // split - mma_n = _mma_n(m) - while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: - ring -= 1 - if a_tmem: - from .m3_mx_gemv_kernel import tmem_cols - - if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: - a_tmem = False - # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into - # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). - ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 - if ar is not None and split <= 4 and n % 2 == 0: - ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar - if ar_mode == 2: - trigger_early = False - if use_pdl is None: - use_pdl = get_env_enable_pdl() - # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. - key = ( + return _core._resolve( + _CORE_TARGET, + m, n, k, - split, - ring, - trigger_early, - contig, - False, - nacc, - a_tmem, - st_async, x_bf16, - ar_mode, - ar_world, - ar_maxt, - mma_n, + trigger_early, + cfg, use_pdl, + ar, x_swiglu, + has_swiglu=True, + has_producer=False, ) - return key, (ar_ws, ar_flags, ar_rank) def _compile(key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple): """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" - import cutlass.cute as cute - - from . import m3_mx_gemv_kernel as kernel - - ( - n, - k, - split, - ring, - trigger_early, - contig, - timed, - nacc, - a_tmem, - st_async, - x_bf16, - ar_mode, - ar_world, - ar_maxt, - mma_n, - use_pdl, - x_swiglu, - ) = key - ar_ws, ar_flags, ar_rank = runtime_ar - with _lock: - fn = _compiled.get(key) - if fn is None: - fn = _compiled[key] = cute.compile( - kernel.m3_mx_gemv, - *args, - m, - n, - k, - split, - ring, - trigger_early, - contig, - timed, - nacc, - a_tmem, - st_async, - x_bf16, - ar_ws, - ar_flags, - ar_rank, - ar_mode, - ar_world, - ar_maxt, - mma_n, - use_pdl, - stream, - *extra, - False, - False, - x_swiglu, - ) - return fn + return _core._compile( + _CORE_TARGET, key, args, m, runtime_ar, stream, extra, has_swiglu=True, has_producer=False + ) def _views( @@ -280,22 +186,18 @@ def _views( x_swiglu: bool = False, ) -> tuple: # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. - k = x.shape[1] // 2 if x_swiglu else x.shape[1] - if x_bf16: - x_bytes = x.view(torch.uint8) - xsf = _placeholder_scales(k, x.device) - x_words = x.view(torch.int32) - else: - x_bytes = x.view(torch.uint8) - x_words = _placeholder_words(x.device) - return ( - _arg(weight.view(torch.uint8)), - _arg(weight_scale.view(torch.uint8).reshape(-1)), - _arg(x_bytes), - _arg(xsf.view(torch.uint8).reshape(-1)), - _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), - _arg(x_words.reshape(-1)), - _arg(_dummy_timers(x.device)), + return _host.gemv_views( + x, + xsf, + weight, + weight_scale, + y, + x_bf16, + x_swiglu, + placeholder_scales=_placeholder_scales, + placeholder_words=_placeholder_words, + arg=_arg, + dummy_timers=_dummy_timers, ) @@ -311,42 +213,21 @@ def _run( ar: Optional[tuple] = None, x_swiglu: bool = False, ) -> Optional[torch.Tensor]: - global last_pushed, last_uncompiled - last_pushed = False - last_uncompiled = None - import cuda.bindings.driver as cuda_driver - - m, k = x.shape - if x_swiglu: - k //= 2 - n = weight.shape[0] - if weight.shape[1] != k: - return None - resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu) - if resolved is None: - return None - key, runtime_ar = resolved - fn = _compiled.get(key) - if fn is None: - if _frozen[0]: - # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller - # reports the miss. - last_uncompiled = key - return None - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "m3_mx_gemv must run once per shape outside CUDA-graph capture first" - ) - y = torch.empty(m, n, dtype=torch.bfloat16, device=x.device) - args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) - extra = _extra_args(x.device) - stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) - if fn is None: - fn = _compile(key, args, m, runtime_ar, stream, extra) - ar_ws, ar_flags, ar_rank = runtime_ar - fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra) - last_pushed = key[11] > 0 - return y + return _core._run( + _CORE_TARGET, + x, + xsf, + weight, + weight_scale, + x_bf16, + trigger_early, + cfg, + use_pdl, + ar, + x_swiglu, + has_swiglu=True, + has_producer=False, + ) def precompile( @@ -361,31 +242,18 @@ def precompile( """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" - import cuda.bindings.driver as cuda_driver - - n, k = weight.shape - push = None if ar is None else (0, 0, 0, *ar) - resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu) - if resolved is None: - return None - key, runtime_ar = resolved - if key in _compiled: - return key - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") - dev = weight.device - x = torch.zeros( + return _core.precompile( + _CORE_TARGET, + weight, + weight_scale, m, - 2 * k if x_swiglu else k, - dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, - device=dev, + x_bf16, + use_pdl, + ar, + x_swiglu, + has_swiglu=True, + has_producer=False, ) - xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) - y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) - args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) - stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) - _compile(key, args, m, runtime_ar, stream, _extra_args(dev)) - return key def freeze(frozen: bool = True) -> None: @@ -403,13 +271,8 @@ def is_compiled(key: tuple) -> bool: def _extra_args(device) -> tuple: - """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" - i32 = _placeholder_words(device) - f32 = i32.view(torch.float32) - u8 = i32.view(torch.uint8) - d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) - r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) - return d + r + "Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off)." + return _host.gemv_extra_args(device, placeholder_words=_placeholder_words, arg=_arg) def m3_mx_gemv( @@ -475,47 +338,25 @@ def m3_mx_gemv_swiglu( def _placeholder_scales(k: int, device) -> torch.Tensor: - key = ("sf", k, device) - t = _PLACEHOLDERS.get(key) - if t is None: - t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) - return t + return _host.gemv_placeholder_scales(k, device, placeholders=_PLACEHOLDERS) def _placeholder_words(device) -> torch.Tensor: - key = ("w", device) - t = _PLACEHOLDERS.get(key) - if t is None: - t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) - return t + return _host.gemv_placeholder_words(device, placeholders=_PLACEHOLDERS) def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: - """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" - if ( - weight.dim() != 2 - or weight.dtype != torch.float8_e4m3fn - or get_config(*weight.shape) is None - ): - return - n, k = weight.shape - if (n, k, use_pdl) in _ensured: - return - _ensured.add((n, k, use_pdl)) - xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) - xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) - m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) - for width in MMA_WIDTHS[1:]: - m3_mx_gemv( - torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), - xsf, - weight, - weight_scale, - use_pdl=use_pdl, - ) - m3_mx_gemv_bf16( - torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + "Compiles the kernel of this weight's shape outside capture (a one-token dummy call)." + return _host.gemv_ensure_compiled( weight, weight_scale, - use_pdl=use_pdl, + use_pdl, + get_config=get_config, + ensured=_ensured, + mma_widths=MMA_WIDTHS, + mx_gemv=m3_mx_gemv, + mx_gemv_bf16=m3_mx_gemv_bf16, ) + + +_CORE_TARGET = cast(_core.GemvTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py index 31c499cab2a6..73019c13bfc4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/norm/m3_embedding_firstnorm.py @@ -9,6 +9,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.decode_embedding import ( + embedding_firstnorm_reference as _embedding_reference, +) from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm @@ -75,13 +78,7 @@ def reference( use_pdl: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" - valid = (ids >= 0) & (ids < table.shape[0]) - rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( - table.dtype - ) - x = rows.double() - rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) - return (x * rstd * weight.double()).to(torch.bfloat16), rows + return _embedding_reference(ids, table, weight, eps) m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py deleted file mode 100644 index be5aa46fe774..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_capture_norm.py +++ /dev/null @@ -1,80 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .m3_chunk_norm import rmsnorm_rows - - -class M3CaptureNorm(OpWrapper): - """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture - buffer slice whose row stride may exceed H. Returns ``dst``.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the - # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. - CELLS: tuple[Cell, ...] = ( - Cell( - why="4 verify rows of one request, layer slot 0, residual added", - spec=dict(m=4, slot=0, residual=True), - ), - Cell( - why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) - ), - Cell( - why="no residual (a capture point without one): plain RMSNorm", - spec=dict(m=8, slot=1, residual=False), - ), - ) - - note = """ - fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional - one-ulp move against torch's statistic. - """ - - def __call__( - self, - h: torch.Tensor, - r: Optional[torch.Tensor], - w: torch.Tensor, - eps: float, - dst: torch.Tensor, - ) -> torch.Tensor: - torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) - return dst - - def reference( - self, - h: torch.Tensor, - r: Optional[torch.Tensor], - w: torch.Tensor, - eps: float, - dst: torch.Tensor, - ) -> torch.Tensor: - """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" - x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h - return rmsnorm_rows(x, w, eps) - - def is_valid( - self, - h: torch.Tensor, - r: Optional[torch.Tensor], - w: torch.Tensor, - eps: float, - dst: torch.Tensor, - ) -> None: - # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read - # past its elements without an error. - assert w.is_contiguous(), "w must be dense [H]" - - -m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py deleted file mode 100644 index 6699a55fc3e6..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_chunk_norm.py +++ /dev/null @@ -1,70 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states -(``m3_spec_tail.chunk_norm``).""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). -MAX_HIDDEN = 8192 - - -def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" - xf = x.float() - r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) - return (xf * r * weight.float()).to(torch.bfloat16) - - -class M3ChunkNorm(OpWrapper): - """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may - have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers - # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. - CELLS: tuple[Cell, ...] = ( - Cell( - why="one request's 4 verify rows, dense [4, 3 x 6144]", - spec=dict(m=4, c=3, h=6144, extra_cols=0), - ), - Cell( - why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", - spec=dict(m=32, c=3, h=6144, extra_cols=6144), - ), - ) - - note = """ - fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then - against torch's statistic; the cells' band states it. - """ - - def __call__( - self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True - ): - torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) - return dst - - def reference( - self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True - ): - """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" - c, h = w.shape - x = src[:, : c * h].reshape(src.shape[0], c, h) - return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) - - def is_valid( - self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True - ) -> None: - c, h = w.shape - assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" - assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" - - -m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py deleted file mode 100644 index 8f0f69ee812b..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_input.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat -(``m3_spec_tail.draft_input``).""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). -MAX_HIDDEN = 8192 - - -def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" - xf = x.float() - r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) - return (xf * r * weight.float()).to(torch.bfloat16) - - -class M3DraftInput(OpWrapper): - """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = - hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of - a vocab-local embedding). Returns ``(out, residual)``.""" - - ARCHS = frozenset({Arch.SM_107}) - - # Every pass of the fused draft loop (modeling.DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the - # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests - # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). - CELLS: tuple[Cell, ...] = ( - *( - Cell( - why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", - spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), - ) - for n in range(1, 9) - ), - *( - Cell( - why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", - spec=dict(m=n, h=6144, gather=True, invalid_id=False), - ) - for n in range(1, 9) - ), - Cell( - why="third pass of 3 requests: one row each, hidden read in place", - spec=dict(m=3, h=6144, gather=False, invalid_id=False), - ), - Cell( - why="one row with an out-of-table id: the embedding half is zero", - spec=dict(m=1, h=6144, gather=False, invalid_id=True), - ), - ) - - note = """ - The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's - summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, - also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the - drafter's embedding is the whole table, so its ids are always in range. - """ - - def __call__( - self, - ids: torch.Tensor, - table: torch.Tensor, - hidden: torch.Tensor, - rows: Optional[torch.Tensor], - w_embed: torch.Tensor, - w_hidden: torch.Tensor, - eps_embed: float, - eps_hidden: float, - out: torch.Tensor, - residual: Optional[torch.Tensor], - pdl: bool = True, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - torch.ops.trtllm.minimax_m3_draft_input( - ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl - ) - return out, residual - - def reference( - self, - ids: torch.Tensor, - table: torch.Tensor, - hidden: torch.Tensor, - rows: Optional[torch.Tensor], - w_embed: torch.Tensor, - w_hidden: torch.Tensor, - eps_embed: float, - eps_hidden: float, - out: torch.Tensor, - residual: Optional[torch.Tensor], - pdl: bool = True, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / - ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" - valid = (ids >= 0) & (ids < table.shape[0]) - emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) - emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) - hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] - ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) - return ref, (hid.clone() if residual is not None else None) - - def is_valid( - self, - ids: torch.Tensor, - table: torch.Tensor, - hidden: torch.Tensor, - rows: Optional[torch.Tensor], - w_embed: torch.Tensor, - w_hidden: torch.Tensor, - eps_embed: float, - eps_hidden: float, - out: torch.Tensor, - residual: Optional[torch.Tensor], - pdl: bool = True, - ) -> None: - h = table.shape[1] - # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. - assert table.is_contiguous() and hidden.is_contiguous(), ( - "table / hidden must be dense [*, H]" - ) - if rows is not None: - assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" - else: - assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" - # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. - assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" - assert residual is None or ( - residual.is_contiguous() and residual.shape == (ids.shape[0], h) - ), "residual must be contiguous [M, H]" - - -m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py index 75626e570244..57fdff8e0ed0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_silu.py @@ -14,7 +14,7 @@ class M3DraftSilu(OpWrapper): """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: - trtllm::minimax_m3_sm107_tp4_silu_and_mul (minimaxM3Sm107Tp4/siluAndMul.cu), bit for bit the Triton + trtllm::minimax_m3_sm107_tp4_silu_and_mul (minimaxM3Sm107/siluAndMul.cu), bit for bit the Triton ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs the Triton op).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py index 5dbb324a7d50..c4cd1da0cc4e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_draft_step.py @@ -7,19 +7,23 @@ from __future__ import annotations import functools -from typing import Optional +import sys +from typing import Optional, cast import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, + true_fp32_matmul as true_fp32_matmul, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + decode as decode_coded, # noqa: F401 — used by the bound family implementation +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec import ( + draft_head_reference as _draft_head_shared, ) -from ..gemm.m3_target_head import decode as decode_coded from .m3_shard_accept import exchange_tokens, local_argmax, make_keys #: Requests (head rows) per pass, and exchange rows per call half of the workspace. @@ -60,18 +64,12 @@ def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). """ - n, k = weight.shape - assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) - q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() - return q.view(n // 16, k // 16, 256) + return _draft_head_shared.bf16_fragments(_DRAFT_HEAD_TARGET, weight) def from_fragments(wf: torch.Tensor) -> torch.Tensor: """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" - tiles, ks, _ = wf.shape - return ( - wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) - ) + return _draft_head_shared.from_fragments(_DRAFT_HEAD_TARGET, wf) def head_weight( @@ -84,10 +82,9 @@ def head_weight( ) -> torch.Tensor: """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the 12-bit coded stream with its exception lists decoded.""" - if exc_off is None: - return from_fragments(wf) - n = (exc_off.numel() - 1) * 16 - return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + return _draft_head_shared.head_weight( + _DRAFT_HEAD_TARGET, wf, k, exc_off, exc_pos, exc_w, exp_base + ) def shard_logits( @@ -100,10 +97,9 @@ def shard_logits( ) -> torch.Tensor: """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, rounded once to BF16 (the BF16 GEMM's output), as fp32.""" - w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) - with true_fp32_matmul(): - logits = x.float() @ w.float().t() - return logits.to(torch.bfloat16).float() + return _draft_head_shared.shard_logits( + _DRAFT_HEAD_TARGET, x, wf, exc_off, exc_pos, exc_w, exp_base + ) class M3DraftStep(OpWrapper): @@ -415,3 +411,5 @@ def is_valid( m3_draft_step = M3DraftStep() + +_DRAFT_HEAD_TARGET = cast(_draft_head_shared.DraftHeadTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py deleted file mode 100644 index b980aeefc8d8..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec.py +++ /dev/null @@ -1,86 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Per-request forced accepted counts of a fractional forced acceptance, drawn as the in-tree worker draws them.""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from ..m3_force_vec import force_vec_kernel - - -class M3ForceVec(OpWrapper): - """Bump ``counter`` (int64 [1]) by one, then for each of the ``batch_size`` requests ``b`` write - ``min(base_total + (pool[(c * counter_stride + b * slot_stride) & pool_mask] < frac), max_total)`` into ``out`` - (int32), ``c`` the bumped counter: the in-tree ``_apply_force_accepted_tokens`` draw over its fixed-seed fp32 - ``pool``. ``block`` is the launch width (a power of two >= batch_size; 16, 32 or 64 are compiled at warmup). - - Only the forced-acceptance testing knob (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` with a fractional part) - runs it. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="one request, fraction 0.5", spec=dict(batch=1, block=16, frac=0.5)), - Cell(why="eight requests, fraction 0.3", spec=dict(batch=8, block=16, frac=0.3)), - ) - - note = """ - Integer arithmetic and one fp32 comparison: `reference` must match exactly. - """ - - def __call__( - self, - pool, - counter, - out, - batch_size, - base_total, - max_total, - frac, - counter_stride, - slot_stride, - pool_mask, - block, - ): - force_vec_kernel[(1,)]( - pool, - counter, - out, - batch_size, - base_total, - max_total, - frac, - counter_stride, - slot_stride, - pool_mask, - BLOCK=block, - ) - return out - - def reference( - self, - pool, - counter, - out, - batch_size, - base_total, - max_total, - frac, - counter_stride, - slot_stride, - pool_mask, - block, - ): - """The counts (a new tensor; ``counter`` and ``out`` are not written).""" - c = int(counter.item()) + 1 - slots = torch.arange(batch_size, dtype=torch.int64, device=pool.device) - idx = (c * counter_stride + slots * slot_stride) & pool_mask - r = pool[idx] - return torch.clamp(base_total + (r < frac).to(torch.int32), max=max_total).to(torch.int32) - - -m3_force_vec = M3ForceVec() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py index 416425f9a01d..ed6f14b240c3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_force_vec_req.py @@ -7,8 +7,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from ..m3_force_vec import force_vec_req_kernel +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.spec.m3_force_vec import ( + force_vec_req_kernel, +) class M3ForceVecReq(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py index 7623f5fe582b..0b833bc4a4da 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_sampled_accept.py @@ -5,14 +5,18 @@ from __future__ import annotations -from typing import Optional +import sys +from typing import Optional, cast import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec import ( + sampling_reference as _sampling_shared, +) -from .m3_shard_accept import strict_accept +from .m3_shard_accept import strict_accept as strict_accept #: Verify rows per call. MAX_ROWS = 48 @@ -25,25 +29,7 @@ def nucleus_probs( ) -> torch.Tensor: """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" - inv = 1.0 / (temperatures.float() + TEMP_EPSILON) - x = logits.float() * inv[:, None] - w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() - z = w.sum(dim=1, keepdim=True) - order = torch.sort(w, dim=1, descending=True) - sw = order.values - # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group - csum = torch.cumsum(sw, dim=1) - sw - first = torch.ones_like(sw, dtype=torch.bool) - first[:, 1:] = sw[:, 1:] != sw[:, :-1] - group = torch.cumsum(first.long(), dim=1) - 1 - starts = torch.zeros_like(csum) - starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) - above = torch.gather(starts, 1, group) - keep_sorted = above <= top_ps.double()[:, None] * z - keep = torch.zeros_like(keep_sorted) - keep.scatter_(1, order.indices, keep_sorted) - p = torch.where(keep, w, torch.zeros_like(w)) - return (p / p.sum(dim=1, keepdim=True)).float() + return _sampling_shared.nucleus_probs(_REFERENCE_TARGET, logits, temperatures, top_ps) class M3SampledAccept(OpWrapper): @@ -215,13 +201,34 @@ def reference( """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" - rows = logits.shape[0] - probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) - - def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: - return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) - - return {"probs": probs, "accept": accept} + return _sampling_shared.sampled_reference( + self, + _REFERENCE_TARGET, + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + top_ks, + min_ps, + ) def is_valid( self, @@ -255,35 +262,36 @@ def is_valid( decided: Optional[torch.Tensor] = None, ilp_walk: bool = False, ) -> None: - rows = logits.shape[0] - # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. - if top_ks is not None: - k = top_ks[:rows] - assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( - "top-k rows must take the generic path" - ) - if min_ps is not None: - assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" - # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the - # wrong rows. - assert rows == draft.shape[0] * (draft.shape[1] + 1), ( - "logits rows must be nreq * (draft_len + 1)" - ) - # Per-row parameters are read at the row index, unchecked. - for t in (temperatures, top_ps, seeds, offsets): - assert t.numel() >= rows and t.is_contiguous(), ( - "per-row parameters must cover every row" - ) - assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( - "fp32 sampling parameters" - ) - assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( - "int64 Philox seeds / offsets" - ) - # The finished-row counter elects the CTA that runs the acceptance. - assert int(counters[0].item()) == 0, ( - "counters[0] (finished rows) must be zero between calls" + return _sampling_shared.validate_sampling( + self, + _REFERENCE_TARGET, + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + top_ks, + min_ps, ) m3_sampled_accept = M3SampledAccept() + +_REFERENCE_TARGET = cast(_sampling_shared.SamplingTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py index 8beba4373b02..fc357c3c4ee3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_shard_accept.py @@ -12,97 +12,18 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from ....ops.spec.acceptance_reference import _LOW32 as _LOW32 +from ....ops.spec.acceptance_reference import exchange_tokens as exchange_tokens +from ....ops.spec.acceptance_reference import key_token as key_token +from ....ops.spec.acceptance_reference import local_argmax as local_argmax +from ....ops.spec.acceptance_reference import make_keys as make_keys +from ....ops.spec.acceptance_reference import ordered_bits as ordered_bits +from ....ops.spec.acceptance_reference import strict_accept as strict_accept +from ....ops.spec.acceptance_reference import unsigned_max as unsigned_max + #: Verify rows per call, and exchange rows per call half of its workspace. SHARD_ROWS = 48 -_LOW32 = 0xFFFFFFFF - - -def ordered_bits(v: torch.Tensor) -> torch.Tensor: - """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" - b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 - b = torch.where(b == 0x80000000, torch.zeros_like(b), b) - return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) - - -def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: - """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" - return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) - - -def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: - """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" - sign = torch.iinfo(torch.int64).min - return (keys ^ sign).max(dim=dim).values ^ sign - - -def key_token(key: torch.Tensor) -> torch.Tensor: - return (_LOW32 - (key & _LOW32)).to(torch.int32) - - -def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: - """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" - best = logits.max(dim=1).values - idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max - return best, idx.to(torch.int64) + vocab_offset - - -def exchange_tokens( - own_keys: torch.Tensor, - ws_uc: torch.Tensor, - counters: torch.Tensor, - rank: int, - world: int, - slot_rows: int, -) -> torch.Tensor: - """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. - - Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A - single-device run stands in for the other ranks by writing their keys there before the call. - """ - m = own_keys.shape[0] - half = int(counters[1].item()) & 1 - slots = ( - ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] - .view(m, world) - .clone() - ) - slots[:, rank] = own_keys - return key_token(unsigned_max(slots, dim=1)) - - -def strict_accept( - tokens: torch.Tensor, - draft: torch.Tensor, - batch_idx: torch.Tensor, - force_total: int, - next_new: Optional[torch.Tensor], - force_vec: Optional[torch.Tensor], -) -> dict[str, torch.Tensor]: - """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). - - accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by - force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and - next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. - """ - nreq, d = draft.shape - k1 = d + 1 - tokens = tokens.to(torch.int32) - accepted = tokens.view(nreq, k1).clone() - match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) - num = 1 + match.cumprod(dim=1).sum(dim=1) - if force_vec is not None: - num = force_vec[:nreq].to(num.dtype) - elif force_total > 0: - num = torch.full_like(num, force_total) - gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 - out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} - if next_new is not None: - nn = next_new.clone() - nn[:, 0] = tokens[gather] - out["next_new"] = nn - return out - class M3ShardAccept(OpWrapper): """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py deleted file mode 100644 index ff07c54de1f0..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_slot_scatter.py +++ /dev/null @@ -1,72 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The spec sampler's per-slot stores of one step's outputs, in one launch (``M3SpecStepFold.scatter`` -> -``m3_spec_tail.slot_scatter``).""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -def padded_rows(src: torch.Tensor, width: int) -> torch.Tensor: - """src rows zero-padded (or cut) to ``width`` columns.""" - out = torch.zeros(src.shape[0], width, dtype=src.dtype, device=src.device) - w = min(width, src.shape[1]) - out[:, :w] = src[:, :w] - return out - - -class M3SlotScatter(OpWrapper): - """For i < len(slots), all int32 2-D strided views: ``da[slots[i]] = sa[i]``, ``db[slots[i]] = sb[i]``, - ``dd[slots[i]] = sd[i]`` (each zero-padded / cut to the destination width) and ``dl[slots[i]] = sl[i]``. The - destinations are the store's transposed [slots, width] views.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The sampler store at draft_len 3 (8 slots): new tokens / next new tokens [slots, 4] (transposed views of - # [4, slots, 1]), next draft tokens [slots, 3]; a worker output narrower than the store (dynamic draft length) - # is zero-padded. - CELLS: tuple[Cell, ...] = ( - Cell( - why="8 rows into shuffled slots, widths equal to the store's", - spec=dict(n=8, slots=8, src_w=4, src_d=3, store_w=4, store_d=3), - ), - Cell( - why="3 rows of a narrower step (runtime draft 2): zero padding to the store width", - spec=dict(n=3, slots=8, src_w=3, src_d=2, store_w=4, store_d=3), - ), - ) - - note = """ - Pure index copies: the reference is exact. Slots are not bounds-checked; is_valid states it. - """ - - def __call__(self, slots, sa, da, sb, db, sl, dl, sd, dd): - torch.ops.trtllm.minimax_m3_slot_scatter(slots, sa, da, sb, db, sl, dl, sd, dd) - - def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): - """The four destinations after the call, on copies.""" - s = slots.long() - out = {} - for name, src, dst in (("da", sa, da), ("db", sb, db), ("dd", sd, dd)): - r = dst.clone() - r[s] = padded_rows(src[: s.numel()], dst.shape[1]) - out[name] = r - lr = dl.clone() - lr[s] = sl[: s.numel()] - out["dl"] = lr - return out - - def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): - assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( - "slots must index the store" - ) - # Two rows into one slot race (which one lands is unordered). - assert slots.unique().numel() == slots.numel(), "slots must be distinct" - assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" - - -m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py index 36e421412f30..3094e2e9965a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_spec_offsets.py @@ -16,7 +16,7 @@ class M3SpecOffsets(OpWrapper): """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths - per pair) in one plain launch: trtllm::minimax_m3_sm107_tp4_spec_offsets (minimaxM3Sm107Tp4/specOffsets.cu), the + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp4_spec_offsets (minimaxM3Sm107/specOffsets.cu), the same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other layouts (the caller runs the ATen adds).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py deleted file mode 100644 index 16c4000199da..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_gather.py +++ /dev/null @@ -1,152 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch -(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3StepGather(OpWrapper): - """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all - int32:: - - ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) - draft[d0 + p*D + d] = next_draft[slots[p], d] - pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) - kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) - - Every other element of ``ids`` / ``draft`` is left as it was. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the - # max batch); the engine's flat buffers sized for the max batch. - CELLS: tuple[Cell, ...] = ( - Cell( - why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", - spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), - ), - Cell( - why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", - spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), - ), - Cell( - why="one request, the store's 2-D view of new tokens (beam dim squeezed)", - spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), - ), - ) - - note = """ - Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid - states them. - """ - - def __call__( - self, - new_tok, - next_draft, - lens, - slots, - ppi, - P, - T, - D, - ids, - ids0, - draft, - d0, - pos_off, - kv_off, - b0, - ): - torch.ops.trtllm.minimax_m3_step_gather( - new_tok, - next_draft, - lens, - slots, - ppi, - P, - T, - D, - ids, - ids0, - draft, - d0, - pos_off, - kv_off, - b0, - ) - - def reference( - self, - new_tok, - next_draft, - lens, - slots, - ppi, - P, - T, - D, - ids, - ids0, - draft, - d0, - pos_off, - kv_off, - b0, - ): - """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" - nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) - s = slots[:P].long() - ids_r, draft_r = ids.clone(), draft.clone() - ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) - draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) - pos_r = torch.zeros_like(pos_off) - pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] - kv_r = torch.zeros_like(kv_off) - kv_r[b0 : b0 + P] = lens[s] - T - return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} - - def is_valid( - self, - new_tok, - next_draft, - lens, - slots, - ppi, - P, - T, - D, - ids, - ids0, - draft, - d0, - pos_off, - kv_off, - b0, - ): - # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or - # writes another buffer. - s = slots[:P] - assert bool( - ( - (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) - ).all() - ), "slots must index the store" - assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( - "ppi must index lens" - ) - assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" - assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" - assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( - "offset buffers too short" - ) - - -m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py index 122d43fb175a..8850057b21b7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/spec/m3_step_prologue.py @@ -14,7 +14,7 @@ class M3StepPrologue(OpWrapper): """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch - (trtllm::minimax_m3_sm107_tp4_step_prologue, minimaxM3Sm107Tp4/stepPrologue.cu). + (trtllm::minimax_m3_sm107_tp4_step_prologue, minimaxM3Sm107/stepPrologue.cu). ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py index 240246d4548e..0f689fa9e06d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode.py @@ -45,15 +45,57 @@ from typing import Optional import torch -import triton import triton.language as tl +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _NVFP4_TUNED_CAPABILITIES as _NVFP4_TUNED_CAPABILITIES, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _TARGET_GRID as _TARGET_GRID, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + MergeMxfp8 as MergeMxfp8, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _check_nvfp4_inputs as _check_nvfp4_inputs, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_launch_options as _sm100f_nvfp4_launch_options, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_merge_launch_options as _sm100f_nvfp4_merge_launch_options, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_num_topk_chunks as _sm100f_nvfp4_num_topk_chunks, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_query_group_size as _sm100f_nvfp4_query_group_size, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_use_linear_softmax as _sm100f_nvfp4_use_linear_softmax, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + resolve_num_topk_chunks as resolve_num_topk_chunks, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( check_decode_span_shape, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers +from ...kernels.attention.triton_sparse_decode import _dequant_nvfp4_rows as _dequant_nvfp4_rows +from ...kernels.attention.triton_sparse_decode import _e2m1x2_to_f16x2 as _e2m1x2_to_f16x2 +from ...kernels.attention.triton_sparse_decode import ( + _gqa_sparse_decode_kernel as _gqa_sparse_decode_kernel, +) +from ...kernels.attention.triton_sparse_decode import ( + _gqa_sparse_decode_pair_kernel as _gqa_sparse_decode_pair_kernel, +) +from ...kernels.attention.triton_sparse_decode import ( + _merge_topk_attn_out_kernel as _merge_topk_attn_out_kernel, +) +from ...kernels.attention.triton_sparse_decode import _rcp_approx_ftz as _rcp_approx_ftz +from ...kernels.attention.triton_sparse_decode import _store_mxfp8_head as _store_mxfp8_head from . import m3_attn_lc, m3_prequant, m3_report # One sparse block is exactly one KV page. @@ -68,7 +110,6 @@ } # Total CTAs the split-K partitioning aims for before it stops splitting. -_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size # NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale # byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. @@ -80,1140 +121,12 @@ MERGE_MXFP8_MIN_TOKENS = 1 # capacity: fixed: a row lower bound -class MergeMxfp8: - """MXFP8 copy of the split-K merge output for the o_proj GEMM. - - From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and swizzled UE8M0 - scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips its own - quantization. One 128-row buffer per width serves every layer: each o_proj consumes it before the next layer's - merge rewrites it, in stream order. - """ - - MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) - - def __init__(self): - self._buffers = {} - - def buffers(self, width: int, device: torch.device): - key = (width, device) - if key not in self._buffers: - self._buffers[key] = ( - torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), - torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), - ) - return self._buffers[key] - - MERGE_MXFP8 = MergeMxfp8() -@triton.jit -def _e2m1x2_to_f16x2(packed): - """Convert one packed E2M1 byte into its low/high FP16 values.""" - return tl.inline_asm_elementwise( - asm="""{ - .reg .b8 src; - .reg .b32 pair; - mov.b32 {src, _, _, _}, $2; - cvt.rn.f16x2.e2m1x2 pair, src; - mov.b32 {$0, $1}, pair; - }""", - constraints="=h,=h,r", - args=(packed,), - dtype=(tl.float16, tl.float16), - is_pure=True, - pack=1, - ) - - -@triton.jit -def _dequant_nvfp4_rows( - packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte - scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements - BLOCK_N: tl.constexpr, - D: tl.constexpr, - SCALE_COLS: tl.constexpr, - SF_VEC: tl.constexpr, -): - """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. - - The caller gathers the block scales in whichever layout its role stores - them, so this handles only the value math and the scale broadcast. The - per-tensor scale is applied by the caller. - """ - # Blackwell converts both E2M1 nibbles with one instruction. Joining the - # low/high results along a new innermost axis and flattening it rebuilds - # the original element order without scalar exponent/sign reconstruction. - value_lo, value_hi = _e2m1x2_to_f16x2(packed) - values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) - - scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) - scales = tl.reshape( - tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) - ) - return values * scales - - -@triton.heuristics( - { - "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), - "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), - } -) -# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary -# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). -@triton.jit( - do_not_specialize=["decode_query_len", "total_q"], - do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], -) -def _gqa_sparse_decode_kernel( - q_ptr, # [total_q, num_heads, head_dim] - k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] - v_ptr, # paged V: same layout as K - kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False - k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] - v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order - k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 - v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 - t_ptr, # topk_idx: [num_kv_heads, total_q, topk] - o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT - lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT - out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused - counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest - block_table_ptr, # [num_reqs, max_blocks] - seq_lens, # [num_reqs] - total_q, - gqa_group_size, - head_dim, - max_topk, - sm_scale, - decode_query_len, - stride_qn, - stride_qh, - stride_qd, - stride_k_blk, - stride_k_h, - stride_k_pos, - stride_k_d, - stride_v_blk, - stride_v_h, - stride_v_pos, - stride_v_d, - stride_ksc_blk, - stride_ksc_h, - stride_vsc_blk, - stride_vsc_h, - stride_th, - stride_tn, - stride_tk, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_bt_b, - num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request - num_pages, # KV pool pages: a page id outside [0, num_pages) is not read - stride_out_n, - stride_out_h, - stride_out_d, - num_kv_heads, - BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_H: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache - WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register - COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math - USE_PDL: tl.constexpr, - KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales - USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence - SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only - DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge - FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place - EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry - IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) -): - sm_scale_log2e = sm_scale * 1.4426950409 - if KV_NVFP4: - # Both per-tensor dequant scales factor out of the inner loop: QK is - # linear in K, and V's scale commutes with the softmax weights. Reading - # them from device memory keeps the launch free of a host sync, so - # CUDA-graph capture stays unconstrained. - sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) - v_global_scale = tl.load(v_gs_ptr) - # Split-K over the topk dimension: pid(0) folds (query token, chunk). - pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) - pid_b = pid_bc % total_q - pid_c = pid_bc // total_q - req_id = pid_b // decode_query_len - q_offset = pid_b - req_id * decode_query_len - pid_h = pid_kh * gqa_group_size - chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_start = pid_c * chunk_size_topk - - if USE_PDL: - tl.extra.cuda.gdc_wait() - if EARLY_TRIGGER: - # Dependents still wait for this grid to complete; triggering here only - # lets them start their input-independent prologue (e.g. weight loads). - tl.extra.cuda.gdc_launch_dependents() - - seq_len = tl.load(seq_lens + req_id) - query_pos = seq_len - decode_query_len + q_offset - # Full-CG padding uses zero-length request rows. Clamp to an empty - # attention range instead of letting padded rows produce negative lengths. - kv_len = tl.maximum(query_pos + 1, 0) - - # Bound the walk by the token's own valid block count rather than a - # sentinel: the selector emits ascending ids and only pads with -1 past - # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. - num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk = tl.minimum(max_topk, num_blocks) - chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) - - off_n = tl.arange(0, BLOCK_SIZE_K) - off_d = tl.arange(0, BLOCK_SIZE_D) - off_h = tl.arange(0, BLOCK_SIZE_H) - if KV_NVFP4: - # Packed byte column and block-scale column. The NVFP4 path runs only at - # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. - off_dp = tl.arange(0, BLOCK_SIZE_D // 2) - off_s = tl.arange(0, SCALE_COLS) - d_mask = off_d < head_dim - h_mask = off_h < gqa_group_size - hd_mask = h_mask[:, None] & d_mask[None, :] - bt_row = block_table_ptr + req_id * stride_bt_b - - m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - if USE_LINEAR_SOFTMAX: - # Keep the linear softmax denominator and reuse the producer's alpha. - # This removes one exp2 plus one log2 from every visited KV page; only - # the final LSE written for split-K merging still needs a logarithm. - l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - else: - lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - q = tl.load( - q_ptr - + pid_b * stride_qn - + (pid_h + off_h[:, None]) * stride_qh - + off_d[None, :] * stride_qd, - mask=hd_mask, - other=0.0, - ) - # Widen before anything reads q.dtype: K, V and p are all cast to it below, - # so leaving q narrow would silently run the whole attention in FP8 and - # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result - # matches a caller that widened q itself before the call. - if WIDEN_Q: - q = q.to(COMPUTE_DTYPE) - kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 - - cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk - for it in tl.range(chunk_start, chunk_end): - if IDENTITY: - # Slot `it` of the identity selection, bounded by this row's own kv extent above. - blk = it - else: - blk = tl.load(cur_idx_ptr) - cur_idx_ptr = cur_idx_ptr + stride_tk - # int64 page offsets: a large cache overflows int32 well before the - # per-page block offsets do. - # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, - # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). - blk_ok = (blk >= 0) & (blk < num_bt_cols) - page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) - blk_ok = blk_ok & (page >= 0) & (page < num_pages) - page = tl.where(blk_ok, page, 0) - pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len - # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block - # would make exp2(-inf - -inf) = NaN). - if blk_ok: - if KV_NVFP4: - # K block scales are token-major linear within the page-head region; - # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. - k_deq = _dequant_nvfp4_rows( - tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[:, None] * stride_k_pos - + off_dp[None, :] * stride_k_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - k_sc_ptr - + page * stride_ksc_blk - + pid_kh * stride_ksc_h - + off_n[:, None] * SCALE_COLS - + off_s[None, :], - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ) - k = tl.trans(k_deq).to(q.dtype) - else: - k = tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[None, :] * stride_k_pos - + off_d[:, None] * stride_k_d, - mask=d_mask[:, None] & pos_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - k = (k * kv_scale).to(q.dtype) - qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) - qk += tl.dot(q, k) * sm_scale_log2e - m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp2(qk - m_ij[:, None]) - l_ij = tl.sum(p, axis=1) - if USE_LINEAR_SOFTMAX: - alpha = tl.exp2(m_i - m_ij) - acc_o = acc_o * alpha[:, None] - else: - acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] - if KV_NVFP4: - # V block scales use vLLM's 4x4 token-quad order rather than K's - # linear one: [token // 4, scale_col, token % 4]. - v = _dequant_nvfp4_rows( - tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_dp[None, :] * stride_v_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - v_sc_ptr - + page * stride_vsc_blk - + pid_kh * stride_vsc_h - + (off_n[:, None] // 4) * (4 * SCALE_COLS) - + 4 * off_s[None, :] - + (off_n[:, None] % 4), - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ).to(q.dtype) - else: - v = tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_d[None, :] * stride_v_d, - mask=pos_mask[:, None] & d_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - v = (v * kv_scale).to(q.dtype) - acc_o += tl.dot(p.to(v.dtype), v) - if USE_LINEAR_SOFTMAX: - l_i = l_i * alpha + l_ij - else: - lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) - m_i = m_ij - - # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. - if USE_LINEAR_SOFTMAX: - scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) - lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) - else: - scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) - lse_out = lse_i - if KV_NVFP4: - # Applying V's scale per partial is exact: the merge is a weighted - # average whose weights sum to one, so a constant survives it. - scale = scale * v_global_scale - acc_o = acc_o * scale[:, None] - o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h - tl.store( - o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, - acc_o.to(o_ptr.dtype.element_ty), - mask=hd_mask, - ) - if not DIRECT_OUTPUT: - lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h - tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) - - if FUSED_MERGE: - # Split-K serial reduction: every chunk publishes its partial, and the - # chunk that arrives last for this (token, KV head) merges all of them - # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. - tl.debug_barrier() - counter = counter_ptr + pid_b * num_kv_heads + pid_kh - arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") - if arrived == NUM_TOPK_CHUNKS - 1: - tl.debug_barrier() - o_rows = ( - o_ptr - + pid_b * stride_o_b - + (pid_h + off_h[:, None]) * stride_o_h - + off_d[None, :] * stride_o_d - ) - l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h - lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - lse_max = tl.maximum(lse_max, lse_c) - # All chunks empty (a CUDA-graph padding row): zero, never NaN. - lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) - denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - denom += tl.exp2(lse_c - lse_max) - denom = tl.where(denom > 0, denom, 1.0) - merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - w_c = tl.exp2(lse_c - lse_max) / denom - o_c = tl.load( - o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" - ).to(tl.float32) - merged += o_c * w_c[:, None] - tl.store( - out_ptr - + pid_b * stride_out_n - + (pid_h + off_h[:, None]) * stride_out_h - + off_d[None, :] * stride_out_d, - merged.to(out_ptr.dtype.element_ty), - mask=hd_mask, - ) - tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") - - # After the stores, never before: this releases either the merge grid or, - # for direct output, the next PDL-enabled consumer of the final tensor. - if USE_PDL and not EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.heuristics( - { - "BLOCK_SIZE_H": lambda args: max( - 16, - triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), - ), - "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), - } -) -@triton.jit( - do_not_specialize=["decode_query_len", "total_q"], - do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], -) -def _gqa_sparse_decode_pair_kernel( - q_ptr, # [total_q, num_heads, head_dim] - k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] - v_ptr, # paged V: same layout as K - kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False - k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] - v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order - k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 - v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 - t_ptr, # topk_idx: [num_kv_heads, total_q, topk] - o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT - lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT - block_table_ptr, # [num_reqs, max_blocks] - seq_lens, # [num_reqs] - total_q, - gqa_group_size, - head_dim, - max_topk, - sm_scale, - decode_query_len, - stride_qn, - stride_qh, - stride_qd, - stride_k_blk, - stride_k_h, - stride_k_pos, - stride_k_d, - stride_v_blk, - stride_v_h, - stride_v_pos, - stride_v_d, - stride_ksc_blk, - stride_ksc_h, - stride_vsc_blk, - stride_vsc_h, - stride_th, - stride_tn, - stride_tk, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_bt_b, - num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request - num_pages, # KV pool pages: a page id outside [0, num_pages) is not read - BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_H: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache - WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register - COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math - USE_PDL: tl.constexpr, - KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales - USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence - QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set - MAX_TOPK: tl.constexpr, - SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only - DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge - IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) -): - sm_scale_log2e = sm_scale * 1.4426950409 - if KV_NVFP4: - # Both per-tensor dequant scales factor out of the inner loop: QK is - # linear in K, and V's scale commutes with the softmax weights. Reading - # them from device memory keeps the launch free of a host sync, so - # CUDA-graph capture stays unconstrained. - sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) - v_global_scale = tl.load(v_gs_ptr) - # Split-K over the topk dimension: pid(0) folds (query group, chunk). - pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) - num_query_groups = total_q // QUERY_GROUP_SIZE - pid_g = pid_bc % num_query_groups - pid_c = pid_bc // num_query_groups - groups_per_req = decode_query_len // QUERY_GROUP_SIZE - req_id = pid_g // groups_per_req - q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE - pid_b = req_id * decode_query_len + q_group_offset - pid_h = pid_kh * gqa_group_size - chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_start = pid_c * chunk_size_topk - - if USE_PDL: - tl.extra.cuda.gdc_wait() - - seq_len = tl.load(seq_lens + req_id) - # Query rows are flattened as [query-in-group, GQA head]. - off_h = tl.arange(0, BLOCK_SIZE_H) - q_lane = off_h // gqa_group_size - h_lane = off_h - q_lane * gqa_group_size - h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE - q_token = pid_b + q_lane - query_pos = seq_len - decode_query_len + q_group_offset + q_lane - # Full-CG padding uses zero-length request rows. Clamp to an empty - # attention range instead of letting padded rows produce negative lengths. - kv_len = tl.maximum(query_pos + 1, 0) - kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) - - # Each query can have a different valid-block count at a page boundary. - # The selector puts those valid entries first and pads only the tail. - kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) - num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk_0 = tl.minimum(max_topk, num_blocks_0) - chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) - if QUERY_GROUP_SIZE == 2: - kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) - num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk_1 = tl.minimum(max_topk, num_blocks_1) - chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) - - off_n = tl.arange(0, BLOCK_SIZE_K) - off_d = tl.arange(0, BLOCK_SIZE_D) - if KV_NVFP4: - # Packed byte column and block-scale column. The NVFP4 path runs only at - # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. - off_dp = tl.arange(0, BLOCK_SIZE_D // 2) - off_s = tl.arange(0, SCALE_COLS) - d_mask = off_d < head_dim - hd_mask = h_mask[:, None] & d_mask[None, :] - bt_row = block_table_ptr + req_id * stride_bt_b - - m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - if USE_LINEAR_SOFTMAX: - # Keep the linear softmax denominator and reuse the producer's alpha. - # This removes one exp2 plus one log2 from every visited KV page; only - # the final LSE written for split-K merging still needs a logarithm. - l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - else: - lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - q = tl.load( - q_ptr - + q_token[:, None] * stride_qn - + (pid_h + h_lane[:, None]) * stride_qh - + off_d[None, :] * stride_qd, - mask=hd_mask, - other=0.0, - ) - # Widen before anything reads q.dtype: K, V and p are all cast to it below, - # so leaving q narrow would silently run the whole attention in FP8 and - # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this - # reproduces the values the caller's standalone widening kernel produced. - if WIDEN_Q: - q = q.to(COMPUTE_DTYPE) - kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 - - idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk - if QUERY_GROUP_SIZE == 2: - idx_ptr_1 = idx_ptr_0 + stride_tn - idx_0 = chunk_start - idx_1 = chunk_start - count_0 = chunk_end_0 - chunk_start - count_1 = chunk_end_1 - chunk_start - chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_offsets = tl.arange(0, chunk_capacity) - sentinel = 2147483647 - if IDENTITY: - blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) - blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) - else: - blocks_0 = tl.load( - idx_ptr_0 + chunk_offsets * stride_tk, - mask=chunk_offsets < count_0, - other=sentinel, - ) - blocks_1 = tl.load( - idx_ptr_1 + chunk_offsets * stride_tk, - mask=chunk_offsets < count_1, - other=sentinel, - ) - chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 - # Equal chunks take the measured shared-page loop. A divergent chunk - # walks the sorted union, bounded by the sum of both row lengths. - num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) - else: - num_candidates = chunk_end_0 - chunk_start - - for cand in tl.range(0, num_candidates): - if QUERY_GROUP_SIZE == 2: - if chunk_equal: - has_block = True - if IDENTITY: - blk = chunk_start + cand - else: - blk = tl.load(idx_ptr_0) - idx_ptr_0 += stride_tk - query_active = h_mask - else: - valid_0 = idx_0 < chunk_end_0 - valid_1 = idx_1 < chunk_end_1 - if IDENTITY: - blk_0 = tl.where(valid_0, idx_0, sentinel) - blk_1 = tl.where(valid_1, idx_1, sentinel) - else: - blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) - blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) - use_0 = valid_0 & (blk_0 <= blk_1) - use_1 = valid_1 & (blk_1 <= blk_0) - has_block = use_0 | use_1 - blk = tl.minimum(blk_0, blk_1) - idx_0 += use_0.to(tl.int32) - idx_1 += use_1.to(tl.int32) - idx_ptr_0 += use_0.to(tl.int64) * stride_tk - idx_ptr_1 += use_1.to(tl.int64) * stride_tk - query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) - else: - has_block = True - if IDENTITY: - blk = chunk_start + cand - else: - blk = tl.load(idx_ptr_0) - idx_ptr_0 += stride_tk - query_active = h_mask - - if has_block: - # int64 page offsets: a large cache overflows int32 well before the - # per-page block offsets do. - # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the - # pool, reads nothing. - blk_ok = (blk >= 0) & (blk < num_bt_cols) - page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) - blk_ok = blk_ok & (page >= 0) & (page < num_pages) - page = tl.where(blk_ok, page, 0) - pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max - q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] - # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block - # would make exp2(-inf - -inf) = NaN). - if blk_ok: - if KV_NVFP4: - # K scales are token-major linear within the page-head region; - # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. - k_deq = _dequant_nvfp4_rows( - tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[:, None] * stride_k_pos - + off_dp[None, :] * stride_k_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - k_sc_ptr - + page * stride_ksc_blk - + pid_kh * stride_ksc_h - + off_n[:, None] * SCALE_COLS - + off_s[None, :], - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ) - k = tl.trans(k_deq).to(q.dtype) - else: - k = tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[None, :] * stride_k_pos - + off_d[:, None] * stride_k_d, - mask=d_mask[:, None] & pos_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - k = (k * kv_scale).to(q.dtype) - active_pos = query_active[:, None] & q_pos_mask - qk = tl.where(active_pos, 0.0, float("-inf")) - qk += tl.dot(q, k) * sm_scale_log2e - m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) - p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) - l_ij = tl.sum(p, axis=1) - alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) - acc_o = acc_o * alpha[:, None] - if KV_NVFP4: - # V scales use vLLM's 4x4 token-quad order rather than K's - # linear one: [token // 4, scale_col, token % 4]. - v = _dequant_nvfp4_rows( - tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_dp[None, :] * stride_v_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - v_sc_ptr - + page * stride_vsc_blk - + pid_kh * stride_vsc_h - + (off_n[:, None] // 4) * (4 * SCALE_COLS) - + 4 * off_s[None, :] - + (off_n[:, None] % 4), - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ).to(q.dtype) - else: - v = tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_d[None, :] * stride_v_d, - mask=pos_mask[:, None] & d_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - v = (v * kv_scale).to(q.dtype) - acc_o += tl.dot(p.to(v.dtype), v) - if USE_LINEAR_SOFTMAX: - l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) - else: - lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) - lse_i = tl.where(query_active, lse_ij, lse_i) - m_i = m_ij - - # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. - if USE_LINEAR_SOFTMAX: - scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) - lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) - else: - scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) - lse_out = lse_i - if KV_NVFP4: - # Applying V's scale per partial is exact: the merge is a weighted - # average whose weights sum to one, so a constant survives it. - scale = scale * v_global_scale - acc_o = acc_o * scale[:, None] - tl.store( - o_ptr - + pid_c * stride_o_c - + q_token[:, None] * stride_o_b - + (pid_h + h_lane[:, None]) * stride_o_h - + off_d[None, :] * stride_o_d, - acc_o.to(o_ptr.dtype.element_ty), - mask=hd_mask, - ) - if not DIRECT_OUTPUT: - tl.store( - lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, - lse_out, - mask=h_mask, - ) - - # After the stores, never before: this releases either the merge grid or, - # for direct output, the next PDL-enabled consumer of the final tensor. - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _rcp_approx_ftz(x): - return tl.inline_asm_elementwise( - "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 - ) - - -@triton.jit -def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): - """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as - torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, - scales in the 128x4-swizzled layout.""" - NUM_SF: tl.constexpr = HEAD_DIM // 32 - xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) - amax = tl.max(tl.abs(xr), axis=1) - sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) - bits = sfv.to(tl.int32, bitcast=True) - e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) - e = tl.minimum(e, 254) - scale = ((254 - e) << 23).to(tl.float32, bitcast=True) - scale = tl.where(amax != 0, scale, 0.0) - q = (xr * scale[:, None]).to(tl.float8e4nv) - off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) - tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) - k_sf = head * NUM_SF + tl.arange(0, NUM_SF) - off_sf = ( - (row // 128) * num_sf_k_tiles * 512 - + (k_sf // 4) * 512 - + (row % 32) * 16 - + ((row % 128) // 32) * 4 - + k_sf % 4 - ) - tl.store(sf_ptr + off_sf, e.to(tl.uint8)) - - -@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) -@triton.jit -def _merge_topk_attn_out_kernel( - o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] - lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] - out_ptr, # merged out: [total_q, num_heads, head_dim] - head_dim, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_out_n, - stride_out_h, - stride_out_d, - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_PDL: tl.constexpr, - EARLY_TRIGGER: tl.constexpr = False, - q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 - sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled - num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) - MXFP8_OUT: tl.constexpr = False, -): - pid_b, pid_h = tl.program_id(0), tl.program_id(1) - - if USE_PDL: - if EARLY_TRIGGER: - # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight - # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits - # for this grid (and therefore the producer) to complete. - tl.extra.cuda.gdc_launch_dependents() - tl.extra.cuda.gdc_wait() - - off_c = tl.arange(0, NUM_TOPK_CHUNKS) - off_d = tl.arange(0, BLOCK_SIZE_D) - o = tl.load( - o_ptr - + pid_b * stride_o_b - + pid_h * stride_o_h - + off_c[:, None] * stride_o_c - + off_d[None, :] * stride_o_d, - mask=off_d[None, :] < head_dim, - other=0.0, - ).to(tl.float32) - # Empty chunks contribute -inf, hence weight 0. - lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) - lse_max = tl.max(lse, axis=0) - # A row whose every chunk is empty (a CUDA-graph padding row) would give - # -inf - -inf here. Zero it instead: the padded output is discarded, but a - # NaN would survive into the residual stream and the all-reduce. - lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) - weights = tl.exp2(lse - lse_max) - denom = tl.sum(weights, axis=0) - weights = weights / tl.where(denom > 0, denom, 1.0) - o_merged = tl.sum(o * weights[:, None], axis=0) - out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d - tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) - if MXFP8_OUT: - _store_mxfp8_head( - o_merged.to(out_ptr.dtype.element_ty), - q8_ptr, - sf_ptr, - pid_b, - pid_h, - num_sf_k_tiles, - BLOCK_SIZE_D, - ) - - if USE_PDL and not EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - -def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: - """Split-K factor over the top-k blocks, as a power of two. - - Depends only on shapes that are fixed for a captured batch size, so the - launch geometry is frozen inside a CUDA graph. - """ - target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) - return 1 << (target.bit_length() - 1) - - # NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. _SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) # Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. -_NVFP4_TUNED_CAPABILITIES = ((10, 3),) - - -def _sm100f_nvfp4_num_topk_chunks( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> Optional[int]: - """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return None - local_batch = total_q // decode_query_len - return { - 2: 8, - 3: 8, - 4: 8, - 5: 16, - 6: 4, - 7: 4, - 8: 2, - 9: 2, - 10: 8, - 11: 4, - 12: 2, - 13: 2, - 14: 2, - }.get(local_batch) - - -def _sm100f_nvfp4_use_linear_softmax( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> bool: - """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads not in (1, 2, 4) - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return False - local_batch = total_q // decode_query_len - # The TP4/DEP32 single-request graph is 8.9% faster with the original - # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral - # or faster with the linear denominator. - return not (num_kv_heads == 1 and local_batch == 1) - - -def _sm100f_nvfp4_query_group_size( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> int: - """Group Eagle queries using the GB300-measured shared-page policy.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return 1 - # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches - # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or - # slower and retain the single-query producer. - return 2 if total_q // decode_query_len in (11, 12, 14) else 1 - - -def _sm100f_nvfp4_launch_options( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> dict[str, int]: - """Select native-NVFP4 launch geometry for the SM100 family. - - The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four - warps/two stages wins through local batch 28, while two warps/one stage - avoids the next residency cliff at local batches 29--32. Leave every - unmeasured architecture, model geometry, and larger batch on Triton's - established defaults. - """ - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - ): - return {} - local_batch = total_q // decode_query_len - if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): - return {"num_warps": 4, "num_stages": 1} - if local_batch == 4: - return {"num_warps": 2, "num_stages": 1} - if 8 <= local_batch <= 28: - return {"num_warps": 4, "num_stages": 2} - if 29 <= local_batch <= 32: - return {"num_warps": 2, "num_stages": 1} - return {} - - -def _sm100f_nvfp4_merge_launch_options( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> dict[str, int]: - """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability in _NVFP4_TUNED_CAPABILITIES - and num_kv_heads == 4 - and gqa_group_size == 16 - and max_topk == 16 - and decode_query_len == 4 - and total_q % decode_query_len == 0 - and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) - ): - return {"num_warps": 1} - return {} - - -def _check_nvfp4_inputs( - k_paged: torch.Tensor, - v_paged: torch.Tensor, - nvfp4_args: tuple[Optional[torch.Tensor], ...], - *, - head_dim: int, - scale_cols: int, -) -> None: - """Validate the packed-NVFP4 cache against what the kernel assumes. - - The kernel indexes block scales by a flat within-page offset and reshapes a - [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows - and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them - rather than let a mismatch read neighboring pages. - """ - k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args - if any(arg is None for arg in nvfp4_args): - raise ValueError( - "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " - "v_block_scale, k_global_scale and v_global_scale." - ) - if head_dim != triton.next_power_of_2(head_dim): - raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") - if v_paged.shape[:3] != k_paged.shape[:3]: - raise ValueError("NVFP4 K and V caches must have matching page, head and token extents.") - for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): - if int(data.shape[3]) * 2 != head_dim: - raise ValueError( - f"NVFP4 {name} must pack two elements per byte: expected " - f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." - ) - for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): - if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): - raise ValueError( - f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " - f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." - ) - if scale.stride(3) != 1 or scale.stride(2) != scale_cols: - raise ValueError( - f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " - f"{tuple(scale.stride())}." - ) - for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): - if gs.dtype != torch.float32 or gs.numel() != 1: - raise ValueError( - f"NVFP4 {name} must be a single-element FP32 tensor; got " - f"{gs.numel()} element(s) of {gs.dtype}." - ) # Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query @@ -1575,7 +488,7 @@ def minimax_m3_sparse_attn_decode( else {} ) if use_mq: - from .triton_sparse_decode_mq import launch_mq_producer + from ...kernels.attention.triton_sparse_decode_mq import launch_mq_producer launch_mq_producer( q, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py deleted file mode 100644 index 34794a95aadd..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/triton_sparse_decode_mq.py +++ /dev/null @@ -1,239 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Multi-query block-sparse decode producer for MiniMax-M3 short-sequence verify steps. - -In short-sequence mode every request attends at most ``topk`` KV blocks and block selection is the identity, so -all ``decode_query_len`` query tokens of a request read the same blocks. The per-token producer -(``_gqa_sparse_decode_kernel``) then dequantizes every block once per query token. This producer runs one program -per (request, block) for all of the request's query tokens: rows are (token, head) pairs, each block is -dequantized once, and the causal limit is a per-row mask. - -It writes the same per-block partials (output scaled by 1/l and the V global scale, log2 LSE) as the per-token -producer with one block per chunk, so ``_merge_topk_attn_out_kernel`` merges them unchanged. -""" - -import triton -import triton.language as tl - -from .triton_sparse_decode import _dequant_nvfp4_rows - - -@triton.jit( - do_not_specialize=["num_reqs"], do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"] -) -def _gqa_sparse_decode_mq_kernel( - q_ptr, # [num_reqs * QLEN, num_heads, head_dim], FP8 or FP16/BF16 - k_ptr, # packed NVFP4 K: [num_pages, 1, page_size, head_dim / 2] uint8 - v_ptr, # packed NVFP4 V: same layout - k_sc_ptr, # K block scales [num_pages, 1, page_size, head_dim / 16], token-linear - v_sc_ptr, # V block scales, 4x4 token-quad order - k_gs_ptr, - v_gs_ptr, - t_ptr, # topk table [1, total_q, topk]; the request's last token's row is read - o_ptr, # partials [topk, total_q, num_heads, head_dim] fp32 - lse_ptr, # partials [topk, total_q, num_heads] fp32 (log2) - block_table_ptr, - seq_lens, - num_reqs, - sm_scale, - stride_qn, - stride_qh, - stride_qd, - stride_k_blk, - stride_k_pos, - stride_v_blk, - stride_v_pos, - stride_ksc_blk, - stride_vsc_blk, - stride_tn, - stride_tk, - stride_o_c, - stride_o_b, - stride_o_h, - stride_l_c, - stride_l_b, - stride_l_h, - stride_bt_b, - QLEN: tl.constexpr, - GROUP: tl.constexpr, # query heads per KV head (16) - BLOCK_R: tl.constexpr, # next power of two of QLEN * GROUP - HEAD_DIM: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - SCALE_COLS: tl.constexpr, - USE_PDL: tl.constexpr, - EARLY_TRIGGER: tl.constexpr, -): - pid_r = tl.program_id(0) - pid_c = tl.program_id(1) # top-k position == block index under identity selection - sm_scale_log2e = sm_scale * 1.4426950409 * tl.load(k_gs_ptr) - v_global_scale = tl.load(v_gs_ptr) - - off_r = tl.arange(0, BLOCK_R) - off_n = tl.arange(0, BLOCK_SIZE_K) - off_d = tl.arange(0, HEAD_DIM) - off_dp = tl.arange(0, HEAD_DIM // 2) - off_s = tl.arange(0, SCALE_COLS) - row_t = off_r // GROUP - row_h = off_r % GROUP - r_mask = off_r < QLEN * GROUP - tok = pid_r * QLEN + row_t - - if USE_PDL: - tl.extra.cuda.gdc_wait() - if EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - seq_len = tl.load(seq_lens + pid_r) - # Token t of the request sits at seq_len - QLEN + t and attends positions [0, kv_len_t). - kv_len = tl.maximum(seq_len - QLEN + row_t + 1, 0) - blk = tl.load(t_ptr + (pid_r * QLEN + QLEN - 1) * stride_tn + pid_c * stride_tk) - blk_ok = blk >= 0 - blk = tl.where(blk_ok, blk, 0) - page = tl.load(block_table_ptr + pid_r * stride_bt_b + blk, mask=blk_ok, other=0).to(tl.int64) - # Rows past the last token's position are never read: their cache bytes may be stale (0 * NaN in the PV dot). - n_valid = tl.where( - blk_ok, tl.minimum(tl.maximum(seq_len - blk * BLOCK_SIZE_K, 0), BLOCK_SIZE_K), 0 - ) - n_mask = off_n < n_valid - - q = tl.load( - q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :] * stride_qd, - mask=r_mask[:, None], - other=0.0, - ).to(tl.float16) - k_deq = _dequant_nvfp4_rows( - tl.load( - k_ptr + page * stride_k_blk + off_n[:, None] * stride_k_pos + off_dp[None, :], - mask=n_mask[:, None], - other=0, - ), - tl.load( - k_sc_ptr + page * stride_ksc_blk + off_n[:, None] * SCALE_COLS + off_s[None, :], - mask=n_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=HEAD_DIM, - SCALE_COLS=SCALE_COLS, - SF_VEC=HEAD_DIM // SCALE_COLS, - ) - valid = ( - (blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None]) & n_mask[None, :] & r_mask[:, None] - ) - qk = tl.dot(q, tl.trans(k_deq).to(tl.float16)) * sm_scale_log2e - qk = tl.where(valid, qk, float("-inf")) - m = tl.max(qk, axis=1) - m_safe = tl.where(m == float("-inf"), 0.0, m) - p = tl.exp2(qk - m_safe[:, None]) - p_sum = tl.sum(p, axis=1) - v = _dequant_nvfp4_rows( - tl.load( - v_ptr + page * stride_v_blk + off_n[:, None] * stride_v_pos + off_dp[None, :], - mask=n_mask[:, None], - other=0, - ), - tl.load( - v_sc_ptr - + page * stride_vsc_blk - + (off_n[:, None] // 4) * (4 * SCALE_COLS) - + 4 * off_s[None, :] - + (off_n[:, None] % 4), - mask=n_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=HEAD_DIM, - SCALE_COLS=SCALE_COLS, - SF_VEC=HEAD_DIM // SCALE_COLS, - ).to(tl.float16) - acc = tl.dot(p.to(tl.float16), v) - # An empty partial stores zero with LSE -inf, so the merge gives it weight 0. - scale = tl.where(p_sum > 0.0, 1.0 / p_sum, 0.0) * v_global_scale - lse = tl.where(p_sum > 0.0, m + tl.log2(p_sum), float("-inf")) - acc = acc * scale[:, None] - tl.store( - o_ptr - + pid_c * stride_o_c - + tok[:, None] * stride_o_b - + row_h[:, None] * stride_o_h - + off_d[None, :], - acc, - mask=r_mask[:, None], - ) - tl.store(lse_ptr + pid_c * stride_l_c + tok * stride_l_b + row_h * stride_l_h, lse, mask=r_mask) - if USE_PDL and not EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - -def launch_mq_producer( - q, - k_paged, - v_paged, - k_block_scale, - v_block_scale, - k_global_scale, - v_global_scale, - topk_idx, - block_table, - seq_lens, - o_arg, - lse_arg, - *, - sm_scale, - decode_query_len, - use_pdl, - early_trigger, - num_warps=8, - num_stages=1, -): - """Launch the multi-query producer; arguments as minimax_m3_sparse_attn_decode after its NVFP4 views.""" - total_q, num_heads, head_dim = q.shape - num_reqs = total_q // decode_query_len - max_topk = int(topk_idx.shape[-1]) - rows = decode_query_len * num_heads - pdl_launch = {"launch_pdl": True} if use_pdl else {} - _gqa_sparse_decode_mq_kernel[(num_reqs, max_topk)]( - q, - k_paged, - v_paged, - k_block_scale, - v_block_scale, - k_global_scale, - v_global_scale, - topk_idx, - o_arg, - lse_arg, - block_table, - seq_lens, - num_reqs, - sm_scale, - q.stride(0), - q.stride(1), - q.stride(2), - k_paged.stride(0), - k_paged.stride(2), - v_paged.stride(0), - v_paged.stride(2), - k_block_scale.stride(0), - v_block_scale.stride(0), - topk_idx.stride(1), - topk_idx.stride(2), - o_arg.stride(0), - o_arg.stride(1), - o_arg.stride(2), - lse_arg.stride(0), - lse_arg.stride(1), - lse_arg.stride(2), - block_table.stride(0), - QLEN=decode_query_len, - GROUP=num_heads, - BLOCK_R=triton.next_power_of_2(rows), - HEAD_DIM=head_dim, - BLOCK_SIZE_K=int(k_paged.shape[2]), - SCALE_COLS=head_dim // 16, - USE_PDL=use_pdl, - EARLY_TRIGGER=early_trigger, - num_warps=num_warps, - num_stages=num_stages, - **pdl_launch, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py index 69ea6f145b80..4c1b7189cc08 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/modeling.py @@ -61,9 +61,9 @@ from __future__ import annotations import atexit -import copy import dataclasses import os +import sys import threading import time from dataclasses import dataclass @@ -79,6 +79,7 @@ Mapping, Optional, Tuple, + cast, ) import torch @@ -106,6 +107,98 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog.quantization.fp4_quantize import ( fp4_quantize, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import ( + modeling_support as _modeling_support, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3 import reporting as _reporting +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _is_vl_config as _is_vl_config, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _moe_layer_ids as _moe_layer_ids, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _sparse_layer_ids as _sparse_layer_ids, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _text_model_config as _text_model_config, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.config import ( + _validate_runtime as _validate_runtime, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.activation.m3_silu_and_mul import ( + m3_silu_and_mul, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + msa_kv_update_hook as _msa_kv_update_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_fp8_indexer_qk_norm_rope import ( + m3_fp8_indexer_qk_norm_rope as m3_fp8_indexer_qk_norm_rope, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qk_norm_rope_fp8 import ( + m3_qk_norm_rope_fp8, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qkv_insert import ( + m3_qkv_insert, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_select_blocks import ( + m3_select_blocks, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( + m3_allreduce, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.m3_embedding import ( + m3_embedding, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_bf16_linear import ( + m3_bf16_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + ALPHA as DENSE_ALPHA, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + BETA as DENSE_BETA, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + LIMIT as DENSE_LIMIT, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( + m3_dense_swiglu, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_fp32_linear import ( + m3_fp32_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( + m3_mxfp8_gemm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_linear import ( + m3_mxfp8_linear as m3_mxfp8_linear, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + MxMode as MxMode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + _mx_mode as _mx_mode, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.mxfp8_state import ( + _Mxfp8GemmState as _Mxfp8GemmState, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_tc import ( + m3_router_tc, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( + m3_capture_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_chunk_norm import ( + m3_chunk_norm, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_draft_input import ( + m3_draft_input, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( + m3_force_vec, + m3_force_vec_req, +) from tensorrt_llm._torch.attention.backends import AttentionMetadata from tensorrt_llm._torch.attention.backends.fmha.msa_prefill import _aligned_nvfp4_dequant_scales from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams @@ -119,22 +212,27 @@ ) from tensorrt_llm._torch.attention.backends.utils import create_attention from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( - FlashInferMXFP8GemmRunner, - MXFP8QuantizeRunner, + FlashInferMXFP8GemmRunner as FlashInferMXFP8GemmRunner, +) +from tensorrt_llm._torch.custom_ops.torch_custom_ops import ( + MXFP8QuantizeRunner as MXFP8QuantizeRunner, ) from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_speculative import ( _ensure_draft_vocab_size, - _get_requested_draft_moe_backend, - _set_draft_kv_cache_quant_algo, _slice_spec_position_ids, ) +from tensorrt_llm._torch.models.modeling_speculative import ( + _get_requested_draft_moe_backend as _get_requested_draft_moe_backend, +) +from tensorrt_llm._torch.models.modeling_speculative import ( + _set_draft_kv_cache_quant_algo as _set_draft_kv_cache_quant_algo, +) from tensorrt_llm._torch.models.modeling_utils import DecoderModel, DecoderModelForCausalLM from tensorrt_llm._torch.modules.embedding import LMHead from tensorrt_llm._torch.modules.linear import ( - MXFP8LinearMethod, TensorParallelMode, flashinfer_mxfp8_autotune_active, flashinfer_mxfp8_decode_graph_capture_active, @@ -152,7 +250,9 @@ _FORCE_ACCEPT_RNG_COUNTER_STRIDE, _FORCE_ACCEPT_RNG_POOL_SIZE, _FORCE_ACCEPT_RNG_SLOT_STRIDE, - should_use_separate_draft_kv_cache, +) +from tensorrt_llm._torch.speculative.interface import ( + should_use_separate_draft_kv_cache as should_use_separate_draft_kv_cache, ) from tensorrt_llm._torch.utils import ActType_TrtllmGen from tensorrt_llm._utils import mpi_rank, prefer_pinned @@ -175,33 +275,19 @@ from .ops import m3_index_score as _index_score_kernel from .ops import m3_sampled_accept as m3_sampled_accept_kernel from .ops import m3_shard_accept as m3_shard_accept_kernel -from .ops import msa_kv_update as _msa_kv_update_kernel from .ops import triton_sparse_decode as _sparse_decode_kernel -from .ops.activation.m3_silu_and_mul import m3_silu_and_mul from .ops.attention.m3_draft_attention import m3_draft_attention -from .ops.attention.m3_fp8_indexer_qk_norm_rope import m3_fp8_indexer_qk_norm_rope from .ops.attention.m3_fullctx_decode import m3_fullctx_decode from .ops.attention.m3_index_score import m3_index_score from .ops.attention.m3_msa_attention import m3_msa_attention from .ops.attention.m3_producer import m3_producer -from .ops.attention.m3_qk_norm_rope_fp8 import m3_qk_norm_rope_fp8 -from .ops.attention.m3_qkv_insert import m3_qkv_insert -from .ops.attention.m3_select_blocks import m3_select_blocks from .ops.attention.m3_select_long import m3_select_long from .ops.attention.m3_select_long_dsl import m3_select_long_dsl from .ops.attention.m3_sparse_decode import m3_sparse_decode from .ops.attention.m3_sparse_decode_lc8 import m3_sparse_decode_lc8 -from .ops.comm.m3_allreduce import m3_allreduce from .ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm -from .ops.embedding.m3_embedding import m3_embedding -from .ops.gemm.m3_bf16_linear import m3_bf16_linear -from .ops.gemm.m3_dense_swiglu import ALPHA as DENSE_ALPHA -from .ops.gemm.m3_dense_swiglu import BETA as DENSE_BETA -from .ops.gemm.m3_dense_swiglu import LIMIT as DENSE_LIMIT -from .ops.gemm.m3_dense_swiglu import m3_dense_swiglu from .ops.gemm.m3_draft_lowm import ROWS as DRAFT_LOWM_TACTIC_ROWS from .ops.gemm.m3_draft_lowm import m3_draft_lowm, tactic_for -from .ops.gemm.m3_fp32_linear import m3_fp32_linear from .ops.gemm.m3_mx_gemv import ( DENSE_DOWN, DENSE_GATE_UP, @@ -213,19 +299,12 @@ m3_mx_gemv, ) from .ops.gemm.m3_mx_gemv_bf16 import m3_mx_gemv_bf16 -from .ops.gemm.m3_mxfp8_gemm import m3_mxfp8_gemm -from .ops.gemm.m3_mxfp8_linear import m3_mxfp8_linear from .ops.long_select import long_select_blocks from .ops.moe.m3_moe_engine import M3MoeState, ar_push_desc, m3_moe_engine -from .ops.moe.m3_router_tc import m3_router_tc from .ops.mx_gemv import op as mx_gemv from .ops.norm.m3_embedding_firstnorm import DecodeEmbedding, m3_embedding_firstnorm -from .ops.spec.m3_capture_norm import m3_capture_norm -from .ops.spec.m3_chunk_norm import m3_chunk_norm -from .ops.spec.m3_draft_input import m3_draft_input from .ops.spec.m3_draft_silu import m3_draft_silu from .ops.spec.m3_draft_step import m3_draft_step -from .ops.spec.m3_force_vec import m3_force_vec, m3_force_vec_req from .ops.spec.m3_sampled_accept import m3_sampled_accept from .ops.spec.m3_shard_accept import m3_shard_accept from .ops.spec.m3_spec_offsets import m3_spec_offsets @@ -363,7 +442,7 @@ class DecodePlan: # position_ids.reshape(-1).contiguous().to(torch.int32) and packed.contiguous() launch a plain torch copy right before # the producer when the dtype or layout differs: also a plain writer, so the contract holds. Bit-identical outputs. # TP8: True runs this target's copy of the producer (trtllm::minimax_m3_sm107_tp8_producer_pw, -# minimaxM3Sm107Tp8/producerPw8.cu, the TP2 producerPw2 at 8 query heads) with the pre-wait reads; False (the kill +# minimaxM3Sm107/producerPw.cu, specialized at 8 query heads) with the pre-wait reads; False (the kill # value) runs the shared minimax_m3_producer_pw at 8 query heads, which reads them after its wait. # Step-input writers added since: the step-input prologue's C++ op (STEP_PACK_CPP_PROLOGUE) and the C++ spec fold tail # (STEP_FOLD_TAIL_CPP) are plain launches like the Triton kernels they replace (their tests assert it on the captured @@ -637,7 +716,9 @@ def covers(self, key: Key) -> bool: return key in self.keys def certified(self) -> FrozenSet[Key]: - return frozenset(k for cell in self.entry.CELLS for k in self.certifies(cell.spec)) + cells_by_tp = getattr(self.entry, "CELLS_BY_TP", None) + cells = self.entry.CELLS if cells_by_tp is None else cells_by_tp[8] + return frozenset(k for cell in cells for k in self.certifies(cell.spec)) def _rows(lo: int, hi: int) -> FrozenSet[int]: @@ -966,7 +1047,7 @@ def moe_path(num_tokens: int) -> Optional[str]: MX_GEMV_PATH = MX_GEMV.path MX_GEMV_BF16_PATH = MX_GEMV_BF16.path # The dense layers' swigluoai activation as this target's C++ op (ops/gemm/m3_dense_swiglu, -# minimaxM3Sm107Tp8/swigluOai.cu, LongAttn c267035d2f: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, +# minimaxM3Sm107/swigluOai.cu, LongAttn c267035d2f: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, # beta 1), PDL-launched with an early release so the down GEMV stages its weights during it; the down GEMV reads the # activation only after its own grid wait. Indexer kdev17 783285 (TP2 chain): -0.69 us per dense layer vs the Triton # kernel; LongAttn gate a 783511 (oracle 23/23). False runs the Triton kernel. @@ -1004,58 +1085,36 @@ class FastPathUnavailable(RuntimeError): _lock = threading.Lock() _require = [False] _logged: set = set() +_POLICY_TARGET = cast(_reporting.ReportingTarget, sys.modules[__name__]) def set_require(require: bool) -> None: """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" - _require[0] = bool(require) + return _reporting.set_require(_POLICY_TARGET, require) def require_mode() -> bool: - return _require[0] + return _reporting.require_mode(_POLICY_TARGET) def _log_once(level: Callable[[str], None], path: str, reason: str, message: str) -> None: - key = (path, reason) - with _lock: - if key in _logged: - return - _logged.add(key) - level(message) + return _reporting._log_once(_POLICY_TARGET, level, path, reason, message) def missed(path: str, rows: Optional[int], reason: str) -> None: """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" - at = "" if rows is None else f" at {rows} rows" - if _require[0]: - raise FastPathUnavailable( - f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: " - f"{reason}" - ) - _log_once( - logger.warning, - path, - reason, - f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", - ) + return _reporting.missed(_POLICY_TARGET, path, rows, reason) def capture_missed(path: str, rows: Optional[int], what: str) -> None: """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" - missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + return _reporting.capture_missed(_POLICY_TARGET, path, rows, what) def taken(path: str, rows: Optional[int], what: str = "") -> None: """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" - at = "" if rows is None else f" at {rows} rows" - detail = f" ({what})" if what else "" - _log_once( - logger.info, - path, - f"taken: {rows}" + (f": {what}" if what else ""), - f"[M3 dispatch] rank {mpi_rank()}: {path} taken{at}{detail}", - ) + return _reporting.taken(_POLICY_TARGET, path, rows, what) def long_select_covers(num_blocks: int) -> bool: @@ -1096,13 +1155,12 @@ def long_select_fallback(num_blocks: int, rows: int) -> None: def uncovered(path: str, reason: str) -> None: """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" - _log_once(logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)") + return _reporting.uncovered(_POLICY_TARGET, path, reason) def reset_logged() -> None: """Forget which (path, reason) pairs were logged.""" - with _lock: - _logged.clear() + return _reporting.reset_logged(_POLICY_TARGET) # Kernel modules report from inside their calls through ops/m3_report: the same policy. @@ -1160,132 +1218,6 @@ def _print_values() -> None: atexit.register(_print_values) -def _is_vl_config(pretrained_config: PretrainedConfig) -> bool: - model_type = getattr(pretrained_config, "model_type", None) - if model_type == "minimax_m3_vl": - return True - if model_type == "minimax_m3": - return False - return any( - "ForConditionalGeneration" in a - for a in (getattr(pretrained_config, "architectures", None) or []) - ) - - -def _text_model_config(model_config: ModelConfig) -> ModelConfig: - """The model config with the language model's ``text_config`` as its pretrained config. - - The checkpoint declares the vision-language wrapper; its ``text_config`` arrives as a plain dict (no typed - sub-config is registered), so it is materialized into a ``PretrainedConfig`` with the wrapper's dtype and - embedding tie carried over. A text config is returned unchanged. - """ - if not _is_vl_config(model_config.pretrained_config): - return model_config - cfg = copy.deepcopy(model_config) - text = getattr(cfg.pretrained_config, "text_config", None) - if text is None: - raise ValueError("MiniMax-M3 VL config has no `text_config`") - if isinstance(text, dict): - out = PretrainedConfig() - for k, v in text.items(): - setattr(out, k, v) - text = out - if ( - getattr(text, "torch_dtype", None) is None - and getattr(cfg.pretrained_config, "torch_dtype", None) is not None - ): - text.torch_dtype = cfg.pretrained_config.torch_dtype - if ( - getattr(text, "tie_word_embeddings", None) is None - and getattr(cfg.pretrained_config, "tie_word_embeddings", None) is not None - ): - text.tie_word_embeddings = cfg.pretrained_config.tie_word_embeddings - return dataclasses.replace(cfg, pretrained_config=text) - - -def _validate_runtime(model_config: ModelConfig) -> None: - """The runtime this target is written for. It is set by engine arguments rather than the checkpoint, so routing - sees only its KV cache algorithm; this check at construction refuses every other runtime. - - The MiniMax-M3 MSA backend (it owns the step metadata and the index-K cache), the fused five-way QKV/index - projection or the separate QKV and index projections, an FP8 indexer and an FP8 or NVFP4 main KV cache. - """ - sparse = model_config.sparse_attention_config - if ( - sparse is None - or getattr(sparse, "algorithm", None) != "minimax_m3" - or sparse.implementation != "msa" - ): - raise ValueError( - "this MiniMax-M3 target requires sparse_attention_config.algorithm='minimax_m3' with implementation='msa'" - ) - if getattr(sparse, "indexer_kv_dtype", None) != "fp8": - raise ValueError( - "this MiniMax-M3 target requires sparse_attention_config.indexer_kv_dtype='fp8'" - ) - mode = model_config.quant_config.quant_mode - if mode is None or not (mode.has_fp8_kv_cache() or mode.has_fp4_kv_cache()): - raise ValueError("this MiniMax-M3 target requires an FP8 or NVFP4 KV cache") - - -def _sparse_layer_ids(cfg: PretrainedConfig) -> List[int]: - freq = cfg.sparse_attention_config["sparse_attention_freq"] - return [i for i, f in enumerate(freq) if int(f) != 0] - - -def _moe_layer_ids(cfg: PretrainedConfig) -> List[int]: - freq = getattr(cfg, "moe_layer_freq", None) - if freq is None: - return list(range(cfg.num_hidden_layers)) - return [i for i, f in enumerate(freq) if int(f) != 0] - - -class _Mxfp8GemmState(nn.Module): - """The MXFP8 GEMM backend state the engine's warmup drives. - - The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM - tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of - this target dispatches identically, so one method object carries that state for all of them; the forward reads it - at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. - """ - - def __init__(self) -> None: - super().__init__() - self.quant_method = MXFP8LinearMethod() - - -@dataclass(frozen=True) -class MxMode: - """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" - - decode_graph: bool # inside a decode CUDA-graph warmup or capture - autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass - tuned_quantize_gemm: ( - bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) - ) - flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation - native_tuning: ( - bool # the native GEMM through the autotuner (warmup, before its tactics are registered) - ) - - -def _mx_mode(state: _Mxfp8GemmState) -> MxMode: - method = state.quant_method - decode_graph = flashinfer_mxfp8_decode_graph_capture_active() - autotune = flashinfer_mxfp8_autotune_active() - return MxMode( - decode_graph=decode_graph, - autotune=autotune, - tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, - flashinfer=method.backend == "flashinfer" - or ( - method.backend == "auto" - and (autotune or (method._flashinfer_autotuned and decode_graph)) - ), - native_tuning=method.needs_native_autotune, - ) - - # -------------------------------------------------------------------------------------------------------------------- # Projections: the MXFP8 GEMM dispatch (``linear_path``) and its hand-offs # -------------------------------------------------------------------------------------------------------------------- @@ -1295,50 +1227,27 @@ def _mxfp8_linear( core, x: torch.Tensor, w: torch.Tensor, s: torch.Tensor, mode: MxMode ) -> torch.Tensor: """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" - _record( - "m3_mxfp8_linear", - tuple(w.shape), - mode.tuned_quantize_gemm, - mode.flashinfer, - mode.native_tuning, - ) - return m3_mxfp8_linear( - x.contiguous(), - w, - s, - core.mx_alpha, - torch.bfloat16, - mode.tuned_quantize_gemm, - mode.autotune, - mode.flashinfer, - mode.native_tuning, - ) + return _modeling_support._mxfp8_linear(_SUPPORT_TARGET, core, x, w, s, mode) _PINNED_RUNNERS: dict = {} def _pinned_runner(cls, dtype): - runner = _PINNED_RUNNERS.get((cls, dtype)) - if runner is None: - runner = _PINNED_RUNNERS[(cls, dtype)] = cls(dtype) - return runner + return _modeling_support._pinned_runner(_SUPPORT_TARGET, cls, dtype) def _pinned_mxfp8_gemm(act, act_scale, weight, weight_scale, out_dtype): """The decode-graph MXFP8 GEMM at the target's fixed tactic (``MXFP8_DECODE_GEMM_TACTIC``), no AutoTuner: the same choice on every rank and every start.""" - _record("pinned_mxfp8_gemm", tuple(weight.shape), act.shape[0], MXFP8_DECODE_GEMM_TACTIC) - runner = _pinned_runner(FlashInferMXFP8GemmRunner, out_dtype) - return runner.forward([act, act_scale, weight, weight_scale], tactic=MXFP8_DECODE_GEMM_TACTIC) + return _modeling_support._pinned_mxfp8_gemm( + _SUPPORT_TARGET, act, act_scale, weight, weight_scale, out_dtype + ) def _pinned_mxfp8_quantize(x): """The decode-graph MXFP8 activation quantize at the target's fixed tactic (``MXFP8_DECODE_QUANT_TACTIC``).""" - _record("pinned_mxfp8_quantize", x.shape[0], MXFP8_DECODE_QUANT_TACTIC) - return _pinned_runner(MXFP8QuantizeRunner, x.dtype).forward( - [x], tactic=MXFP8_DECODE_QUANT_TACTIC - ) + return _modeling_support._pinned_mxfp8_quantize(_SUPPORT_TARGET, x) _MX_TAC_ENGAGED: set = set() @@ -1665,8 +1574,7 @@ def _routed_experts(core, i: int, x, x_sf, router_logits, do_finalize: bool): def _split_chunk(rows: int, num_chunks: int) -> List[int]: """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" - base, extra = divmod(rows, num_chunks) - return [base + 1] * extra + [base] * (num_chunks - extra) + return _modeling_support._split_chunk(_SUPPORT_TARGET, rows, num_chunks) # -------------------------------------------------------------------------------------------------------------------- @@ -1675,7 +1583,7 @@ def _split_chunk(rows: int, num_chunks: int) -> List[int]: def _nvfp4_layer(manager, i: int) -> bool: - return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + return _modeling_support._nvfp4_layer(_SUPPORT_TARGET, manager, i) def _producer_covered(core) -> bool: @@ -1702,31 +1610,7 @@ def _producer_layout(kv_data, kv_scale) -> bool: def _nvfp4_fused_write_layout(data: torch.Tensor, scales: torch.Tensor, num_heads: int) -> bool: """Reject invalid HND pools; select the fused producer only for its packed layout.""" - if ( - data.dim() != 5 - or scales.dim() != 5 - or data.shape[0] <= 0 - or data.shape[0] != scales.shape[0] - or tuple(data.shape[1:]) != (2, num_heads, 128, 64) - or tuple(scales.shape[1:]) != (2, num_heads, 128, 8) - or data.dtype not in (torch.int8, torch.uint8) - or scales.dtype != torch.uint8 - ): - raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") - for cache in (data, scales): - if ( - cache.stride(4) != 1 - or cache.stride(3) != cache.shape[4] - or cache.stride(2) < 128 * cache.stride(3) - or cache.stride(1) < num_heads * cache.stride(2) - or cache.stride(0) < 2 * cache.stride(1) - ): - raise ValueError( - "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" - ) - if scales.stride(2) != 128 * 8: - raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") - return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and data.stride(1) % 2 == 0 + return _modeling_support._nvfp4_fused_write_layout(_SUPPORT_TARGET, data, scales, num_heads) #: The step phases ("decode" / "prefill") whose producer engagement has been logged, once per process. @@ -2055,14 +1939,14 @@ def _dense_qk_norm_rope_fp8( def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: - """The dense MLP's SwiGLU-OAI activation: this target's C++ op (``m3_dense_swiglu``, minimaxM3Sm107Tp8/ + """The dense MLP's SwiGLU-OAI activation: this target's C++ op (``m3_dense_swiglu``, minimaxM3Sm107/ swigluOai.cu: bit for bit the Triton silu_and_mul at limit 7, alpha 1.702, beta 1, PDL-launched with an early release so the down GEMV stages its weights during it) at the rows its cells certify, in any graph or eager step; otherwise the Triton kernel (``m3_silu_and_mul``). Indexer kdev17 (TP2 chain): -0.69 us per dense layer vs the Triton kernel. Every decline is logged once.""" path = DENSE_SWIGLU_CPP_PATH rows = gu.shape[0] - max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS) + max_rows = max(cell.spec["rows"] for cell in m3_dense_swiglu.CELLS_BY_TP[8]) params = (core.swiglu_alpha, core.swiglu_limit, core.swiglu_beta) act = None if rows > max_rows: @@ -2071,7 +1955,12 @@ def _dense_swiglu(core, gu: torch.Tensor) -> torch.Tensor: uncovered(path, f"activation parameters {params}") else: act = m3_dense_swiglu( - gu, core.swiglu_limit, core.swiglu_alpha, core.swiglu_beta, get_env_enable_pdl() + gu, + core.swiglu_limit, + core.swiglu_alpha, + core.swiglu_beta, + get_env_enable_pdl(), + tp_size=8, ) if act is None: uncovered(path, "the gate_up layout") @@ -2089,31 +1978,9 @@ def _fp8_index_qk_norm_rope( """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` are the calling target's ``step_args`` entries.""" - if position_ids is None or idx_qk.dtype != torch.bfloat16: - raise NotImplementedError( - "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." - ) - index_k_cache = msa_idx_k_cache(i) - if index_k_cache.dtype != torch.float8_e4m3fn: - raise ValueError( - "MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, " - f"got {index_k_cache.dtype}." - ) - rows = int(idx_qk.shape[0]) - _record("m3_fp8_indexer_qk_norm_rope", rows) - return m3_fp8_indexer_qk_norm_rope( - idx_qk.contiguous(), - index_k_cache, - msa_out_cache_loc[:rows], - core.index_heads, - core.index_dim, - core.rotary_dim, - core.rms_eps, - core.w[f"l{i}_idx_q_norm"], - core.w[f"l{i}_idx_k_norm"], - core.rope_theta, - position_ids.reshape(-1).contiguous().to(torch.int32), - ).flatten(1) + return _modeling_support._fp8_index_qk_norm_rope( + _SUPPORT_TARGET, core, i, idx_qk, position_ids, msa_idx_k_cache, msa_out_cache_loc + ) # -------------------------------------------------------------------------------------------------------------------- @@ -3941,6 +3808,7 @@ def _index_scores(self, i: int, idx_q, step, num_tokens: int, qlen: int): qlen, split_k, early, + runner=_index_score_kernel.IndexScoreRunner, ) taken(INDEX_SCORE_PATH, rows) else: @@ -4528,7 +4396,9 @@ def _lowm_linear(self, weight: torch.Tensor, x: torch.Tensor) -> Optional[torch. def prebuild_linears(self) -> int: """Compile (and run once) every covered row count of every drafter projection, outside capture, on the current stream, before the executor's warmup. Returns the number of calls compiled.""" - from .ops.m3_lowm_splitk import compile_push + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_lowm_splitk import ( + compile_push, + ) use_pdl = get_env_enable_pdl() built = 0 @@ -4611,7 +4481,7 @@ def attention(self, x: torch.Tensor, md, reduce: bool) -> torch.Tensor: def mlp(self, h: torch.Tensor, reduce: bool, cpp_silu: bool = False) -> torch.Tensor: """gate_up, the plain SwiGLU, down (all-reduced unless ``reduce`` is False). ``cpp_silu`` (the fused draft - loop): the SwiGLU on this target's C++ op (``ops/spec/m3_draft_silu``, minimaxM3Sm107Tp8/siluAndMul.cu), bit + loop): the SwiGLU on this target's C++ op (``ops/spec/m3_draft_silu``, minimaxM3Sm107/siluAndMul.cu), bit for bit the in-tree Triton silu_and_mul in one launch either way (no Triton on the GEN decode path).""" gu = self.linear(h, "gu") act = None @@ -4679,41 +4549,14 @@ def _eagle3_capture(spec_metadata, layer_id: int, hidden_states, residual, captu in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; otherwise it is the in-tree capture. """ - if capture_norm is None: - spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) - return - for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): - if captured_layer_id == layer_id: - assert spec_metadata.hidden_states is not None - num_tokens = spec_metadata.num_capture_tokens - if num_tokens > spec_metadata.hidden_states.shape[0]: - raise RuntimeError( - "EAGLE3 hidden-state capture token count exceeds buffer capacity: " - f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" - ) - if num_tokens > hidden_states.shape[0]: - raise RuntimeError( - "EAGLE3 hidden-state capture token count exceeds available hidden states: " - f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" - ) - hidden_size = spec_metadata.hidden_size - m3_spec_tail.capture_norm( - hidden_states[:num_tokens], - residual[:num_tokens] if residual is not None else None, - capture_norm[0][i], - capture_norm[1], - spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], - ) - break + return _modeling_support._eagle3_capture( + _SUPPORT_TARGET, spec_metadata, layer_id, hidden_states, residual, capture_norm + ) def _sharded_head(lm_head: nn.Module) -> bool: """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" - return ( - getattr(lm_head, "tp_mode", None) == TensorParallelMode.COLUMN - and getattr(lm_head, "padding_size", 0) == 0 - and getattr(lm_head, "tp_size", 1) > 1 - ) + return _modeling_support._sharded_head(_SUPPORT_TARGET, lm_head) class M3Eagle3Worker(Eagle3OneModelWorker): @@ -4919,6 +4762,7 @@ def _forward_impl( max_draft_len=runtime_draft_len, ) + next_new_tokens = None if self._uses_external_shared_target_kv: next_draft_tokens = self._forward_external_shared_target_kv_draft_loop( position_ids=position_ids, @@ -4932,6 +4776,8 @@ def _forward_impl( batch_size=batch_size, ) else: + # Retain the pre-draft tensor before metadata is replaced, as in the parent worker. + pre_draft_seq_lens = attn_metadata._seq_lens_cuda # Save the old attn_metadata and spec_metadata self._prepare_attn_metadata_for_spec_dec(attn_metadata) @@ -4954,7 +4800,7 @@ def _forward_impl( # Get the draft KV cache manager if using separate layouts draft_kv_cache_manager = self.get_draft_kv_cache_manager(resource_manager) - next_draft_tokens = self._forward_draft_loop( + next_draft_tokens, next_new_tokens = self._forward_draft_loop( inputs, attn_metadata, spec_metadata, @@ -4966,6 +4812,8 @@ def _forward_impl( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ) # restore attn_metadata to support cuda graph self._restore_attn_metadata_from_spec_dec(attn_metadata) @@ -4978,7 +4826,7 @@ def _forward_impl( if self._m3_next_new_tokens is not None: # Filled in place by the fused verify or the seed (column 0) and the fused draft passes (columns 1..). next_new_tokens = self._m3_next_new_tokens - else: + elif next_new_tokens is None: next_new_tokens = self._prepare_next_new_tokens( accepted_tokens, next_draft_tokens, @@ -5701,10 +5549,12 @@ def _forward_draft_loop( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ): """The fused draft loop on a step it covers (``_m3_fused_loop_ok``), else the in-tree loop.""" if self._m3_fused_loop_ok(attn_metadata, spec_metadata, batch_size, draft_model): - return self._m3_fused_draft_loop( + next_draft_tokens = self._m3_fused_draft_loop( inputs, attn_metadata, spec_metadata, @@ -5713,6 +5563,7 @@ def _forward_draft_loop( batch_size, num_accepted_tokens, ) + return next_draft_tokens, None return super()._forward_draft_loop( inputs, attn_metadata, @@ -5725,6 +5576,8 @@ def _forward_draft_loop( num_accepted_tokens, original_all_rank_num_tokens, resource_manager, + accepted_tokens, + pre_draft_seq_lens, ) def _m3_fused_draft_loop( @@ -5851,40 +5704,7 @@ def _m3_fused_draft_loop( def _build_speculation(shell, model_config: ModelConfig): """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for the llama3 drafter architecture).""" - spec = model_config.spec_config - if ( - not spec.spec_dec_mode.use_one_engine() - or not spec.spec_dec_mode.is_eagle3_one_model() - or spec.eagle3_model_arch != "llama3" - or getattr(spec, "use_dynamic_tree", False) - ): - raise ValueError( - "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" - ) - draft_config = ModelConfig.from_pretrained( - spec.speculative_model, - trust_remote_code=True, - attn_backend=model_config.attn_backend, - moe_backend=_get_requested_draft_moe_backend(model_config, spec), - mapping=model_config.mapping, - spec_config=model_config.spec_config, - max_num_tokens=model_config.max_num_tokens, - moe_max_num_tokens=model_config.moe_max_num_tokens, - ) - _set_draft_kv_cache_quant_algo(draft_config, model_config) - draft_config.extra_attrs = model_config.extra_attrs - use_separate = should_use_separate_draft_kv_cache(spec) - draft = M3Eagle3Draft( - shell.model, - draft_config, - model_config.pretrained_config.num_hidden_layers, - spec.num_capture_layers, - ) - worker = M3Eagle3Worker(spec, model_config.mapping, use_separate_draft_kv_cache=use_separate) - worker.set_draft_model(draft) - shell.epilogue.append(draft) - shell.epilogue.append(worker) - return draft_config, draft, worker, use_separate + return _modeling_support._build_speculation(_SUPPORT_TARGET, shell, model_config) # -------------------------------------------------------------------------------------------------------------------- @@ -6138,3 +5958,6 @@ def forward( # at its tail (one launch inside the step graph). spec_fold.store_in_forward(outputs) return outputs + + +_SUPPORT_TARGET = cast(_modeling_support.PinnedModelingSupport, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py deleted file mode 100644 index e7ceefc5bdab..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/activation/m3_silu_and_mul.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The gated MLPs' fused SwiGLU: ``silu_and_mul`` with MiniMax-M3's SwiGLU-OAI gain, offset and clamp.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3SiluAndMul(OpWrapper): - """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with - ``g = min(gate, limit)`` and ``u = clamp(up, -limit, limit)``; ``limit``, ``alpha`` and ``beta`` None give plain - SwiGLU ``silu(gate) * up``. - - The MXFP8 dense MLPs (layers 0-2, intermediate 6144 per rank) and the MoE's shared expert (intermediate 1536 per - rank) run it at MiniMax-M3's (limit 7, alpha 1.702, beta 1); the BF16 Eagle3 drafter's MLP (intermediate 9216 per - rank) runs it plain. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="shared expert, one decode token", spec=dict(m=1, inter=384, oai=True)), - Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=384, oai=True)), - Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=1536, oai=True)), - Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=1536, oai=True)), - Cell( - why="drafter MLP (plain SwiGLU), eight decode rows", - spec=dict(m=8, inter=2304, oai=False), - ), - Cell( - why="drafter MLP (plain SwiGLU), a prefill chunk", - spec=dict(m=1024, inter=2304, oai=False), - ), - ) - - note = """ - The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two - differ only by the fast sigmoid's last bits (the test's band). - """ - - def __call__( - self, - x: torch.Tensor, - limit: Optional[float], - alpha: Optional[float], - beta: Optional[float], - ) -> torch.Tensor: - return torch.ops.trtllm.silu_and_mul( - x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta - ) - - def reference( - self, - x: torch.Tensor, - limit: Optional[float], - alpha: Optional[float], - beta: Optional[float], - ) -> torch.Tensor: - gate, up = x.float().chunk(2, dim=-1) - if limit is not None: - gate = gate.clamp(max=limit) - up = up.clamp(min=-limit, max=limit) - a = 1.0 if alpha is None else alpha - b = 0.0 if beta is None else beta - return (gate * torch.sigmoid(a * gate) * (up + b)).to(x.dtype) - - def is_valid( - self, - x: torch.Tensor, - limit: Optional[float], - alpha: Optional[float], - beta: Optional[float], - ) -> None: - assert x.stride(-1) == 1, "the gate/up halves are read from contiguous rows" - - -m3_silu_and_mul = M3SiluAndMul() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py deleted file mode 100644 index 789ae79db88e..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_dense_decode.py +++ /dev/null @@ -1,190 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The dense layers' decode attention (trtllm-gen over the full page table) with the KV split this target sizes per -request count (modeling's ``DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH``).""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( - fp8_e4m3_saturate, -) -from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.trtllm_gen_dense_decode import ( # noqa: E501 - minimax_m3_trtllm_gen_dense_decode, -) - -__all__ = ["m3_dense_decode"] - -#: The TP4 dense layers' geometry: 16 query heads and one KV head of dim 128 per rank, an FP8 KV pool. -HEADS = 16 -KV_HEADS = 1 -HEAD_DIM = 128 -#: The split this target hands trtllm-gen at one request (modeling's DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH), and a model -#: maximum the in-tree step passes otherwise. -CAPPED = 16384 -MODEL_MAX = 1 << 20 - - -def dense_attention_reference( - q: torch.Tensor, - keys: torch.Tensor, - values: torch.Tensor, - seq_lens: torch.Tensor, - decode_query_len: int, - sm_scale: float, -) -> torch.Tensor: - """fp32 causal decode attention: query row ``b * qlen + i`` of request b attends its first - ``seq_lens[b] - qlen + i + 1`` keys. ``keys`` / ``values``: [batch, kv_heads, max_len, head_dim] (the request's - tokens in order). Returns fp32 [total_q, heads, head_dim].""" - batch, kv_heads, max_len, _ = keys.shape - heads = q.shape[1] - group = heads // kv_heads - out = torch.zeros(q.shape, dtype=torch.float32, device=q.device) - pos = torch.arange(max_len, device=q.device) - with true_fp32_matmul(): - for b in range(batch): - full = int(seq_lens[b]) - rows = q[b * decode_query_len : (b + 1) * decode_query_len].float() # [qlen, heads, d] - limit = full - decode_query_len + 1 + torch.arange(decode_query_len, device=q.device) - valid = pos[None, :] < limit[:, None] # [qlen, max_len] - for h in range(kv_heads): - qh = rows[:, h * group : (h + 1) * group] # [qlen, group, d] - logits = torch.einsum("qgd,kd->qgk", qh, keys[b, h].float()) * sm_scale - logits = logits.masked_fill(~valid[:, None, :], float("-inf")) - probs = torch.softmax(logits, dim=-1) - out[ - b * decode_query_len : (b + 1) * decode_query_len, h * group : (h + 1) * group - ] = torch.einsum("qgk,kd->qgd", probs, values[b, h].float()) - return out - - -def gather_paged_kv( - kv_cache_manager, layer_idx: int, block_table: torch.Tensor, seq_lens: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - """The K and V of each request in token order, [batch, kv_heads, max_len, head_dim], from the flat sub-page pool - (``get_kv_subpage_pool``: K of slot s at sub-page s * subpages_per_slot, V one sub-page later).""" - pool, subpages = kv_cache_manager.get_kv_subpage_pool(layer_idx, "HND") - page = pool.shape[2] - pages = (int(seq_lens.max()) + page - 1) // page - slots = block_table[:, :pages].long() - k = pool[slots * subpages] # [batch, pages, kv_heads, page, d] - v = pool[slots * subpages + 1] - batch, _, kv_heads, _, d = k.shape - - def order(t): - return t.permute(0, 2, 1, 3, 4).reshape(batch, kv_heads, pages * page, d) - - return order(k), order(v) - - -class M3DenseDecode(OpWrapper): - """``minimax_m3_trtllm_gen_dense_decode`` as ``MsaDecodeFmha`` calls it for a dense layer, with the max_seq_len - ``msa_decode.dense_decode_max_seq_len`` resolves: this target's per-request-count cap, else the step's bound. - - Decode attention of every query row over its request's full causal context in the paged FP8 pool, written into - ``output`` (BF16 [T, heads, head_dim]) in place. ``max_seq_len`` only sizes trtllm-gen's multi-CTA KV split: - every CTA still covers its share of each request's actual length, which may exceed it. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="one request verifying 1 + 3 rows over a 20000-token context, the split sized for 16384", - spec=dict(kv_lens=(20000,), qlen=4, max_seq_len=CAPPED, capped=True), - ), - Cell( - why="one request, one row, over a 65536-token context (an AgentX-length cache) at the capped split", - spec=dict(kv_lens=(65536,), qlen=1, max_seq_len=CAPPED, capped=True), - ), - Cell( - why="one request verifying 1 + 3 rows over a 3000-token context: the cap above the length", - spec=dict(kv_lens=(3000,), qlen=4, max_seq_len=CAPPED, capped=True), - ), - Cell( - why="two requests keep the model maximum (no cap at two requests)", - spec=dict(kv_lens=(20000, 5000), qlen=4, max_seq_len=MODEL_MAX, capped=False), - ), - ) - - note = """ - The split only reorders the merge of the CTAs' partial softmax sums, so the capped split's output is the uncapped - one's to a few BF16 ulps (the test holds it there), and both are held to the fp32 reference within the FP8-KV - band of the in-tree dense decode test. - """ - - def __call__( - self, - q: torch.Tensor, - kv_cache_manager, - layer_idx: int, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - sm_scale: float, - output: torch.Tensor, - decode_query_len: int, - max_seq_len: int, - max_num_requests: int, - ) -> torch.Tensor: - minimax_m3_trtllm_gen_dense_decode( - q, - kv_cache_manager, - layer_idx, - block_table, - seq_lens, - sm_scale=sm_scale, - output=output, - decode_query_len=decode_query_len, - max_seq_len=max_seq_len, - max_num_requests=max_num_requests, - ) - return output - - def reference( - self, - q: torch.Tensor, - kv_cache_manager, - layer_idx: int, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - sm_scale: float, - output: torch.Tensor, - decode_query_len: int, - max_seq_len: int, - max_num_requests: int, - ) -> torch.Tensor: - """fp32 attention over the pool's K / V (q rounded to the pool's dtype first, as the kernel runs it). - - ``output``, ``max_seq_len`` and ``max_num_requests`` do not change the result.""" - keys, values = gather_paged_kv(kv_cache_manager, layer_idx, block_table, seq_lens) - if keys.dtype == torch.float8_e4m3fn: - q = fp8_e4m3_saturate(q) - return dense_attention_reference(q, keys, values, seq_lens, decode_query_len, sm_scale) - - def is_valid( - self, - q: torch.Tensor, - kv_cache_manager, - layer_idx: int, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - sm_scale: float, - output: torch.Tensor, - decode_query_len: int, - max_seq_len: int, - max_num_requests: int, - ) -> None: - # The multi-CTA KV counters are sized for max_num_requests: more requests overrun them. - assert seq_lens.shape[0] <= max_num_requests, "more requests than the counters cover" - # Each verify row attends a causal prefix of at least one key. - assert bool((seq_lens >= decode_query_len).all()), "a request shorter than its query rows" - - -m3_dense_decode = M3DenseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py index c2b23aa18750..2989104e372c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_draft_attention.py @@ -4,17 +4,24 @@ from __future__ import annotations -import torch +import torch as torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.draft_attention import ( + GroupedDraftAttentionBase, +) +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs as AttentionForwardArgs, +) from tensorrt_llm._torch.attention.backends.interface import ( - AttentionForwardArgs, - PredefinedAttentionMask, + PredefinedAttentionMask as PredefinedAttentionMask, ) -from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType +from tensorrt_llm.functional import PositionEmbeddingType as PositionEmbeddingType +from tensorrt_llm.functional import RotaryScalingType as RotaryScalingType -class M3DraftAttention(OpWrapper): +class M3DraftAttention(GroupedDraftAttentionBase): """The drafter layer's attention over its fused QKV rows ``qkv`` [T, (q + 2 kv) heads x 128] (bf16): the backend applies the drafter's RoPE (fused), appends K / V to the drafter layer's paged cache (FP8, the target's NVFP4 cache keeps its draft layer in FP8) and runs causal attention over the step's live rows (``metadata`` @@ -38,94 +45,5 @@ class M3DraftAttention(OpWrapper): kernel's input precision, its P rounding and accumulation order), which the test bounds. """ - def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: - num_tokens = metadata.num_tokens - out = backend.forward( - qkv[:num_tokens, :], - None, - None, - metadata, - forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), - ) - return out[0] if isinstance(out, tuple) else out - - def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: - """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. - - Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and - attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from - the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated - after the call (whose roped, quantized append it then reads back; the test checks that append itself). - Mirrors grouped KV heads (query head h reads KV head h // (heads / kv heads); two per rank at tp2) and - un-scaled NeoX RoPE, the drafter's configuration. - """ - manager = metadata.kv_cache_manager - layer = backend.layer_idx - heads, dim = backend.num_heads, backend.head_dim - kv_heads = int(backend.num_kv_heads) - assert heads % kv_heads == 0 - group = heads // kv_heads - rope = backend.rope_params - assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) - assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 - page = int(manager.tokens_per_block) - sm_scale = dim**-0.5 / float(backend.q_scaling) - kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) - num_seqs = int(metadata.num_seqs) - qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] - cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] - total = int(metadata.num_tokens) - assert sum(qlens) == total - device = qkv.device - - rot = int(rope.dim) - inv_freq = float(rope.theta) ** ( - -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot - ) - pool = manager.get_buffers(layer, "HND") - q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) - out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) - first = 0 - for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): - n, length = qlens[row], cached[row] + qlens[row] - pages = torch.tensor( - manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ - : (length + page - 1) // page - ], - dtype=torch.int64, - device=device, - ) - # [kv heads, tokens, dim]: each KV head's rows in token order across the pages. - keys = ( - pool[pages, 0] - .to(torch.float64) - .transpose(0, 1) - .reshape(kv_heads, -1, dim)[:, :length] - * kv_scale - ) - values = ( - pool[pages, 1] - .to(torch.float64) - .transpose(0, 1) - .reshape(kv_heads, -1, dim)[:, :length] - * kv_scale - ) - keys = keys.repeat_interleave(group, dim=0) # [heads, tokens, dim] - values = values.repeat_interleave(group, dim=0) - qpos = cached[row] + torch.arange(n, device=device) - angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] - cos, sin = angles.cos()[:, None, :], angles.sin()[:, None, :] - rows = q[first : first + n] - x1, x2 = rows[..., : rot // 2], rows[..., rot // 2 : rot] - rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) - allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] - logits = torch.einsum("thd,hkd->thk", rows, keys) * sm_scale - logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) - out[first : first + n] = torch.einsum( - "thk,hkd->thd", torch.softmax(logits, dim=-1), values - ) - first += n - return out.reshape(total, heads * dim).to(qkv.dtype) - m3_draft_attention = M3DraftAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py deleted file mode 100644 index 2fecad0e9afa..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fp8_indexer_qk_norm_rope.py +++ /dev/null @@ -1,114 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The unfused sparse layers' index branch: per-head Gemma RMSNorm + partial NeoX RoPE of [index-Q | index-K], index-Q -out in FP8 E4M3 and index-K inserted into the FP8 index-K cache, in one kernel.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( - fused_qk_norm_rope, -) - - -class M3Fp8IndexerQkNormRope(OpWrapper): - """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and - writes index-K (E4M3) at ``out_cache_loc`` of ``index_k_cache`` [pages, 1, 128, 128]. Gemma RMSNorm of each head, - then RoPE on its leading ``rotary_dim`` channels (NeoX). - - The unfused sparse layers (``fuse_qkv_index_projection=False``) run it on the index branch with an FP8 indexer; - the fused layers' producer writes the index cache itself. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="one decode token, two index heads per rank", spec=dict(m=1, heads=2)), - Cell(why="an 8-row decode verify", spec=dict(m=8, heads=2)), - Cell(why="a prefill chunk", spec=dict(m=2048, heads=2)), - ) - - note = """ - `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 - before the cast); its index-K is what the kernel writes to the cache. The two kernels may round a handful of - values differently: the test allows fewer than 0.1% of the bytes to differ, as the in-tree test does. - """ - - def __call__( - self, - qk: torch.Tensor, - index_k_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - num_heads_q: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - base: float, - position_ids: torch.Tensor, - ) -> torch.Tensor: - return torch.ops.trtllm.minimax_m3_fp8_indexer_qk_norm_rope( - qk, - index_k_cache, - out_cache_loc, - num_heads_q, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - position_ids, - ) - - def reference( - self, - qk: torch.Tensor, - index_k_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - num_heads_q: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_weight: torch.Tensor, - k_weight: torch.Tensor, - base: float, - position_ids: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """(index-Q [T, heads, head_dim] E4M3, index-K [T, head_dim] E4M3): the catalog's in-place kernel on a copy, - then the E4M3 cast.""" - out = fused_qk_norm_rope( - qk.clone(), - num_heads_q, - 1, - 0, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - True, - position_ids, - 1.0, - 0.0, - 0.0, - 1.0, - True, - True, - False, - 0, - 0, - ) - q, k = out.split([num_heads_q * head_dim, head_dim], dim=-1) - return ( - q.reshape(q.shape[0], num_heads_q, head_dim).to(torch.float8_e4m3fn), - k.to(torch.float8_e4m3fn), - ) - - -m3_fp8_indexer_qk_norm_rope = M3Fp8IndexerQkNormRope() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py index 88ce9728f40b..d5ab5a6e0c8d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_fullctx_decode.py @@ -8,12 +8,15 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_dense_decode import ( + dense_attention_reference, + gather_paged_kv, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( fp8_e4m3_saturate, ) from .. import m3_fullctx_decode as kernel -from .m3_dense_decode import dense_attention_reference, gather_paged_kv __all__ = ["m3_fullctx_decode"] diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py index 5eb42546612e..900a62503ac1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_index_score.py @@ -8,13 +8,9 @@ from __future__ import annotations -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_index_score import ( + M3IndexScoreBase, ) __all__ = ["m3_index_score"] @@ -23,17 +19,7 @@ HEAD_DIM = 128 -def compiled_scorer(key: tuple): - """The compiled kernel of ``key`` (dtype, heads, query length, head dim, split_k, early), compiling it on a miss - (the target compiles every key it can take before warmup: ``layers.prebuild_index_scorer``).""" - from ..m3_index_score import IndexScoreRunner - - if key not in IndexScoreRunner.kernel_cache: - IndexScoreRunner._compile(*key) - return IndexScoreRunner.kernel_cache[key] - - -class M3IndexScore(OpWrapper): +class M3IndexScore(M3IndexScoreBase): """``score[h, t, j]`` (fp32, a transposed view of the [heads, blocks, tokens] selector buffer) = the max over the positions ``p`` of block ``j`` of request ``b = t // qlen`` with ``p <= seq_lens[b] - qlen + t % qlen`` of ``idx_q[t, h] . index_k[block_table[b, j], p % 128]``, for ``j < ceil(seq_lens[b] / 128)`` (no other entry is @@ -90,78 +76,5 @@ class M3IndexScore(OpWrapper): differs from the reference matmul by a few fp32 ulps; it is held to the in-tree scorer bit for bit. """ - def __call__( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> torch.Tensor: - _, num_heads, head_dim = idx_q.shape - key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) - compiled_scorer(key)(idx_q, index_k, block_table, score, seq_lens) - return score - - def reference( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> torch.Tensor: - """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and - ``early`` do not change them.""" - total_q, heads, _ = idx_q.shape - qlen = max_decode_query_len - out = torch.full( - (heads, total_q, score.shape[2]), - -float("inf"), - dtype=torch.float32, - device=idx_q.device, - ) - with true_fp32_matmul(): - for b, seq in enumerate(seq_lens.tolist()): - nblocks = (seq + PAGE - 1) // PAGE - if nblocks == 0: - continue - t0 = b * qlen - q = idx_q[t0 : t0 + qlen].float() - k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() - s = torch.einsum("qhd,kd->hqk", q, k) - q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) - k_pos = torch.arange(k.shape[0], device=idx_q.device) - s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) - out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) - return out - - def is_valid( - self, - idx_q: torch.Tensor, - index_k: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - score: torch.Tensor, - max_decode_query_len: int, - split_k: int, - early: bool, - ) -> None: - # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. - nblocks = (seq_lens + PAGE - 1) // PAGE - assert bool((nblocks <= block_table.shape[1]).all()), ( - "a request runs past its block-table row" - ) - assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" - assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( - "idx_q must hold max_decode_query_len rows per request" - ) - m3_index_score = M3IndexScore() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py index f1a7f210a590..ffe49916508a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_msa_attention.py @@ -4,20 +4,29 @@ from __future__ import annotations -from typing import Optional +from typing import Optional as Optional import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa import ( + GroupedMsaAttentionBase, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa import ( + _selected_blocks as _selected_blocks, +) from tensorrt_llm._torch.attention.backends.interface import ( - AttentionForwardArgs, - SparseBackendForwardArgs, + AttentionForwardArgs as AttentionForwardArgs, +) +from tensorrt_llm._torch.attention.backends.interface import ( + SparseBackendForwardArgs as SparseBackendForwardArgs, ) -from .m3_sparse_decode_lc8 import dequant_nvfp4_pages +from .m3_sparse_decode_lc8 import dequant_nvfp4_pages as dequant_nvfp4_pages -class M3MsaAttention(OpWrapper): +class M3MsaAttention(GroupedMsaAttentionBase): """One layer's attention core on the MSA backend (``MiniMaxM3MsaSparseAttention``), written into ``output`` [T, heads * 128] (bf16, the step's live rows). @@ -31,6 +40,14 @@ class M3MsaAttention(OpWrapper): objects for this layer and step. """ + @staticmethod + def _dequant_pages(*args, **kwargs) -> torch.Tensor: + return dequant_nvfp4_pages(*args, **kwargs) + + @staticmethod + def _select_blocks(*args, **kwargs) -> torch.Tensor: + return _selected_blocks(*args, **kwargs) + ARCHS = frozenset({Arch.SM_107}) CELLS: tuple[Cell, ...] = ( @@ -62,188 +79,5 @@ class M3MsaAttention(OpWrapper): rounding and accumulation order), which the test bounds. """ - def __call__( - self, - backend, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - idx_q: Optional[torch.Tensor], - metadata, - output: torch.Tensor, - sparse: bool, - kv_scale_orig_quant: Optional[torch.Tensor], - kv_scale_quant_orig: Optional[torch.Tensor], - ) -> torch.Tensor: - if k is not None: - backend.write_layer_caches( - k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant - ) - if sparse: - blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) - args = AttentionForwardArgs( - output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) - ) - else: - args = AttentionForwardArgs(output=output) - args.kv_scale_orig_quant = kv_scale_orig_quant - args.kv_scale_quant_orig = kv_scale_quant_orig - backend.forward(q, None, None, metadata, forward_args=args) - return output - - def reference( - self, - backend, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - idx_q: Optional[torch.Tensor], - metadata, - output: torch.Tensor, - sparse: bool, - kv_scale_orig_quant: Optional[torch.Tensor], - kv_scale_quant_orig: Optional[torch.Tensor], - ) -> torch.Tensor: - """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. - - Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the - metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of - ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the - test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and - attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors grouped KV heads (query - head h reads KV head h // (heads / kv heads); two per rank at tp2), each KV head selecting its own top-k - blocks from its own index heads against the rank's single index-K head. - """ - manager = metadata.kv_cache_manager - layer = backend.layer_idx - heads, dim = backend.num_heads, backend.head_dim - kv_heads = int(backend.num_kv_heads) - assert heads % kv_heads == 0 - group = heads // kv_heads - page = int(manager.tokens_per_block) - sm_scale = dim**-0.5 / float(backend.q_scaling) - num_seqs = int(metadata.num_seqs) - qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] - cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] - total = sum(qlens) - assert q.shape[0] == total and output.shape[0] == total - - pool = manager.get_buffers(layer, "HND") - nvfp4 = manager.is_nvfp4_layer(layer) - if nvfp4: - scale_pool = manager.get_block_scale_buffers(layer, "HND") - k_global = float(kv_scale_quant_orig.reshape(-1)[1]) - v_global = float(kv_scale_quant_orig.reshape(-1)[2]) - q_rows = q.reshape(total, heads, dim) - if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: - # The FP8 FMHA shares one dtype across Q / K / V: Q is cast to E4M3, unscaled. - q_rows = q_rows.to(torch.float8_e4m3fn) - q64 = q_rows.to(torch.float64) - if sparse: - config = backend.m3_config - idx_cache = metadata.msa_idx_k_cache(layer) - idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( - torch.float64 - ) - assert config.num_index_heads % kv_heads == 0 - idx_group = config.num_index_heads // kv_heads - - out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) - first = 0 - for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): - n, length = qlens[row], cached[row] + qlens[row] - blocks = (length + page - 1) // page - pages = torch.tensor( - manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], - dtype=torch.int64, - device=q.device, - ) - qpos = cached[row] + torch.arange(n, device=q.device) - causal = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] - for g in range(kv_heads): - if nvfp4: - keys = dequant_nvfp4_pages( - pool[:, 0].view(torch.uint8), - scale_pool[:, 0], - pages, - quad_scales=False, - kv_head=g, - ) - values = dequant_nvfp4_pages( - pool[:, 1].view(torch.uint8), - scale_pool[:, 1], - pages, - quad_scales=True, - kv_head=g, - ) - keys, values = keys * k_global, values * v_global - else: - keys = pool[pages, 0, g].to(torch.float64) - values = pool[pages, 1, g].to(torch.float64) - keys = keys.reshape(-1, dim)[:length] - values = values.reshape(-1, dim)[:length] - allowed = causal - if sparse: - # One index-K head per rank, shared by every index head (the cache is [pages, 1, page, dim]). - idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] - chosen = _selected_blocks( - idx64[first : first + n, g * idx_group : (g + 1) * idx_group], - idx_keys, - causal, - qpos, - page, - config, - ) - allowed = causal & chosen.repeat_interleave(page, dim=1)[:, :length] - h0, h1 = g * group, (g + 1) * group - for lo in range(0, n, 256): - hi = min(n, lo + 256) - logits = ( - torch.einsum("thd,kd->thk", q64[first + lo : first + hi, h0:h1], keys) - * sm_scale - ) - logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) - out[first + lo : first + hi, h0:h1] = torch.softmax(logits, dim=-1) @ values - first += n - return out.reshape(total, heads * dim).to(output.dtype) - - -def _selected_blocks( - idx_q: torch.Tensor, - idx_keys: torch.Tensor, - allowed: torch.Tensor, - qpos: torch.Tensor, - page: int, - config, -) -> torch.Tensor: - """The blocks each query token of one request attends, [n, blocks] bool. - - A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of - ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last - ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the - backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the - ranking, so none is applied. - """ - n, length = allowed.shape - blocks = (length + page - 1) // page - scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) - scores = scores.masked_fill(~allowed, float("-inf")) - padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) - padded[:, :length] = scores - block_scores = padded.view(n, blocks, page).amax(dim=-1) - ids = torch.arange(blocks, device=scores.device)[None, :] - n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] - block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) - forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( - (ids >= n_valid - config.local_blocks) & (ids < n_valid) - ) - block_scores = block_scores.masked_fill(forced, float("inf")) - order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] - keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid - chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) - rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) - chosen[rows[keep], order[keep]] = True - return chosen - m3_msa_attention = M3MsaAttention() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py index ff34b8eac631..c09475aaacc9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_producer.py @@ -5,10 +5,16 @@ from __future__ import annotations +import sys +from typing import cast + import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + producer_reference as _producer_shared, +) __all__ = ["m3_producer"] @@ -27,29 +33,20 @@ def _norm_rope( ): """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. x [T, 128] fp32; cos / sin [T, 32].""" - y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) - lo, hi, rest = y[:, :32], y[:, 32:64], y[:, 64:] - return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + return _producer_shared._norm_rope(_REFERENCE_TARGET, x, weight, cos, sin, eps) def decode_nvfp4_token( data: torch.Tensor, scales: torch.Tensor, quant_scale: float ) -> torch.Tensor: """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" - table = torch.tensor(E2M1_VALUES, dtype=torch.float32, device=data.device) - nib = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(-1).long() - vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) - s = scales.view(torch.float8_e4m3fn).float() / quant_scale - return vals * s.repeat_interleave(FP4_BLOCK) + return _producer_shared.decode_nvfp4_token(_REFERENCE_TARGET, data, scales, quant_scale) def scale_index(within_page: int, plane: int) -> torch.Tensor: """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token swizzle for V.""" - cols = torch.arange(8) - if plane == 0: - return within_page * 8 + cols - return (within_page // 4) * 32 + cols * 4 + within_page % 4 + return _producer_shared.scale_index(_REFERENCE_TARGET, within_page, plane) class M3Producer(OpWrapper): @@ -108,7 +105,7 @@ def __call__( prewait_step_inputs: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if prewait_step_inputs: - # This target's copy (producerPw8): the same outputs bit for bit, with the position ids, slots and RoPE + # The TP8 producer specialization: same outputs bit for bit, with position ids, slots and RoPE # rows read before its grid wait (none of them may be written by a PDL predecessor). return tuple( torch.ops.trtllm.minimax_m3_sm107_tp8_producer_pw( @@ -173,20 +170,27 @@ def reference( """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change them.""" - t = packed.shape[0] - x = packed.float().view(t, heads_q + 4, HEAD_DIM) - rope = rotary_cos_sin[position_ids[:t].long()] - cos, sin = rope[:, 0], rope[:, 1] - q = torch.stack( - [_norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + return _producer_shared.producer_reference_single_kv( + self, + _REFERENCE_TARGET, + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + prewait_step_inputs, ) - return { - "q": q, - "k": _norm_rope(x[:, heads_q], k_weight, cos, sin, eps), - "v": x[:, heads_q + 1], - "index_q": _norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), - "index_k": _norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), - } def is_valid( self, @@ -208,16 +212,29 @@ def is_valid( early: bool = True, prewait_step_inputs: bool = False, ) -> None: - t = packed.shape[0] - # The RoPE row is read at the position with no bound. - pos = position_ids[:t] - assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( - "position_ids must index rotary_cos_sin" + return _producer_shared.validate_producer( + self, + _REFERENCE_TARGET, + packed, + kv_data, + kv_scale, + index_k, + out_cache_loc, + kv_quant_scale, + eps, + q_weight, + k_weight, + index_q_weight, + index_k_weight, + rotary_cos_sin, + position_ids, + heads_q, + pdl, + early, + prewait_step_inputs, ) - # Two tokens writing one slot race. - slots = out_cache_loc[:t] - live = slots[(slots >= 0) & (slots < kv_data.shape[0] * PAGE)] - assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" m3_producer = M3Producer() + +_REFERENCE_TARGET = cast(_producer_shared.ProducerTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py deleted file mode 100644 index 88180c2f1a1f..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qk_norm_rope_fp8.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The dense layers' per-head Gemma RMSNorm + partial NeoX RoPE of the packed [Q | K | V], into a fresh FP8 copy.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch._experimental.modeling_v2.catalog.attention.fused_qk_norm_rope import ( - fused_qk_norm_rope, -) - - -class M3QkNormRopeFp8(OpWrapper): - """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded - once to E4M3 (saturating at +/-448), V cast, ``qkv`` left untouched. Same arguments as the catalog's in-place - entry. - - The dense layers (0-2) feed the FP8 main KV cache with it: on every step of an FP8 cache, and on the decode steps - of a hybrid NVFP4 cache (whose dense layers sit in its FP8 pool). - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="one decode token, tp8 heads (8 q, 1 kv)", spec=dict(m=1, q=8, kv=1)), - Cell(why="a 48-row decode verify", spec=dict(m=48, q=8, kv=1)), - Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=8, kv=1)), - ) - - note = """ - `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from - fp32 where the composition rounds through bf16 first, so they are within one E4M3 step of it. - """ - - def __call__( - self, - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) -> torch.Tensor: - return torch.ops.trtllm.fused_qk_norm_rope_to_fp8( - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) - - def reference( - self, - qkv, - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) -> torch.Tensor: - """The catalog's in-place kernel on a copy, then the E4M3 cast (V exact; Q and K within one E4M3 step).""" - out = fused_qk_norm_rope( - qkv.clone(), - num_heads_q, - num_heads_k, - num_heads_v, - head_dim, - rotary_dim, - eps, - q_weight, - k_weight, - base, - is_neox, - position_ids, - factor, - low, - high, - attention_factor, - is_qk_norm, - use_gemma, - use_mrope, - mrope_section1, - mrope_section2, - ) - return out.to(torch.float8_e4m3fn) - - -m3_qk_norm_rope_fp8 = M3QkNormRopeFp8() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py deleted file mode 100644 index c6f4be426ba3..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_qkv_insert.py +++ /dev/null @@ -1,134 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The sparse layers' generic horizontal producer: packed QKV/index rows -> FP8 Q / index-Q, K / V / index-K into -the paged caches.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3QkvInsert(OpWrapper): - """The packed ``[Q | K | V | index-Q | index-K]`` bf16 rows of a sparse layer: per-head Gemma RMSNorm and the - partial NeoX RoPE of Q, K, index-Q and index-K, K / V written to the main paged cache at ``out_cache_loc`` and - index-K to the FP8 index cache; returns (Q, index-Q) in FP8 E4M3 [M, heads, 128]. - - ``kv_scale`` given: the main cache is NVFP4 (packed E2M1 data [pages, 2, kv, 128, 64] uint8 + E4M3 scale pool - [pages, 2, kv, 128, 8], ``kv_quant_scale`` the fp32 (Q, K, V) quantization scales); None: FP8 E4M3 - [pages, 2, kv, 128, 128]. Token-major and step-type agnostic. ``m3_producer`` takes the NVFP4 cache at the tp8 - heads; this entry takes the FP8 cache and NVFP4 pools outside that layout. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="FP8 main cache, a prefill chunk", spec=dict(m=1024, cache="fp8")), - Cell(why="FP8 main cache, one decode token", spec=dict(m=1, cache="fp8")), - ) - - note = """ - The arithmetic is the CUDA producer's (fp32 norm and RoPE with bf16 roundings, then the E4M3 casts); `reference` - returns the fp32 values, and the test holds the FP8 outputs and cache rows within one E4M3 step of them. - """ - - def __call__( - self, - packed: torch.Tensor, - kv_cache: torch.Tensor, - kv_scale: Optional[torch.Tensor], - index_k_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - kv_quant_scale: Optional[torch.Tensor], - num_heads: int, - num_kv_heads: int, - num_index_heads: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_norm: torch.Tensor, - k_norm: torch.Tensor, - index_q_norm: torch.Tensor, - index_k_norm: torch.Tensor, - rotary_cos_sin: torch.Tensor, - position_ids: torch.Tensor, - ): - common = ( - num_heads, - num_kv_heads, - num_index_heads, - head_dim, - rotary_dim, - eps, - q_norm, - k_norm, - index_q_norm, - index_k_norm, - rotary_cos_sin, - position_ids, - ) - if kv_scale is not None: - return torch.ops.trtllm.minimax_m3_nvfp4_qkv_indexer_norm_rope_kv_insert( - packed, kv_cache, kv_scale, index_k_cache, out_cache_loc, kv_quant_scale, *common - ) - return torch.ops.trtllm.minimax_m3_fp8_qkv_indexer_norm_rope_kv_insert( - packed, kv_cache, index_k_cache, out_cache_loc, *common - ) - - def reference( - self, - packed: torch.Tensor, - kv_cache: torch.Tensor, - kv_scale: Optional[torch.Tensor], - index_k_cache: torch.Tensor, - out_cache_loc: torch.Tensor, - kv_quant_scale: Optional[torch.Tensor], - num_heads: int, - num_kv_heads: int, - num_index_heads: int, - head_dim: int, - rotary_dim: int, - eps: float, - q_norm: torch.Tensor, - k_norm: torch.Tensor, - index_q_norm: torch.Tensor, - index_k_norm: torch.Tensor, - rotary_cos_sin: torch.Tensor, - position_ids: torch.Tensor, - ): - """The fp32 values the producer rounds (a new dict, nothing written): Gemma RMSNorm + NeoX partial RoPE of - Q [M, heads, D], index-Q [M, index heads, D], K [M, kv heads, D] and index-K [M, D]; V [M, kv heads, D] as - packed. The cache rows at ``out_cache_loc`` and the returned FP8 Q / index-Q are these values cast.""" - m = packed.shape[0] - sizes = [ - num_heads * head_dim, - num_kv_heads * head_dim, - num_kv_heads * head_dim, - num_index_heads * head_dim, - head_dim, - ] - q, k, v, iq, ik = packed.float().split(sizes, dim=-1) - cos, sin = rotary_cos_sin[position_ids.long()].unbind(1) - half = rotary_dim // 2 - - def norm_rope(x, w, heads): - x = x.reshape(m, heads, head_dim) - x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * (1.0 + w.float()) - x1, x2, rest = x[..., :half], x[..., half:rotary_dim], x[..., rotary_dim:] - c, s = cos.unsqueeze(1), sin.unsqueeze(1) - return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s, rest], dim=-1) - - return { - "q": norm_rope(q, q_norm, num_heads), - "index_q": norm_rope(iq, index_q_norm, num_index_heads), - "k": norm_rope(k, k_norm, num_kv_heads), - "v": v.reshape(m, num_kv_heads, head_dim), - "index_k": norm_rope(ik, index_k_norm, 1)[:, 0], - } - - -m3_qkv_insert = M3QkvInsert() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py deleted file mode 100644 index 76841cca732d..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_blocks.py +++ /dev/null @@ -1,82 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MSA top-k block selection from per-KV-head block scores (the built-in selector).""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3SelectBlocks(OpWrapper): - """Per query row the ``topk`` highest-scoring KV blocks of ``max_score_kv`` [kv_heads, n_blocks, total_q] (fp32), - with the first ``init_blocks`` and the last ``local_blocks`` valid blocks forced in and blocks at or beyond the - row's ``n_valid_blocks`` excluded; [total_q, kv_heads, topk] int32 ascending ids, -1 padded. ``head_major_output`` - backs the result head-major so ``permute(1, 0, 2)`` is contiguous. - - The decode dispatch runs it where the long-row selector does not apply: requests of 17..128 KV blocks. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="eight decode rows of 17..128 blocks", - spec=dict(total_q=8, n_blocks=128, mixed=True), - ), - Cell( - why="one decode row at the 128-block edge", - spec=dict(total_q=1, n_blocks=128, mixed=False), - ), - Cell(why="a 4-row verify at 40 blocks", spec=dict(total_q=4, n_blocks=40, mixed=False)), - ) - - note = """ - Exact: ties between equal scores break toward the lower block id; `reference` sorts with the same rule. - """ - - def __call__( - self, - max_score_kv: torch.Tensor, - topk: int, - n_valid_blocks: torch.Tensor, - init_blocks: int, - local_blocks: int, - head_major_output: bool, - ) -> torch.Tensor: - nvb = n_valid_blocks.to(device=max_score_kv.device, dtype=torch.int32).contiguous() - return torch.ops.trtllm.minimax_m3_select_blocks( - max_score_kv, nvb, topk, init_blocks, local_blocks, head_major_output - ) - - def reference( - self, - max_score_kv: torch.Tensor, - topk: int, - n_valid_blocks: torch.Tensor, - init_blocks: int, - local_blocks: int, - head_major_output: bool, - ) -> torch.Tensor: - """Plain torch selection; ``head_major_output`` changes the backing only.""" - heads, n_blocks, total_q = max_score_kv.shape - out = torch.full((total_q, heads, topk), -1, dtype=torch.int32, device=max_score_kv.device) - blocks = torch.arange(n_blocks, device=max_score_kv.device) - for t in range(total_q): - nv = int(n_valid_blocks[t]) - for h in range(heads): - s = max_score_kv[h, :, t].clone().double() - s[blocks >= nv] = float("-inf") - forced = (blocks < min(init_blocks, nv)) | ( - (blocks >= nv - local_blocks) & (blocks < nv) - ) - s[forced] = float("inf") - order = sorted(range(nv), key=lambda b: (-float(s[b]), b))[: min(topk, nv)] - ids = torch.tensor(sorted(order), dtype=torch.int32, device=out.device) - out[t, h, : ids.numel()] = ids - return out - - -m3_select_blocks = M3SelectBlocks() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py index 9a92050c6abb..4c35f54175d2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_select_long.py @@ -4,10 +4,19 @@ from __future__ import annotations +import sys +from typing import cast + import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.select_reference import ( + SelectReferenceTarget, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.select_reference import ( + select_reference as _shared_select_reference, +) #: Blocks per decode row the op covers (9 keys per thread of one 1024-thread CTA). MAX_BLOCKS = 9216 @@ -135,32 +144,18 @@ def reference( (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and ``trigger_phase`` only say where and how the op writes; they do not enter the result. """ - num_kv_heads, num_blocks, total_q = scores.shape - device = scores.device - blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) - raw = n_valid_blocks[:total_q].to(torch.int64) - valid = raw.clamp(0, num_blocks) - local_start = (raw - local_blocks).clamp_min(0) - result = torch.full((total_q, num_kv_heads, TOPK), -1, dtype=torch.int32, device=device) - for head in range(num_kv_heads): - forced = scores[head].float().t().clone() # [total_q, num_blocks] - forced = torch.where(blocks[None, :] < init_blocks, INIT_SCORE, forced) - forced = torch.where(blocks[None, :] >= local_start[:, None], LOCAL_SCORE, forced) - bits = forced.view(torch.int32).to(torch.int64) & 0xFFFFFFFF - key = torch.where(bits >= 0x80000000, bits, (~bits) & 0x7FFFFFFF) - in_row = blocks[None, :] < valid[:, None] - packed = torch.where( - in_row, (key << 16) | blocks[None, :], torch.iinfo(torch.int64).max - ) - k = min(TOPK, num_blocks) - picked = torch.topk(packed, k, dim=1, largest=False).values - picked_key = picked >> 16 - keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 0xFF800000) - ids = torch.where(keep, picked & 0xFFFF, torch.iinfo(torch.int64).max) - ids = ids.sort(dim=1).values - ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) - result[:, head, :k] = ids.to(torch.int32) - return result + return _shared_select_reference( + self, + _SELECT_REFERENCE_TARGET, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, + ) def is_valid( self, @@ -191,3 +186,5 @@ def is_valid( m3_select_long = M3SelectLong() + +_SELECT_REFERENCE_TARGET = cast(SelectReferenceTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py index fc770d010643..b718d9a515cf 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode.py @@ -9,6 +9,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_reference import ( + validate_sparse_selection, +) from ..triton_sparse_decode import minimax_m3_sparse_attn_decode from .m3_sparse_decode_lc8 import PAGE, sparse_attention_reference @@ -183,33 +186,9 @@ def is_valid( num_topk_chunks: Optional[int] = None, identity_blocks: bool = False, ) -> None: - """The inputs this call takes and answers wrongly. - - Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` - table entries and indexes the block table with them unchecked, so they - must be real blocks of the request. ``identity_blocks`` is a promise the - multi-query producer relies on: every token of a request reads the same - blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the - host: a test-only cost. - """ - total_q = q.shape[0] - seq = seq_lens.tolist() - max_topk = topk_idx.shape[-1] - for tok in range(total_q): - req, offset = divmod(tok, decode_query_len) - kv_len = max(seq[req] - decode_query_len + offset + 1, 0) - nblocks = (kv_len + PAGE - 1) // PAGE - ids = topk_idx[:, tok, : min(max_topk, nblocks)] - assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( - f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " - "the block table is indexed with them unchecked" - ) - if identity_blocks: - assert nblocks <= max_topk and bool( - (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() - ), ( - f"token {tok}: identity_blocks promises the table holds every block of the request in order" - ) + validate_sparse_selection( + q, topk_idx, block_table, seq_lens, decode_query_len, identity_blocks, PAGE + ) m3_sparse_decode = M3SparseDecode() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py index 76a99144322a..d52757921a0e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/m3_sparse_decode_lc8.py @@ -16,6 +16,13 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from ....ops.attention.sparse_reference import _E2M1 as _E2M1 +from ....ops.attention.sparse_reference import dequant_nvfp4_pages as _dequant_nvfp4_pages +from ....ops.attention.sparse_reference import mxfp8_handoff_reference as _mxfp8_handoff_reference +from ....ops.attention.sparse_reference import ( + sparse_attention_reference as _sparse_attention_reference, +) + #: The TP8 geometry the kernel is compiled for: 8 query heads on one KV head. NUM_HEADS = 8 HEAD_DIM = 128 @@ -27,8 +34,6 @@ MX_BLOCK = 32 WIDTH = NUM_HEADS * HEAD_DIM -_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) - def dequant_nvfp4_pages( packed: torch.Tensor, @@ -37,26 +42,8 @@ def dequant_nvfp4_pages( quad_scales: bool, kv_head: int = 0, ) -> torch.Tensor: - """Dequantize KV head ``kv_head`` of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. - - ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low - nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, - HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token - ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at - ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor - global scale is not applied. Every product is exact in float64. - """ - table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) - rows = packed[pages, kv_head].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] - nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) - cols = HEAD_DIM // SF_VEC - region = block_scale[pages, kv_head].reshape(rows.shape[0], PAGE * cols) - t = torch.arange(PAGE, device=packed.device)[:, None] - c = torch.arange(cols, device=packed.device)[None, :] - index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c - scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) - scales = scales.view(torch.float8_e4m3fn).to(torch.float64) - return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + """Use the shared reference with this target's geometry.""" + return _dequant_nvfp4_pages(packed, block_scale, pages, quad_scales, kv_head) def sparse_attention_reference( @@ -73,87 +60,28 @@ def sparse_attention_reference( decode_query_len: int, sm_scale: float, ) -> torch.Tensor: - """Block-sparse decode attention over one KV head of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. - - Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset - ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / - PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at - ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that - attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or - its byte view); softmax and both products are float64. - """ - q64 = ( - q.view(torch.float8_e4m3fn).to(torch.float64) - if q.element_size() == 1 - else q.to(torch.float64) + """Use the shared reference with this target's geometry.""" + return _sparse_attention_reference( + q, + k, + v, + k_block_scale, + v_block_scale, + k_global_scale, + v_global_scale, + topk, + block_table, + seq_lens, + decode_query_len, + sm_scale, ) - total_q, num_heads, _ = q64.shape - kgs = float(k_global_scale.reshape(-1)[0]) - vgs = float(v_global_scale.reshape(-1)[0]) - out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) - seq = seq_lens.tolist() - positions = torch.arange(PAGE, device=q.device) - for tok in range(total_q): - req, offset = divmod(tok, decode_query_len) - kv_len = max(seq[req] - decode_query_len + offset + 1, 0) - real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) - blocks = topk[tok, :real_topk] - blocks = blocks[blocks >= 0].to(torch.int64) - if blocks.numel() == 0: - continue - pages = block_table[req, blocks].to(torch.int64) - keys = ( - dequant_nvfp4_pages(k, k_block_scale, pages, quad_scales=False).reshape(-1, HEAD_DIM) - * kgs - ) - values = ( - dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True).reshape(-1, HEAD_DIM) - * vgs - ) - valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) - logits = (q64[tok] @ keys.T) * sm_scale - logits = logits.masked_fill(~valid[None, :], float("-inf")) - values = values.masked_fill(~valid[:, None], 0) - out[tok] = torch.softmax(logits, dim=-1) @ values - return out def mxfp8_handoff_reference( out: torch.Tensor, num_rows_alloc: int ) -> Tuple[torch.Tensor, torch.Tensor]: - """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, WIDTH]: ``(q8 [T, WIDTH] uint8, - sf [num_rows_alloc * WIDTH // 32] uint8)``. - - Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` - rounded up when its mantissa is non-zero (capped at 254), the elements are - scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands - in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` - for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a - comparison against it is exact. - """ - total_q = out.shape[0] - x = out.float().reshape(total_q, WIDTH // MX_BLOCK, MX_BLOCK) - amax = x.abs().amax(dim=-1) - sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) - bits = sfv.view(torch.int32) - e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) - e = e.clamp(max=254) - qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) - xq = (x * qs[..., None]).clamp(-448.0, 448.0) - q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, WIDTH) - num_sf_k_tiles = WIDTH // 128 - sf = torch.zeros(num_rows_alloc * (WIDTH // MX_BLOCK), dtype=torch.uint8, device=out.device) - tok = torch.arange(total_q, device=out.device)[:, None] - k_sf = torch.arange(WIDTH // MX_BLOCK, device=out.device)[None, :] - offset = ( - (tok // 128) * num_sf_k_tiles * 512 - + (k_sf // 4) * 512 - + (tok % 32) * 16 - + ((tok % 128) // 32) * 4 - + k_sf % 4 - ) - sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) - return q8, sf + """Use the shared reference with this target's geometry.""" + return _mxfp8_handoff_reference(out, num_rows_alloc, width=WIDTH) class M3SparseDecodeLc8(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py deleted file mode 100644 index 43166b472eca..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/attention/msa_kv_update.py +++ /dev/null @@ -1,200 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The MSA KV-length correction in one launch (this target's ``MsaKvUpdate`` hook).""" - -from __future__ import annotations - -from typing import Optional, Tuple - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: The MSA pool's page: one sparse block is one 128-token KV page. -TOKENS_PER_BLOCK = 128 - - -class MsaKvUpdate(OpWrapper): - """Re-derive the MSA decode state after the overlap scheduler moved ``kv_lens``. - - ``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` calls this with its - own buffers at the head of every captured speculative decode step. One - single-CTA launch writes, for ``batch`` requests and ``total_q`` query - tokens (all int32): - - * ``msa_seq_lens[b] = min(kv_lens[b], kv_lens_staged[b])`` and, when given, - ``cu_kv_lens`` as its exclusive prefix sum (``batch + 1`` entries); - * ``out_cache_loc[t]``, the KV write slot of the position token ``t`` - attends up to, through ``block_table``; - * ``n_valid_blocks[t]``, when given, the block count top-k selection is - bounded by. - - It returns None; the cells read the four buffers. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="decode without speculation at one request: qo = 1, every buffer written", - spec=dict(batch=1, query_len=1, accept="full", cu_kv=True, n_valid=True), - ), - Cell( - why="decode without speculation at the captured batch 8", - spec=dict(batch=8, query_len=1, accept="full", cu_kv=True, n_valid=True), - ), - Cell( - why="Eagle3 verify (draft 3): qo = 4 per request, partial acceptance shortens kv_lens below staged", - spec=dict(batch=8, query_len=4, accept="partial", cu_kv=True, n_valid=True), - ), - Cell( - why="the metadata without cu_kv_lens (no extend_ctx promotion): the optional prefix sum is skipped", - spec=dict(batch=4, query_len=4, accept="partial", cu_kv=False, n_valid=True), - ), - Cell( - why="CUDA-graph padding rows (kv_lens 0, pos -1) and no n_valid_blocks buffer", - spec=dict(batch=8, query_len=4, accept="padded", cu_kv=True, n_valid=False), - ), - Cell( - why="positions crossing a page boundary: the slot moves to the previous page of the request", - spec=dict(batch=3, query_len=4, accept="page_edge", cu_kv=True, n_valid=True), - ), - ) - - note = """ - The kernel covers any batch up to the block size's loop (it strides over - 256 threads), and any page size; only `tokens_per_block = 128`, the MSA - pool's page, is certified. - - The op does not bounds-check the request index a token reads - (`q_batch_row`) or the block-table column it derives: both come from the - metadata's own staging, which keeps them in range by construction, and an - out-of-range value reads whatever memory sits there. `is_valid` states both. - """ - - def __call__( - self, - kv_lens: torch.Tensor, - kv_lens_staged: torch.Tensor, - seq_lens: torch.Tensor, - q_batch_row: torch.Tensor, - q_intra: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - batch: int, - total_q: int, - msa_seq_lens: torch.Tensor, - cu_kv_lens: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - n_valid_blocks: Optional[torch.Tensor], - ) -> None: - torch.ops.trtllm.minimax_m3_msa_kv_update( - kv_lens, - kv_lens_staged, - seq_lens, - q_batch_row, - q_intra, - block_table, - tokens_per_block, - batch, - total_q, - msa_seq_lens, - cu_kv_lens, - out_cache_loc, - n_valid_blocks, - ) - - def reference( - self, - kv_lens: torch.Tensor, - kv_lens_staged: torch.Tensor, - seq_lens: torch.Tensor, - q_batch_row: torch.Tensor, - q_intra: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - batch: int, - total_q: int, - msa_seq_lens: torch.Tensor, - cu_kv_lens: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - n_valid_blocks: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, Optional[torch.Tensor]]: - """``(msa_seq_lens[:batch], cu_kv_lens[:batch + 1], out_cache_loc[:total_q], - n_valid_blocks[:total_q])`` as the op is supposed to write them. - - The in-tree metadata's own torch re-derivation - (``MiniMaxM3MsaSparseAttentionMetadata.on_update_kv_lens`` past its hook - call), statement for statement, returning new tensors rather than - writing the buffers. The buffer arguments only say which outputs exist: - a ``None`` buffer gives a ``None`` result. - """ - page = int(tokens_per_block) - kv_true = torch.minimum(kv_lens[:batch], kv_lens_staged[:batch]) - cu_kv = None - if cu_kv_lens is not None: - cu_kv = torch.zeros(batch + 1, dtype=torch.int32, device=kv_lens.device) - cu_kv[1:] = torch.cumsum(kv_true, 0) - qbr = q_batch_row[:total_q].to(torch.long) - qo_dev = seq_lens[:batch] - pos = kv_true[qbr] - qo_dev[qbr] + q_intra[:total_q] - pos_slot = pos.clamp_min(0) - block_col = torch.div(pos_slot, page, rounding_mode="floor").clamp( - max=int(block_table.shape[1]) - 1 - ) - slots = block_table[qbr, block_col.to(torch.long)] * page + torch.remainder(pos_slot, page) - n_valid = None - if n_valid_blocks is not None: - n_valid = torch.div( - (pos + 1).clamp_min(0) + (page - 1), page, rounding_mode="floor" - ).to(torch.int32) - return kv_true.to(torch.int32), cu_kv, slots.to(torch.int32), n_valid - - def is_valid( - self, - kv_lens: torch.Tensor, - kv_lens_staged: torch.Tensor, - seq_lens: torch.Tensor, - q_batch_row: torch.Tensor, - q_intra: torch.Tensor, - block_table: torch.Tensor, - tokens_per_block: int, - batch: int, - total_q: int, - msa_seq_lens: torch.Tensor, - cu_kv_lens: Optional[torch.Tensor], - out_cache_loc: torch.Tensor, - n_valid_blocks: Optional[torch.Tensor], - ) -> None: - """The inputs this op takes and answers wrongly. - - Every operand but the block table is read as a dense run from its first - element, so a strided view reads the wrong elements. The request index of - each token and the table row it selects are not bounds-checked. Reads - ``q_batch_row`` back to the host: a test-only cost. - """ - for name, t in ( - ("kv_lens", kv_lens), - ("kv_lens_staged", kv_lens_staged), - ("seq_lens", seq_lens), - ("q_batch_row", q_batch_row), - ("q_intra", q_intra), - ("msa_seq_lens", msa_seq_lens), - ("out_cache_loc", out_cache_loc), - ("cu_kv_lens", cu_kv_lens), - ("n_valid_blocks", n_valid_blocks), - ): - assert t is None or t.is_contiguous(), ( - f"{name} must be contiguous: the kernel indexes it densely from data_ptr and ignores strides" - ) - rows = q_batch_row[:total_q] - assert total_q == 0 or ( - int(rows.min()) >= 0 and int(rows.max()) < min(batch, block_table.shape[0]) - ), ( - "q_batch_row must index a request below batch (and a block-table row): the op reads the " - "request's lengths and table row unchecked" - ) - - -msa_kv_update = MsaKvUpdate() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py deleted file mode 100644 index 8daa2d6b9540..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_allreduce.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The built-in TP all-reduce, fused with the residual add and the next RMSNorm.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams - - -class M3AllReduce(OpWrapper): - """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the - plain sum over the ranks (the drafter's o_proj and down projection). With ``RESIDUAL_RMS_NORM`` params: - ``updated = residual + sum_ranks(x)``, ``normed = rmsnorm(updated) * norm_weight``; - returns (normed, updated). The strategy (``AllReduceStrategy.AUTO``: one-shot / two-shot Lamport, NCCL) is the - communicator's per call. - - The layer boundaries the multicast all-reduce does not take run it: prefill and mixed steps, and generation steps - beyond the multicast kernel's rows. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # No cell names a world size: the result does not depend on it, and the certifying matrix runs at every world - # size one node can form (the target's two-tray group of 8 cannot be formed on one node). - CELLS: tuple[Cell, ...] = ( - Cell(why="a prefill chunk", spec=dict(m=2048, hidden=6144)), - Cell(why="a 64-row generation step", spec=dict(m=64, hidden=6144)), - Cell(why="plain sum of eight drafter rows", spec=dict(m=8, hidden=6144, plain=True)), - ) - - note = """ - The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then - adds the residual and normalizes (one bf16 rounding of the sum and one of the norm). A strategy may round the sum - to bf16 before the residual add, so the residual sum agrees within one bf16 ulp of its operands' magnitude. - """ - - def __call__(self, x: torch.Tensor, comm: AllReduce, params: Optional[AllReduceParams]): - return comm(x, all_reduce_params=params) - - def reference(self, xs, comm: AllReduce, params: Optional[AllReduceParams]): - """``xs``: every rank's ``x`` in rank order (a collective's reference takes the group's inputs).""" - total = torch.stack([x.float() for x in xs]).sum(0) - if params is None: - return total.to(xs[0].dtype) - updated = (params.residual.float() + total).to(xs[0].dtype) - h = updated.float() - normed = ( - h - * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + params.eps) - * params.norm_weight.float() - ) - return normed.to(xs[0].dtype), updated - - -m3_allreduce = M3AllReduce() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py index cf0a4be970d9..4fabf0fc0f87 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_allreduce_norm.py @@ -4,85 +4,23 @@ from __future__ import annotations -from typing import Optional, Sequence +from typing import Optional import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper -_E4M3_MAX = 448.0 +from ....ops.comm.reference import _E4M3_MAX as _E4M3_MAX +from ....ops.comm.reference import McAllreduceNormReference +from ....ops.comm.reference import boundary as boundary +from ....ops.comm.reference import mxfp8_rows as mxfp8_rows +from ....ops.comm.reference import sf_written as sf_written +from ....ops.comm.reference import swizzle_sf as swizzle_sf -def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. - - One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), - data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the - op leaves them (the comparison masks them out). - """ - m, h = normed.shape - blocks = normed.float().view(m, h // 32, 32) - amax = blocks.abs().amax(-1) - bits = (amax / _E4M3_MAX).view(torch.int32) - byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) - scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) - data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) - data = data.view(m, h).to(torch.float8_e4m3fn) - return data, swizzle_sf(byte.to(torch.uint8)) - - -def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: - """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" - rows, cols = byte.shape - padded_rows = (rows + 127) // 128 * 128 - r = torch.arange(rows, device=byte.device).view(-1, 1) - c = torch.arange(cols, device=byte.device).view(1, -1) - index = ( - (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) - ) - out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) - out[index.reshape(-1)] = byte.reshape(-1) - return out - - -def sf_written(rows: int, cols: int, device) -> torch.Tensor: - """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" - return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() - - -def boundary( - partials: Sequence[torch.Tensor], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - drafter: bool = False, -) -> tuple[torch.Tensor, ...]: - """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). - - ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer - fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; - only the stored residual is rounded to BF16. - - The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are - summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic - and the normalization are done in float64 and rounded to BF16 once. - """ - total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) - for p in partials: - total = total + p.float() - pre = total.to(torch.bfloat16).float() + residual.float() - updated = pre.to(torch.bfloat16) - u = pre.double() if drafter else updated.double() - inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) - normed = (u * inv * norm_w.double()).to(torch.bfloat16) - if not quant: - return normed, updated - return (normed, updated, *mxfp8_rows(normed)) - - -class M3McAllreduceNorm(OpWrapper): +class M3McAllreduceNorm(McAllreduceNormReference): """One TP group's boundary: ``updated = bf16(residual + bf16(sum_r bf16(x_r + x2_r)))``, ``normed = RMSNorm(updated)``: bitwise the stock all-reduce + fused residual RMSNorm (+ mxfp8_quantize) of SM107 TP2. @@ -232,29 +170,6 @@ def __call__( *((cap_out, cap_w, cap_eps) if cap_out is not None else ()), ) - def reference( - self, - per_rank_x: Sequence[torch.Tensor], - per_rank_x2: Optional[Sequence[torch.Tensor]], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - drafter: bool = False, - ) -> tuple[torch.Tensor, ...]: - """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). - - The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling - rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. - """ - partials = [] - for r, x in enumerate(per_rank_x): - if per_rank_x2 is None: - partials.append(x) - else: - partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) - return boundary(partials, residual, norm_w, eps, quant, drafter) - def is_valid( self, x: torch.Tensor, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py index 8b74c9974397..b7f8ab5a67db 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/comm/m3_mc_finalize_allreduce_norm.py @@ -4,40 +4,20 @@ from __future__ import annotations -from typing import Optional, Sequence +from typing import Optional import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper as OpWrapper -from .m3_mc_allreduce_norm import boundary +from ....ops.comm.reference import McFinalizeAllreduceNormReference +from ....ops.comm.reference import boundary as boundary +from ....ops.comm.reference import finalize_partial as finalize_partial -def finalize_partial( - fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] -) -> torch.Tensor: - """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k - order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). - - The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two - FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. - """ - m, topk = perm.shape - acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) - sc = scales.reshape(m, topk).float() - for k in range(topk): - rows = perm[:, k].long() - keep = (rows >= 0).unsqueeze(-1) - gathered = fc2[rows.clamp_min(0)].double() - step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() - acc = torch.where(keep, step, acc) - if shared is not None: - acc = acc.to(torch.bfloat16).float() + shared.float() - return acc.to(torch.bfloat16) - - -class M3McFinalizeAllreduceNorm(OpWrapper): +class M3McFinalizeAllreduceNorm(McFinalizeAllreduceNormReference): """The MoE boundary with the routed experts' finalize folded in. Per rank the partial is ``bf16(bf16(sum_k scales[t, k] * fc2[perm[t, k]]) + shared[t])`` (``perm`` -1 = dropped, @@ -146,30 +126,6 @@ def __call__( *((cap_out, cap_w, cap_eps) if cap_out is not None else ()), ) - def reference( - self, - per_rank_fc2: Sequence[torch.Tensor], - per_rank_perm: Sequence[torch.Tensor], - per_rank_scales: Sequence[torch.Tensor], - per_rank_shared: Optional[Sequence[torch.Tensor]], - residual: torch.Tensor, - norm_w: torch.Tensor, - eps: float, - quant: bool, - ) -> tuple[torch.Tensor, ...]: - """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the - collective exception, as in ``m3_mc_allreduce_norm.reference``).""" - partials = [ - finalize_partial( - fc2, - per_rank_perm[r], - per_rank_scales[r], - None if per_rank_shared is None else per_rank_shared[r], - ) - for r, fc2 in enumerate(per_rank_fc2) - ] - return boundary(partials, residual, norm_w, eps, quant) - def is_valid( self, fc2: torch.Tensor, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py deleted file mode 100644 index 71fce6b253cb..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding/m3_embedding.py +++ /dev/null @@ -1,48 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The generic token-embedding gather over the replicated table (an id outside the table gives a zero row).""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper -from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops - - -class M3Embedding(OpWrapper): - """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). - - The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode - embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) - ), - Cell( - why="an 8-row decode verify, int64 ids", - spec=dict(rows=8, id_dtype=torch.int64, pad=False), - ), - Cell( - why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) - ), - ) - - note = """ - An exact gather: `reference` must match bit for bit. - """ - - def __call__(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: - return pre_comm_embedding_ops(ids, table, 1, 0, None, 0, table.shape[0], False, 0) - - def reference(self, ids: torch.Tensor, table: torch.Tensor) -> torch.Tensor: - valid = (ids >= 0) & (ids < table.shape[0]) - rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] - return rows * valid.unsqueeze(-1).to(table.dtype) - - -m3_embedding = M3Embedding() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py deleted file mode 100644 index 78563b4db40d..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/m3_embedding_firstnorm_kernel.py +++ /dev/null @@ -1,128 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Fixed M3 H6144 gather/raw residual/RMS, reusing pinned K3 norm kernel. - -128 threads x 8 BF16 x 6 iterations covers the row. Unlike the original -K3 wrapper this target does not promise FlashInfer's bit-identical reduction. -""" - -import cuda.bindings.driver as cuda_driver -import cutlass -import cutlass.cute as cute -from cutlass.experimental import primitives as prims - -NORM_THREADS = 128 -NORM_VEC = 8 -NORM_WARPS = 4 # capacity: fixed: the CTA's warps - - -@cute.kernel -def m3_embed_norm_kernel( - ids: cute.Tensor, # int32 or int64 [N] - table: cute.Tensor, # bf16 [V, H] - weight: cute.Tensor, # bf16 [H] - raw: cute.Tensor, # bf16 [N, H], out: the rows - out: cute.Tensor, # bf16 [N, H], out: the normed rows - vocab: cutlass.Int64, - eps: cutlass.Float32, - tv_layout: cute.Layout, - tiler_mn: cute.Shape, - hidden: cutlass.Constexpr[int], -): - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - prims.griddepcontrol(prims.GridDepAction.LAUNCH_DEPENDENTS) - prims.griddepcontrol(prims.GridDepAction.WAIT) - - smem = cutlass.utils.SmemAllocator() - reduction_buffer = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((1, NORM_WARPS)), byte_alignment=4 - ) - - tok = cutlass.Int64(ids[bidx]) - valid = (tok >= cutlass.Int64(0)) & (tok < vocab) - row = cutlass.Int64(cutlass.select_(valid, tok, cutlass.Int64(0))) - - gX = cute.local_tile(table, tiler_mn, (row, 0)) - gR = cute.local_tile(raw, tiler_mn, (bidx, 0)) - gY = cute.local_tile(out, tiler_mn, (bidx, 0)) - w_layout = cute.prepend(weight.layout, cute.make_layout((tiler_mn[0],), stride=(0,))) - gW = cute.local_tile(cute.make_tensor(weight.iterator, w_layout), tiler_mn, (0, 0)) - - copy_atom_load = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), table.element_type, num_bits_per_copy=128 - ) - copy_atom_store = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), out.element_type, num_bits_per_copy=128 - ) - thr_copy_X = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) - thr_copy_W = cute.make_tiled_copy(copy_atom_load, tv_layout, tiler_mn).get_slice(tidx) - thr_copy_O = cute.make_tiled_copy(copy_atom_store, tv_layout, tiler_mn).get_slice(tidx) - - tXgX = thr_copy_X.partition_S(gX) - tXrX = cute.make_fragment_like(tXgX) - tWgW = thr_copy_W.partition_S(gW) - tWrW = cute.make_fragment_like(tWgW) - tXrW = thr_copy_X.retile(tWrW) - tXgR = thr_copy_O.partition_D(gR) - tXgO = thr_copy_O.partition_D(gY) - tXrO = cute.make_fragment_like(tXgO) - - # The row (zero for an id outside [0, V)), straight to the snapshot slot; the weight. - tXrX.store(cute.zeros_like(tXrX, dtype=table.element_type)) - if valid: - cute.copy(copy_atom_load, tXgX, tXrX) - cute.copy(copy_atom_load, tWgW, tWrW) - cute.copy(copy_atom_store, tXrX, tXgR) - - # flashinfer's RMSNormKernel arithmetic, in its order (see the header). - x = tXrX.load().to(cutlass.Float32) - x_sq = x * x - sum_sq = x_sq.reduce(cute.ReductionOp.ADD, init_val=cutlass.Float32(0.0), reduction_profile=0) - for i in cutlass.range_constexpr(5): - sum_sq = sum_sq + cute.arch.shuffle_sync_bfly(sum_sq, offset=1 << i) - lane = cute.arch.lane_idx() - warp = cute.arch.warp_idx() - if lane == 0: - reduction_buffer[0, warp] = sum_sq - cute.arch.barrier() - total = cutlass.Float32(0.0) - if lane < NORM_WARPS: - total = reduction_buffer[0, lane] - for i in cutlass.range_constexpr(5): - total = total + cute.arch.shuffle_sync_bfly(total, offset=1 << i) - mean_sq = total / cutlass.Float32(hidden) - rstd = cute.math.rsqrt(mean_sq + eps, fastmath=True) - w = tXrW.load().to(cutlass.Float32) - y = x * rstd * (w + cutlass.Float32(0.0)) - tXrO.store(y.to(out.element_type)) - cute.copy(copy_atom_store, tXrO, tXgO) - - -@cute.jit -def m3_embed_norm( - ids: cute.Tensor, - table: cute.Tensor, - weight: cute.Tensor, - raw: cute.Tensor, - out: cute.Tensor, - vocab: cutlass.Int64, - eps: cutlass.Float32, - n_tokens: cutlass.Constexpr[int], - hidden: cutlass.Constexpr[int], - use_pdl: cutlass.Constexpr[bool], - stream: cuda_driver.CUstream, -) -> None: - # flashinfer's thread-value layout for one row: thread t, value (v, k) -> column 8 t + v + 1024 k. - blocks = hidden // (NORM_VEC * NORM_THREADS) - tv_layout = cute.make_layout(((NORM_THREADS, 1), (NORM_VEC, blocks)), - stride=((NORM_VEC, 1), (1, NORM_VEC * NORM_THREADS))) # fmt: skip - m3_embed_norm_kernel( - ids, table, weight, raw, out, vocab, eps, tv_layout, (1, hidden), hidden - ).launch( - grid=[n_tokens, 1, 1], - block=[NORM_THREADS, 1, 1], - smem=NORM_WARPS * 4, - stream=stream, - use_pdl=use_pdl, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py index 276f8c32023b..f123d2f97db2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/embedding_firstnorm/op.py @@ -7,15 +7,22 @@ (:class:`DecodeEmbedding` compiles at construction). """ +import sys import threading -from typing import Dict, Tuple +from typing import Dict, Tuple, cast -import cuda.bindings.driver as cuda_driver -import cutlass.cute as cute +import cuda.bindings.driver as cuda_driver # noqa: F401 — used by the bound family implementation +import cutlass.cute as cute # noqa: F401 — used by the bound family implementation import torch -from cutlass.cute.runtime import from_dlpack +from cutlass.cute.runtime import from_dlpack as from_dlpack -from .m3_embedding_firstnorm_kernel import m3_embed_norm as _m3_embed_norm_jit +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding import ( + decode_embedding as _embedding_shared, +) + +from ....kernels.embedding_firstnorm.m3_embedding_firstnorm_kernel import ( # noqa: F401 + m3_embed_norm as _m3_embed_norm_jit, +) HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size MAX_ROWS = 8 # capacity: fixed: the kernel's built rows (decode steps) @@ -25,12 +32,7 @@ def arg(tensor): - return from_dlpack( - tensor.detach(), - assumed_align=tensor.element_size() - if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) - else 16, - ) + return _embedding_shared.arg(_EMBED_TARGET, tensor) def _check(ids, table, weight, raw, out) -> None: @@ -57,32 +59,7 @@ def _check(ids, table, weight, raw, out) -> None: def _compile(ids, table, weight, eps, raw, out, use_pdl: bool): - key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) - fn = _compiled.get(key) - if fn is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "m3_embed_norm compiles on its first call per shape: call it once before capture" - ) - with _lock: - fn = _compiled.get(key) - if fn is None: - stream = cuda_driver.CUstream(torch.cuda.current_stream(table.device).cuda_stream) - fn = _compiled[key] = cute.compile( - _m3_embed_norm_jit, - arg(ids), - arg(table), - arg(weight), - arg(raw), - arg(out), - int(table.shape[0]), - float(eps), - ids.shape[0], - HIDDEN, - bool(use_pdl), - stream, - ) - return fn + return _embedding_shared._compile(_EMBED_TARGET, ids, table, weight, eps, raw, out, use_pdl) def m3_embed_norm( @@ -96,20 +73,9 @@ def m3_embed_norm( ) -> Tuple[torch.Tensor, torch.Tensor]: """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / ``eps`` into ``out``; returns ``(out, raw)``.""" - _check(ids, table, weight, raw, out) - fn = _compile(ids, table, weight, eps, raw, out, use_pdl) - stream = cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) - fn( - arg(ids), - arg(table), - arg(weight), - arg(raw), - arg(out), - int(table.shape[0]), - float(eps), - stream, + return _embedding_shared.m3_embed_norm( + _EMBED_TARGET, ids, table, weight, eps, raw, out, use_pdl ) - return out, raw class DecodeEmbedding: @@ -118,25 +84,12 @@ class DecodeEmbedding: ``eps``. Built (and compiled) before CUDA-graph capture.""" def __init__(self, table, weight, eps, rows, id_dtype, use_pdl=True): - assert not torch.cuda.is_current_stream_capturing() - assert table.dtype == torch.bfloat16 and table.is_contiguous() and table.shape[1] == HIDDEN - assert weight.dtype == torch.bfloat16 and weight.shape == (HIDDEN,) - assert 1 <= rows <= MAX_ROWS and id_dtype in (torch.int32, torch.int64) - self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( - table, - weight, - rows, - id_dtype, - use_pdl, + _embedding_shared.init_decode_embedding( + self, _EMBED_TARGET, table, weight, eps, rows, id_dtype, use_pdl ) - self.eps = float(eps) - self.raw = torch.empty(rows, HIDDEN, device=table.device, dtype=torch.bfloat16) - self.output = torch.empty_like(self.raw) - ids = torch.empty(rows, device=table.device, dtype=id_dtype) - _compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) def __call__(self, ids): - assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() - return m3_embed_norm( - ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl - ) + return _embedding_shared.call_decode_embedding(self, _EMBED_TARGET, ids) + + +_EMBED_TARGET = cast(_embedding_shared.EmbeddingTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py deleted file mode 100644 index b973b175d8b4..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_bf16_linear.py +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The BF16 Eagle3 drafter's projections: ``x @ w^T`` through torch's GEMM.""" - -from __future__ import annotations - -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3Bf16Linear(OpWrapper): - """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. - - The drafter's fc over the captured target layers (K = 3 x 6144), its fused QKV (K = 2 x 6144, this rank's heads), - o_proj (this rank's input columns), gate_up and down projections run it. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), - Cell(why="drafter QKV, nine draft rows", spec=dict(m=9, n=1280, k=12288)), - Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=1024)), - Cell(why="drafter gate_up, 16 rows", spec=dict(m=16, n=4608, k=6144)), - Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=2304)), - ) - - note = """ - cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. - """ - - def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - return F.linear(x, w) - - def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - return (x.float() @ w.float().t()).to(x.dtype) - - -m3_bf16_linear = M3Bf16Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py index 8d4a56e3b437..2cbf500c414c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_draft_lowm.py @@ -99,7 +99,11 @@ def __call__( elif own_splitk: # This target's copy of the split-K kernel (ops/m3_lowm_splitk): bit for bit FlashInfer's; with ``ar`` # it also pushes the rows into a drafter boundary all-reduce (modeling's DRAFT_AR_PUSH). - from ..m3_lowm_splitk import SplitKTactic, run_splitk_dense, run_splitk_dense_push + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_lowm_splitk import ( + SplitKTactic, + run_splitk_dense, + run_splitk_dense_push, + ) if ar is not None: run_splitk_dense_push(x, weight.t(), out, use_pdl, SplitKTactic(*tactic), ar) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py deleted file mode 100644 index 9702cc0778a4..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_fp32_linear.py +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The router's fp32 projection beyond the decode routers' rows: ``float(x) @ w^T`` through torch's GEMM.""" - -from __future__ import annotations - -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) - - -class M3Fp32Linear(OpWrapper): - """fp32 router logits ``F.linear(float(x[M, 6144]), w[128, 6144])`` for prefill and steps above 48 rows (the - fp32 GEMV and the tensor-core router take 1..48). Whatever fp32 matmul precision the process runs with applies - (the engine leaves torch's default).""" - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell(why="a prefill chunk", spec=dict(m=2048, e=128, k=6144)), - Cell( - why="a generation step above the tensor-core router's 48 rows", - spec=dict(m=64, e=128, k=6144), - ), - ) - - note = """ - Under torch's default fp32 matmul precision the product may be TF32: the test's band covers that rounding against - the true fp32 reference. - """ - - def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - return F.linear(x.to(torch.float32), w) - - def reference(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - with true_fp32_matmul(): - return x.float() @ w.t() - - -m3_fp32_linear = M3Fp32Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py deleted file mode 100644 index 3f07fe2fb256..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_linear.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""One MXFP8 projection on the engine-tuned path: dynamic MXFP8 quantize of a bf16 activation + block-scaled GEMM.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) -from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_quantize_gemm_autotuned - -_BLOCK = 32 - - -def _unswizzle(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: - """[rows, cols] UE8M0 bytes from the 128x4-swizzled layout (rows padded to 128, columns to 4).""" - pad_r = (rows + 127) // 128 * 128 - pad_c = (cols + 3) // 4 * 4 - s = scale.view(torch.uint8).reshape(pad_r // 128, pad_c // 4, 32, 4, 4) - return s.permute(0, 3, 2, 1, 4).reshape(pad_r, pad_c)[:rows, :cols] - - -def _dequant(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: - rows, k = x.shape - s = _unswizzle(scale, rows, k // _BLOCK).to(torch.int32) - factor = torch.pow(2.0, (s - 127).to(torch.float32)) - return ( - x.view(torch.float8_e4m3fn).float().view(rows, k // _BLOCK, _BLOCK) * factor.unsqueeze(-1) - ).view(rows, k) - - -def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: - """fp32 product of two dequantized MXFP8 operands (swizzled UE8M0 scales).""" - with true_fp32_matmul(): - return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() - - -class M3Mxfp8Linear(OpWrapper): - """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 - out, on the backend the engine selected: - - * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize - and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); - * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or - the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by - ``alpha`` (fp32 [1], 1.0 for these checkpoints). - - Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it - outside the decode-step kernels. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( - Cell( - why="sparse-layer QKV, a prefill chunk, native GEMM", - spec=dict(m=2048, n=1536, k=6144, path="native"), - ), - Cell( - why="o_proj, a prefill chunk, native GEMM", - spec=dict(m=2048, n=6144, k=1024, path="native"), - ), - Cell( - why="dense down projection, eager decode, native GEMM", - spec=dict(m=4, n=6144, k=1536, path="native"), - ), - Cell( - why="shared gate_up, FlashInfer GEMM", - spec=dict(m=64, n=768, k=6144, path="flashinfer"), - ), - Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=1024, path="tuned")), - ) - - note = """ - Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` - quantizes x with the same block rule (UE8M0 scale = the block's power-of-two ceiling over e4m3's 448) and - multiplies dequantized operands in fp32; the test's band covers the accumulation order. - """ - - def __call__( - self, - x: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - alpha: torch.Tensor, - output_dtype: torch.dtype, - tuned_quantize_gemm: bool, - tune: bool, - flashinfer: bool, - native_tuning: bool, - ) -> torch.Tensor: - if tuned_quantize_gemm: - return mxfp8_quantize_gemm_autotuned(x, weight, weight_scale, output_dtype, tune=tune) - act, act_scale = torch.ops.trtllm.mxfp8_quantize(x, True) - if flashinfer: - return torch.ops.trtllm.flashinfer_mm_mxfp8( - act, act_scale, weight, weight_scale, output_dtype - ) - gemm = ( - torch.ops.trtllm.mxfp8_mxfp8_gemm_autotuned - if native_tuning - else torch.ops.trtllm.mxfp8_mxfp8_gemm - ) - return gemm(act, act_scale, weight, weight_scale, alpha, output_dtype) - - def reference( - self, - x: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - alpha: torch.Tensor, - output_dtype: torch.dtype, - tuned_quantize_gemm: bool, - tune: bool, - flashinfer: bool, - native_tuning: bool, - ) -> torch.Tensor: - """The MXFP8 round trip of ``x`` times the dequantized weight, times ``alpha``; the backend flags are - ignored.""" - m, k = x.shape - blocks = x.float().view(m, k // _BLOCK, _BLOCK) - amax = blocks.abs().amax(-1, keepdim=True) - exp = torch.ceil(torch.log2(torch.where(amax > 0, amax / 448.0, torch.ones_like(amax)))) - scale = torch.pow(2.0, exp) - q = (blocks / scale).to(torch.float8_e4m3fn).float() * scale - with true_fp32_matmul(): - out = q.view(m, k) @ _dequant(weight, weight_scale).t() - return (out * alpha.float()).to(output_dtype) - - -m3_mxfp8_linear = M3Mxfp8Linear() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py deleted file mode 100644 index f8deb858c843..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_target_head.py +++ /dev/null @@ -1,346 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The target lm_head shard GEMV of a verify step of at most 8 rows over a lossless 12-bit coding of its BF16 weight, -and that coding (``encode_chunked`` / ``decode_reference`` / ``round_trip_ok``). - -The coding: the BF16 exponents of a trained weight sit almost all in a window of 15 consecutive values -``[base, base + 14]``. Each weight is one byte ``sign << 7 | mantissa`` plus a 4-bit code ``exponent - base``; code 15 -marks an exception, which decodes to +0 (it adds nothing to the MMA sums); its exact product ``x * w`` is added from a -per-tile list. The bytes are laid out in ``mma.m16n8k16`` A-fragment order: -per 16-row tile and unit of 4 k16 steps, each lane reads the sign/mantissa bytes of steps 0-1, of steps 2-3, then 4 -code words (``[tiles][K / 64][3][32 lanes][16 bytes]``). A coded weight is used only after ``round_trip_ok`` proved -it decodes to the BF16 weight bit for bit. -""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) - -__all__ = ["m3_target_head"] - -#: Verify rows per call (requests x (draft length + 1)). -MAX_ROWS = 8 -#: Warps per CTA (the split-tile scratch holds 2 x 4 x 32 fp32 per warp). -CTA_WARPS = 16 -#: The target lm_head's TP4 shard (200064-id vocabulary) and the hidden size. -TARGET_SHARD = 50016 -HIDDEN = 6144 -#: Codes per window: 0..14 decode, 15 marks an exception. -WINDOW = 15 -#: The largest exponent base whose window top plus one (base + 15) is still a finite BF16 exponent. -MAX_BASE = 254 - WINDOW -#: Rows per block of encode_chunked / round_trip_ok (bounded temporaries). -ROWS_PER_CHUNK = 4096 - - -def _pad_rows(weight: torch.Tensor) -> torch.Tensor: - n = weight.shape[0] - tiles = (n + 15) // 16 - if tiles * 16 == n: - return weight - return torch.cat([weight, weight.new_zeros(tiles * 16 - n, weight.shape[1])]) - - -def _fragment_view(w: torch.Tensor) -> torch.Tensor: - """[tiles * 16, K] -> [tiles, K / 16, 32 lanes, 8 values] in mma A-fragment order (a permutation of the values).""" - n, k = w.shape - # rows: (tile, rh, g); cols: (ks, ch, j, e) -> (tile, ks, g, j, ch, rh, e) - return ( - w.view(n // 16, 2, 8, k // 16, 2, 4, 2) - .permute(0, 3, 2, 5, 4, 1, 6) - .reshape(n // 16, k // 16, 32, 8) - ) - - -def _exponents(weight: torch.Tensor) -> torch.Tensor: - return (weight.contiguous().view(torch.int16).to(torch.int32) >> 7) & 0xFF - - -def _window_base(hist: torch.Tensor) -> int: - """The base of the 15-exponent window holding the most weights (the lowest such base).""" - window = torch.stack([hist[b : b + WINDOW].sum() for b in range(0, MAX_BASE + 1)]) - return int(window.argmax()) - - -def encode(weight: torch.Tensor, base: Optional[int] = None): - """Lossless 12-bit coding of a BF16 [N, K] weight (see the module docstring); ``base`` defaults to this weight's - own window. - - Returns (packed uint8 [tiles * K / 64 * 1536], base, exc_off int32 [tiles + 1], exc_pos int32 [E] (row in tile - << 16 | k), exc_w fp32 [E, 2] ({w, 0})). - """ - w = _pad_rows(weight) - n, k = w.shape - tiles, ks = n // 16, k // 16 - assert ks % 4 == 0, "K must be a multiple of 64" - bits = w.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF - expo = (bits >> 7) & 0xFF - if base is None: - base = _window_base(torch.bincount(expo.flatten(), minlength=256)) - code = expo - base - exc = (code < 0) | (code >= WINDOW) - code = torch.where(exc, torch.full_like(code, WINDOW), code) - sm = ((bits >> 8) & 0x80) | (bits & 0x7F) - sm = torch.where(exc, torch.zeros_like(sm), sm) - # Exceptions: (tile, row in tile, k) in tile-major order; code 15 decodes to +0. - rows, cols = exc.nonzero(as_tuple=True) - order = torch.argsort((rows // 16).long() * (16 * k) + (rows % 16).long() * k + cols.long()) - rows, cols = rows[order], cols[order] - exc_w = torch.stack( - [ - w[rows, cols].float(), - torch.zeros((rows.numel(),), device=w.device), - ], - dim=1, - ) - exc_pos = (((rows % 16) << 16) | cols).to(torch.int32) - exc_off = torch.zeros(tiles + 1, dtype=torch.int32, device=w.device) - exc_off[1:] = torch.cumsum(torch.bincount((rows // 16).long(), minlength=tiles), 0).to( - torch.int32 - ) - # Fragment order, then per (tile, group of 4 k16 steps, lane): 32 sign/mantissa bytes and 4 code words. - smf = _fragment_view(sm.to(torch.uint8).view(n, k)).view(tiles, ks // 4, 4, 32, 8) - cf = ( - _fragment_view(code.to(torch.uint8).view(n, k)) - .view(tiles, ks // 4, 4, 32, 8) - .to(torch.int32) - ) - # Code word of a k16 step: value 2p -> bits 4p..4p+3, value 2p+1 -> bits 16+4p.. - cw = torch.zeros(tiles, ks // 4, 4, 32, dtype=torch.int64, device=w.device) - for p in range(4): - cw |= (cf[..., 2 * p].long() << (4 * p)) | (cf[..., 2 * p + 1].long() << (16 + 4 * p)) - cw = cw.to(torch.int32) # wraps for bit 31; the kernel reads it as uint32 - # Word 0 of a lane = the bytes of steps 0 and 1; word 1 = steps 2 and 3; word 2 = the code words of steps 0..3. - seg01 = ( - smf.permute(0, 1, 3, 2, 4).reshape(tiles, ks // 4, 32, 2, 16).permute(0, 1, 3, 2, 4) - ) # [.., 2, 32, 16] - seg2 = cw.permute(0, 1, 3, 2).contiguous().view(torch.uint8).view(tiles, ks // 4, 1, 32, 16) - packed = torch.cat([seg01.contiguous(), seg2], dim=2) # [tiles, ks / 4, 3, 32, 16] - return packed.contiguous().view(-1), base, exc_off, exc_pos, exc_w.contiguous() - - -def encode_chunked(weight: torch.Tensor, rows_per_chunk: int = ROWS_PER_CHUNK): - """``encode`` over row blocks (bounded temporaries), with the exponent window chosen on the whole weight.""" - n, _ = weight.shape - assert rows_per_chunk % 16 == 0 - hist = torch.zeros(256, dtype=torch.int64, device=weight.device) - for r in range(0, n, rows_per_chunk): - hist += torch.bincount(_exponents(weight[r : r + rows_per_chunk]).flatten(), minlength=256) - base = _window_base(hist) - packed, offs, poss, ws = [], [], [], [] - n_exc = 0 - for r in range(0, n, rows_per_chunk): - p, _, off, pos, ew = encode(weight[r : r + rows_per_chunk], base=base) - packed.append(p) - offs.append(off[:-1] + n_exc) - poss.append(pos) - ws.append(ew) - n_exc += pos.numel() - exc_off = torch.cat(offs + [torch.tensor([n_exc], dtype=torch.int32, device=weight.device)]) - return torch.cat(packed), base, exc_off, torch.cat(poss), torch.cat(ws) - - -def decode_reference(packed: torch.Tensor, base: int, n: int, k: int) -> torch.Tensor: - """BF16 [N, K] from the coded stream with the in-kernel rule (exceptions, code 15, decode to +0).""" - tiles, ks = (n + 15) // 16, k // 16 - p = packed.view(tiles, ks // 4, 3, 32, 16) - sm = ( - p[:, :, :2] - .permute(0, 1, 3, 2, 4) - .reshape(tiles, ks // 4, 32, 4, 8) - .permute(0, 1, 3, 2, 4) - .to(torch.int32) - ) - cw = ( - p[:, :, 2] - .contiguous() - .view(torch.int32) - .view(tiles, ks // 4, 32, 4) - .permute(0, 1, 3, 2) - .long() - & 0xFFFFFFFF - ) - codes = torch.stack([(cw >> (4 * (i // 2) + 16 * (i % 2))) & 0xF for i in range(8)], dim=-1).to( - torch.int32 - ) - bits = ((sm & 0x80) << 8) | ((codes + base) << 7) | (sm & 0x7F) - bits = torch.where(codes == WINDOW, torch.zeros_like(bits), bits) - frag = bits.reshape(tiles, ks, 32, 8) - # Inverse of _fragment_view: (tile, ks, g, j, ch, rh, e) -> (tile, rh, g, ks, ch, j, e). - w = frag.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, k) - return w.to(torch.int16).view(torch.bfloat16)[:n] - - -def decode( - packed: torch.Tensor, - base: int, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - n: int, - k: int, - row0: int = 0, - row1: Optional[int] = None, -) -> torch.Tensor: - """BF16 rows [row0, row1) of the coded weight [N, K] with the exceptions applied (what the kernel multiplies by). - ``row0`` must be a multiple of 16.""" - row1 = n if row1 is None else row1 - assert row0 % 16 == 0 - t0, t1 = row0 // 16, (row1 + 15) // 16 - unit_bytes = (k // 64) * 3 * 32 * 16 - dec = decode_reference(packed[t0 * unit_bytes : t1 * unit_bytes], base, row1 - row0, k).clone() - counts = (exc_off[1:] - exc_off[:-1]).long() - tile_of = torch.repeat_interleave(torch.arange(counts.numel(), device=packed.device), counts) - rows = tile_of * 16 + (exc_pos >> 16).long() - cols = (exc_pos & 0xFFFF).long() - sel = (rows >= row0) & (rows < row1) - dec[rows[sel] - row0, cols[sel]] = exc_w.view(-1, 2)[sel, 0].to(torch.bfloat16) - return dec - - -def round_trip_ok( - packed: torch.Tensor, - base: int, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - weight: torch.Tensor, - rows_per_chunk: int = ROWS_PER_CHUNK, -) -> bool: - """Whether the coded weight (code values + exception list) decodes to ``weight`` bit for bit, with its exponent - window below the BF16 infinity exponent.""" - n, k = weight.shape - if not 0 <= base <= MAX_BASE or k % 64: - return False - if ( - packed.numel() != ((n + 15) // 16) * 16 * k * 3 // 2 - or exc_off.numel() != (n + 15) // 16 + 1 - ): - return False - for r0 in range(0, n, rows_per_chunk): - r1 = min(n, r0 + rows_per_chunk) - dec = decode(packed, base, exc_off, exc_pos, exc_w, n, k, r0, r1) - if not torch.equal(dec.view(torch.int16), weight[r0:r1].contiguous().view(torch.int16)): - return False - return True - - -class M3TargetHead(OpWrapper): - """``out[:M, :n]`` (BF16) = x [M, K] (BF16, M <= 8) times the coded BF16 shard W [n, K] transposed: each logit - the fp32 sum of the exact BF16 products, rounded once to BF16 (the in-tree BF16 head GEMM's arithmetic, in - another summation order). - - ``w`` / ``exp_base`` / ``exc_off`` / ``exc_pos`` / ``exc_w`` are ``encode_chunked``'s outputs for W; ``part`` - (fp32 [grid * 16 * 256]) and ``counter`` (int32 [ceil(n / 16)], zero between calls, re-armed by the kernel) are - the split-tile scratch; ``grid`` is the number of 16-warp CTAs (one per SM). - """ - - ARCHS = frozenset({Arch.SM_107}) - - # The verify of a fused-verify step at most 8 rows wide: 1 + 3 drafts for one or two requests (4 / 8 rows), and - # every other row count the kernel's one n8 MMA tile serves; one cell with an exception in many tiles, one with - # a shard whose rows are not a multiple of 16. - CELLS: tuple[Cell, ...] = ( - *( - Cell( - why=f"{m} verify row(s) of the TP4 target head shard", - spec=dict(m=m, n=TARGET_SHARD, k=HIDDEN), - ) - for m in range(1, MAX_ROWS + 1) - ), - Cell( - why="4 rows with outliers and exact zeros in many tiles (the exception lists)", - spec=dict(m=4, n=TARGET_SHARD, k=HIDDEN, exceptions=True), - ), - Cell( - why="2 rows over a shard of 4000 rows (250 tiles), K 1024", - spec=dict(m=2, n=4000, k=1024, exceptions=True), - ), - ) - - note = """ - The weight the kernel multiplies by is the checkpoint's BF16 weight bit for bit (round_trip_ok); an exception - adds nothing to the MMA sums and its exact product x * w is added at its tile's flush. A logit therefore differs - from the in-tree head's only by the fp32 summation order before its single BF16 rounding (ulp-level); - integer-valued data, whose sums are exact in any order, is equal bit for bit. Split tiles are summed in warp - order: the result is deterministic. - """ - - def __call__( - self, - w: torch.Tensor, - x: torch.Tensor, - out: torch.Tensor, - part: torch.Tensor, - counter: torch.Tensor, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - exp_base: int, - n: int, - grid: int, - pdl: bool = True, - ) -> torch.Tensor: - torch.ops.trtllm.minimax_m3_target_head( - w, x, out, part, counter, exc_off, exc_pos, exc_w, exp_base, n, grid, pdl - ) - return out - - def reference( - self, - w: torch.Tensor, - x: torch.Tensor, - out: torch.Tensor, - part: torch.Tensor, - counter: torch.Tensor, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - exp_base: int, - n: int, - grid: int, - pdl: bool = True, - ) -> torch.Tensor: - """The BF16 logits [M, n] of x against the decoded weight: the fp32 product, rounded once to BF16. - - ``out``, ``part``, ``counter``, ``grid`` and ``pdl`` do not change the result.""" - weight = decode(w, exp_base, exc_off, exc_pos, exc_w, n, x.shape[1]) - with true_fp32_matmul(): - logits = x.float() @ weight.float().t() - return logits.to(torch.bfloat16) - - def is_valid( - self, - w: torch.Tensor, - x: torch.Tensor, - out: torch.Tensor, - part: torch.Tensor, - counter: torch.Tensor, - exc_off: torch.Tensor, - exc_pos: torch.Tensor, - exc_w: torch.Tensor, - exp_base: int, - n: int, - grid: int, - pdl: bool = True, - ) -> None: - # A split tile's last finisher is the CTA whose arrival makes the count reach its owners: a stale count - # stores a partial sum. - assert not bool(counter[: (n + 15) // 16].any()), "counter must be zero between calls" - # The exceptions are read per tile in list order with no bound against the coded stream. - assert int(exc_off[-1].item()) == exc_pos.numel(), "exc_off must end at the exception count" - assert bool((exc_off[1:] >= exc_off[:-1]).all()), "exc_off must be non-decreasing" - assert out.data_ptr() != x.data_ptr(), "out must not alias x" - - -m3_target_head = M3TargetHead() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py index c977362e92b8..8fa99fade5c2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_draft_step.py @@ -27,9 +27,12 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + encode_chunked, + round_trip_ok, +) from tensorrt_llm.logger import logger -from ..ops.gemm.m3_target_head import encode_chunked, round_trip_ok from ..ops.spec.m3_draft_step import MAX_TILES_PER_CTA, MAX_TOKENS, bf16_fragments, max_tile_pieces from ..ops.spec.m3_draft_step import m3_draft_step as _draft_step_op from . import m3_report diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py index e582e2309e61..30f5d8c26700 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_drafter_attn.py @@ -20,27 +20,50 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Mapping, Optional, Tuple +import sys +from typing import TYPE_CHECKING as TYPE_CHECKING +from typing import Mapping as Mapping +from typing import Optional as Optional +from typing import Tuple as Tuple +from typing import cast import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + drafter as _drafter_shared, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + DrafterFullCtxBase, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + drafter_rope_applies as _shared_rope_applies, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + initialize_drafter_state as _shared_initialize, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.drafter import ( + install as _shared_install, +) from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( - FlashInferTrtllmGenFmha, + FlashInferTrtllmGenFmha as FlashInferTrtllmGenFmha, ) -from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase -from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase as FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams as FmhaParams from tensorrt_llm._torch.attention.backends.fmha.utils import ( - get_attention_chunk_size, - get_bmm1_scale, + get_attention_chunk_size as get_attention_chunk_size, +) +from tensorrt_llm._torch.attention.backends.fmha.utils import get_bmm1_scale as get_bmm1_scale +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionInputType as AttentionInputType, ) -from tensorrt_llm._torch.attention.backends.interface import AttentionInputType from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( - fp8_e4m3_saturate, + fp8_e4m3_saturate as fp8_e4m3_saturate, ) -from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor -from tensorrt_llm.bindings.internal import thop -from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.quantization.mode import QuantMode +from tensorrt_llm._utils import TensorWrapper as TensorWrapper +from tensorrt_llm._utils import convert_to_torch_tensor as convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop as thop +from tensorrt_llm.functional import PositionEmbeddingType as PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode as QuantMode from . import m3_fullctx_decode as fullctx @@ -57,9 +80,15 @@ def fwd_has_no_mrope(fwd) -> bool: from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention -class M3DrafterFullCtxFmha(FlashInferTrtllmGenFmha): +class M3DrafterFullCtxFmha(DrafterFullCtxBase): """Generation phase of the drafter layer: in-tree preprocess, PrimTS split-KV decode.""" + FULLCTX = fullctx + GEOMETRY_ERROR = "not the TP8 FP8 drafter geometry" + + def _launch_rope(self, *args) -> None: + torch.ops.trtllm.minimax_m3_sm107tp8_drafter_rope_fp8_kv(*args, ROPE_CARVEOUT) + def __init__( self, attn: "TrtllmAttention", @@ -74,285 +103,25 @@ def __init__( on_rope_taken=None, ) -> None: super().__init__(attn) - self._recipes = recipes - self._max_kv_len = max_kv_len - self._own_reducer = own_reducer - self._fused_rope_mode = fused_rope_mode - self._pdl = pdl - self._on_rope_taken = on_rope_taken - self._on_taken = on_taken - self._on_uncovered = on_uncovered - self._on_missed = on_missed - - def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: - nreq = int(metadata.num_generations) - if nreq <= 0 or int(q.shape[0]) % nreq: - return None - return int(q.shape[0]) // nreq, nreq - - def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: - if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: - return False - if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: - return False - attn = self.attn - if ( - attn.num_heads != fullctx.HEADS - or attn.num_kv_heads != fullctx.KV_HEADS - or attn.head_dim != fullctx.HEAD_DIM - or not QuantMode(attn.quant_mode).has_fp8_kv_cache() - ): - self._on_uncovered("not the TP8 FP8 drafter geometry") - return False - # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's - # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. - if ( - forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant - or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig - ): - self._on_uncovered("per-step KV cache scales") - return False - if int(metadata.max_seq_len) > self._max_kv_len: - self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") - return False - shape = self._shape(q, metadata) - if shape is None or shape not in self._recipes: - self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") - return False - # Compiled before warmup (``install``); the forward never compiles one. - if fullctx.plan_for(q.device, *shape, self._own_reducer) is None: - self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") - return False - return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) - - def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: - attn = params.attn - width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim - rotary_dim = int(attn.rope_params.dim) - return ( - self._fused_rope_mode is not None - and params.qkv_input is not None - and attention_input.dim() == 2 - and attention_input.dtype == torch.bfloat16 - and attention_input.is_contiguous() - and int(attention_input.shape[1]) == width - and not params.is_cross - and int(attn.position_embedding_type) == int(PositionEmbeddingType.rope_gpt_neox) - and attn.rotary_cos_sin is not None - and attn.rotary_cos_sin.dtype == torch.float32 - and fwd_has_no_mrope(params.fwd) - and 0 < rotary_dim <= attn.head_dim - and rotary_dim % 16 == 0 - and int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length) - # A one-request multi-token pass: the in-tree preprocess does not build its token map (buildDecoderInfo - # runs for several requests only) and reads the workspace's, so it keeps that path. - and not (self._multi_token(params) and int(params.batch_size) == 1) - ) - - @staticmethod - def _multi_token(params: FmhaParams) -> bool: - return ( - params.spec_decoding_generation_lengths is not None - and int(params.attn.predicted_tokens_per_seq) > 1 + _shared_initialize( + self, + _DRAFTER_BINDINGS, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode, + pdl, + on_rope_taken, + has_one_request_multi=False, ) - def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: - """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" - if not self._fused_rope_applies(params, attention_input): - return False - attn = params.attn - meta = params.meta - qlen, nreq = int(params.input_seq_length), int(params.batch_size) - plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) - if plan is None: - return False - # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary - # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the - # block table is the pool's [requests, 2, blocks] K / V sub-page indices. - pool_index, layer_in_pool = ( - int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] - ) - tpb = int(params.tokens_per_block) - sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim - base = ( - int(meta.host_kv_cache_pool_pointers[pool_index, 0]) - + layer_in_pool * int(params.kv_factor) * sub_page_bytes - ) - pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 - kv_pages = convert_to_torch_tensor( - TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) - ) - tables = meta.kv_cache_block_offsets[ - pool_index, params.seq_offset : params.seq_offset + nreq - ] - multi = self._multi_token(params) - if not torch.cuda.is_current_stream_capturing(): - # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess - # path checks; the op itself writes V at the table's V sub-page. - tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb - cols = torch.arange(tables.shape[-1], device=tables.device) - in_use = cols[None, :] < tpb_held[:, None] - offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() - if offsets.numel() != 1 or int(offsets) != 1: - raise RuntimeError( - f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" - ) - q = torch.empty( - int(params.num_tokens), - attn.num_heads, - attn.head_dim, - dtype=torch.float8_e4m3fn, - device=attention_input.device, - ) - torch.ops.trtllm.minimax_m3_sm107tp8_drafter_rope_fp8_kv( - attention_input, - attn.rotary_cos_sin, - params.sequence_lengths, - params.spec_decoding_position_offsets if multi else None, - kv_pages, - tables, - q, - qlen, - attn.num_heads, - attn.num_kv_heads, - int(attn.rope_params.dim), - tpb, - self._fused_rope_mode, - self._pdl, - ROPE_CARVEOUT, - ) - plan.run( - q, - kv_pages[:-1], - kv_pages[1:], - tables[:, 0], - params.sequence_lengths[:nreq], - get_bmm1_scale(attn), - params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), - ) - if self._on_rope_taken is not None: - self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") - self._on_taken(params.num_tokens) - return True - - def run_generation(self, params: FmhaParams) -> None: - attention_input = params.qkv_input if params.qkv_input is not None else params.query_input - if attention_input is None: - raise RuntimeError("the drafter decode requires QKV or query input.") - attn = params.attn - meta = params.meta - fwd = params.fwd - rope_params = attn.rope_params - if self._run_fused_rope(params, attention_input): - return - # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 - # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). - fp8_context_fmha = False - ( - q_processed, - kv_pool, - block_tables, - _kv_scale_pool, - _bmm1_scale, - _bmm2_scale, - _fmha_workspace, - _cu_seqlens, - _max_q_len, - _max_kv_len, - _window_left, - is_multi_token_gen, - ) = thop.trtllm_gen_generation_preprocess( - attention_input, - params.workspace, - params.sequence_lengths, - params.spec_decoding_generation_lengths, - params.spec_decoding_position_offsets, - meta.kv_cache_block_offsets, - meta.host_kv_cache_pool_pointers, - meta.host_kv_cache_pool_mapping, - fwd.kv_scale_orig_quant, - fwd.kv_scale_quant_orig, - fwd.out_scale, - attn.rotary_inv_freq, - attn.rotary_cos_sin, - fwd.mrope_position_deltas, - attn.local_layer_idx, - params.seq_offset, - attn.num_heads, - attn.num_kv_heads, - attn.head_dim, - params.tokens_per_block, - attn.quant_mode, - params.max_attention_window_size, - params.cyclic_attention_window_size, - params.num_tokens, - params.batch_size, - params.input_seq_length, - params.max_past_kv_length, - rope_params.dim, - rope_params.theta, - int(rope_params.scale_type), - rope_params.scale, - rope_params.max_positions, - attn.position_embedding_type, - get_bmm1_scale(attn), - 1.0, - fp8_context_fmha, - attn.predicted_tokens_per_seq, - get_attention_chunk_size(attn), - self._multi_processor_count, - params.total_num_blocks, - params.kv_factor, - True, - params.is_cross, - ) - qlen, nreq = params.input_seq_length, params.batch_size - plan = fullctx.plan_for(attention_input.device, qlen, nreq, self._own_reducer) - # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the - # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row - # pass is never multi-token. - if plan is None or (is_multi_token_gen and qlen == 1): - raise RuntimeError( - f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" - ) - rows = params.num_tokens * attn.num_heads * attn.head_dim - q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) - q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) - # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 - # of the pool (the two views have the same page count, as the kernel requires). - if not torch.cuda.is_current_stream_capturing(): - # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. - lengths = params.spec_decoding_generation_lengths - if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): - raise RuntimeError( - f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" - ) - # Only the blocks a request holds are laid out; the table's tail entries are unset. - tpb = int(params.tokens_per_block) - held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb - cols = torch.arange(block_tables.shape[-1], device=block_tables.device) - in_use = cols[None, :] < held[:, None] - offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() - if offsets.numel() != 1 or int(offsets) != 1: - raise RuntimeError( - f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" - ) - # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. - # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. - pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 - kv_pages = convert_to_torch_tensor( - TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) - ) - plan.run( - q, - kv_pages[:-1], - kv_pages[1:], - block_tables[:nreq, 0], - params.sequence_lengths[:nreq], - get_bmm1_scale(attn), - params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + def _fused_rope_applies(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + return _shared_rope_applies( + self, _DRAFTER_BINDINGS, params, attention_input, has_one_request_multi=False ) - self._on_taken(params.num_tokens) def install( @@ -370,35 +139,20 @@ def install( """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it was installed on.""" - from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention - - count = 0 - device = next(draft_model.parameters()).device - # The backend is not a module: each attention module of the drafter holds it as ``attn``. - for module in draft_model.modules(): - backend = getattr(module, "attn", None) - if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): - continue - if not bool((backend.kv_cache_scaling_factor == 1.0).all()): - on_uncovered("a scaled FP8 KV cache") - continue - fullctx.prebuild(device, recipes, max_kv_len, own_reducer, pdl) - libs = backend._fmha_manager.fmha_libs - if not any(isinstance(lib, M3DrafterFullCtxFmha) for lib in libs): - libs.insert( - 0, - M3DrafterFullCtxFmha( - backend, - recipes, - max_kv_len, - own_reducer, - on_taken, - on_uncovered, - on_missed, - fused_rope_mode=fused_rope_mode, - pdl=pdl, - on_rope_taken=on_rope_taken, - ), - ) - count += 1 - return count + return _shared_install( + draft_model=draft_model, + recipes=recipes, + max_kv_len=max_kv_len, + own_reducer=own_reducer, + pdl=pdl, + on_taken=on_taken, + on_uncovered=on_uncovered, + on_missed=on_missed, + fused_rope_mode=fused_rope_mode, + on_rope_taken=on_rope_taken, + prebuild=fullctx.prebuild, + fmha_class=M3DrafterFullCtxFmha, + ) + + +_DRAFTER_BINDINGS = cast(_drafter_shared.DrafterBindings, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py index 292d1df264b1..51f517f5433b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fold_tail.py @@ -14,70 +14,27 @@ Integer copies only: the values are the ones the sampler's and the engine's own copies write. With ``cpp`` the same writes run as this target's C++ op (trtllm::minimax_m3_sm107_tp8_fold_tail, -minimaxM3Sm107Tp8/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. +minimaxM3Sm107/foldTail.cu: one 32-lane CTA per row) instead of the Triton kernel, for widths up to 32. """ import torch -import triton -import triton.language as tl +import triton as triton +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.spec.m3_fold_tail import ( + _fold_tail_kernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.fold_tail import ( + fold_tail as _shared_fold_tail, +) -@triton.jit -def _fold_tail_kernel( - ctrl, - slots, - src_new, - sn_s0, - sn_s1, - SN_W: tl.constexpr, - dst_new, - dn_s0, - dn_s1, - DN_W: tl.constexpr, - src_next, - sx_s0, - sx_s1, - SX_W: tl.constexpr, - dst_next, - dx_s0, - dx_s1, - DX_W: tl.constexpr, - src_draft, - sd_s0, - sd_s1, - SD_W: tl.constexpr, - dst_draft, - dd_s0, - dd_s1, - DD_W: tl.constexpr, - src_lens, - dst_lens, - ids, - draft_ids, - pos_off, - kv_off, - T: tl.constexpr, - D: tl.constexpr, - BLOCK: tl.constexpr, -): - i = tl.program_id(0) - n = tl.load(ctrl) - if i < n: - s = tl.load(slots + i).to(tl.int64) - c = tl.arange(0, BLOCK) - new = tl.load(src_new + i * sn_s0 + c * sn_s1, mask=c < SN_W, other=0) - tl.store(dst_new + s * dn_s0 + c * dn_s1, new, mask=c < DN_W) - nxt = tl.load(src_next + i * sx_s0 + c * sx_s1, mask=c < SX_W, other=0) - tl.store(dst_next + s * dx_s0 + c * dx_s1, nxt, mask=c < DX_W) - drf = tl.load(src_draft + i * sd_s0 + c * sd_s1, mask=c < SD_W, other=0) - tl.store(dst_draft + s * dd_s0 + c * dd_s1, drf, mask=c < DD_W) - ln = tl.load(src_lens + i) - tl.store(dst_lens + s, ln) - # The next step's inputs read the store's padded rows: zero past each source width. - tl.store(ids + i * T + c, tl.where(c < DX_W, nxt, 0), mask=c < T) - tl.store(draft_ids + i * D + c, tl.where(c < DD_W, drf, 0), mask=c < D) - tl.store(pos_off + i * T + c, tl.zeros([BLOCK], dtype=tl.int32) + ln, mask=c < T) - tl.store(kv_off + i, ln - T) + +def _launch_cpp(*args: object, one_warp: bool) -> None: + op = ( + torch.ops.trtllm.minimax_m3_sm107_tp8_fold_tail_one_warp + if one_warp + else torch.ops.trtllm.minimax_m3_sm107_tp8_fold_tail + ) + op(*args) def fold_tail( @@ -103,78 +60,25 @@ def fold_tail( """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width]. ``one_warp`` (with ``cpp``): the in-kernel form (one warp looping over the rows, foldTailWarp in - minimaxM3Sm107Tp8/foldTail.h) that a caller folding the tail into its own kernel runs; the same writes.""" - two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) - one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) - if any(t.dim() != 2 for t in two) or any(t.dim() != 1 or t.stride(0) != 1 for t in one): - return False - if ( - any(t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,)) - or slots.dtype != torch.long - ): - return False - t_step = tokens_per_step - d_step = t_step - 1 - if ( - rows <= 0 - or slots.numel() < rows - or any(t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens)) - or ids.numel() < rows * t_step - or pos_off.numel() < rows * t_step - or draft_ids.numel() < rows * d_step - or kv_off.numel() < rows - ): - return False - widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) - widths += (src_draft.shape[1], dst_draft.shape[1], t_step) - if cpp and max(widths) <= 32: - op = ( - torch.ops.trtllm.minimax_m3_sm107_tp8_fold_tail_one_warp - if one_warp - else torch.ops.trtllm.minimax_m3_sm107_tp8_fold_tail - ) - op( - ctrl, slots, rows, src_new, dst_new, src_next, dst_next, src_draft, dst_draft, - src_lens, dst_lens, ids, draft_ids, pos_off, kv_off, t_step, - ) # fmt: skip - return True - block = triton.next_power_of_2(max(widths)) - _fold_tail_kernel[(rows,)]( - ctrl, - slots, - src_new, - src_new.stride(0), - src_new.stride(1), - src_new.shape[1], - dst_new, - dst_new.stride(0), - dst_new.stride(1), - dst_new.shape[1], - src_next, - src_next.stride(0), - src_next.stride(1), - src_next.shape[1], - dst_next, - dst_next.stride(0), - dst_next.stride(1), - dst_next.shape[1], - src_draft, - src_draft.stride(0), - src_draft.stride(1), - src_draft.shape[1], - dst_draft, - dst_draft.stride(0), - dst_draft.stride(1), - dst_draft.shape[1], - src_lens, - dst_lens, - ids, - draft_ids, - pos_off, - kv_off, - t_step, - d_step, - block, - num_warps=1, + minimaxM3Sm107/foldTail.h) that a caller folding the tail into its own kernel runs; the same writes.""" + return _shared_fold_tail( + ctrl=ctrl, + slots=slots, + rows=rows, + src_new=src_new, + dst_new=dst_new, + src_next=src_next, + dst_next=dst_next, + src_draft=src_draft, + dst_draft=dst_draft, + src_lens=src_lens, + dst_lens=dst_lens, + ids=ids, + draft_ids=draft_ids, + pos_off=pos_off, + kv_off=kv_off, + tokens_per_step=tokens_per_step, + cpp=cpp, + native_launch=lambda *args: _launch_cpp(*args, one_warp=one_warp), + triton_kernel=_fold_tail_kernel, ) - return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py deleted file mode 100644 index cc0ce28d9639..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_force_vec.py +++ /dev/null @@ -1,63 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Per-request forced accepted counts of a fractional forced acceptance (the in-tree speculative worker's -``force_num_accepted_tokens``), in one launch. Reached through the ``ops/spec`` entries ``m3_force_vec`` and -``m3_force_vec_req``. - -Triton. ``force_vec_kernel``: one CTA bumps the shared step counter and draws each request's count from the fixed-seed -pool with the in-tree (counter, slot) hash, exactly as ``_apply_force_accepted_tokens`` does. ``force_vec_req_kernel``: -one draw per request from its own key (its KV length at this verify), no shared counter. -""" - -import triton -import triton.language as tl - - -# One compiled variant for every batch and accept count: specializing them (== 1, divisible by 16) would compile -# mid-serving the first time a step reaches a new one, past the worker's warm-up. -@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) -def force_vec_kernel( - pool, - counter, - out, - batch_size, - base_total, - max_total, - frac, - counter_stride, - slot_stride, - pool_mask, - BLOCK: tl.constexpr, -): - c = tl.load(counter) + 1 - slots = tl.arange(0, BLOCK) - tl.store(counter, c) - idx = (c * counter_stride + slots.to(tl.int64) * slot_stride) & pool_mask - live = slots < batch_size - r = tl.load(pool + idx, mask=live, other=1.0) - v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) - tl.store(out + slots, v, mask=live) - - -@triton.jit(do_not_specialize=["batch_size", "base_total", "max_total"]) -def force_vec_req_kernel( - pool, - keys, - out, - batch_size, - base_total, - max_total, - frac, - key_stride, - pool_mask, - BLOCK: tl.constexpr, -): - # One draw per request from its own key (its KV length at this verify): no shared counter, so a request's forced - # counts depend only on its own prompt and progress, not on what else the server ran before or beside it. - slots = tl.arange(0, BLOCK) - live = slots < batch_size - k = tl.load(keys + slots, mask=live, other=0).to(tl.int64) - idx = (k * key_stride) & pool_mask - r = tl.load(pool + idx, mask=live, other=1.0) - v = tl.minimum(base_total + (r < frac).to(tl.int32), max_total) - tl.store(out + slots, v, mask=live) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py index de9f6d6666f3..0330a12bbde7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_fullctx_decode.py @@ -27,6 +27,18 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + FullCtxGeometry, + FullCtxPlanBase, + check_layout, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + device_index as _index, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.fullctx import ( + kv_page_views as kv_page_views, +) + # The TP8 per-rank geometry: 8 query heads over 1 KV head. HEADS = 8 KV_HEADS = 1 # capacity: fixed: the tp8 KV heads per rank @@ -43,116 +55,10 @@ _lock = threading.Lock() -class FullCtxPlan: +class FullCtxPlan(FullCtxPlanBase): """One compiled PrimTS split-KV decode at a pinned recipe, its workspace and its FP16 staging output.""" - def __init__( - self, - device: torch.device, - qlen: int, - nreq: int, - max_kv_len: int, - recipe: Mapping, - own_reducer: bool, - pdl: bool, - ): - from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec - from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( - get_max_active_clusters_for_cluster_size, - make_decode_config, - ) - - self._tsdec = tsdec - self.qlen = qlen - self.nreq = nreq - self.pdl = pdl - self.recipe = dict(recipe) - splits = int(self.recipe["splits"]) - fp8 = tsdec._cutlass_dtype("float8_e4m3fn") - # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer - # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents - # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an - # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the - # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own - # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge - # (minimax_m3_sm107_tp2_fullctx_reduce) waits as its first instruction. - cfg = make_decode_config( - headdim=HEAD_DIM, - args={"use_pdl": pdl, **self.recipe.get("config", {})}, - seq_len_q=qlen, - seq_len_kv=min(max_kv_len, POLICY_KV_LEN), - batch_size=nreq, - num_heads_q=HEADS, - num_heads_kv=KV_HEADS, - q_dtype=fp8, - k_dtype=fp8, - v_dtype=fp8, - o_dtype=tsdec._cutlass_dtype("float16"), - qkv_layout="pagedKv", - num_tokens_per_page=PAGE, - storage_tokens_per_page=PAGE, - split_kv_mode="gmem_reduction_with_separate_kernel", - splits_kv=splits, - max_splits_kv=splits, - min_loop_iters_per_split=1, - mask_type="causal" if qlen > 1 else "dense", - auto_tuner=True, - split_kv=True, - ) - _check_layout(cfg, qlen, splits) - with torch.cuda.device(device): - max_active = get_max_active_clusters_for_cluster_size(1) - spec = tsdec._decode_launch_spec_from_config( - cfg, - batch_size=nreq, - num_qo_heads=HEADS, - num_kv_heads=KV_HEADS, - head_dim=HEAD_DIM, - seq_len_q=qlen, - max_active_clusters=int(max_active), - ) - compile_spec = tsdec._make_decode_compile_spec( - spec, - device_index=_index(device), - num_qo_heads=HEADS, - num_kv_heads=KV_HEADS, - head_dim=HEAD_DIM, - page_size=PAGE, - max_kv_len=max_kv_len, - seq_len_q=qlen, - q_dtype_key="float8_e4m3fn", - k_dtype_key="float8_e4m3fn", - v_dtype_key="float8_e4m3fn", - output_dtype_key="float16", - use_packed_q=False, - kv_prefix_mode="dynamic", - kv_lengths_mode="dynamic", - ) - self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) - layout = tsdec._make_decode_workspace_layout( - spec.scratch_shapes, - torch.float16, - use_separate_reduction_kernel=True, - use_split_kv=True, - ) - self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) - self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) - self.workspace.split_kv_counter.zero_() - self.workspace.cu_seqlens_q.zero_() - self.workspace.attention_sinks.zero_() - rows = HEADS_PER_KV * qlen - if tuple(self.workspace.partial_o.shape) != (nreq, KV_HEADS, splits, rows, HEAD_DIM): - raise RuntimeError( - f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " - f"[requests, KV heads, splits, {rows} rows, {HEAD_DIM}] layout the BF16 merge reads" - ) - self.own_reducer = own_reducer - self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) - self.tile_kv = int(cfg.tile_size_kv) - self.insts_kv = int(cfg.num_insts_kv) - self.causal = qlen > 1 - shape = (nreq, qlen, HEADS, HEAD_DIM) if qlen > 1 else (nreq, HEADS, HEAD_DIM) - self.staging = torch.empty(shape, dtype=torch.float16, device=device) + GEOMETRY = FullCtxGeometry(HEADS, KV_HEADS, HEAD_DIM, PAGE, POLICY_KV_LEN, HEADS_PER_KV) def run( self, @@ -240,32 +146,7 @@ def run( def _check_layout(cfg, qlen: int, splits: int) -> None: - """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. - - Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is - ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding - window, attention sinks or packed Q. - """ - problems = [] - if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): - problems.append("not split-KV with the separate GMEM reducer") - if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: - problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") - if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: - problems.append("cluster-SMEM reduction or a persistent scheduler") - if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: - problems.append("packed Q, a sliding window or attention sinks") - if int(cfg.heads_q_per_kv) not in (0, HEADS_PER_KV): - problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") - if qlen > 1: - if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): - problems.append("multi-token Q without the grouped [token, head] tile layout") - elif int(cfg.q_tokens_per_cta) * HEADS_PER_KV != int(cfg.tile_size_q): - problems.append( - f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" - ) - if problems: - raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + return check_layout(cfg, qlen, splits, heads_per_kv=HEADS_PER_KV) # Select the TP8 copy of the BF16 reducer; both copies use the same per-row arithmetic. @@ -317,12 +198,3 @@ def prebuild( _plans[key] = FullCtxPlan(device, qlen, nreq, max_kv_len, recipe, own_reducer, pdl) built += 1 return built - - -def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: - """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" - return pool[0::subpages_per_slot], pool[1::subpages_per_slot] - - -def _index(device: torch.device) -> int: - return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py index 1467852d67b0..71c8d79b4b64 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_index_score.py @@ -26,27 +26,18 @@ and called through this target's ``m3_index_score`` entry. """ -import cutlass -import torch -from cuda.bindings.driver import CUstream -from cutlass import Float8E4M3FN, Float16, Float32, Int32, Int64, Uint32, cute -from cutlass._mlir.dialects import llvm -from cutlass.cute.nvgpu import cpasync, warp -from cutlass.cutlass_dsl import dsl_user_op -from cutlass.utils.smem_allocator import SmemAllocator - -from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import ( - EVICT_FIRST, - fp8x4_to_fp16x4, - mma_sync, - simple_tma_copy, +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.index_score_runtime import ( + IndexScoreRunner as SharedIndexScoreRunner, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + IndexDecodeScoreKernel as IndexDecodeScoreKernel, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + _fp8_to_f16_mma_fragments as _fp8_to_f16_mma_fragments, ) -from tensorrt_llm._torch.kernels.blackwell.utils import ( - TRTLLM_ENABLE_PDL, - griddepcontrol_launch_dependents, - griddepcontrol_wait, +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.m3_index_score import ( + _l2_prefetch_bulk as _l2_prefetch_bulk, ) -from tensorrt_llm.logger import logger # With the early issue, a CTA whose first block holds only tokens older than this step also TMA-loads that block's # index-K tile before the grid wait (bit-identical scores). @@ -58,635 +49,9 @@ __all__ = ["IndexDecodeScoreKernel", "IndexScoreRunner"] -@dsl_user_op -def _l2_prefetch_bulk(addr: Int64, size: Int32, *, loc=None, ip=None) -> None: - """Prefetch ``size`` contiguous bytes (multiple of 16) at global ``addr`` into L2.""" - llvm.inline_asm( - res=None, - operands_=[Int64(addr).ir_value(loc=loc, ip=ip), Int32(size).ir_value(loc=loc, ip=ip)], - asm_string="cp.async.bulk.prefetch.L2.global [$0], $1;", - constraints="l,r", - has_side_effects=True, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@cute.jit -def _fp8_to_f16_mma_fragments(src: cute.Tensor): - """Split an FP8 ldmatrix fragment into the two FP16 k-fragments MMA wants.""" - src_elems = cute.size(src) - src_u32 = cute.recast_tensor(src, Uint32) - src_f16 = cute.make_rmem_tensor(src_elems, Float16) - src_f16_u32 = cute.recast_tensor(src_f16, Uint32) - # Packed conversion; faster and fewer SASS instructions than - # src.load().to(Float16). - for i in cutlass.range_constexpr(src_elems // 4): - converted = fp8x4_to_fp16x4(src_u32[i]) - src_f16_u32[i * 2] = converted[0] - src_f16_u32[i * 2 + 1] = converted[1] - lower = cute.make_rmem_tensor(src_elems // 2, Float16) - upper = cute.make_rmem_tensor(src_elems // 2, Float16) - - # FP8 ldmatrix yields four consecutive values along K. Split each group - # into the lower two and upper two for the two FP16 MMA k-fragments. - for i in cutlass.range_constexpr(src_elems // 2): - lower[i] = src_f16[(i // 2) * 4 + i % 2] - upper[i] = src_f16[(i // 2) * 4 + 2 + i % 2] - return lower, upper - - -class IndexDecodeScoreKernel: - """Per-block max index score for one decode step. - - Grid is (batch, split_k): CTA (b, s) walks KV blocks - s, s + split_k, s + 2 * split_k, ... of request b and writes each block's - score directly, so no cross-CTA reduction is needed. CTAs whose split_id - exceeds the request's block count exit immediately. - """ - - BLOCK_K = 128 # capacity: fixed: the MSA block size of the checkpoint - BAR_MMA = 1 - num_stages = 2 - - def __init__( - self, - dtype: type[cutlass.Numeric], - num_heads: int, - max_decode_query_len: int, - split_k: int, - head_dim: int = 128, - prefetch_old: bool = False, - early_issue: bool = False, - prefetch_depth: int = 0, - prestage_old: bool = False, - early_trigger: bool = False, - ): - # prefetch_old: before the grid dependency wait, the TMA warp L2-prefetches this CTA's index-K pages that - # hold only tokens older than this step (the predecessor writes only this step's tokens). - self.prefetch_old = prefetch_old - # early_issue: the TMA warp reads a seq_lens hint and this CTA's first page id before the grid dependency - # wait, then issues the Q tile and the first K tile right after it, overlapping the post-wait seq_lens load. - # The tiles are read after the wait, so their contents are final; the block table is only written between - # graph replays, so the early page id is the right one whenever the block is valid. A CTA whose early tiles - # turn out to lie past the request's true length drains them before exiting. - self.early_issue = early_issue - # prefetch_depth: blocks per CTA that prefetch_old covers (0 = all of them). Prefetching every older page at - # large batches floods L2; the first few per CTA are the ones whose latency the pipeline cannot hide. - self.prefetch_depth = prefetch_depth - # prestage_old (with early_issue): when this CTA's first block holds only tokens older than this step, its K - # tile is TMA-loaded into its stage BEFORE the grid dependency wait (the predecessor writes only this step's - # tokens; the block table and the seq_lens hint are final before the grid starts). Q stays after the wait. - self.prestage_old = prestage_old and early_issue - # early_trigger: release the dependent block selector right after this CTA's setup (see the module docstring). - self.early_trigger = early_trigger - self.dtype = dtype - self.num_heads = num_heads - self.max_decode_query_len = max_decode_query_len - self.split_k = split_k - self.head_dim = head_dim - - def _swizzle_elems(self) -> int: - """Elements spanned by one 128-byte swizzle atom.""" - return 128 * 8 // self.dtype.width - - def _sq_layout(self): - """Composed SMEM layout for the Q tile, shared by the descriptor and the kernel.""" - elems = self._swizzle_elems() - head_dim = self.head_dim - block_q = self.num_heads * self.max_decode_query_len - layout = cute.make_layout( - (self.max_decode_query_len, self.num_heads, (elems, head_dim // elems)), - stride=(elems, self.max_decode_query_len * elems, (1, block_q * elems)), - ) - return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) - - def _sk_layout(self): - """Composed SMEM layout for the pipelined K tiles.""" - elems = self._swizzle_elems() - head_dim = self.head_dim - block_k = self.BLOCK_K - layout = cute.make_layout( - (1, block_k, (elems, head_dim // elems), self.num_stages), - stride=(0, elems, (1, block_k * elems), block_k * head_dim), - ) - return cute.make_composed_layout(cute.make_swizzle(3, 4, 3), 0, layout) - - @cute.jit - def __call__( - self, - gQ: cute.Tensor, # [bs * runtime_decode_query_len, num_heads, head_dim] - gK_cache: cute.Tensor, # [num_pages, page_size, head_dim] - block_table: cute.Tensor, # [bs, max_pages] - score: cute.Tensor, # [num_heads, bs * runtime_decode_query_len, max_pages] - seq_lens: cute.Tensor, # [bs] - stream: CUstream, - ): - num_heads = self.num_heads - head_dim = self.head_dim - MAX_DQL = self.max_decode_query_len - assert num_heads * MAX_DQL <= 32 - - batch = seq_lens.shape[0] - decode_query_len = gQ.shape[0] // batch - grid = (batch, self.split_k, 1) - block = (32 * 5, 1, 1) - - tma_g2s = cpasync.CopyBulkTensorTileG2SOp() - elems = self._swizzle_elems() - - q_tma_atom, q_tma_tensor = cpasync.make_tiled_tma_atom( - tma_g2s, - cute.logical_divide(gQ, (None, None, elems)), - self._sq_layout(), - (MAX_DQL, num_heads, head_dim), - ) - k_tma_atom, k_tma_tensor = cpasync.make_tiled_tma_atom( - tma_g2s, - cute.logical_divide(gK_cache, (None, None, elems)), - self._sk_layout(), - (1, self.BLOCK_K, head_dim), - ) - - self.kernel( - q_tma_atom, - q_tma_tensor, - k_tma_atom, - k_tma_tensor, - gK_cache, - block_table, - score, - seq_lens, - decode_query_len, - ).launch(grid=grid, block=block, stream=stream, use_pdl=TRTLLM_ENABLE_PDL) - - @cute.kernel - def kernel( - self, - q_tma_atom: cute.CopyAtom, - q_tma_tensor: cute.Tensor, - k_tma_atom: cute.CopyAtom, - k_tma_tensor: cute.Tensor, - gK_cache: cute.Tensor, - block_table: cute.Tensor, - score: cute.Tensor, - seq_lens: cute.Tensor, - decode_query_len, - ): - batch_id, split_id, _ = cute.arch.block_idx() - _, split_k, _ = cute.arch.grid_dim() - warp_id = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane_id = cute.arch.lane_idx() - - NUM_HEADS = self.num_heads - MAX_DQL = self.max_decode_query_len - BLOCK_Q = NUM_HEADS * MAX_DQL - BLOCK_K = self.BLOCK_K - head_dim = self.head_dim - dtype = self.dtype - MMA_N = 8 - num_stages = self.num_stages - Q_TILES = cute.ceil_div(BLOCK_Q, MMA_N) - EPI_Q = Q_TILES * MMA_N - - sq_layout = self._sq_layout() - sk_layout = self._sk_layout() - - smem = SmemAllocator() - sK = smem.allocate_tensor( - dtype, - sk_layout.outer, - byte_alignment=128, - swizzle=sk_layout.inner, - )[0, None, None, None] - # sQ aliases the first K stage: Q is consumed into registers before the - # first K tile is needed, so the two never overlap in time. - sQ_tma = cute.make_tensor(sK[None, None, 0].iterator, layout=sq_layout.outer) - # TMA sees Q as (query, head, dim) while ldmatrix consumes a flattened Q - # column mode. The target profile keeps the rank-2 view even for - # degenerate shapes such as MAX_DQL == 1. - q_tma_elems = self._swizzle_elems() - sQ = cute.coalesce( - cute.group_modes(sQ_tma, 0, 2), - target_profile=(BLOCK_Q, (q_tma_elems, head_dim // q_tma_elems)), - ) - epi_buffer = smem.allocate_tensor(Float32, cute.make_layout((EPI_Q, 4))) - - tma_full_mbar = smem.allocate_array(Int64, num_stages) - tma_empty_mbar = smem.allocate_array(Int64, num_stages) - - # Barrier init and descriptor prefetch do not depend on the predecessor - # grid, so they are the work PDL can overlap. Every CTA runs them, even - # the ones that own no blocks, since ownership is not known until - # seq_lens is read below. - if warp_id == 0: - with cute.arch.elect_one(): - for i in cutlass.range_constexpr(num_stages): - cute.arch.mbarrier_init(tma_full_mbar + i, 1) - cute.arch.mbarrier_init(tma_empty_mbar + i, 128) - cute.arch.mbarrier_init_fence() - elif warp_id == 1: - cpasync.prefetch_descriptor(q_tma_atom) - cpasync.prefetch_descriptor(k_tma_atom) - cute.arch.sync_threads() - if cutlass.const_expr(self.early_trigger): - griddepcontrol_launch_dependents() - - if cutlass.const_expr(self.prefetch_old): - if warp_id == 4: - with cute.arch.elect_one(): - seqlen0 = seq_lens[batch_id] - num_blocks0 = cute.ceil_div(seqlen0, BLOCK_K) - stable = (seqlen0 - decode_query_len) // BLOCK_K - if stable > num_blocks0: - stable = num_blocks0 - if cutlass.const_expr(self.prefetch_depth > 0): - if stable > split_id + self.prefetch_depth * split_k: - stable = split_id + self.prefetch_depth * split_k - for block_id in range(split_id, stable, split_k): - page = block_table[batch_id, block_id] - _l2_prefetch_bulk( - gK_cache[page, None, None].iterator.toint().to(Int64), - Int32(BLOCK_K * head_dim * (dtype.width // 8)), - ) - - # Without early_trigger the dependent block selector is released as the CTAs exit. - early_q = False - early_k = False - early_page = Int32(0) - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - seqlen_hint = seq_lens[batch_id] - hint_blocks = cute.ceil_div(seqlen_hint, BLOCK_K) - if hint_blocks > block_table.shape[1]: - hint_blocks = block_table.shape[1] - if split_id < hint_blocks: - early_q = True - early_page = block_table[batch_id, split_id] - if cutlass.const_expr(self.prestage_old): - # Stage the first K tile now when every token of the block predates this step. - if split_id < (seqlen_hint - decode_query_len) // BLOCK_K: - early_k = True - cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - tma_full_mbar + 1 % num_stages, - BLOCK_K * head_dim * (dtype.width // 8), - ) - simple_tma_copy( - k_tma_atom, - k_tma_tensor[early_page, None, None], - sK[None, None, 1 % num_stages], - tma_full_mbar + 1 % num_stages, - cache_policy=EVICT_FIRST, - ) - - griddepcontrol_wait() - - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - if early_q: - gQ_early = cute.local_tile( - cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), - tiler=(MAX_DQL, NUM_HEADS, head_dim), - coord=(0, 0, 0), - ) - cute.arch.mbarrier_wait(tma_empty_mbar, 1) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - tma_full_mbar, BLOCK_Q * head_dim * (dtype.width // 8) - ) - simple_tma_copy(q_tma_atom, gQ_early, sQ_tma, tma_full_mbar) - if not early_k: - cute.arch.mbarrier_wait(tma_empty_mbar + 1 % num_stages, 1) - with cute.arch.elect_one(): - cute.arch.mbarrier_arrive_and_expect_tx( - tma_full_mbar + 1 % num_stages, - BLOCK_K * head_dim * (dtype.width // 8), - ) - simple_tma_copy( - k_tma_atom, - k_tma_tensor[early_page, None, None], - sK[None, None, 1 % num_stages], - tma_full_mbar + 1 % num_stages, - cache_policy=EVICT_FIRST, - ) - - # The length that bounds this grid's work is read after the wait (the - # pre-wait reads above are prefetch / early-issue hints only, clamped to - # the table width). seq_lens is a graph-stable buffer that on_update_kv_lens - # patches every step, and num_blocks bounds an unmasked block_table load - # and an unmasked score store, so a stale length is an out-of-bounds write - # rather than just a wrong answer. - seqlen = seq_lens[batch_id] - num_blocks = cute.ceil_div(seqlen, BLOCK_K) - - if split_id < num_blocks: - if warp_id == 4: - # TMA warp - tma_stage = 0 - tma_parity = 1 - first_block = split_id - - if not early_q: - gQ_tile = cute.local_tile( - cute.domain_offset((batch_id * decode_query_len, 0, 0), q_tma_tensor), - tiler=(MAX_DQL, NUM_HEADS, head_dim), - coord=(0, 0, 0), - ) - cute.arch.mbarrier_wait(tma_empty_mbar, tma_parity) - with cute.arch.elect_one(): - Q_size = BLOCK_Q * head_dim * (dtype.width // 8) - cute.arch.mbarrier_arrive_and_expect_tx(tma_full_mbar, Q_size) - # TMA bounds-checks rows when the runtime decode_query_len is - # below MAX_DQL; padded Q columns are masked before the stores. - simple_tma_copy(q_tma_atom, gQ_tile, sQ_tma, tma_full_mbar) - - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - if early_q: - # The first K tile was issued right after the wait. - first_block = split_id + split_k - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - for block_id in range(first_block, num_blocks, split_k): - page_id = block_table[batch_id, block_id] - gK_tile = k_tma_tensor[page_id, None, None] - k_mbar = tma_full_mbar + tma_stage - - cute.arch.mbarrier_wait(tma_empty_mbar + tma_stage, tma_parity) - with cute.arch.elect_one(): - K_size = BLOCK_K * head_dim * (dtype.width // 8) - cute.arch.mbarrier_arrive_and_expect_tx(k_mbar, K_size) - simple_tma_copy( - k_tma_atom, - gK_tile, - sK[None, None, tma_stage], - k_mbar, - cache_policy=EVICT_FIRST, - ) - - tma_stage = (tma_stage + 1) % num_stages - if tma_stage == 0: - tma_parity ^= 1 - - else: - # MMA warps; each handles K[32, head_dim] @ Q[BLOCK_Q, head_dim].T - sK_warp = cute.local_tile(sK, (32, head_dim, num_stages), (warp_id, 0, 0)) - q_start = seqlen - decode_query_len - - elems = 128 // dtype.width # 16B - MMA_K = 32 * 8 // dtype.width # 32B - - # Pre-compute ldmatrix addresses. - # sK loads a [16 x 16B] tile: - # ((16, (16B, 2), 1), (32 / 16, head_dim / 32B, num_stages)) - # sQ loads an [8 x 32B] tile: - # ((8, (16B, 4)), (BLOCK_Q / MMA_N, head_dim / 64B)) - sK_ldsm = cute.zipped_divide(sK_warp, (16, cute.make_layout((elems, 2)), 1)) - sQ_ldsm = cute.zipped_divide(sQ, (MMA_N, cute.make_layout((elems, 4)))) - - # sK: (16B, (32 / 16, head_dim / 32B, num_stages)) - # sQ: (16B, (BLOCK_Q / MMA_N, head_dim / 64B)) - sK_ldsm = sK_ldsm[(lane_id % 16, (None, lane_id // 16), 0), None] - sQ_ldsm = sQ_ldsm[(lane_id % MMA_N, (None, lane_id // 8)), None] - - ldsm_op = warp.LdMatrix8x8x16bOp(num_matrices=4) - ldsm_atom = cute.make_copy_atom(ldsm_op, dtype) - - rQ = cute.make_rmem_tensor( - ((elems // 2, 2), head_dim // (MMA_K * 2), Q_TILES), dtype - ) - rK = cute.make_rmem_tensor((elems, 2, head_dim // MMA_K), dtype) - rC = cute.make_rmem_tensor((4, 2, Q_TILES), Float32) - - if warp_id == 0: - cute.arch.mbarrier_wait(tma_full_mbar, 0) - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - for q in cutlass.range_constexpr(Q_TILES): - cute.copy(ldsm_atom, sQ_ldsm[None, (q, None)], rQ[None, None, q]) - cute.arch.mbarrier_arrive(tma_empty_mbar) - - tma_stage = 1 % self.num_stages - tma_parity = 0 - if tma_stage == 0: - tma_parity ^= 1 - - if cutlass.const_expr(dtype is Float8E4M3FN): - rQ_f16 = cute.make_rmem_tensor((4, head_dim // MMA_K, Q_TILES, 2), Float16) - q_lower, q_upper = _fp8_to_f16_mma_fragments(rQ) - rQ_f16[None, None, None, 0].store(q_lower.load()) - rQ_f16[None, None, None, 1].store(q_upper.load()) - - for block_id in range(split_id, num_blocks, split_k): - rC.fill(0.0) - - if warp_id == 0: - cute.arch.mbarrier_wait(tma_full_mbar + tma_stage, tma_parity) - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - - for k in cutlass.range_constexpr(head_dim // MMA_K): - cute.copy( - ldsm_atom, - sK_ldsm[None, (None, k, tma_stage)], - rK[None, None, k], - ) - for m in cutlass.range_constexpr(2): - if cutlass.const_expr(dtype is Float8E4M3FN): - rK_lower, rK_upper = _fp8_to_f16_mma_fragments(rK[None, m, k]) - for n in cutlass.range_constexpr(Q_TILES): - rC[None, m, n] = mma_sync( - rK_lower, - rQ_f16[None, k, n, 0], - rC[None, m, n], - ) - rC[None, m, n] = mma_sync( - rK_upper, - rQ_f16[None, k, n, 1], - rC[None, m, n], - ) - else: - for n in cutlass.range_constexpr(Q_TILES): - rC[None, m, n] = mma_sync( - rK[None, m, k], - rQ[(None, k % 2), k // 2, n], - rC[None, m, n], - ) - - cute.arch.mbarrier_arrive(tma_empty_mbar + tma_stage) - - k_start = block_id * BLOCK_K + warp_id * 32 - - for q in cutlass.range_constexpr(Q_TILES): - for i in cutlass.range_constexpr(4): - for j in cutlass.range_constexpr(2): - col = q * 8 + (lane_id % 4) * 2 + j - q_local_pos = col % MAX_DQL - q_pos = q_start + q_local_pos - k_pos = k_start + i * 8 + lane_id // 4 - rC[q * 8 + i * 2 + j] = ( - rC[q * 8 + i * 2 + j] if q_pos >= k_pos else float("-inf") - ) - - for q in cutlass.range_constexpr(Q_TILES): - # Thread-local reduction along the BLOCK_K dim. - rScore = cute.make_rmem_tensor(2, Float32) - rScore.fill(float("-inf")) - for i in cutlass.range_constexpr(4): - rScore[0] = cute.arch.fmax(rScore[0], rC[i * 2 + 0 + q * 8]) - rScore[1] = cute.arch.fmax(rScore[1], rC[i * 2 + 1 + q * 8]) - - # Warp reduction among lanes 0, 4, 8, 12, ... - for i in cutlass.range_constexpr(3): - offset = 4 << i - other0 = cute.arch.shuffle_sync_bfly( - rScore[0], offset=offset, mask=-1, mask_and_clamp=31 - ) - other1 = cute.arch.shuffle_sync_bfly( - rScore[1], offset=offset, mask=-1, mask_and_clamp=31 - ) - rScore[0] = cute.arch.fmax(rScore[0], other0) - rScore[1] = cute.arch.fmax(rScore[1], other1) - - if lane_id * 2 < MMA_N: - epi_buffer[q * MMA_N + lane_id * 2 + 0, warp_id] = rScore[0] - epi_buffer[q * MMA_N + lane_id * 2 + 1, warp_id] = rScore[1] - cute.arch.barrier(barrier_id=self.BAR_MMA, number_of_threads=128) - - head_id = lane_id // MAX_DQL - q_local_pos = lane_id - head_id * MAX_DQL - valid_q = head_id < NUM_HEADS and q_local_pos < decode_query_len - if lane_id < BLOCK_Q and valid_q: - final_score = epi_buffer[lane_id, 0] - for i in cutlass.range_constexpr(1, 4): - final_score = cute.arch.fmax(final_score, epi_buffer[lane_id, i]) - - t = batch_id * decode_query_len + q_local_pos - score[head_id, t, block_id] = final_score - - tma_stage = (tma_stage + 1) % self.num_stages - if tma_stage == 0: - tma_parity ^= 1 - else: - if cutlass.const_expr(self.early_issue): - if warp_id == 4: - if early_q: - # Early tiles past the request's true length: let them land before the CTA exits. - cute.arch.mbarrier_wait(tma_full_mbar, 0) - cute.arch.mbarrier_wait(tma_full_mbar + 1 % num_stages, 0) - - -class IndexScoreRunner: - """Compiled ``IndexDecodeScoreKernel`` variants, keyed by (index dtype, index heads, decode query length, head dim, - split_k, early issue). Batch, query-token count, page count and block-table width stay symbolic, so one compile - serves every decode step of its key. ``prefetch_old`` is always on; ``early`` also caps the prefetch at one block - per CTA; ``early_trigger`` is ``INDEX_SCORE_SET`` for every key of the process.""" +class IndexScoreRunner(SharedIndexScoreRunner): + """TP8 schedule and private compilation cache.""" kernel_cache: dict = {} - SUPPORTED_HEAD_DIM = 128 # capacity: fixed: the checkpoint's index head dim - SUPPORTED_PAGE_SIZE = IndexDecodeScoreKernel.BLOCK_K - # BLOCK_Q must fit one warp's worth of epilogue lanes. - MAX_BLOCK_Q = 32 # capacity: fixed: the scorer's query tile - EARLY_PREFETCH_DEPTH = 1 - - @classmethod - def is_supported( - cls, - *, - q_dtype: torch.dtype, - num_heads: int, - head_dim: int, - page_size: int, - max_decode_query_len: int, - ) -> bool: - from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE - from tensorrt_llm._utils import is_sm_100f - - return ( - is_sm_100f() - and q_dtype in TORCH_TO_CUTE_DTYPE - and head_dim == cls.SUPPORTED_HEAD_DIM - and page_size == cls.SUPPORTED_PAGE_SIZE - and num_heads * max_decode_query_len <= cls.MAX_BLOCK_Q - and max_decode_query_len >= 1 - ) - - @classmethod - def _compile( - cls, - q_dtype: torch.dtype, - num_heads: int, - max_decode_query_len: int, - head_dim: int, - split_k: int, - early: bool, - ) -> None: - from tensorrt_llm._torch.kernels.blackwell.cute_ptx_utils import TORCH_TO_CUTE_DTYPE - - key = (q_dtype, num_heads, max_decode_query_len, head_dim, split_k, early) - if key in cls.kernel_cache: - return - cute_dtype = TORCH_TO_CUTE_DTYPE[q_dtype] - page_size = cls.SUPPORTED_PAGE_SIZE - sym_total_tokens = cute.sym_int() - sym_batch = cute.sym_int() - - # 16-element divisibility on every non-innermost stride is what the TMA descriptors for Q and K assume. - def _sym_stride(): - return cute.sym_int64(divisibility=16) - - q_fake = cute.runtime.make_fake_tensor( - cute_dtype, - (sym_total_tokens, num_heads, head_dim), - stride=(_sym_stride(), _sym_stride(), 1), - ) - # The index-K pool may be coalesced with the main K/V cache (per-page stride above page_size * head_dim). - k_fake = cute.runtime.make_fake_tensor( - cute_dtype, - (cute.sym_int(), page_size, head_dim), - stride=(_sym_stride(), _sym_stride(), 1), - ) - bt_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (sym_batch, cute.sym_int()), stride_order=(1, 0) - ) - # The target passes a transposed view of the [heads, blocks, tokens] selector buffer. - score_fake = cute.runtime.make_fake_tensor( - cutlass.Float32, - (num_heads, sym_total_tokens, cute.sym_int()), - stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), - ) - sl_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (sym_batch,), stride_order=(0,) - ) - fake_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - kernel = IndexDecodeScoreKernel( - cute_dtype, - num_heads, - max_decode_query_len, - split_k, - head_dim, - prefetch_old=True, - early_issue=early, - prefetch_depth=cls.EARLY_PREFETCH_DEPTH if early else 0, - prestage_old=early and INDEX_SCORE_PRESTAGE_OLD, - early_trigger=INDEX_SCORE_SET, - ) - if INDEX_SCORE_SET and not cls.kernel_cache: - logger.info( - "tp8 index scorer: the block selector is released after the scorer's setup (INDEX_SCORE_SET)" - ) - cls.kernel_cache[key] = cute.compile( - kernel, - q_fake, - k_fake, - bt_fake, - score_fake, - sl_fake, - fake_stream, - options="--enable-tvm-ffi", - ) + PRESTAGE_OLD = INDEX_SCORE_PRESTAGE_OLD + EARLY_TRIGGER = INDEX_SCORE_SET diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py index 3e88c7fb1ce2..6af262dafe82 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_locality.py @@ -1,183 +1,19 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Locality domains ("uGPU") of a multi-die GPU (Rubin VR200: 2 domains of 100 SMs + 12 remainder SMs) for the -streaming decode kernels of this target: the SM -> domain map, per-domain weight placement, and the CTA -> SM -assignment a persistent kernel uses to read only its own die's memory. - -A kernel launched with one CTA per SM reads ``%smid`` once and takes its logical block index from a host-built -permutation (:meth:`LocalityMap.logical_map`), so every per-CTA structure (work plan row, epochs, roles) follows its -die. Weights that are split by output rows into per-domain halves (:meth:`LocalityMap.split_rows`) live in one virtual -range whose halves are physically in the two domains. Placement only: the arithmetic is unchanged, so results are -bitwise those of the unplaced weights. - -The SM -> domain map is attested by the driver: the SMs are split by locality domain -(``cuDevSmResourceSplit`` with ``CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID``), and a ``%smid`` probe runs on a green -context of each split. Allocations use driver VMM with ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN`` and are checked -with ``CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL``; any failure raises (fail closed).""" - -from __future__ import annotations +"""Target-owned locality-map cache with shared placement implementation.""" import threading -from typing import Dict, List, Sequence, Tuple +from typing import Dict import torch +from ...ops.gemm.locality import LocalityMap as _LocalityMap + _LOCK = threading.Lock() _MAPS: Dict[int, "LocalityMap"] = {} -_PROBE_SRC = r""" -extern "C" __global__ void smid_probe(int* out) { - if (threadIdx.x == 0) { - unsigned s; - asm volatile("mov.u32 %0, %%smid;" : "=r"(s)); - out[blockIdx.x] = (int)s; - } -} -""" - - -def _drv(): - from cuda.bindings import driver - - return driver - - -def _ck(res, what: str): - """cuda-python returns (CUresult, *values): raise on an error, return the values.""" - err = res[0] if isinstance(res, tuple) else res - if int(err) != 0: - raise RuntimeError(f"locality: {what} failed: {err}") - if isinstance(res, tuple): - return res[1] if len(res) == 2 else res[1:] - return None - - -class _Mapping: - """One virtual range backed by per-domain physical allocations, exposed to torch through - ``__cuda_array_interface__``. Lives as long as a tensor that views it (the tensor holds this object).""" - - def __init__(self, device: int, pieces: Sequence[Tuple[int, int]]): - """``pieces``: (bytes, domain) in address order; each a multiple of the allocation granularity.""" - d = _drv() - self.device = device - self.size = sum(n for n, _ in pieces) - self.va = _ck(d.cuMemAddressReserve(self.size, 0, 0, 0), "cuMemAddressReserve") - self.handles = [] - off = 0 - try: - for nbytes, dom in pieces: - prop = _alloc_prop(device, dom) - h = _ck(d.cuMemCreate(nbytes, prop, 0), f"cuMemCreate(domain {dom})") - self.handles.append((h, nbytes)) - _ck(d.cuMemMap(int(self.va) + off, nbytes, 0, h, 0), "cuMemMap") - off += nbytes - acc = d.CUmemAccessDesc() - acc.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - acc.location.id = device - acc.flags = d.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE - _ck(d.cuMemSetAccess(self.va, self.size, [acc], 1), "cuMemSetAccess") - except Exception: - self.close() - raise - self.__cuda_array_interface__ = { - "shape": (self.size,), - "typestr": "|u1", - "data": (int(self.va), False), - "version": 3, - "strides": None, - "stream": None, - } - - def close(self): - d = _drv() - off = 0 - for h, n in self.handles: - d.cuMemUnmap(int(self.va) + off, n) - d.cuMemRelease(h) - off += n - self.handles = [] - if self.va is not None: - d.cuMemAddressFree(self.va, self.size) - self.va = None - - def __del__(self): - try: - self.close() - except Exception: - pass - - -_LIBCUDA = None - - -def _libcuda(): - global _LIBCUDA - if _LIBCUDA is None: - import ctypes - _LIBCUDA = ctypes.CDLL("libcuda.so.1") - return _LIBCUDA - - -def _alloc_prop(device: int, domain: int): - d = _drv() - prop = d.CUmemAllocationProp() - prop.type = d.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED - prop.requestedHandleTypes = d.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE - prop.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN - prop.location.localized.deviceId = device - prop.location.localized.localityDomainId = domain - return prop - - -class LocalityMap: - """The locality domains of one device. Build with :meth:`get` before CUDA-graph capture.""" - - def __init__(self, device: int): - d = _drv() - _ck(d.cuInit(0), "cuInit") - dev = _ck(d.cuDeviceGet(device), "cuDeviceGet") - attr = d.CUdevice_attribute - self.device = device - self.n_domains = int( - _ck( - d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, dev), - "domain count", - ) - ) - self.sms_per_domain = int( - _ck( - d.cuDeviceGetAttribute( - attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_MULTIPROCESSOR_COUNT, dev - ), - "SMs per domain", - ) - ) - self.num_sms = int( - _ck(d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev), "SMs") - ) - if self.n_domains != 2: - raise RuntimeError( - f"locality: {self.n_domains} locality domains (this layer handles 2)" - ) - self._domain_sms = self._probe(d, dev) - dom = [-1] * self.num_sms - for k, sms in enumerate(self._domain_sms): - for s in sms: - if dom[s] != -1: - raise RuntimeError(f"locality: SM {s} probed in two domains") - dom[s] = k - self.sm_domain: List[int] = dom - self.remainder: List[int] = [s for s in range(self.num_sms) if dom[s] < 0] - rt = _alloc_prop(device, 0) - self.granularity = int( - _ck( - d.cuMemGetAllocationGranularity( - rt, d.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM - ), - "granularity", - ) - ) +class LocalityMap(_LocalityMap): @classmethod def get(cls, device) -> "LocalityMap": idx = torch.device(device).index if not isinstance(device, int) else device @@ -191,258 +27,3 @@ def get(cls, device) -> "LocalityMap": ) m = _MAPS[idx] = cls(idx) return m - - def sms(self, domain: int) -> List[int]: - return list(self._domain_sms[domain]) - - @property - def sm_table(self) -> torch.Tensor: - """int32 [num_sms] on the device: (domain << 16) | rank within the domain's sorted SMs; -1 = remainder SM. - For kernels that pick work by %smid themselves.""" - t = getattr(self, "_sm_table", None) - if t is None: - v = [-1] * self.num_sms - for k, sms in enumerate(self._domain_sms): - for i, s in enumerate(sms): - v[s] = (k << 16) | i - t = self._sm_table = torch.tensor(v, dtype=torch.int32, device=f"cuda:{self.device}") - return t - - def _probe(self, d, dev) -> List[List[int]]: - """Split the SMs by locality domain (strict, no backfill) and record %smid of a probe kernel launched on a - green context of each part: the driver's own SM -> domain assignment.""" - from cuda.bindings import nvrtc - - full = _ck( - d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), - "dev resource", - ) - params = [] - for k in range(self.n_domains): - p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() - p.smCount = self.sms_per_domain - p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID - p.localityDomainId = k - params.append(p) - res, _rem = _ck( - d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" - ) - # the probe module, compiled for this device - major = _ck( - d.cuDeviceGetAttribute( - d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev - ), - "cc", - ) - minor = _ck( - d.cuDeviceGetAttribute( - d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev - ), - "cc", - ) - prog = _ck( - nvrtc.nvrtcCreateProgram(_PROBE_SRC.encode(), b"smid_probe.cu", 0, [], []), - "nvrtcCreateProgram", - ) - opts = [f"--gpu-architecture=sm_{major}{minor}".encode()] - err = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)[0] - if int(err) != 0: - raise RuntimeError(f"locality: probe compile failed: {err}") - n = _ck(nvrtc.nvrtcGetCUBINSize(prog), "cubin size") - cubin = b" " * n - _ck(nvrtc.nvrtcGetCUBIN(prog, cubin), "cubin") - out: List[List[int]] = [] - nb = 4 * self.sms_per_domain # more blocks than SMs: every SM of the part shows up - torch.cuda.set_device(self.device) - buf = torch.full((nb,), -1, dtype=torch.int32, device=f"cuda:{self.device}") - for k in range(self.n_domains): - desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") - g = _ck( - d.cuGreenCtxCreate(desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM), - "green ctx", - ) - st = _ck( - d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), - "green stream", - ) - ctx = _ck(d.cuCtxFromGreenCtx(g), "ctx from green ctx") - _ck(d.cuCtxPushCurrent(ctx), "push ctx") - try: - mod = _ck(d.cuModuleLoadData(cubin), "probe module") - fn = _ck(d.cuModuleGetFunction(mod, b"smid_probe"), "probe fn") - buf.fill_(-1) - torch.cuda.synchronize(self.device) - import ctypes - - arg = ctypes.c_void_p(buf.data_ptr()) - _ck( - d.cuLaunchKernel( - fn, nb, 1, 1, 32, 1, 1, 0, st, ((arg,), (ctypes.c_void_p,)), 0 - ), - "probe launch", - ) - _ck(d.cuStreamSynchronize(st), "probe sync") - sms = sorted(set(int(x) for x in buf.tolist() if x >= 0)) - d.cuModuleUnload(mod) - finally: - d.cuCtxPopCurrent() - d.cuStreamDestroy(st) - d.cuGreenCtxDestroy(g) - if len(sms) != self.sms_per_domain: - raise RuntimeError( - f"locality: domain {k} probed {len(sms)} SMs, expected {self.sms_per_domain}" - ) - out.append(sms) - return out - - def logical_map(self, per_domain: int, extra: int) -> Tuple[torch.Tensor, List[int]]: - """The SM -> logical block index permutation for a kernel of ``2 * per_domain + extra`` CTAs, one per SM: - logical blocks [d * per_domain, (d + 1) * per_domain) run on domain d's first ``per_domain`` SMs; logical - blocks 2 * per_domain .. run on the rest (the leftover domain SMs, alternating domains, then the remainder - SMs). Returns (int32 [num_sms] on the device, -1 = no block; the domain of each extra block, -1 remainder).""" - if 2 * per_domain + extra != self.num_sms: - raise ValueError( - f"locality: {2 * per_domain} + {extra} CTAs must cover the {self.num_sms} SMs" - ) - if per_domain > self.sms_per_domain: - raise ValueError(f"locality: {per_domain} CTAs per domain > {self.sms_per_domain} SMs") - lmap = [-1] * self.num_sms - for k in range(2): - for i, s in enumerate(self._domain_sms[k][:per_domain]): - lmap[s] = k * per_domain + i - left = [self._domain_sms[k][per_domain:] for k in range(2)] - rest = [x for pair in zip(*left) for x in pair] + self.remainder - ext_dom = [] - for i, s in enumerate(rest): - lmap[s] = 2 * per_domain + i - ext_dom.append(self.sm_domain[s]) - assert sorted(lmap) == list(range(self.num_sms)) - return torch.tensor(lmap, dtype=torch.int32, device=f"cuda:{self.device}"), ext_dom - - def domain_streams(self) -> List[torch.cuda.ExternalStream]: - """One non-blocking stream per domain on a green context holding exactly that domain's SMs (kept alive with - this map): work launched on stream d runs only on domain d's SMs. Fork / join with events from the caller's - stream; graph capture follows the forks. Built once, before capture.""" - st = getattr(self, "_streams", None) - if st is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("locality.domain_streams: build before CUDA-graph capture") - d = _drv() - dev = _ck(d.cuDeviceGet(self.device), "cuDeviceGet") - full = _ck( - d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), - "dev resource", - ) - params = [] - for k in range(self.n_domains): - p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() - p.smCount = self.sms_per_domain - p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID - p.localityDomainId = k - params.append(p) - res, _rem = _ck( - d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" - ) - self._green = [] - st = [] - for k in range(self.n_domains): - desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") - g = _ck( - d.cuGreenCtxCreate( - desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM - ), - "green", - ) - s_ = _ck( - d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), - "green stream", - ) - self._green.append((g, s_)) - st.append(torch.cuda.ExternalStream(int(s_), device=f"cuda:{self.device}")) - self._streams = st - return st - - def domain_of(self, ptr: int) -> int: - """The locality domain ordinal of a device address (-1: not placed in one domain). Through libcuda - directly: cuda-python's attribute helper does not know this attribute.""" - import ctypes - - d = _drv() - lib = _libcuda() - out = ctypes.c_int(-9) - attr = int(d.CUpointer_attribute.CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL) - err = lib.cuPointerGetAttribute(ctypes.byref(out), ctypes.c_int(attr), ctypes.c_uint64(ptr)) - if err != 0: - raise RuntimeError( - f"locality: cuPointerGetAttribute(LOCALITY_DOMAIN_ORDINAL) failed: {err}" - ) - return int(out.value) - - def place_split(self, t: torch.Tensor, split_row: int) -> torch.Tensor: - """Split-mapped placement: a copy of ``t`` (contiguous, rows along dim 0) with the SAME shape and layout, in one - virtual range whose rows [0, split_row) are physically in domain 0 and rows [split_row, N) in domain 1. The - split must fall on an allocation granule (rows [0, split_row) a multiple of 2 MiB on VR200); the tail is padded - to a granule in domain 1. Stock consumers read it as a plain tensor; no extra memory once ``t`` is dropped.""" - if not t.is_contiguous(): - raise ValueError("locality.place_split: needs a contiguous tensor") - row_bytes = t[0].numel() * t.element_size() if t.dim() > 1 else t.element_size() - head = split_row * row_bytes - total = t.numel() * t.element_size() - if not 0 < head < total or head % self.granularity: - raise ValueError( - f"locality.place_split: rows [0, {split_row}) = {head} B must be a granule multiple " - f"({self.granularity} B) inside the {total} B tensor" - ) - tail = -(-(total - head) // self.granularity) * self.granularity - m = _Mapping(self.device, [(head, 0), (tail, 1)]) - flat = torch.as_tensor(m, device=f"cuda:{self.device}") - out = flat[:total].view(t.dtype).view(t.shape) - out.copy_(t) - out._locality_mapping = m - for q, k in ( - (out.data_ptr(), 0), - (out.data_ptr() + head - 1, 0), - (out.data_ptr() + head, 1), - (out.data_ptr() + total - 1, 1), - ): - if self.domain_of(q) != k: - raise RuntimeError(f"locality.place_split: byte {q:#x} is not in domain {k}") - return out - - def replace_param_(self, param: torch.nn.Parameter, split_row: int) -> torch.nn.Parameter: - """Re-place a parameter in place with :meth:`place_split` (swaps ``param.data``; the old block returns to the - caching allocator, so call ``torch.cuda.empty_cache()`` between large re-placements).""" - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("locality.replace_param_: not during CUDA-graph capture") - param.data = self.place_split(param.data, split_row) - return param - - def split_rows(self, t: torch.Tensor, dim: int = 1) -> torch.Tensor: - """Place ``t`` (contiguous, 1-byte elements, rows along ``dim`` split into two equal halves) as - [2, *t.shape[:dim], t.shape[dim] // 2, *t.shape[dim + 1:]]: half k (rows [k h, (k + 1) h) of every leading - index) in domain k's memory, one virtual range. Each half's byte size must be a multiple of the allocation - granularity (2 MiB on VR200). The copy is exact; the caller drops ``t`` afterwards.""" - if t.element_size() != 1 or not t.is_contiguous(): - raise ValueError("locality.split_rows: needs a contiguous 1-byte tensor") - n = t.shape[dim] - if n % 2: - raise ValueError(f"locality.split_rows: {n} rows don't split in halves") - half_bytes = t.numel() // 2 - if half_bytes % self.granularity: - raise ValueError( - f"locality.split_rows: a half is {half_bytes} B, not a multiple of the {self.granularity} B granularity" - ) - m = _Mapping(self.device, [(half_bytes, 0), (half_bytes, 1)]) - flat = torch.as_tensor(m, device=f"cuda:{self.device}") - shape = (2,) + tuple(t.shape[:dim]) + (n // 2,) + tuple(t.shape[dim + 1 :]) - out = flat.view(torch.uint8).view(shape) - src = t.view(torch.uint8) - for k in range(2): - out[k].copy_(src.narrow(dim, k * (n // 2), n // 2)) - out._locality_mapping = m # the tensor keeps its mapping alive - for k in range(2): - p = out[k].data_ptr() - for q in (p, p + half_bytes - 1): - if self.domain_of(q) != k: - raise RuntimeError(f"locality.split_rows: byte {q:#x} is not in domain {k}") - return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py index eb025bc1694a..5c71ec93fe34 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_mc_allreduce.py @@ -29,11 +29,13 @@ """ import functools -from dataclasses import dataclass -from typing import Any import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.mc_workspace import ( + McAllReduceWorkspace as McAllReduceWorkspace, +) + from ..ops.comm.m3_mc_allreduce_norm import m3_mc_allreduce_norm as _allreduce_norm_op from ..ops.comm.m3_mc_finalize_allreduce_norm import ( m3_mc_finalize_allreduce_norm as _finalize_allreduce_norm_op, @@ -73,74 +75,6 @@ def _ar_opts(m: int, bucket: int, finalize: bool = False) -> int: return opts | _OPT_CLUSTER_BARRIER -@dataclass(eq=False) -class McAllReduceWorkspace: - """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call - counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" - - uc: torch.Tensor - mc: torch.Tensor - flags: dict - rank: int - world_size: int - hidden: int - handle: Any - comm: Any - max_tokens: int = 8 - - @classmethod - def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": - from tensorrt_llm._torch.distributed.ops import ( - _make_mnnvl_mcast_buffer, - _mnnvl_workspace_all_succeeded, - ) - from tensorrt_llm._utils import mpi_comm - - session = mpi_comm() - session_group = session.Get_group() - group = session_group.Incl(mapping.tp_group) - session_group.Free() - comm = session.Create_group(group) - group.Free() - words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves - ok = not torch.cuda.is_current_stream_capturing() - if not _mnnvl_workspace_all_succeeded(comm, ok): - comm.Free() - raise RuntimeError( - "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" - ) - handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) - uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) - mc = handle.get_mc_buffer((words,), torch.int32, 0) - uc.fill_(EMPTY_WORD) - torch.cuda.synchronize() - if not _mnnvl_workspace_all_succeeded(comm, True): - raise RuntimeError("McAllReduceWorkspace.create failed on another rank") - return cls( - uc=uc, - mc=mc, - flags={}, - rank=mapping.tp_rank, - world_size=mapping.tp_size, - hidden=hidden, - handle=handle, - comm=comm, - max_tokens=max_tokens, - ) - - def flags_for(self, cluster: int): - """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. - get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) - that None never happens and cannot differ between ranks.""" - flags = self.flags.get(cluster) - if flags is None: - if torch.cuda.is_current_stream_capturing(): - return None - flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) - self.flags[cluster] = flags - return flags - - class McBuckets: """One workspace per token bucket (8 / 16 / 48): calls of different sizes never share exchange slots.""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py deleted file mode 100644 index d1730239463c..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/m3_moe_kernel.py +++ /dev/null @@ -1,3841 +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. -"""MiniMax-M3 routed experts for up to 8 decode tokens on one TP rank: ``m3_moe`` (group-phased weight stream). - -One persistent kernel replaces TRTLLM-Gen's MiniMax2 routing, FC1 (NVFP4 x NVFP4 + SwiGLU-OAI + NVFP4 requantization), -FC2 and finalize, optionally with the router GEMV and the activation's NVFP4 quantization folded in (FUSED_ROUTER). It -writes the rank's routed partial [M, 6144] bf16, bit for bit what the TRTLLM-Gen op returns. - -- Routing (warps 4..7, token t on warp 4 + t % 4): TRTLLM-Gen's MiniMax2 top-4, operation for operation; warp 4 then - makes the distinct experts slots in (token, top-k) order (match.any), slot s in group s / 4. -- Work: the groups' FC1 stages and FC2 tiles stream through one ring per CTA in phases: phase p = FC1 of group p (a - stream-K range of the group's unit-major stages), then FC2 of group p - LAG (this CTA's 32-row output blocks, 4 - slots' slices stacked in one 128-row MMA, B row j M + t). LAG phases of streaming hide the hand-off of group - p - LAG's intermediate between CTAs. -- FC1 stage: a 32-row unit's rows at 4 consecutive 256-K tiles (rows 32 q + r) against resident B rows QS q + t. - FC2 stage: 4 slots' 32 rows at one 256-K tile. Both carry 8 KB of raw scale atoms, spliced into the 4 instruction - atoms by warp 6 with one formula (word [r][q] <- raw atom 4 q + kb, word [r][m1]). -- FC1 epilogue: k-quarters summed in order, SwiGLU-OAI, NVFP4 requant into tagged 16-byte records (3 words of 3 data - bytes + the call's tag); a unit split between two CTAs is summed by its tail's holder (Lamport slot). -- Warps 5 and 8.. (XW warps) build each group's FC2 B (tiles + scale atoms) from the records once they carry the - call's tag; warp 5 copies the scales to TMEM. -- Finalize: bf16(fc2_alpha * acc) per (slot, token) into its top-k position, then per token the fma chain over its - top-4 in order, bf16 out. With an all-reduce descriptor, every output word is also pushed to the exchange - (AR_PRE: call parities and descriptor loaded during the last group's FC2). - -Configuration (trace-time constants) via M3_CONFIG = {"m_max", "num_ctas", "stages", "nb2_cap", "fused_router", "lag", -"xw", ...}; the plan (op.make_phase_plan) must be built for the same lag. -Launched with PDL: barrier setup, TMEM allocation, tensormap prefetch, bias / plan / router weight loads before -griddepcontrol.wait; every producer-output read and every global write after it. -""" - -from __future__ import annotations - -import cuda.bindings.driver as cuda_driver -import cutlass -import cutlass.cute as cute -import cutlass.experimental.cuda as cuda -from cutlass._mlir.dialects import llvm as _llvm -from cutlass._mlir.extras import types as _T -from cutlass.cutlass_dsl import dsl_user_op -from cutlass.experimental import primitives as prims - -_CFG = globals().get("M3_CONFIG") or {} - - -def _cfg(key: str, default): - return type(default)(_CFG.get(key, default)) - - -# ============================================================================= -# Problem shape (trace-time constants). -# ============================================================================= -H = 6144 # capacity: fixed: the checkpoint's hidden size -# routed intermediate per rank: 768 at TP4, 384 at TP8 (FC2 K = 1.5 256-K tiles: the last tile's A box is half out -# of bounds, zero-filled by TMA, and only its first KB2_LAST K-64 MMAs / scale atoms are issued; no padding, no extra -# bytes) -I_TP = _cfg("i_tp", 768) -TWO_I = 2 * I_TP -assert I_TP % 128 == 0 and 128 <= I_TP <= 768, I_TP -E_LOCAL = 128 -TOP_K = 4 # capacity: fixed: the checkpoint's top-k -ROUTE_SCALE = 2.0 -SUM_EPS = 1e-20 -M_MAX = _cfg("m_max", 1) -NCTA = _cfg("num_ctas", 152) -FUSED_ROUTER = bool(_cfg("fused_router", 0)) -NROUTE = E_LOCAL if FUSED_ROUTER else 0 # experts whose logits the router GEMV computes -NGEMV = (NROUTE + NCTA - 1) // NCTA # experts per CTA: bx, bx + NCTA, .. -TIMING = bool(_cfg("timing", 0)) -HANG_DBG = bool(_cfg("hang_dbg", 0)) -XVDUMP = bool( - _cfg("xvdump", 0) -) # debug only (needs a plan without split units): FC1 sums into `partials` -YDUMP = bool( - _cfg("ydump", 0) -) # debug only: FC2 outputs per (token, k) into `partials` (breaks its protocol) -LAG = _cfg("lag", 2) # phase p runs FC1 of group p and FC2 of group p - LAG -G_CAP = TOP_K * M_MAX # distinct experts per call -NG_CAP = M_MAX # groups of 4 slots -XW = _cfg("xw", 5 if M_MAX > 4 else (4 if M_MAX > 2 else 1)) # FC2 B builder warps: 5, then 8 .. -THREADS = 256 + 32 * (XW - 1) -GEMV_THREADS = 256 # capacity: fixed: the CTA size -# SHARED: the MoE's shared expert (MXFP8 gate_up [1536, 6144] -> SwiGLU-OAI -> MXFP8 down [6144, 768]) on NSH more -# CTAs (warps 0-3), its output added in finalize: out = bf16(shared + routed) (shared_output.add_(routed_output)). -# The math is mxfp8_low_m_gemm's (native MXFP8 activation quant, mma.sync m16n8k32 e4m3 per 32-K block, block scales -# as one fma, its partial-sum order, bf16) and the Triton silu_and_mul's, operation for operation. Shared CTA s owns -# intermediate positions 32 s .. 32 s + 31 (gate and up rows: MMA tiles = warps 0-3) and down rows 256 s .. + 255; -# 16-row x 768-byte weight chunks (+ their scale atoms) stream through NBW buffers per warp, the first NBW before -# the grid dependency. h crosses the shared CTAs, y_sh reaches the routed CTAs' finalize (both tagged words). -SHARED = bool(_cfg("shared", 0)) -# LATE_TRIGGER: release the dependents at teardown, not after the grid dependency. An early release lets the next -# kernels' CTAs (the all-reduce, then the attention it releases at entry) become resident during this kernel, packed -# onto the SMs it leaves free, where the attention then runs on a fraction of the GPU. -LATE_TRIGGER = bool(_cfg("late_trigger", 1)) -# AR_PRE: the finalize's all-reduce push uses the call parities and descriptor fields loaded while the last group's -# FC2 runs (not two dependent global loads per push on the kernel's tail). -AR_PRE = bool(_cfg("ar_pre", 1)) -# EARLY_AR (scheduling only): when the finalize pushes into the multicast all-reduce, EVERY CTA releases the -# dependents right after the grid dependency, so the all-reduce (which reads only the pushed tagged words, without a -# grid wait) and, through its own early release, the next layer's QKV GEMV become resident on each SM this kernel -# frees, instead of after the last CTA's teardown. The arithmetic is unchanged. -EARLY_AR = bool(_cfg("early_ar", 1)) -# XSTAGE (fused router): the hidden states land in the (still idle) ring by one bulk copy per token right after the -# grid dependency; the NVFP4 quantization of B1 reads them from shared memory (one HBM / L2 round trip, not ~8). -XSTAGE = bool(_cfg("xstage", 1)) and bool(_cfg("fused_router", 0)) -NSH = 24 if SHARED else 0 -# the shared intermediate per rank: 768 (TP4), 384 (TP8). A shared CTA owns SH_ROWS intermediate positions: its -# gate_up MMA tile is SH_SPLIT K slices x (SH_ROWS gate + SH_ROWS up rows) = 128 rows (TP4: 2 x (32 + 32), the order -# of m3_mx_gemv split 2; TP8: 4 x (16 + 16), m3_mx_gemv split 4), and 6144 / NSH = 256 down rows over the full K. -SH_I = _cfg("sh_i", 768) -SH_ROWS = SH_I // 24 -# gate_up K slices (= m3_mx_gemv's cluster ranks, contiguous K ranges) -SH_SPLIT = 128 // (2 * SH_ROWS) -SH_NB = 8 * SH_SPLIT # gate_up B rows (8 q + t: token t's slice-q k-tile) = its MMA N -SH_KB = H // 32 # capacity: fixed: a kernel tile constant (gate_up K blocks) -SH_DKB = SH_I // 32 # down K blocks -SH_DKT = SH_I // 128 # down k-tiles -SH_GT = SH_KB // 4 // SH_SPLIT # gate_up k-tiles (128 K) per slice -SH_STAGES_GU = SH_GT -SH_STAGES = SH_GT + 2 * SH_DKT # + 2 down tiles x SH_DKT k-tiles -SHX_BAR = 6 -# SH_PRETILE: the shared weights come pre-tiled (op._pretile_shared): per shared CTA its SH_STAGES stages, each the -# 16 KB shared-memory image of its 128-row MMA tile (128-byte swizzle) followed by the stage's 512-byte instruction -# scale atom, so a stage is two contiguous bulk copies (not 128 row segments of 128 bytes) and needs no splicing. -SH_PRETILE = bool( - _cfg("sh_pretile", 1 if M_MAX <= 2 or SH_SPLIT != 2 else 0) -) # (in-chain: a win at M <= 2, not 4..8) -if SHARED: - assert SH_I in (384, 768), SH_I - assert SH_PRETILE or SH_SPLIT == 2, ( - "the un-pretiled shared stream is the 32-row (TP4) layout only" - ) -SH_STAGE_BYTES = 16384 + 512 -MMA_M = 128 -KT_BYTES = 128 -INC_K = 2 -NUM_KB = 4 -# FC1 -N1 = 16 if M_MAX <= 4 else 32 # B rows: QS q + t (token t at the stage's k-tile 4 j + q) -QS = N1 // 4 -R1 = 32 -C1 = 4 -U1 = TWO_I // R1 # 48 units per slot -K1_TILES = H // 256 -S1 = K1_TILES // C1 # capacity: fixed: pipeline stages (6 stages per unit) -NUM_BYTES_A = MMA_M * KT_BYTES # capacity: fixed: a stage's operand bytes -NUM_BYTES_BOX1 = R1 * KT_BYTES -SF1_KA = H // 64 -SF1_RB = TWO_I // 128 -NUM_BYTES_SFA_RAW = 16 * 512 -NUM_WORDS_SFA_RAW = NUM_BYTES_SFA_RAW // 4 -NUM_BYTES_SFA = NUM_KB * 512 # capacity: fixed: a stage's scale bytes -NUM_WORDS_SFA = NUM_BYTES_SFA // 4 -NUM_BYTES_B1 = N1 * KT_BYTES -NUM_SFB1_ATOMS = S1 * NUM_KB -# FC2 -R2 = 32 -NB2 = H // R2 # 192 -KT2 = (I_TP + 255) // 256 # 3 (TP4), 2 (TP8) -# K-64 instructions of the last FC2 k-tile: 4 (TP4), 2 (TP8) -KB2_LAST = (I_TP - 256 * (KT2 - 1)) // 64 - - -def _nkb2(t2: int) -> int: - """K-64 MMA instructions (= scale atoms) of FC2 k-tile t2.""" - return NUM_KB if t2 < KT2 - 1 else KB2_LAST - - -SF2_KA = I_TP // 64 -SF2_RB = H // 128 -N2 = max(8, (4 * M_MAX + 7) // 8 * 8) # B rows j M + t -NUM_BYTES_B2 = N2 * KT_BYTES -NB2_CAP = _cfg("nb2_cap", 2) -UNITS_ROW = TWO_I // R1 # 48 records per (slot, token) row -H_ROW = UNITS_ROW * 16 -# HB_REPL: the tagged FC1 -> FC2 records written to HB_REPL copies (HB_STRIDE int32 words apart); CTA bx polls copy -# bx % HB_REPL. Every routed CTA polls every record, so with one copy the hand-off latency depends on where the few -# hot lines land in L2 (TP8 M=1: 1.5 KB of records, +-1 us across allocations); copies spread the polls. -HB_REPL = _cfg("hb_repl", 1) -assert 1 <= HB_REPL <= 32, HB_REPL # (copy i is written by lane i) -HB_STRIDE = (G_CAP * M_MAX * UNITS_ROW * 4 + 255) // 256 * 256 + 32 -PARTIAL_EMPTY = -1 # capacity: fixed: a sentinel, not a size -EPI_THREADS = 128 -EPI_BAR = 1 -ROUTE_BAR = 2 -GEMV_BAR = 3 -B2_BAR = 5 -TD_STRIDE = 64 # tdbg columns per CTA -TOPK_BAR = 4 -B1_AFTER_TOPK = bool( - _cfg("b1_after_topk", 0) -) # (experiment) the B1 build waits for the routing warps' top-k -TOPK2_BAR = 7 -EARLY0 = bool( - _cfg("early0", 1) -) # warp 4 streams group 0 (token 0's experts) right after its own top-k -EARLY0_N = _cfg("early0_n", 99) # stages issued early (capped at STAGES) -# routing warps: 4..7, then the FC2 B builder warps 8.. (idle until the routing is done) up to one warp per token -RW_EXTRA = max(0, min(M_MAX - 4, (THREADS // 32) - 8)) if bool(_cfg("rw_extra", 1)) else 0 -NRW0 = ( - min(4, M_MAX) if bool(_cfg("rw_trim", 1)) else 4 -) # warps 4.. with a token (no idle pollers in the barrier) -NRW = NRW0 + RW_EXTRA -LG_REPL = _cfg("lg_repl", 32 if M_MAX > 4 else 128) -GEMV_SMEM = bool( - _cfg("gemv_smem", 1) -) # (XSTAGE) the router GEMV reads the staged hidden states from smem -LG_STRIDE = M_MAX * E_LOCAL + 32 -POLL_NS = _cfg("poll_ns", 150) -B2_POLL_NS = _cfg("b2_poll_ns", 200) -# B2IMG (needs SFA_CP: warp 6 is free): the FC1 epilogue writes each finished unit's FP4 bytes and scale byte straight -# into a per-group global image laid out as the FC2 B tile + its scale atoms (exactly what the builder warps assembled -# from the tagged records); per group, warp 6 joins the epilogue (fc1_done), fences the async proxy and does one -# release add of this CTA's finished units to the group counter; warp 5 acquires the counter (every group gains -# CNT_PER_CALL per call, CTA 0 pads absent slots), fences, and bulk-copies the image into the B2 buffer. -B2IMG = bool(_cfg("b2img", 0)) and bool(_cfg("sfa_cp", 0)) -IMG_B2 = KT2 * NUM_BYTES_B2 -IMG_SF = KT2 * NUM_KB * 512 -IMG_BYTES = IMG_B2 + IMG_SF -CNT_PER_CALL = 4 * (TWO_I // R1) -B2IMG_NS = _cfg("b2img_ns", 16) # (B2IMG) counter poll back-off -B2_DLY = _cfg( - "b2_dly", 0 -) # (stress) the release warp spins up to this many ns (CTA- and group-dependent) first -W_WARP = 4 # capacity: fixed: a warp role index -X_WARP = 5 # capacity: fixed: a warp role index -S_WARP = 6 # capacity: fixed: a warp role index -MMA_WARP = 7 # capacity: fixed: a warp role index -# E0_M4: at M = 4 the MMA warp routes -# token 3 (NRW0 = 4), so the early group-0 stages go to FC2 B builder warp 8 instead (idle until the routing is done; -# RW_EXTRA = 0 at M = 4). Issue order only: exact. -E0_M4 = ( - bool(_cfg("e0_m4", 1)) and NRW0 == 4 and RW_EXTRA == 0 and THREADS // 32 > 8 -) # M == 4 only: at M > 4 warp 8 routes (RW_EXTRA) -E0_WARP = ( - 8 if E0_M4 else MMA_WARP -) # (M <= 3) the early stages are issued by the idle MMA warp, not by warp 4 (routing's critical path) -E0_BAR = 8 -E1_BAR = 9 -E0_OFFLOAD = EARLY0 and (NRW0 <= 3 or E0_M4) and bool(_cfg("e0_off", 1)) -# B1_AFTER_TOPK holds warps 0..3 (the B1 build, which arrives x_bar[1] when XSTAGE) on TOPK2_BAR until every routing -# warp arrives; with EARLY0 and no E0 offload, warp 4 first waits on x_bar[1] to fill the ring before it arrives, so the -# two wait on each other. Named barriers have no HANG_NS trap, so the kernel would spin forever: refuse the config. -if B1_AFTER_TOPK and XSTAGE and EARLY0 and not E0_OFFLOAD: - raise ValueError( - "m3_moe: b1_after_topk=1 with xstage, early0 and no early-stage offload deadlocks (warp 4 waits on x_bar[1] " - "before TOPK2_BAR, the B1 build waits on TOPK2_BAR before x_bar[1]); set b1_after_topk=0 or early0=0" - ) -EVICT_FIRST = 0x12F0000000000000 -_LOG2E = 1.4426950408889634 -LD_POW2 = ( - 1 if M_MAX == 1 else 2 if M_MAX == 2 else 4 if M_MAX <= 4 else 8 -) # tcgen05.ld .num covering M columns - -_FIXED_SMEM = ( - S1 * NUM_BYTES_B1 - + 2 * KT2 * NUM_BYTES_B2 - + 2 * KT2 * NUM_KB * 512 - + NB2_CAP * M_MAX * TOP_K * R2 * 4 - + 2 * M_MAX * C1 * R1 * 4 - + (M_MAX * (H // 16) if FUSED_ROUTER else 0) - + 6 * 1024 -) -# SFA_CP: the A scale words go by 4-byte cp.async straight into the instruction atom layout (no raw scale box, no -# splicer warp; the MMA warp tcgen05.cp's them), completing on the stage's full barrier: 18 KB stages instead of 26 KB -SFA_CP = bool(_cfg("sfa_cp", 0)) -_RAW_STAGE = 0 if SFA_CP else NUM_BYTES_SFA_RAW -_STAGE_SMEM = NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA -_STAGES_FIT = (227 * 1024 - _FIXED_SMEM - (128 if SHARED else 0)) // _STAGE_SMEM -ACC1_COL = 0 -ACC2_COL = 2 * N1 -SFB1_COL = ACC2_COL + 2 * N2 -SFB2_COL = SFB1_COL + NUM_SFB1_ATOMS * 4 -SFA1_COL = SFB2_COL + 2 * KT2 * NUM_KB * 4 -STAGES = min(_cfg("stages", 12 if SFA_CP else 8), _STAGES_FIT, (512 - SFA1_COL) // 16) -TMEM_COLS = 512 # capacity: fixed: the SM's tensor-memory columns -assert 1 <= M_MAX <= 8, M_MAX -assert STAGES >= 3, (STAGES, _STAGES_FIT, SFA1_COL) -assert M_MAX * H * 2 <= STAGES * (NUM_BYTES_A + _RAW_STAGE + NUM_BYTES_SFA), ( - "x staging must fit the ring" -) -# one shared-memory region for the ring and the FC1 / FC2 B operands (a shared CTA reuses all of it) -_OFF_SFAR = STAGES * NUM_BYTES_A -_OFF_SFA = _OFF_SFAR + STAGES * _RAW_STAGE -_OFF_B1 = _OFF_SFA + STAGES * NUM_BYTES_SFA -_OFF_B2 = _OFF_B1 + S1 * NUM_BYTES_B1 -_OFF_SFB2 = _OFF_B2 + 2 * KT2 * NUM_BYTES_B2 -BIG = _OFF_SFB2 + 2 * KT2 * NUM_KB * 512 -# shared CTA layout in BIG: weight buffers, x quarters (2), x scales [M, 192], gate / up [4, 16, M], h, h scales -# shared CTA layout in BIG: A ring (16 KB stages), raw scale atoms (4 per stage), spliced atoms, B (x: [stage][16 rows] -# [128 B]; h: [6][8][128 B]), x / h scale codes, gate_up partial exchange -_SB_A = 0 -_SB_RAW = None -# slices 1.. gate / up partials [slice - 1][u][p][8 tokens] f32 -_SH_XCH_BYTES = (SH_SPLIT - 1) * 2 * SH_ROWS * 8 * 4 -_SH_FIX = ( - SH_GT * SH_NB * 128 - + SH_DKT * 8 * 128 - + M_MAX * SH_KB - + M_MAX * SH_DKB - + _SH_XCH_BYTES - + 2 * SH_ROWS * 8 * 4 - + 1024 -) -SH_NST = min(8, (BIG - _SH_FIX) // (NUM_BYTES_A + 4 * 512 + 512)) if SHARED else 1 -_SB_RAW = SH_NST * NUM_BYTES_A -_SB_SPL = _SB_RAW + SH_NST * 2048 -_SB_BX = (_SB_SPL + SH_NST * 512 + 1023) // 1024 * 1024 -_SB_BH = _SB_BX + SH_GT * SH_NB * 128 -_SB_XC = _SB_BH + SH_DKT * 8 * 128 -_SB_HC = _SB_XC + M_MAX * SH_KB -_SB_XCH = (_SB_HC + M_MAX * SH_DKB + 15) // 16 * 16 -_SB_GU = _SB_XCH + _SH_XCH_BYTES -_SB_END = _SB_GU + 2 * SH_ROWS * 8 * 4 -# shared CTA TMEM columns: gate_up accumulator (N SH_NB), down accumulators (2 x N 8), A scales per stage, B scales -SH_ACC_GU = 0 -SH_ACC_DN = SH_NB -SH_SFA = SH_ACC_DN + 16 -SH_SFB_GU = SH_SFA + 4 * 8 -SH_SFB_DN = SH_SFB_GU + 4 * SH_GT -SH_TMEM = 256 # capacity: fixed: the shared CTAs' TMEM columns -if SHARED: - assert FUSED_ROUTER, "shared expert: fused router only" - assert SH_STAGES - SH_NST <= 64, SH_STAGES - assert SH_NST >= 3 and _SB_END <= BIG and SH_SFB_DN + 24 <= SH_TMEM, (SH_NST, _SB_END, BIG) - - -# ============================================================================= -# Inline-PTX helpers. -# ============================================================================= -def _asm(ret, args, ptx, cons, side=False, loc=None, ip=None): - return _llvm.inline_asm( - ret, - [a.ir_value(loc=loc, ip=ip) for a in args], - ptx, - cons, - has_side_effects=side, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -# ---- all-reduce push (the MoE boundary's multicast exchange, set A): -# arp = int64 [enable, mc, flags, rank, world, maxt] -@dsl_user_op -def _ld_volatile_s32(addr, *, loc=None, ip=None): - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Int64(addr).ir_value(loc=loc, ip=ip)], "ld.volatile.global.s32 $0, [$1];", "=r,l", - has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -@dsl_user_op -def _st_global_b32_if(addr, value, pred, *, loc=None, ip=None): - """st.global.b32 when pred != 0.""" - _llvm.inline_asm( - None, [cutlass.Int64(addr).ir_value(loc=loc, ip=ip), cutlass.Int32(value).ir_value(loc=loc, ip=ip), - cutlass.Int32(pred).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .pred p;\n\tsetp.ne.s32 p, $2, 0;\n\t@p st.global.b32 [$0], $1;\n\t}", "l,r,r", - has_side_effects=True, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) # fmt: skip - - -@dsl_user_op -def _bf16x2_word(lo, hi, *, loc=None, ip=None): - """bf16(lo) | bf16(hi) << 16 (RN), with -0.0 halves turned into +0.0 (0x80000000 marks an empty exchange word).""" - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), [cutlass.Float32(hi).ir_value(loc=loc, ip=ip), cutlass.Float32(lo).ir_value(loc=loc, ip=ip)], - "{\n\t.reg .b32 w, l, h;\n\t.reg .pred pl, ph;\n\tcvt.rn.bf16x2.f32 w, $1, $2;\n\t" - "and.b32 l, w, 0xFFFF;\n\tsetp.eq.u32 pl, l, 0x8000;\n\t@pl and.b32 w, w, 0xFFFF0000;\n\t" - "shr.u32 h, w, 16;\n\tsetp.eq.u32 ph, h, 0x8000;\n\t@ph and.b32 w, w, 0x0000FFFF;\n\t" - "mov.b32 $0, w;\n\t}", "=r,f,f", - has_side_effects=False, is_align_stack=False, asm_dialect=_llvm.AsmDialect.AD_ATT, loc=loc, ip=ip, - ) - ) # fmt: skip - - -def _ar_push(arp, v, low, t, cl): - """Push bf16 channel pair (cl, cl + 1) of token t (v = channel cl's bf16 value on this lane, its neighbour's on - lane ^ 8) into this rank's set-A slot of the exchange; the slot half is the all-reduce CTA (t, cl // 768)'s call - parity. Only lanes with ``low`` store; every lane of the warp must call it (shuffle).""" - hi = cute.arch.shuffle_sync_bfly(v, 8) - pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) - cs = cutlass.select_(low, cl, cutlass.Int32(0)) - world = arp.load(idx=4) - parity = cutlass.Int64( - _ld_volatile_s32( - arp.load(idx=2) - + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cs // cutlass.Int32(H // 8)) - ) - & cutlass.Int32(1) - ) - off = ( - parity * arp.load(idx=5) * world * cutlass.Int64(H // 2) - + (cutlass.Int64(t) * world + arp.load(idx=3)) * (cutlass.Int64(H // 2)) - + cutlass.Int64(cs // cutlass.Int32(2)) - ) - _st_global_b32_if(arp.load(idx=1) + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) - - -def _ar_parity(arp, t, cl): - """The call parity _ar_push reads for token t's channel cl (fixed between this kernel's grid dependency and the - next all-reduce, which this kernel releases at teardown): loaded ahead of the finalize (AR_PRE).""" - return _ld_volatile_s32( - arp.load(idx=2) - + cutlass.Int64(4) * cutlass.Int64(t * cutlass.Int32(8) + cl // cutlass.Int32(H // 8)) - ) - - -def _ar_push_pre(ard, parity32, v, low, t, cl): - """_ar_push with the parity (_ar_parity) and the descriptor fields ard = (mc, rank, world, maxt) preloaded.""" - hi = cute.arch.shuffle_sync_bfly(v, 8) - pred = cutlass.select_(low, cutlass.Int32(1), cutlass.Int32(0)) - cs = cutlass.select_(low, cl, cutlass.Int32(0)) - mc, rk, world, maxt = ard - parity = cutlass.Int64(parity32 & cutlass.Int32(1)) - off = ( - parity * maxt * world * cutlass.Int64(H // 2) - + (cutlass.Int64(t) * world + rk) * (cutlass.Int64(H // 2)) - + cutlass.Int64(cs // cutlass.Int32(2)) - ) - _st_global_b32_if(mc + off * cutlass.Int64(4), _bf16x2_word(v, hi), pred) - - -@dsl_user_op -def _bar_arrive(bar_id: int, threads: int, *, loc=None, ip=None): - """bar.arrive: signal a named barrier without waiting (its memory ordering is bar.sync's).""" - _llvm.inline_asm( - None, - [], - f"bar.arrive {bar_id}, {threads};", - "", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def _st_relaxed_u32(addr, val, *, loc=None, ip=None): - _asm( - None, - [cutlass.Int64(addr), cutlass.Int32(val)], - "st.relaxed.gpu.global.u32 [$0], $1;", - "l,r", - True, - loc, - ip, - ) - - -@dsl_user_op -def _ld_relaxed_u32(addr, *, loc=None, ip=None): - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Int64(addr)], - "ld.relaxed.gpu.global.u32 $0, [$1];", - "=r,l", - True, - loc, - ip, - ) - ) - - -@dsl_user_op -def _st_relaxed_u64(addr, val, *, loc=None, ip=None): - _asm( - None, - [cutlass.Int64(addr), cutlass.Int64(val)], - "st.relaxed.gpu.global.u64 [$0], $1;", - "l,l", - True, - loc, - ip, - ) - - -@dsl_user_op -def _ld_relaxed_u64(addr, *, loc=None, ip=None): - return cutlass.Int64( - _asm( - _T.i64(), - [cutlass.Int64(addr)], - "ld.relaxed.gpu.global.u64 $0, [$1];", - "=l,l", - True, - loc, - ip, - ) - ) - - -@dsl_user_op -def _nanosleep(ns: int, *, loc=None, ip=None): - _llvm.inline_asm( - None, - [], - f"nanosleep.u32 {ns};", - "", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -@dsl_user_op -def _gtime(*, loc=None, ip=None): - return cutlass.Int64( - _llvm.inline_asm( - _T.i64(), - [], - "mov.u64 $0, %globaltimer;", - "=l", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) - - -@dsl_user_op -def _smid(*, loc=None, ip=None): - return cutlass.Int32( - _llvm.inline_asm( - _T.i32(), - [], - "mov.u32 $0, %smid;", - "=r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - ) - - -@cute.jit -def _stamp(tdbg: cutlass.Array, bx, i: cutlass.Constexpr): - if cutlass.const_expr(TIMING): - tdbg.store(_gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(i)) - - -@dsl_user_op -def _rcp(a, *, loc=None, ip=None): - """rcp.approx.ftz.f32 (the reciprocal TRT-LLM's FP4 quantization uses).""" - return cutlass.Float32( - _asm(_T.f32(), [cutlass.Float32(a)], "rcp.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) - ) - - -@dsl_user_op -def _div_rn(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "div.rn.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _fma_rn(a, b, c, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b), cutlass.Float32(c)], - "fma.rn.f32 $0, $1, $2, $3;", - "=f,f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _add_rn(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "add.rn.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _sub_rn(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "sub.rn.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _mul_rn(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "mul.rn.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _e4m3_bits(x, *, loc=None, ip=None): - """e4m3 (round to nearest even, saturating) of x as an int 0..255.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(x)], - "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, 0f00000000, $1;\ncvt.u32.u16 $0, t;\n}", - "=r,f", - False, - loc, - ip, - ) - ) & cutlass.Int32(0xFF) - - -@dsl_user_op -def _e4m3_to_f32(b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Int32(b)], - "{\n.reg .b16 t, l, u;\n.reg .b32 h;\ncvt.u16.u32 t, $1;\ncvt.rn.f16x2.e4m3x2 h, t;\nmov.b32 {l, u}, h;\n" - "cvt.f32.f16 $0, l;\n}", - "=f,r", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _e2m1x2(hi, lo, *, loc=None, ip=None): - """One byte of two E2M1 values (lo in the low nibble), round to nearest even, saturating.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(hi), cutlass.Float32(lo)], - "{\n.reg .b8 t;\ncvt.rn.satfinite.e2m1x2.f32 t, $1, $2;\ncvt.u32.u8 $0, t;\n}", - "=r,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _ex2(x, *, loc=None, ip=None): - return cutlass.Float32( - _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.ftz.f32 $0, $1;", "=f,f", False, loc, ip) - ) - - -@dsl_user_op -def _bf16_round(x, *, loc=None, ip=None): - """x rounded to bf16 (nearest even), as f32.""" - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(x)], - "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\nmov.b32 $0, {0, t};\n}", - "=f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _redux_max_u32(v, *, loc=None, ip=None): - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Int32(v)], - "redux.sync.max.u32 $0, $1, 0xffffffff;", - "=r,r", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _bf16_bits(x, *, loc=None, ip=None): - """bf16 (nearest even) bits of x in the low 16 bits.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(x)], - "{\n.reg .b16 t;\ncvt.rn.bf16.f32 t, $1;\ncvt.u32.u16 $0, t;\n}", - "=r,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _ue8m0_rp(x, *, loc=None, ip=None): - """UE8M0 code of x rounded toward +inf, saturating (__nv_cvt_float_to_e8m0(x, SATFINITE, RoundPosInf)).""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(x)], - "{\n.reg .b16 t;\ncvt.rp.satfinite.ue8m0x2.f32 t, $1, $1;\ncvt.u32.u16 $0, t;\n}", - "=r,f", - False, - loc, - ip, - ) - ) & cutlass.Int32(0xFF) - - -@dsl_user_op -def _e4m3x2(hi, lo, *, loc=None, ip=None): - """Two e4m3 (nearest even, saturating) in 16 bits, lo in the low byte.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Float32(hi), cutlass.Float32(lo)], - "{\n.reg .b16 t;\ncvt.rn.satfinite.e4m3x2.f32 t, $1, $2;\ncvt.u32.u16 $0, t;\n}", - "=r,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _min_xorsign_abs(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "min.xorsign.abs.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _div_full(a, b, *, loc=None, ip=None): - return cutlass.Float32( - _asm( - _T.f32(), - [cutlass.Float32(a), cutlass.Float32(b)], - "div.full.f32 $0, $1, $2;", - "=f,f,f", - False, - loc, - ip, - ) - ) - - -@dsl_user_op -def _ex2_full(x, *, loc=None, ip=None): - """ex2.approx.f32 (no flush to zero), Triton's tl.exp.""" - return cutlass.Float32( - _asm(_T.f32(), [cutlass.Float32(x)], "ex2.approx.f32 $0, $1;", "=f,f", False, loc, ip) - ) - - -@dsl_user_op -def _l2_prefetch(addr, nbytes, *, loc=None, ip=None): - """cp.async.bulk.prefetch.L2 of nbytes (a multiple of 16) at global address addr (16-byte aligned).""" - _llvm.inline_asm( - None, - [ - cutlass.Int64(addr).ir_value(loc=loc, ip=ip), - cutlass.Int32(nbytes).ir_value(loc=loc, ip=ip), - ], - "cp.async.bulk.prefetch.L2.global [$0], $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=_llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) - - -def _mx_osc(amax): - """mxfp8 quantization of a block with absolute max amax: (UE8M0 code, the element multiplier).""" - code = _ue8m0_rp(_mul_rn(amax, _rcp(cutlass.Float32(448.0)))) - cf = cutlass.Int32( - cutlass.select_( - code == cutlass.Int32(255), - cutlass.Int32(0x7FFFFFFF), - cutlass.select_( - code == cutlass.Int32(0), cutlass.Int32(0x00400000), code << cutlass.Int32(23) - ), - ) - ).bitcast(cutlass.Float32) - return code, cutlass.select_(amax != cutlass.Float32(0.0), _rcp(cf), cutlass.Float32(0.0)) - - -def _mx_words(v, osc): - """e4m3 words of len(v) (a multiple of 4) values times osc, in element order.""" - words = [] - for t in range(len(v) // 4): - lo = _e4m3x2(_mul_rn(v[4 * t + 1], osc), _mul_rn(v[4 * t], osc)) - hi = _e4m3x2(_mul_rn(v[4 * t + 3], osc), _mul_rn(v[4 * t + 2], osc)) - words.append(lo | (hi << cutlass.Int32(16))) - return words - - -def _twiddle(score): - """Float bits mapped so that unsigned order is float order (cub TwiddleIn).""" - bits = ( - score.bitcast(cutlass.Int32) - if hasattr(score, "bitcast") - else cutlass.Float32(score).bitcast(cutlass.Int32) - ) - return cutlass.select_(bits < cutlass.Int32(0), ~bits, bits | cutlass.Int32(-2147483648)) - - -@cute.jit -def _wait(bar: cutlass.Array, parity): - while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): - pass - cute.arch.sync_warp() - - -def _tptr(base_row, col, dtype=cutlass.Int32): - return cutlass.inttoptr((base_row << 16) | col, 6, dtype) - - -@dsl_user_op -def _match_any(value, *, loc=None, ip=None): - """The mask of the warp's lanes holding the same 32-bit value.""" - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Int32(value)], - "match.any.sync.b32 $0, $1, 0xffffffff;", - "=r,r", - False, - loc, - ip, - ) - ) - - -@cute.jit -def _waitd(bar: cutlass.Array, parity, tdbg: cutlass.Array, site: cutlass.Constexpr): - """_wait, or (HANG_DBG) give up after ~2^22 polls and record the site in tdbg[bx, 24 + warp].""" - if cutlass.const_expr(HANG_DBG): - n = cutlass.Int32(0) - while (not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity)) & ( - n < cutlass.Int32(1 << 22) - ): - n = n + cutlass.Int32(1) - if n >= cutlass.Int32(1 << 22): - bx, _, _ = cute.arch.block_idx() - w = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - tdbg.store( - cutlass.Int64(site), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + w - ) - else: - while not cute.arch.mbarrier_try_wait(bar.data_ptr(), parity): - pass - cute.arch.sync_warp() - - -@cute.jit -def _adv(stage, phase): - """The next ring position's (stage, phase).""" - stage = stage + cutlass.Int32(1) - wrap = stage == cutlass.Int32(STAGES) - return cutlass.select_(wrap, cutlass.Int32(0), stage), cutlass.select_( - wrap, phase ^ cutlass.Int32(1), phase - ) - - -@cute.jit -def _seg(i0, i1, k): - """Segment k of the FC1 item range [i0, i1) (items = unit-major stages): (unit, first stage, end stage). Order: - the last unit (its head, the part a split unit's next CTA does not cover), then the first unit, then the rest.""" - uf = i0 // cutlass.Int32(S1) - ul = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - unit = uf + k - cutlass.Int32(1) - unit = cutlass.select_(k == cutlass.Int32(0), ul, unit) - unit = cutlass.select_(ul == uf, uf, unit) - j0 = i0 - unit * cutlass.Int32(S1) - j0 = cutlass.select_(j0 < cutlass.Int32(0), cutlass.Int32(0), j0) - j1 = i1 - unit * cutlass.Int32(S1) - j1 = cutlass.select_(j1 > cutlass.Int32(S1), cutlass.Int32(S1), j1) - return unit, j0, j1 - - -@cute.jit -def _nsegs(i0, i1): - n = (i1 - cutlass.Int32(1)) // cutlass.Int32(S1) - i0 // cutlass.Int32(S1) + cutlass.Int32(1) - return cutlass.select_(i1 > i0, n, cutlass.Int32(0)) - - -@cute.jit -def _sfb1_to_tmem(warp, lane, base_row, base_col, src: cutlass.Array, b1_ready: cutlass.Array): - """B1 scales to TMEM: instruction (j, kb)'s column holds, at lane n = QS q + t, token t's 4 scales of k-tile - 4 j + q, K 64 kb .. (src: int32 words [M, 96] of the token scales); each warp writes its sub-partition.""" - q = lane // cutlass.Int32(QS) - tl = lane - q * cutlass.Int32(QS) - sf_ok = (lane < cutlass.Int32(N1)) & (tl < cutlass.Int32(M_MAX)) - sfv = [] - for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): - j = a // NUM_KB - kb = a % NUM_KB - wsf = src.load( - idx=cutlass.select_( - sf_ok, - tl * cutlass.Int32(H // 64) - + (cutlass.Int32(4 * j) + q) * cutlass.Int32(4) - + cutlass.Int32(kb), - cutlass.Int32(0), - ) - ) - sfv.append(cutlass.select_(sf_ok, cutlass.Int32(wsf), cutlass.Int32(0))) - for a in cutlass.range_constexpr(NUM_SFB1_ATOMS): - prims.tcgen05_st( - "32x32b", - _tptr(base_row + warp * cutlass.Int32(32), base_col + SFB1_COL + a * 4), - sfv[a], - ) - prims.tcgen05_wait(prims.Tcgen05Wait.STORE) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - prims.mbarrier_arrive(b1_ready) - - -@cute.jit -def _b1_addr(t, kb16): - """Word index in the resident B tiles of the 8 FP4 bytes of token t's 16-value block kb16.""" - kt = kb16 >> cutlass.Int32(4) - j = kt >> cutlass.Int32(2) - n = (kt & cutlass.Int32(3)) * cutlass.Int32(QS) + t - cc = (kb16 >> cutlass.Int32(1)) & cutlass.Int32(7) - return ( - j * cutlass.Int32(NUM_BYTES_B1) - + n * cutlass.Int32(KT_BYTES) - + (cc ^ (n & cutlass.Int32(7))) * cutlass.Int32(16) - + (kb16 & cutlass.Int32(1)) * cutlass.Int32(8) - ) // cutlass.Int32(4) - - -@cute.jit -def _build_b1( - tidx, - warp, - lane, - base_row, - base_col, - x32: cutlass.Array, - xsf32: cutlass.Array, - sB1: cutlass.Array, - b1_ready: cutlass.Array, -): - """B (the M tokens' whole K) into the resident tiles from the NVFP4 activation, its scales to TMEM.""" - sB1_32 = cutlass.Array( - sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 - ) - for r in cutlass.range_constexpr((M_MAX * (H // 16) + EPI_THREADS - 1) // EPI_THREADS): - b = tidx + cutlass.Int32(r * EPI_THREADS) - if b < cutlass.Int32(M_MAX * (H // 16)): - t = b // cutlass.Int32(H // 16) - kb = b - t * cutlass.Int32(H // 16) - v2 = x32.load(idx=b * cutlass.Int32(2), vector_size=2, alignment=8) - sB1_32.store( - (cutlass.Int32(v2[0]), cutlass.Int32(v2[1])), idx=_b1_addr(t, kb), alignment=8 - ) - _sfb1_to_tmem(warp, lane, base_row, base_col, xsf32, b1_ready) - - -@cute.jit -def _build_b1_fused( - tidx, - warp, - lane, - base_row, - base_col, - xb32: cutlass.Array, - sfsc: cutlass.Array, - sB1: cutlass.Array, - s_xsf: cutlass.Array, - b1_ready: cutlass.Array, -): - """_build_b1 from the bf16 hidden states, quantized exactly as fp4_quantize(x, sf_scale, 16).""" - sB1_32 = cutlass.Array( - sB1.data_ptr(0), shape=(S1 * NUM_BYTES_B1 // 4,), dtype=cutlass.Int32, alignment=16 - ) - s_xsf8 = cutlass.Array( - s_xsf.data_ptr(0), shape=(M_MAX * (H // 16),), dtype=cutlass.Int8, alignment=16 - ) - sc = cutlass.Float32(sfsc.load(idx=0)) - rsc = _rcp(sc) - rsix = _rcp(cutlass.Float32(6.0)) - NBLK = M_MAX * (H // 16) // EPI_THREADS - QB = ( - 3 if NBLK % 3 == 0 else (2 if NBLK % 2 == 0 else 1) - ) # blocks whose loads are in flight together - # a rolled loop over groups of QB blocks (a fully unrolled one is ~50 KB of SASS at M = 8: instruction fetch bound) - for rg in cutlass.range(NBLK // QB, unroll=1): - wall = [] - for rr in cutlass.range_constexpr(QB): - bq = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) - for h in cutlass.range_constexpr(2): - wall.append( - xb32.load( - idx=bq * cutlass.Int32(8) + cutlass.Int32(4 * h), - vector_size=4, - alignment=16, - ) - ) - for rr in cutlass.range_constexpr(QB): - b = tidx + (rg * cutlass.Int32(QB) + cutlass.Int32(rr)) * cutlass.Int32(EPI_THREADS) - t = b // cutlass.Int32(H // 16) - kb = b - t * cutlass.Int32(H // 16) - w = [] - for h in cutlass.range_constexpr(2): - v4 = wall[rr * 2 + h] - for z in cutlass.range_constexpr(4): - w.append(cutlass.Int32(v4[z])) - xs = [] - for z in cutlass.range_constexpr(8): - xs.append((w[z] << cutlass.Int32(16)).bitcast(cutlass.Float32)) - xs.append((w[z] & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) - vmax = cute.math.abs(xs[0]) - for z in cutlass.range_constexpr(1, 16): - vmax = cute.arch.fmax(vmax, cute.math.abs(xs[z])) - sfb = _e4m3_bits(_mul_rn(sc, _mul_rn(vmax, rsix))) - osc = cutlass.select_( - vmax != cutlass.Float32(0.0), - _rcp(_mul_rn(_e4m3_to_f32(sfb), rsc)), - cutlass.Float32(0.0), - ) - by = [ - _e2m1x2(_mul_rn(xs[2 * z + 1], osc), _mul_rn(xs[2 * z], osc)) & cutlass.Int32(0xFF) - for z in range(8) - ] - lo = ( - by[0] - | (by[1] << cutlass.Int32(8)) - | (by[2] << cutlass.Int32(16)) - | (by[3] << cutlass.Int32(24)) - ) - hi = ( - by[4] - | (by[5] << cutlass.Int32(8)) - | (by[6] << cutlass.Int32(16)) - | (by[7] << cutlass.Int32(24)) - ) - sB1_32.store((lo, hi), idx=_b1_addr(t, kb), alignment=8) - s_xsf8.store(cutlass.Int8(sfb), idx=b) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) - _sfb1_to_tmem(warp, lane, base_row, base_col, s_xsf, b1_ready) - - -@cute.jit -def _swiglu_quant(xv, gate_raw, is_up, a, c, alpha, beta, lim): - """TRTLLM-Gen's FC1 epilogue (the swiGlu Bmm E2m1 output), operation for operation: lane value xv is the up row, - gate_raw its gate row (fp32 sums). Returns (q, the scale's e4m3 bits) of the 16-row NVFP4 block (lanes r, r ^ 16 - and r ^ 8 cover it; amax over the warp's up lanes of a half). a = fc31_alpha, c = fc31_scale_c.""" - neg_max = cutlass.Float32(-3.40282346638528859812e38) - gc = cutlass.select_(gate_raw > lim, lim, gate_raw) - gc = cutlass.select_(gate_raw < neg_max, neg_max, gc) - uc = cutlass.select_(xv > lim, lim, xv) - uc = cutlass.select_(xv < -lim, -lim, uc) - c1 = _mul_rn(a, _mul_rn(alpha, cutlass.Float32(1.4426950216293334961))) - sgm = _rcp(_add_rn(_ex2(-_mul_rn(gc, c1)), cutlass.Float32(1.0))) - gs = _mul_rn(gc, sgm) - h = _mul_rn(_fma_rn(uc, a, _mul_rn(a, beta)), gs) - h = cutlass.select_(is_up, h, cutlass.Float32(0.0)) - amax = cute.arch.fmax(_mul_rn(cute.math.abs(h), c), cutlass.Float32(0.0)) - for off in (16, 8, 4, 2, 1): - amax = cute.arch.fmax(amax, cute.arch.shuffle_sync_bfly(amax, off)) - sfb = _e4m3_bits(_mul_rn(amax, cutlass.Float32(0.16666667163372039795))) - sff = _e4m3_to_f32(sfb) - osc = cutlass.select_(sff != cutlass.Float32(0.0), _rcp(sff), cutlass.Float32(0.0)) - return _mul_rn(_mul_rn(h, osc), c), sfb - - -@cute.jit -def _plan_range(s_plan: cutlass.Array, p): - w = s_plan.load(idx=p) - return w & cutlass.Int32(0xFFFF), w >> cutlass.Int32(16) - - -@cute.jit -def _sh_issue( - lane, - sid, - g, - s_big: cutlass.Array, - sh_full: cutlass.Array, - tma_gu, - tma_dn, - gusf: cutlass.Array, - dnsf: cutlass.Array, -): - """Stage g of the shared CTA's stream into ring slot g % SH_NST (warp 0). g < SH_GT: gate_up k-tile g of both - slices: quadrant Q = 2 q + u (q slice, u 0 gate / 1 up) = rows 32 sid .. (+ 768 for up) at k-tile SH_GT q + g, its - raw scale atom at raw slot Q. Else down tile j = (g - SH_GT) / 6 (rows 256 sid + 128 j ..), k-tile - (g - SH_GT) % 6.""" - st = g % cutlass.Int32(SH_NST) - bar = sh_full.subview(st) - if cutlass.const_expr(SH_PRETILE): - if lane == 0: - src = (sid * cutlass.Int32(SH_STAGES) + g) * cutlass.Int32(SH_STAGE_BYTES) - prims.mbarrier_arrive_expect_tx(bar, SH_STAGE_BYTES) - prims.cp_async_bulk_shared_cluster_global( - s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), gusf.subview(src), bar, 16384 - ) - prims.cp_async_bulk_shared_cluster_global( - s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), - gusf.subview(src + cutlass.Int32(16384)), - bar, - 512, - ) - cute.arch.sync_warp() - else: - if g < cutlass.Int32(SH_GT): - if lane == 0: - prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 4 * 512) - cute.arch.sync_warp() - if lane < 4: - q = lane >> cutlass.Int32(1) - u = lane & cutlass.Int32(1) - row0 = sid * cutlass.Int32(32) + u * cutlass.Int32(SH_I) - kt = q * cutlass.Int32(SH_GT) + g - prims.cp_async_bulk_tensor_shared_cta_global( - s_big.subview(st * cutlass.Int32(NUM_BYTES_A) + lane * cutlass.Int32(4096)), - tma_gu, - (kt * cutlass.Int32(128), row0), - bar, - l2_cache_hint=EVICT_FIRST, - ) - prims.cp_async_bulk_shared_cluster_global( - s_big.subview( - cutlass.Int32(_SB_RAW) - + st * cutlass.Int32(2048) - + lane * cutlass.Int32(512) - ), - gusf.subview( - ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_KB // 4) + kt) - * cutlass.Int32(512) - ), - bar, - 512, - ) - else: - gd = g - cutlass.Int32(SH_GT) - jt = gd // cutlass.Int32(SH_DKT) - kt = gd - jt * cutlass.Int32(SH_DKT) - row0 = sid * cutlass.Int32(256) + jt * cutlass.Int32(128) - if lane == 0: - prims.mbarrier_arrive_expect_tx(bar, NUM_BYTES_A + 512) - prims.cp_async_bulk_tensor_shared_cta_global( - s_big.subview(st * cutlass.Int32(NUM_BYTES_A)), - tma_dn, - (kt * cutlass.Int32(128), row0), - bar, - l2_cache_hint=EVICT_FIRST, - ) - prims.cp_async_bulk_shared_cluster_global( - s_big.subview(cutlass.Int32(_SB_RAW) + st * cutlass.Int32(2048)), - dnsf.subview( - ((row0 >> cutlass.Int32(7)) * cutlass.Int32(SH_DKB // 4) + kt) - * cutlass.Int32(512) - ), - bar, - 512, - ) - cute.arch.sync_warp() - - -@dsl_user_op -def _red_release_add(addr, val, *, loc=None, ip=None): - _asm( - None, - [cutlass.Int64(addr), cutlass.Int32(val)], - "red.release.gpu.global.add.u32 [$0], $1;", - "l,r", - True, - loc, - ip, - ) - - -@dsl_user_op -def _ld_acquire_u32(addr, *, loc=None, ip=None): - return cutlass.Int32( - _asm( - _T.i32(), - [cutlass.Int64(addr)], - "ld.acquire.gpu.global.u32 $0, [$1];", - "=r,l", - True, - loc, - ip, - ) - ) - - -@dsl_user_op -def _fence_proxy_async_global(*, loc=None, ip=None): - _asm(None, [], "fence.proxy.async.global;", "", True, loc, ip) - - -@cute.jit -def _sfa_to_tmem(sSFA: cutlass.Array, stage, base_row, base_col): - """(SFA_CP, elected MMA thread) the stage's 4 scale atoms smem -> TMEM (what the splicer warp does otherwise).""" - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - sd0 = prims.Tcgen05SmemDesc.build( - sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - for kb in cutlass.range_constexpr(NUM_KB): - prims.tcgen05_cp( - s2t_shape, - _tptr(base_row, base_col + SFA1_COL + stage * cutlass.Int32(16) + kb * 4), - sd0 + (kb * 512 >> 4), - multicast=s2t_mc, - ) - - -@cute.jit -def _sf1_cp( - lane, stage, sf_line, c, sSFA: cutlass.Array, s1w: cutlass.Array, ab_full: cutlass.Array -): - """(SFA_CP) FC1 stage scale atoms: word [r][q] of atom kb <- raw atom 4 q + kb (raw line sf_line = 4 x atom) - word [r][c % 4] (what the splicer copies out of the raw box), 4-byte cp.async by the warp's 32 lanes; their - arrivals (noinc) complete the stage's full barrier together with the A boxes' bytes.""" - src0 = ( - (sf_line >> cutlass.Int32(2)) * cutlass.Int32(128) - + lane * cutlass.Int32(4) - + (c & cutlass.Int32(3)) - ) - dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) - for q in cutlass.range_constexpr(C1): - for kb in cutlass.range_constexpr(NUM_KB): - prims.cp_async_shared_global( - sSFA.subview(dst0 + cutlass.Int32(kb * 128 + q)), - s1w.subview(src0 + cutlass.Int32((4 * q + kb) * 128)), - 4, - "ca", - ) - prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) - cute.arch.sync_warp() - - -@cute.jit -def _sf2_cp( - lane, - stage, - g, - blk, - t2: cutlass.Constexpr, - n_here, - s_el: cutlass.Array, - sSFA: cutlass.Array, - s2w: cutlass.Array, - ab_full: cutlass.Array, -): - """(SFA_CP) FC2 k-tile t2 scale atoms: word [r][jj] of atom kb <- raw atom (rows / 128, K block 4 t2 + kb) of slot - jj's expert, word [r][blk % 4].""" - dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) - for jj in cutlass.range_constexpr(4): - if cutlass.Int32(jj) < n_here: - e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) - src0 = ( - ( - (e * cutlass.Int32(SF2_RB) + (blk >> cutlass.Int32(2))) * cutlass.Int32(SF2_KA) - + cutlass.Int32(4 * t2) - ) - * cutlass.Int32(128) - + lane * cutlass.Int32(4) - + (blk & cutlass.Int32(3)) - ) - # (a partial last k-tile has no atoms past I_TP) - for kb in cutlass.range_constexpr(_nkb2(t2)): - prims.cp_async_shared_global( - sSFA.subview(dst0 + cutlass.Int32(kb * 128 + jj)), - s2w.subview(src0 + cutlass.Int32(kb * 128)), - 4, - "ca", - ) - prims.cp_async_mbarrier_arrive(ab_full.subview(stage), noinc=True) - cute.arch.sync_warp() - - -@cute.jit -def _shared_cta( - tidx, - warp, - lane, - sid, - s_big: cutlass.Array, - sh_bar: cutlass.Array, - epochs: cutlass.Array, - tdbg: cutlass.Array, - bx, - xb32: cutlass.Array, - tma_gu, - tma_dn, - gusf: cutlass.Array, - dnsf: cutlass.Array, - hx: cutlass.Array, - dnx: cutlass.Array, - arp: cutlass.Array, - tmem_ptr_i32: cutlass.Array, - gu8: cutlass.Array, - dn8: cutlass.Array, -): - """Shared CTA sid (see SHARED): tcgen05 MXF8 (E4M3 x E4M3, UE8M0 per 32 K), m3_mx_gemv's tile order. gate_up: one - 128-row MMA tile = 2 K slices x (32 gate + 32 up rows), B rows 8 q + t = token t's slice-q k-tile; down: 2 tiles - of 128 rows, full K. Warps: 0 weights (TMA), 1 TMEM + scale splice + MMA, 4-7 epilogue; 2.. activation quant.""" - big32 = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Int32, alignment=16) - bigf = cutlass.Array(s_big.data_ptr(0), shape=(BIG // 4,), dtype=cutlass.Float32, alignment=16) - big8 = cutlass.Array(s_big.data_ptr(0), shape=(BIG,), dtype=cutlass.Int8, alignment=16) - sh_full = sh_bar # [0, SH_NST): stage full - sh_empty = cutlass.Array(sh_bar.data_ptr(8), shape=(8,), dtype=cutlass.Int64, alignment=8) - b_gu_rdy = sh_bar.subview(16) - b_dn_rdy = sh_bar.subview(17) - acc_gu = sh_bar.subview(18) - acc_dn = sh_bar.subview(19) - if warp == 0: - prims.prefetch_tensormap(tma_gu) - prims.prefetch_tensormap(tma_dn) - if lane < SH_NST: - prims.mbarrier_init(sh_full.subview(lane), 1) - prims.mbarrier_init( - sh_empty.subview(lane), 33 - ) # the MMAs' commit + each splicer lane's reads done - if lane == 0: - prims.mbarrier_init(b_gu_rdy, 1) - prims.mbarrier_init(b_dn_rdy, 1) - prims.mbarrier_init(acc_gu, 1) - prims.mbarrier_init(acc_dn, 1) - if warp == 1: - prims.tcgen05_alloc(tmem_ptr_i32, SH_TMEM) - prims.tcgen05_relinquish_alloc_permit() - prims.fence_mbarrier_init() - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - prims.barrier_cta_sync(0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - tmem_raw = tmem_ptr_i32.load() - base_col = tmem_raw & 0xFFFF - base_row = tmem_raw >> 16 - if warp == 0: - for g in cutlass.range_constexpr(SH_NST): # layer constants: before the grid dependency - _sh_issue(lane, sid, cutlass.Int32(g), s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) - # the rest of this CTA's weights into L2 (64 gate / up rows of 6 KB, 256 down rows of 768 B) - if cutlass.const_expr(SH_PRETILE): - # this CTA's stages beyond the ring, 4 per lane - for r in cutlass.range_constexpr(2): - pg = cutlass.Int32(SH_NST) + lane * cutlass.Int32(2) + cutlass.Int32(r) - if pg < cutlass.Int32(SH_STAGES): - _l2_prefetch( - gusf.subview( - (sid * cutlass.Int32(SH_STAGES) + pg) * cutlass.Int32(SH_STAGE_BYTES) - ) - .data_ptr() - .toint(), - SH_STAGE_BYTES, - ) - else: - for r in cutlass.range_constexpr(2): - prw = sid * cutlass.Int32(32) + cutlass.Int32(r * SH_I) + lane - _l2_prefetch(gu8.subview(prw * cutlass.Int32(H)).data_ptr().toint(), H) - if lane == 0: - _l2_prefetch( - dn8.subview(sid * cutlass.Int32(256 * SH_I)).data_ptr().toint(), 256 * SH_I - ) - cute.arch.griddepcontrol_wait() - if cutlass.const_expr(not LATE_TRIGGER): - if arp.load(idx=0) == cutlass.Int64(0): - cute.arch.griddepcontrol_launch_dependents() - if tidx == 0: - _stamp(tdbg, bx, 1) - ep = epochs.load(idx=bx) - tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) - NQ = THREADS - 64 # activation quant threads (warps 2 ..) - if warp == 0: - for g in cutlass.range(SH_NST, SH_STAGES, unroll=1): - _wait( - sh_empty.subview(g % cutlass.Int32(SH_NST)), - ((g // cutlass.Int32(SH_NST)) - cutlass.Int32(1)) & cutlass.Int32(1), - ) - _sh_issue(lane, sid, g, s_big, sh_full, tma_gu, tma_dn, gusf, dnsf) - elif warp == 1: - # ---- MMA: gate_up (scale atoms spliced to the tile's row order), then down - idesc_gu = prims.Tcgen05MxInstrDesc.build( - a_dtype=cutlass.Float8E4M3FN, - b_dtype=cutlass.Float8E4M3FN, - scale_format=1, - n_dim=SH_NB, - m_dim=128, - ) - idesc_dn = prims.Tcgen05MxInstrDesc.build( - a_dtype=cutlass.Float8E4M3FN, - b_dtype=cutlass.Float8E4M3FN, - scale_format=1, - n_dim=8, - m_dim=128, - ) - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - desc_a0 = prims.Tcgen05SmemDesc.build( - s_big.subview(0), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - desc_bx = prims.Tcgen05SmemDesc.build( - s_big.subview(_SB_BX), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - desc_bh = prims.Tcgen05SmemDesc.build( - s_big.subview(_SB_BH), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - tcol = (base_row << 16) | base_col - sq = sid & cutlass.Int32(3) - _wait(b_gu_rdy, 0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - for g in cutlass.range(SH_STAGES, unroll=1): - st = g % cutlass.Int32(SH_NST) - if g == cutlass.Int32(SH_GT): - _wait(b_dn_rdy, 0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - _wait(sh_full.subview(st), (g // cutlass.Int32(SH_NST)) & cutlass.Int32(1)) - raw = cutlass.Int32(_SB_RAW // 4) + st * cutlass.Int32(512) - if (g < cutlass.Int32(SH_GT)) & cutlass.Boolean(not SH_PRETILE): - # MMA row 32 Q + i = quadrant Q's row i: word 4 i + Q <- raw atom Q's word 4 i + sid % 4 - w4 = [ - big32.load(idx=raw + cutlass.Int32(128 * qq) + lane * cutlass.Int32(4) + sq) - for qq in range(4) - ] - big32.store( - ( - cutlass.Int32(w4[0]), - cutlass.Int32(w4[1]), - cutlass.Int32(w4[2]), - cutlass.Int32(w4[3]), - ), - idx=cutlass.Int32(_SB_SPL // 4) - + st * cutlass.Int32(128) - + lane * cutlass.Int32(4), - alignment=16, - ) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - cute.arch.sync_warp() - prims.mbarrier_arrive( - sh_empty.subview(st) - ) # this lane's generic reads of the slot are done - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - if prims.elect_sync(): - if g < cutlass.Int32(SH_GT): - dsf = prims.Tcgen05SmemDesc.build( - s_big.subview(cutlass.Int32(_SB_SPL) + st * cutlass.Int32(512)), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - prims.tcgen05_cp( - s2t_shape, - _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), - dsf, - multicast=s2t_mc, - ) - da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) - db = desc_bx + cutlass.Int64(g * cutlass.Int32(SH_NB * 128 >> 4)) - for kb in cutlass.range_constexpr(4): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF8F6F4, - prims.CTAGroup.CTA_1, - cutlass.inttoptr(tcol + SH_ACC_GU, 6, cutlass.Float32), - da + 2 * kb, - db + 2 * kb, - idesc_gu.set_sf_ids(a_sf_id=kb, b_sf_id=kb), - (g > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), - cutlass.inttoptr( - tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 - ), - cutlass.inttoptr( - tcol + SH_SFB_GU + g * cutlass.Int32(4), 6, cutlass.Int32 - ), - ) - prims.tcgen05_commit(sh_empty.subview(st)) - if g == cutlass.Int32(SH_GT - 1): - prims.tcgen05_commit(acc_gu) - else: - gd = g - cutlass.Int32(SH_GT) - jt = gd // cutlass.Int32(SH_DKT) - kt = gd - jt * cutlass.Int32(SH_DKT) - dsf = prims.Tcgen05SmemDesc.build( - s_big.subview( - cutlass.Int32(_SB_SPL if SH_PRETILE else _SB_RAW) - + st * cutlass.Int32(512 if SH_PRETILE else 2048) - ), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - prims.tcgen05_cp( - s2t_shape, - _tptr(base_row, base_col + SH_SFA + st * cutlass.Int32(4)), - dsf, - multicast=s2t_mc, - ) - da = desc_a0 + cutlass.Int64(st * cutlass.Int32(NUM_BYTES_A >> 4)) - db = desc_bh + cutlass.Int64(kt * cutlass.Int32(1024 >> 4)) - for kb in cutlass.range_constexpr(4): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF8F6F4, - prims.CTAGroup.CTA_1, - cutlass.inttoptr( - tcol + SH_ACC_DN + jt * cutlass.Int32(8), 6, cutlass.Float32 - ), - da + 2 * kb, - db + 2 * kb, - idesc_dn.set_sf_ids(a_sf_id=kb, b_sf_id=kb), - (kt > cutlass.Int32(0)) | cutlass.Boolean(kb > 0), - cutlass.inttoptr( - tcol + SH_SFA + st * cutlass.Int32(4), 6, cutlass.Int32 - ), - cutlass.inttoptr( - tcol + SH_SFB_DN + kt * cutlass.Int32(4), 6, cutlass.Int32 - ), - ) - prims.tcgen05_commit(sh_empty.subview(st)) - if g == cutlass.Int32(SH_STAGES - 1): - prims.tcgen05_commit(acc_dn) - cute.arch.sync_warp() - if warp >= 2: - # ---- the activation, MXFP8 per 32-K block (native mxfp8_quantize), into B (row 8 q + t of stage tile i = - # token t's k-tile SH_GT q + i, 128-byte swizzle) and its codes - qt = tidx - cutlass.Int32(64) - for r in cutlass.range_constexpr((M_MAX * SH_KB + NQ - 1) // NQ): - q_ = qt + cutlass.Int32(r * NQ) - if q_ < cutlass.Int32(M_MAX * SH_KB): - t = q_ // cutlass.Int32(SH_KB) - b = q_ - t * cutlass.Int32(SH_KB) - wv = [ - xb32.load( - idx=t * cutlass.Int32(H // 2) - + b * cutlass.Int32(16) - + cutlass.Int32(4 * z), - vector_size=4, - alignment=16, - ) - for z in range(4) - ] - v = [] - for z in cutlass.range_constexpr(4): - for zz in cutlass.range_constexpr(4): - w = cutlass.Int32(wv[z][zz]) - v.append((w << cutlass.Int32(16)).bitcast(cutlass.Float32)) - v.append((w & cutlass.Int32(-65536)).bitcast(cutlass.Float32)) - amax = cutlass.Float32(0.0) - for z in cutlass.range_constexpr(32): - amax = cute.arch.fmax(amax, cute.math.abs(v[z])) - code, osc = _mx_osc(amax) - words = _mx_words(v, osc) - kt = b >> cutlass.Int32(2) - sl = kt // cutlass.Int32(SH_GT) - it = kt - sl * cutlass.Int32(SH_GT) - n = sl * cutlass.Int32(8) + t - c0 = (b & cutlass.Int32(3)) * cutlass.Int32(2) - rw = ( - cutlass.Int32(_SB_BX) + it * cutlass.Int32(SH_NB * 128) + n * cutlass.Int32(128) - ) >> cutlass.Int32(2) - big32.store( - (words[0], words[1], words[2], words[3]), - idx=rw + ((c0 ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), - alignment=16, - ) - big32.store( - (words[4], words[5], words[6], words[7]), - idx=rw - + (((c0 + cutlass.Int32(1)) ^ (n & cutlass.Int32(7))) << cutlass.Int32(2)), - alignment=16, - ) - big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_XC) + q_) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR, number_of_threads=NQ) - if warp >= 4: - if warp < 8: - qd = warp - cutlass.Int32(4) - # B scales of gate_up to TMEM: column of stage i, lane n = 8 q + t: token t's codes of k-tile SH_GT q + i - n = lane - t = n & cutlass.Int32(7) - ok = (n < cutlass.Int32(SH_NB)) & (t < cutlass.Int32(M_MAX)) - for i in cutlass.range_constexpr(SH_GT): - kt = (n >> cutlass.Int32(3)) * cutlass.Int32(SH_GT) + cutlass.Int32(i) - wsf = big32.load( - idx=( - cutlass.Int32(_SB_XC) - + cutlass.select_(ok, t, cutlass.Int32(0)) * cutlass.Int32(SH_KB) - + kt * cutlass.Int32(4) - ) - >> cutlass.Int32(2) - ) - prims.tcgen05_st( - "32x32b", - _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_GU + i * 4), - cutlass.select_(ok, cutlass.Int32(wsf), cutlass.Int32(0)), - ) - prims.tcgen05_wait(prims.Tcgen05Wait.STORE) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - et = tidx - cutlass.Int32(128) - if et == 0: - prims.mbarrier_arrive(b_gu_rdy) - _stamp(tdbg, bx, 54) - # ---- gate_up out: quadrant qd lane l = MMA row 32 qd + l; slice q = qd / 2 holds token t at column 8 q + t - _wait(acc_gu, 0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - if cutlass.const_expr(SH_SPLIT == 2): - av = prims.tcgen05_ld( - "32x32b", - _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_ACC_GU, cutlass.Float32), - num=16, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - if qd >= cutlass.Int32(2): - for tt in cutlass.range_constexpr(8): - bigf.store( - cutlass.Float32(av[8 + tt]), - idx=cutlass.Int32(_SB_XCH // 4) - + ((qd - cutlass.Int32(2)) * cutlass.Int32(32) + lane) - * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - if qd < cutlass.Int32(2): - # m3_mx_gemv's owner sum: (0 + rank 0) + rank 1, then bf16 - for tt in cutlass.range_constexpr(8): - tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) - tot = _add_rn( - tot, - bigf.load( - idx=cutlass.Int32(_SB_XCH // 4) - + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) - + cutlass.Int32(tt) - ), - ) - bigf.store( - _bf16_round(tot), - idx=cutlass.Int32(_SB_GU // 4) - + (qd * cutlass.Int32(32) + lane) * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - else: - # 16-row octants: warp qd = slice qd, lane l = (u = l / 16 gate / up, position l % 16); its tokens are - # accumulator columns 8 qd + t - av = prims.tcgen05_ld( - "32x32b", - _tptr( - base_row + qd * cutlass.Int32(32), - base_col + SH_ACC_GU + qd * cutlass.Int32(8), - cutlass.Float32, - ), - num=8, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - if qd >= cutlass.Int32(1): - for tt in cutlass.range_constexpr(8): - bigf.store( - cutlass.Float32(av[tt]), - idx=cutlass.Int32(_SB_XCH // 4) - + ((qd - cutlass.Int32(1)) * cutlass.Int32(32) + lane) - * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - if qd == cutlass.Int32(0): - # m3_mx_gemv's owner sum: (((0 + rank 0) + rank 1) + rank 2) + rank 3, then bf16 - for tt in cutlass.range_constexpr(8): - tot = _add_rn(cutlass.Float32(0.0), cutlass.Float32(av[tt])) - for s_ in cutlass.range_constexpr(1, SH_SPLIT): - tot = _add_rn( - tot, - bigf.load( - idx=cutlass.Int32(_SB_XCH // 4) - + (cutlass.Int32((s_ - 1) * 32) + lane) * cutlass.Int32(8) - + cutlass.Int32(tt) - ), - ) - bigf.store( - _bf16_round(tot), - idx=cutlass.Int32(_SB_GU // 4) - + lane * cutlass.Int32(8) - + cutlass.Int32(tt), - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - # h = silu_and_mul(gate, up) of positions 32 sid + p, published under the call's tag - for r in cutlass.range_constexpr((SH_ROWS * M_MAX + 127) // 128): - pq = et + cutlass.Int32(r * 128) - if pq < cutlass.Int32(SH_ROWS * M_MAX): - p = pq // cutlass.Int32(M_MAX) - tt = pq - p * cutlass.Int32(M_MAX) - ga = bigf.load(idx=cutlass.Int32(_SB_GU // 4) + p * cutlass.Int32(8) + tt) - up = bigf.load( - idx=cutlass.Int32(_SB_GU // 4) - + (cutlass.Int32(SH_ROWS) + p) * cutlass.Int32(8) - + tt - ) - q_a = cute.arch.fmin(cutlass.Float32(ga), cutlass.Float32(7.0)) - q_bu = _min_xorsign_abs(cutlass.Float32(up), cutlass.Float32(7.0)) - q_ub = _add_rn(q_bu, cutlass.Float32(1.0)) - q_e2 = _ex2_full( - _mul_rn( - _mul_rn(q_a, cutlass.Float32(-1.70200002193450927734)), - cutlass.Float32(1.4426950216293334961), - ) - ) - q_sg = _div_full(cutlass.Float32(1.0), _add_rn(q_e2, cutlass.Float32(1.0))) - q_h = _mul_rn(q_ub, _mul_rn(q_sg, q_a)) - _st_relaxed_u32( - hx.subview(tt * cutlass.Int32(SH_I) + sid * cutlass.Int32(SH_ROWS) + p) - .data_ptr() - .toint(), - (tag16 << cutlass.Int32(16)) | _bf16_bits(q_h), - ) - if et == 0: - _stamp(tdbg, bx, 50) - # every shared CTA's h: (token, K block) per thread, 32 tagged words, MXFP8 into B (h) and its codes - for rr in cutlass.range_constexpr((M_MAX * SH_DKB + 127) // 128): - q_ = et + cutlass.Int32(rr * 128) - if q_ < cutlass.Int32(M_MAX * SH_DKB): - t = q_ // cutlass.Int32(SH_DKB) - kb = q_ - t * cutlass.Int32(SH_DKB) - base = t * cutlass.Int32(SH_I) + kb * cutlass.Int32(32) - hv = [cutlass.Int32(0) for _ in range(32)] - pend = cutlass.Int32(1) - while pend != cutlass.Int32(0): - for q4 in cutlass.range_constexpr(8): - v4 = prims.load_ext( - hx.subview(base + cutlass.Int32(4 * q4)), - dtype=cutlass.Int32, - count=4, - order="relaxed", - scope="gpu", - ) - for z in cutlass.range_constexpr(4): - hv[4 * q4 + z] = cutlass.Int32(v4[z]) - okw = cutlass.Boolean(True) - for z in cutlass.range_constexpr(32): - okw = okw & ( - ((hv[z] >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) == tag16 - ) - pend = cutlass.select_(okw, cutlass.Int32(0), cutlass.Int32(1)) - vals = [ - (hv[z] << cutlass.Int32(16)).bitcast(cutlass.Float32) for z in range(32) - ] - amax = cutlass.Float32(0.0) - for z in cutlass.range_constexpr(32): - amax = cute.arch.fmax(amax, cute.math.abs(vals[z])) - code, osc = _mx_osc(amax) - words = _mx_words(vals, osc) - kt = kb >> cutlass.Int32(2) - c0 = (kb & cutlass.Int32(3)) * cutlass.Int32(2) - rw = ( - cutlass.Int32(_SB_BH) + kt * cutlass.Int32(1024) + t * cutlass.Int32(128) - ) >> cutlass.Int32(2) - big32.store( - (words[0], words[1], words[2], words[3]), - idx=rw + ((c0 ^ t) << cutlass.Int32(2)), - alignment=16, - ) - big32.store( - (words[4], words[5], words[6], words[7]), - idx=rw + (((c0 + cutlass.Int32(1)) ^ t) << cutlass.Int32(2)), - alignment=16, - ) - big8.store(cutlass.Int8(code), idx=cutlass.Int32(_SB_HC) + q_) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - okh = lane < cutlass.Int32(M_MAX) - for kt in cutlass.range_constexpr(SH_DKT): - wsf = big32.load( - idx=( - cutlass.Int32(_SB_HC) - + cutlass.select_(okh, lane, cutlass.Int32(0)) * cutlass.Int32(SH_DKB) - + cutlass.Int32(4 * kt) - ) - >> cutlass.Int32(2) - ) - prims.tcgen05_st( - "32x32b", - _tptr(base_row + qd * cutlass.Int32(32), base_col + SH_SFB_DN + kt * 4), - cutlass.select_(okh, cutlass.Int32(wsf), cutlass.Int32(0)), - ) - prims.tcgen05_wait(prims.Tcgen05Wait.STORE) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=SHX_BAR + 1, number_of_threads=128) - if et == 0: - prims.mbarrier_arrive(b_dn_rdy) - _stamp(tdbg, bx, 51) - # ---- down out: tile j row 32 qd + l -> y_sh row 256 sid + 128 j + 32 qd + l - _wait(acc_dn, 0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - for jt in cutlass.range_constexpr(2): - dv = prims.tcgen05_ld( - "32x32b", - _tptr( - base_row + qd * cutlass.Int32(32), - base_col + SH_ACC_DN + jt * 8, - cutlass.Float32, - ), - num=8, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - rowo = ( - sid * cutlass.Int32(256) - + cutlass.Int32(128 * jt) - + qd * cutlass.Int32(32) - + lane - ) - for tt in cutlass.range_constexpr(M_MAX): - _st_relaxed_u32( - dnx.subview(cutlass.Int32(tt * H) + rowo).data_ptr().toint(), - (tag16 << cutlass.Int32(16)) | _bf16_bits(cutlass.Float32(dv[tt])), - ) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - if et == 0: - _stamp(tdbg, bx, 52) - cute.arch.sync_warp() - prims.barrier_cta_sync(0) - if warp == 1: - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - prims.tcgen05_dealloc(cutlass.inttoptr(tmem_raw, 6, cutlass.Float32), SH_TMEM) - if tidx == 0: - epochs.store(ep + cutlass.Int32(1), idx=bx) - - -@cute.kernel -def m3_moe_kernel( - tma_a1_desc: cutlass.GridConstant[cuda.TensorMap], - tma_a2_desc: cutlass.GridConstant[cuda.TensorMap], - tma_s1_desc: cutlass.GridConstant[cuda.TensorMap], # w3_w1_weight_scale as 128-byte lines - tma_s2_desc: cutlass.GridConstant[cuda.TensorMap], # w2_weight_scale as 128-byte lines - tma_gu_desc: cutlass.GridConstant[ - cuda.TensorMap - ], # SHARED: gate_up e4m3 [1536, 6144], box 128 B x 32 rows - tma_dn_desc: cutlass.GridConstant[ - cuda.TensorMap - ], # SHARED: down e4m3 [6144, 768], box 128 B x 128 rows - logits: cutlass.Array, # f32 [M, 128] (not FUSED_ROUTER) - bias: cutlass.Array, # f32 [128] - x32: cutlass.Array, # int32 words of x_fp4 [M, 3072 B] (not FUSED_ROUTER) - xsf32: cutlass.Array, # int32 words of x_sf [M, 384 B] (not FUSED_ROUTER) - escal: cutlass.Array, # f32 [6, 128]: fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha - plan: cutlass.Array, # int32 [4 M + 1, NCTA, NG_CAP + 1] by G: per group i0 | i1 << 16, then b0 | nb << 16 - partials: cutlass.Array, # int32 bits of f32 [G_CAP * 48, M, 32]: split units' head sums - hbuf32: cutlass.Array, # int32 [G_CAP * M, 48, 4]: tagged intermediate records - epochs: cutlass.Array, # int32 [NCTA] - out: cutlass.Array, # bf16 [M, 6144] - route_dbg: cutlass.Array, # int32 [M, 8]: ids, weight bits (CTA 0) - tdbg: cutlass.Array, # int64 [NCTA, 32] timing stamps (TIMING builds) - xb32: cutlass.Array, # FUSED_ROUTER: int32 words (bf16 pairs) of the hidden states [M, 6144] - rw: cutlass.Array, # FUSED_ROUTER: f32 router weight [128, 6144] - sfsc: cutlass.Array, # FUSED_ROUTER: f32 [1] fc31_input_scale - lgbuf: cutlass.Array, # FUSED_ROUTER: int64 [LG_REPL, LG_STRIDE] logits: tag << 32 | f32 bits - guw: cutlass.Array, # SHARED: uint8 gate_up e4m3 [1536, 6144] - gusf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales - dnw: cutlass.Array, # SHARED: uint8 down e4m3 [6144, 768] - dnsf: cutlass.Array, # SHARED: uint8 its swizzled UE8M0 scales - hx: cutlass.Array, # SHARED: int32 [M, 768]: h as (call tag << 16 | bf16), tag 0 = never written - dnx: cutlass.Array, # SHARED: int32 [M, 6144]: y_sh as (call tag << 16 | bf16), tag 0 = never written - arp: cutlass.Array, # int64 [8]: all-reduce push (enable, mc, flags, rank, world, maxt); enable 0 = off - s1w: cutlass.Array, # SFA_CP: int32 words of w3_w1_weight_scale - s2w: cutlass.Array, # SFA_CP: int32 words of w2_weight_scale - img32: cutlass.Array, # B2IMG: int32 [NG_CAP, IMG_BYTES / 4] per-group FC2 B image (tile + scale atoms) - b2cnt: cutlass.Array, # B2IMG: int32 [NG_CAP * 8 + 8] per-group unit counters (index 8 g) -): - tidx, _, _ = cute.arch.thread_idx() - bx, _, _ = cute.arch.block_idx() - warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - lane = tidx % 32 - if tidx == 0: - _stamp(tdbg, bx, 0) - if cutlass.const_expr(TIMING): - if tidx == 0: - tdbg.store( - cutlass.Int64(_smid()), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(63) - ) - - s_big = cutlass.Array(cutlass.Int8, BIG, space=cutlass.AddressSpace.smem, alignment=1024) - sA = cutlass.Array( - s_big.data_ptr(0), shape=(NUM_BYTES_A * STAGES,), dtype=cutlass.Int8, alignment=1024 - ) - sSFAraw = cutlass.Array( - s_big.data_ptr(_OFF_SFAR), - shape=(max(1, _RAW_STAGE // 4) * STAGES,), - dtype=cutlass.Int32, - alignment=1024, - ) - sSFA = cutlass.Array( - s_big.data_ptr(_OFF_SFA), - shape=(NUM_WORDS_SFA * STAGES,), - dtype=cutlass.Int32, - alignment=1024, - ) - sB1 = cutlass.Array( - s_big.data_ptr(_OFF_B1), shape=(NUM_BYTES_B1 * S1,), dtype=cutlass.Int8, alignment=1024 - ) - sB2 = cutlass.Array( - s_big.data_ptr(_OFF_B2), shape=(2 * KT2 * NUM_BYTES_B2,), dtype=cutlass.Int8, alignment=1024 - ) - sSFB2 = cutlass.Array( - s_big.data_ptr(_OFF_SFB2), - shape=(2 * KT2 * NUM_KB * 128,), - dtype=cutlass.Int32, - alignment=1024, - ) - ys = cutlass.Array( - cutlass.Float32, NB2_CAP * M_MAX * TOP_K * R2, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_part = cutlass.Array( - cutlass.Float32, 2 * M_MAX * C1 * R1, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_bias = cutlass.Array(cutlass.Float32, E_LOCAL, space=cutlass.AddressSpace.smem, alignment=16) - s_red = cutlass.Array(cutlass.Float32, 9 * M_MAX, space=cutlass.AddressSpace.smem, alignment=16) - s_xsf = cutlass.Array( - cutlass.Int32, - M_MAX * (H // 64) if FUSED_ROUTER else 4, - space=cutlass.AddressSpace.smem, - alignment=16, - ) - ab_full = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) - ab_empty = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) - sfa_ready = cutlass.Array(cutlass.Int64, STAGES, space=cutlass.AddressSpace.smem, alignment=8) - acc_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - acc_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - acc2_full = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - acc2_empty = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - b2_ready = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - b2_free = cutlass.Array(cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8) - # this CTA's FC1 share of group g is out: one barrier per group parity (a single barrier could complete twice - # while warp 5 still works on an earlier group, and its parity wait would then miss them) - fc1_done = cutlass.Array(cutlass.Int64, 4, space=cutlass.AddressSpace.smem, alignment=8) - b1_ready = cutlass.Array(cutlass.Int64, 1, space=cutlass.AddressSpace.smem, alignment=8) - x_bar = cutlass.Array( - cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 - ) # x landed, x consumed - img_bar = cutlass.Array( - cutlass.Int64, 2, space=cutlass.AddressSpace.smem, alignment=8 - ) # (B2IMG) image landed - tmem_ptr_i32 = cutlass.Array(cutlass.Int32, 1, space=cutlass.AddressSpace.smem) - s_el = cutlass.Array( - cutlass.Int32, G_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16 - ) # slots, then G - s_kpos = cutlass.Array( - cutlass.Int32, G_CAP * M_MAX, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_w = cutlass.Array( - cutlass.Float32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_ids = cutlass.Array( - cutlass.Int32, M_MAX * TOP_K, space=cutlass.AddressSpace.smem, alignment=16 - ) - s_plan = cutlass.Array(cutlass.Int32, NG_CAP + 1, space=cutlass.AddressSpace.smem, alignment=16) - # this CTA's plan rows for every G (static data: loaded before the grid dependency) - s_misc = cutlass.Array(cutlass.Int32, 4, space=cutlass.AddressSpace.smem, alignment=16) - s_plan_all = cutlass.Array( - cutlass.Int32, (G_CAP + 1) * (NG_CAP + 1), space=cutlass.AddressSpace.smem, alignment=16 - ) - s_esc = cutlass.Array( - cutlass.Float32, 8 * G_CAP, space=cutlass.AddressSpace.smem, alignment=32 - ) # slot s at 8 s - - is_routed = cutlass.Boolean(True) - if cutlass.const_expr(SHARED): - sh_full = cutlass.Array(cutlass.Int64, 32, space=cutlass.AddressSpace.smem, alignment=8) - is_routed = bx < cutlass.Int32(NCTA) - if bx >= cutlass.Int32(NCTA): - _shared_cta( - tidx, - warp, - lane, - bx - cutlass.Int32(NCTA), - s_big, - sh_full, - epochs, - tdbg, - bx, - xb32, - tma_gu_desc.get_ptr(), - tma_dn_desc.get_ptr(), - gusf, - dnsf, - hx, - dnx, - arp, - tmem_ptr_i32, - guw, - dnw, - ) - if is_routed: - # ---- before the grid dependency - if warp == 0: - if tidx < STAGES: - prims.mbarrier_init(ab_full.subview(tidx), 33 if SFA_CP else 1) - prims.mbarrier_init(ab_empty.subview(tidx), 1) - prims.mbarrier_init(sfa_ready.subview(tidx), 1) - if tidx < 2: - prims.mbarrier_init(acc_full.subview(tidx), 1) - prims.mbarrier_init(acc_empty.subview(tidx), EPI_THREADS) - prims.mbarrier_init(acc2_full.subview(tidx), 1) - prims.mbarrier_init(acc2_empty.subview(tidx), EPI_THREADS) - prims.mbarrier_init(b2_ready.subview(tidx), 1) - prims.mbarrier_init(b2_free.subview(tidx), EPI_THREADS) - if cutlass.const_expr(B2IMG): - prims.mbarrier_init(img_bar.subview(tidx), 1) - if tidx < 4: - prims.mbarrier_init(fc1_done.subview(tidx), EPI_THREADS) - if tidx == 0: - prims.mbarrier_init(b1_ready.subview(0), EPI_THREADS) - if cutlass.const_expr(XSTAGE): - prims.mbarrier_init(x_bar.subview(0), 1) - prims.mbarrier_init(x_bar.subview(1), EPI_THREADS) - if warp == MMA_WARP: - prims.tcgen05_alloc(tmem_ptr_i32, TMEM_COLS) - prims.tcgen05_relinquish_alloc_permit() - if warp == W_WARP: - prims.prefetch_tensormap(tma_a1_desc.get_ptr()) - prims.prefetch_tensormap(tma_a2_desc.get_ptr()) - prims.prefetch_tensormap(tma_s1_desc.get_ptr()) - prims.prefetch_tensormap(tma_s2_desc.get_ptr()) - for i in cutlass.range_constexpr(E_LOCAL // 32): - s_bias.store(cutlass.Float32(bias.load(idx=lane + i * 32)), idx=lane + i * 32) - for i in cutlass.range_constexpr(((G_CAP + 1) * (NG_CAP + 1) + 31) // 32): - q = lane + cutlass.Int32(32 * i) - if q < cutlass.Int32((G_CAP + 1) * (NG_CAP + 1)): - gq = q // cutlass.Int32(NG_CAP + 1) - s_plan_all.store( - plan.load( - idx=(gq * cutlass.Int32(NCTA) + bx) * cutlass.Int32(NG_CAP + 1) - + q - - gq * cutlass.Int32(NG_CAP + 1) - ), - idx=q, - ) - rwv = [] - if cutlass.const_expr(FUSED_ROUTER): - tg = cutlass.select_(tidx < cutlass.Int32(GEMV_THREADS), tidx, cutlass.Int32(0)) - for rg in cutlass.range_constexpr(NGEMV): - eg = bx + cutlass.Int32(rg * NCTA) - rbx = cutlass.select_(eg < cutlass.Int32(NROUTE), eg, cutlass.Int32(0)) - rwv.append( - [ - rw.load( - idx=rbx * cutlass.Int32(H) - + tg * cutlass.Int32(4) - + cutlass.Int32(1024 * j), - vector_size=4, - alignment=16, - ) - for j in range(H // 1024) - ] - ) - if cutlass.const_expr(TIMING): - if tidx == 0: - s_el.store(cutlass.Int32(-7), idx=G_CAP) - prims.fence_mbarrier_init() - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - prims.barrier_cta_sync(0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - tmem_raw = tmem_ptr_i32.load() - base_col = tmem_raw & 0xFFFF - base_row = tmem_raw >> 16 - - cute.arch.griddepcontrol_wait() - if cutlass.const_expr(XSTAGE): - if warp == 0: - if lane == 0: - prims.mbarrier_arrive_expect_tx(x_bar.subview(0), M_MAX * H * 2) - cute.arch.sync_warp() - if lane < M_MAX: - prims.cp_async_bulk_shared_cluster_global( - s_big.subview(lane * cutlass.Int32(H * 2)), - xb32.subview(lane * cutlass.Int32(H // 2)), - x_bar.subview(0), - H * 2, - ) - ar_on = arp.load(idx=0) != cutlass.Int64(0) - ar_off = arp.load(idx=0) == cutlass.Int64(0) - if cutlass.const_expr(not LATE_TRIGGER): - if ar_off: - cute.arch.griddepcontrol_launch_dependents() - if cutlass.const_expr(EARLY_AR): - if ar_on: - cute.arch.griddepcontrol_launch_dependents() - if tidx == 0: - _stamp(tdbg, bx, 1) - - ep = epochs.load(idx=bx) - tagv = ((ep & cutlass.Int32(0x7F)) | cutlass.Int32(0x80)) << cutlass.Int32(24) - tag16 = (ep & cutlass.Int32(0x7FFF)) | cutlass.Int32(0x8000) - ltag = cutlass.Int64(ep + cutlass.Int32(1)) << cutlass.Int64(32) - if cutlass.const_expr(FUSED_ROUTER): - # ---- router GEMV (experts bx, bx + NCTA, .., threads 0..255), m3_router_gemv's order exactly - for rg in cutlass.range_constexpr(NGEMV): - eg = bx + cutlass.Int32(rg * NCTA) - if (eg < cutlass.Int32(NROUTE)) & (tidx < cutlass.Int32(GEMV_THREADS)): - accs = [cutlass.Float32(0.0) for _ in range(M_MAX)] - xsrc = xb32 - if cutlass.const_expr(GEMV_SMEM and XSTAGE): - _wait(x_bar.subview(0), 0) - xsrc = cutlass.Array( - s_big.data_ptr(0), - shape=(M_MAX * H // 2,), - dtype=cutlass.Int32, - alignment=16, - ) - for j in cutlass.range_constexpr(H // 1024): - for t in cutlass.range_constexpr(M_MAX): - xw = xsrc.load( - idx=cutlass.Int32(t * (H // 2) + 512 * j) + tidx * cutlass.Int32(2), - vector_size=2, - alignment=8, - ) - x0 = cutlass.Int32(xw[0]) - x1 = cutlass.Int32(xw[1]) - a = accs[t] - a = _fma_rn( - (x0 << cutlass.Int32(16)).bitcast(cutlass.Float32), - cutlass.Float32(rwv[rg][j][0]), - a, - ) - a = _fma_rn( - (x0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), - cutlass.Float32(rwv[rg][j][1]), - a, - ) - a = _fma_rn( - (x1 << cutlass.Int32(16)).bitcast(cutlass.Float32), - cutlass.Float32(rwv[rg][j][2]), - a, - ) - a = _fma_rn( - (x1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), - cutlass.Float32(rwv[rg][j][3]), - a, - ) - accs[t] = a - for t in cutlass.range_constexpr(M_MAX): - v = accs[t] - for off in (16, 8, 4, 2, 1): - v = _add_rn(v, cute.arch.shuffle_sync_bfly(v, off)) - if lane == 0: - s_red.store(v, idx=warp * cutlass.Int32(M_MAX) + cutlass.Int32(t)) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) - if tidx < M_MAX: - sacc = cutlass.Float32(0.0) - for w in cutlass.range_constexpr(8): - sacc = _add_rn(sacc, s_red.load(idx=cutlass.Int32(w * M_MAX) + tidx)) - s_red.store(sacc, idx=cutlass.Int32(8 * M_MAX) + tidx) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS) - if tidx < LG_REPL: - for t in cutlass.range_constexpr(M_MAX): - lw = ltag | ( - cutlass.Int64( - cutlass.Float32(s_red.load(idx=8 * M_MAX + t)).bitcast( - cutlass.Int32 - ) - ) - & cutlass.Int64(0xFFFFFFFF) - ) - _st_relaxed_u64( - lgbuf.subview( - tidx * cutlass.Int32(LG_STRIDE) - + cutlass.Int32(t * E_LOCAL) - + eg - ) - .data_ptr() - .toint(), - lw, - ) - if cutlass.const_expr(NGEMV > 1): - cute.arch.sync_warp() - cute.arch.barrier( - barrier_id=GEMV_BAR, number_of_threads=GEMV_THREADS - ) # s_red reused next - if tidx == 0: - _stamp(tdbg, bx, 21) - - # ---- routing (warps 4..7, token t on warp 4 + t % 4); the epilogue warps build B meanwhile - if ((warp >= W_WARP) & (warp < cutlass.Int32(W_WARP + NRW0))) | ( - (warp >= cutlass.Int32(8)) & (warp < cutlass.Int32(8 + RW_EXTRA)) - ): - rw_i = cutlass.select_( - warp < cutlass.Int32(8), warp - cutlass.Int32(W_WARP), warp - cutlass.Int32(4) - ) - # ILP: poll every token's logits first, then run the independent top-k chains back to back - lgv_all = [] - t_all = [] - for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): - t = rw_i + cutlass.Int32(NRW * tr) - if cutlass.const_expr(M_MAX % NRW != 0): - t = cutlass.select_(t < cutlass.Int32(M_MAX), t, cutlass.Int32(M_MAX - 1)) - t_all.append(t) - lgv = [] - if cutlass.const_expr(FUSED_ROUTER): - want = ltag >> cutlass.Int64(32) - rbase = ( - (bx % cutlass.Int32(LG_REPL)) * cutlass.Int32(LG_STRIDE) - + t * cutlass.Int32(E_LOCAL) - + lane - ) - done = cutlass.Int32(0) - res = [cutlass.Int64(0) for _ in range(4)] - while done == cutlass.Int32(0): - ok = cutlass.Boolean(True) - for z in cutlass.range_constexpr(4): - res[z] = _ld_relaxed_u64( - lgbuf.subview(cutlass.Int32(z * 32) + rbase).data_ptr().toint() - ) - ok = ok & ((res[z] >> cutlass.Int64(32)) == want) - if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32( - -1 - ): - done = cutlass.Int32(1) - else: - _nanosleep(POLL_NS) - for z in cutlass.range_constexpr(4): - lgv.append(cutlass.Int32(res[z]).bitcast(cutlass.Float32)) - if (lane == 0) & (warp == cutlass.Int32(W_WARP)): - _stamp(tdbg, bx, 22) - else: - for i in cutlass.range_constexpr(4): - lgv.append( - cutlass.Float32( - logits.load( - idx=t * cutlass.Int32(E_LOCAL) + cutlass.Int32(i * 32) + lane - ) - ) - ) - lgv_all.append(lgv) - if (lane == 0) & (warp == cutlass.Int32(W_WARP)): - _stamp(tdbg, bx, 61) - res_all = [] - for tr in cutlass.range_constexpr((M_MAX + NRW - 1) // NRW): - t = t_all[tr] - lgv = lgv_all[tr] - own = (rw_i + cutlass.Int32(NRW * tr)) < cutlass.Int32(M_MAX) - # top-4 (keys sigmoid + bias, ties to the lower id) - keys = [] - for i in cutlass.range_constexpr(4): - sg = _fma_rn( - cutlass.Float32(0.5), - cute.math.tanh(_mul_rn(cutlass.Float32(0.5), lgv[i])), - cutlass.Float32(0.5), - ) - keys.append(_add_rn(sg, s_bias.load(idx=i * 32 + lane))) - hi = [_twiddle(keys[i]) for i in range(4)] - taken = [cutlass.Boolean(False) for _ in range(4)] - ids_t = [] - kt_t = [] - for k in cutlass.range_constexpr(TOP_K): - bh = cutlass.Int32(0) - bl = cutlass.Int32(0) - for i in cutlass.range_constexpr(4): - lo_i = cutlass.Int32(65535 - i * 32) - lane - hv = cutlass.select_(taken[i], cutlass.Int32(0), hi[i]) - lv = cutlass.select_(taken[i], cutlass.Int32(0), lo_i) - better = (cutlass.Uint32(hv) > cutlass.Uint32(bh)) | ( - (hv == bh) & (cutlass.Uint32(lv) > cutlass.Uint32(bl)) - ) - bh = cutlass.select_(better, hv, bh) - bl = cutlass.select_(better, lv, bl) - mh = _redux_max_u32(bh) - ml = _redux_max_u32(cutlass.select_(bh == mh, bl, cutlass.Int32(0))) - e_k = cutlass.Int32(65535) - ml - for i in cutlass.range_constexpr(4): - taken[i] = taken[i] | (e_k == cutlass.Int32(i * 32) + lane) - ids_t.append(e_k) - kv = cutlass.Float32(0.0) - for i in cutlass.range_constexpr(4): - kv = cutlass.select_(e_k == cutlass.Int32(i * 32) + lane, keys[i], kv) - kt_t.append(cute.arch.shuffle_sync(kv, e_k % cutlass.Int32(32))) - sig = cutlass.Float32(0.0) - e_l = cutlass.Int32(0) - for k in cutlass.range_constexpr(TOP_K): - sk = _sub_rn(kt_t[k], s_bias.load(idx=ids_t[k])) - sig = cutlass.select_(lane == cutlass.Int32(k), sk, sig) - e_l = cutlass.select_(lane == cutlass.Int32(k), ids_t[k], e_l) - total = sig - for off in (16, 8, 4, 2, 1): - total = _add_rn(total, cute.arch.shuffle_sync_bfly(total, off)) - wv = _bf16_round( - _div_rn( - _mul_rn(sig, cutlass.Float32(ROUTE_SCALE)), - _add_rn(total, cutlass.Float32(SUM_EPS)), - ) - ) - res_all.append((t, own, wv, e_l)) - if (lane == 0) & (warp == cutlass.Int32(W_WARP)): - _stamp(tdbg, bx, 62) - for tr in cutlass.range_constexpr( - (M_MAX + NRW - 1) // NRW - ): # stores after every chain (one branch-free block) - t, own, wv, e_l = res_all[tr] - if (lane < TOP_K) & own: - s_w.store(wv, idx=t * cutlass.Int32(TOP_K) + lane) - s_ids.store(e_l, idx=t * cutlass.Int32(TOP_K) + lane) - if bx == 0: - route_dbg.store( - wv.bitcast(cutlass.Int32), - idx=t * cutlass.Int32(8) + cutlass.Int32(4) + lane, - ) - cute.arch.sync_warp() - if cutlass.const_expr(EARLY0 and not E0_OFFLOAD): - if warp == cutlass.Int32(W_WARP): - # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same - # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now - if cutlass.const_expr(XSTAGE): - _wait(x_bar.subview(1), 0) - if lane == 0: - _stamp(tdbg, bx, 56) - n0 = cutlass.Int32(0) - w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) - ei0 = w0 & cutlass.Int32(0xFFFF) - ei1 = w0 >> cutlass.Int32(16) - for k in range(_nsegs(ei0, ei1)): - ug, j0, j1 = _seg(ei0, ei1, k) - c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - e = s_ids.load(idx=ug // cutlass.Int32(U1)) - sf_line = ( - (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) - * cutlass.Int32(SF1_KA) - + j0 * cutlass.Int32(C1 * NUM_KB) - ) * cutlass.Int32(4) - kx = j0 * cutlass.Int32(C1 * KT_BYTES) - row = c * cutlass.Int32(R1) - for j in cutlass.range(j0, j1, unroll=1): - if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx( - ab_full.subview(n0), - NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, - ) - for q in cutlass.range_constexpr(C1): - prims.cp_async_bulk_tensor_shared_cta_global( - sA.subview( - n0 * cutlass.Int32(NUM_BYTES_A) - + cutlass.Int32(q * NUM_BYTES_BOX1) - ), - tma_a1_desc.get_ptr(), - (kx + cutlass.Int32(q * KT_BYTES), row, e), - ab_full.subview(n0), - l2_cache_hint=EVICT_FIRST, - ) - if cutlass.const_expr(not SFA_CP): - prims.cp_async_bulk_tensor_shared_cta_global( - sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), - tma_s1_desc.get_ptr(), - (cutlass.Int32(0), sf_line), - ab_full.subview(n0), - ) - if cutlass.const_expr(SFA_CP): - _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) - n0 = n0 + cutlass.Int32(1) - kx = kx + cutlass.Int32(C1 * KT_BYTES) - sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) - if lane == 0: - s_misc.store(n0, idx=0) - cute.arch.sync_warp() - if cutlass.const_expr(E0_OFFLOAD): - if warp == cutlass.Int32(W_WARP): - _bar_arrive( - E0_BAR, 64 - ) # token 0's experts are in s_ids: warp 7 issues the early stages - if cutlass.const_expr(B1_AFTER_TOPK): - _bar_arrive(TOPK2_BAR, 128 + 32 * NRW) - cute.arch.barrier(barrier_id=TOPK_BAR, number_of_threads=32 * NRW) - if cutlass.const_expr(E0_OFFLOAD): - if warp == cutlass.Int32(E0_WARP): - cute.arch.barrier(barrier_id=E0_BAR, number_of_threads=64) - # slots 0..3 are always token 0's four (distinct) experts and group 0's stage ranges are the same - # for every G (the plan guarantees it): fill the ring with this CTA's first group-0 stages now - if cutlass.const_expr(XSTAGE): - _wait(x_bar.subview(1), 0) - if lane == 0: - _stamp(tdbg, bx, 56) - n0 = cutlass.Int32(0) - w0 = s_plan_all.load(idx=G_CAP * (NG_CAP + 1)) - ei0 = w0 & cutlass.Int32(0xFFFF) - ei1 = w0 >> cutlass.Int32(16) - for k in range(_nsegs(ei0, ei1)): - ug, j0, j1 = _seg(ei0, ei1, k) - c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - e = s_ids.load(idx=ug // cutlass.Int32(U1)) - sf_line = ( - (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) - * cutlass.Int32(SF1_KA) - + j0 * cutlass.Int32(C1 * NUM_KB) - ) * cutlass.Int32(4) - kx = j0 * cutlass.Int32(C1 * KT_BYTES) - row = c * cutlass.Int32(R1) - for j in cutlass.range(j0, j1, unroll=1): - if n0 < cutlass.Int32(min(STAGES, EARLY0_N)): - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx( - ab_full.subview(n0), - NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, - ) - for q in cutlass.range_constexpr(C1): - prims.cp_async_bulk_tensor_shared_cta_global( - sA.subview( - n0 * cutlass.Int32(NUM_BYTES_A) - + cutlass.Int32(q * NUM_BYTES_BOX1) - ), - tma_a1_desc.get_ptr(), - (kx + cutlass.Int32(q * KT_BYTES), row, e), - ab_full.subview(n0), - l2_cache_hint=EVICT_FIRST, - ) - if cutlass.const_expr(not SFA_CP): - prims.cp_async_bulk_tensor_shared_cta_global( - sSFAraw.subview(n0 * cutlass.Int32(NUM_WORDS_SFA_RAW)), - tma_s1_desc.get_ptr(), - (cutlass.Int32(0), sf_line), - ab_full.subview(n0), - ) - if cutlass.const_expr(SFA_CP): - _sf1_cp(lane, n0, sf_line, c, sSFA, s1w, ab_full) - n0 = n0 + cutlass.Int32(1) - kx = kx + cutlass.Int32(C1 * KT_BYTES) - sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) - if lane == 0: - s_misc.store(n0, idx=0) - cute.arch.sync_warp() - _bar_arrive(E1_BAR, 64) - if warp == W_WARP: - if lane == 0: - _stamp(tdbg, bx, 20) - # slots: lane l < 4 M is (token l / 4, top-k l % 4); the distinct experts in lane order (match.any) - valid = lane < cutlass.Int32(TOP_K * M_MAX) - my_e = cutlass.select_( - valid, - s_ids.load(idx=cutlass.select_(valid, lane, cutlass.Int32(0))), - cutlass.Int32(-1) - lane, - ) - same = _match_any(my_e) - leader = cute.arch.popc((same & (cutlass.Int32(0) - same)) - cutlass.Int32(1)) - first = valid & (lane == leader) - bal = prims.vote_sync(0xFFFFFFFF, first, prims.VoteSync.BALLOT) - rank = cute.arch.popc(bal & cutlass.Int32(cute.arch.lanemask_lt())) - slot_l = cute.arch.shuffle_sync(rank, leader) - for z in cutlass.range_constexpr((G_CAP * M_MAX + 31) // 32): - if lane + cutlass.Int32(32 * z) < cutlass.Int32(G_CAP * M_MAX): - s_kpos.store(cutlass.Int32(-1), idx=lane + cutlass.Int32(32 * z)) - cute.arch.sync_warp() - if first: - s_el.store(my_e, idx=rank) - if valid: - s_kpos.store( - lane % cutlass.Int32(TOP_K), - idx=slot_l * cutlass.Int32(M_MAX) + lane // cutlass.Int32(TOP_K), - ) - if bx == 0: - route_dbg.store( - my_e, - idx=(lane // cutlass.Int32(TOP_K)) * cutlass.Int32(8) - + lane % cutlass.Int32(TOP_K), - ) - gcnt = cute.arch.popc(bal) - if lane == 0: - s_el.store(gcnt, idx=G_CAP) - if lane < cutlass.Int32(NG_CAP + 1): # the plan row of this G (preloaded) - s_plan.store(s_plan_all.load(idx=gcnt * cutlass.Int32(NG_CAP + 1) + lane), idx=lane) - cute.arch.sync_warp() - if cutlass.const_expr(B2IMG): - # CTA 0 pads every group counter to exactly CNT_PER_CALL for this call (absent slots, unused groups) - if (bx == 0) & (lane < cutlass.Int32(NG_CAP)): - nh = gcnt - lane * cutlass.Int32(4) - nh = cutlass.select_(nh > cutlass.Int32(4), cutlass.Int32(4), nh) - nh = cutlass.select_(nh < cutlass.Int32(0), cutlass.Int32(0), nh) - pad = cutlass.Int32(CNT_PER_CALL) - nh * cutlass.Int32(UNITS_ROW) - if pad > cutlass.Int32(0): - _red_release_add( - b2cnt.subview(lane * cutlass.Int32(8)).data_ptr().toint(), pad - ) - cute.arch.sync_warp() - if lane == 0: - _stamp(tdbg, bx, 24) - cute.arch.sync_warp() - _bar_arrive(ROUTE_BAR, THREADS) - if warp < 4: - if cutlass.const_expr(B1_AFTER_TOPK): - cute.arch.barrier(barrier_id=TOPK2_BAR, number_of_threads=128 + 32 * NRW) - if cutlass.const_expr(FUSED_ROUTER): - if tidx == 0: - _stamp(tdbg, bx, 57) - if cutlass.const_expr(XSTAGE): - _wait(x_bar.subview(0), 0) - if tidx == 0: - _stamp(tdbg, bx, 58) - xs32 = cutlass.Array( - s_big.data_ptr(0), - shape=(M_MAX * H // 2,), - dtype=cutlass.Int32, - alignment=16, - ) - _build_b1_fused( - tidx, warp, lane, base_row, base_col, xs32, sfsc, sB1, s_xsf, b1_ready - ) - prims.mbarrier_arrive(x_bar.subview(1)) # the ring is free for the weights - else: - _build_b1_fused( - tidx, warp, lane, base_row, base_col, xb32, sfsc, sB1, s_xsf, b1_ready - ) - if tidx == 0: - _stamp(tdbg, bx, 23) - else: - _build_b1(tidx, warp, lane, base_row, base_col, x32, xsf32, sB1, b1_ready) - if warp != W_WARP: # one sync site for every thread but warp 4's (which arrived) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=ROUTE_BAR, number_of_threads=THREADS) - if tidx == 0: - _stamp(tdbg, bx, 2) - if cutlass.const_expr(TIMING): - tdbg.store( - cutlass.Int64(s_el.load(idx=G_CAP)), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(28), - ) - - g_n = s_el.load(idx=G_CAP) - ng = (g_n + cutlass.Int32(3)) >> cutlass.Int32(2) - pr2 = s_plan.load(idx=NG_CAP) - blk0 = pr2 & cutlass.Int32(0xFFFF) - nb = pr2 >> cutlass.Int32(16) - - # ---- weights producer (warp 4): phase p = FC1 stages of group p, then FC2 tiles of group p - LAG - if warp == W_WARP: - if cutlass.const_expr(XSTAGE): - _wait(x_bar.subview(1), 0) # the staged hidden states are quantized - stage = cutlass.Int32(0) - phase = cutlass.Int32(1) - skip = cutlass.Int32(0) - if cutlass.const_expr(E0_OFFLOAD): - cute.arch.barrier(barrier_id=E1_BAR, number_of_threads=64) - if cutlass.const_expr(EARLY0): - skip = s_misc.load(idx=0) # group-0 stages already in the ring - stage = cutlass.select_(skip == cutlass.Int32(STAGES), cutlass.Int32(0), skip) - phase = cutlass.select_( - skip == cutlass.Int32(STAGES), cutlass.Int32(0), cutlass.Int32(1) - ) - for p in range(ng + cutlass.Int32(LAG)): - if p < ng: - i0, i1 = _plan_range(s_plan, p) - for k in range(_nsegs(i0, i1)): - ug, j0, j1 = _seg(i0, i1, k) - slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) - c = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - e = s_el.load(idx=slot) - sf_line = ( - (e * cutlass.Int32(SF1_RB) + (c >> cutlass.Int32(2))) - * cutlass.Int32(SF1_KA) - + j0 * cutlass.Int32(C1 * NUM_KB) - ) * cutlass.Int32(4) - kx = j0 * cutlass.Int32(C1 * KT_BYTES) - row = c * cutlass.Int32(R1) - for j in cutlass.range(j0, j1, unroll=1): - if skip > cutlass.Int32(0): - skip = skip - cutlass.Int32(1) - else: - _waitd(ab_empty.subview(stage), phase, tdbg, 1) - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx( - ab_full.subview(stage), - NUM_BYTES_A if SFA_CP else NUM_BYTES_A + NUM_BYTES_SFA_RAW, - ) - for q in cutlass.range_constexpr(C1): - prims.cp_async_bulk_tensor_shared_cta_global( - sA.subview( - stage * cutlass.Int32(NUM_BYTES_A) - + cutlass.Int32(q * NUM_BYTES_BOX1) - ), - tma_a1_desc.get_ptr(), - (kx + cutlass.Int32(q * KT_BYTES), row, e), - ab_full.subview(stage), - l2_cache_hint=EVICT_FIRST, - ) - if cutlass.const_expr(not SFA_CP): - prims.cp_async_bulk_tensor_shared_cta_global( - sSFAraw.subview( - stage * cutlass.Int32(NUM_WORDS_SFA_RAW) - ), - tma_s1_desc.get_ptr(), - (cutlass.Int32(0), sf_line), - ab_full.subview(stage), - ) - if cutlass.const_expr(SFA_CP): - _sf1_cp(lane, stage, sf_line, c, sSFA, s1w, ab_full) - stage, phase = _adv(stage, phase) - kx = kx + cutlass.Int32(C1 * KT_BYTES) - sf_line = sf_line + cutlass.Int32(NUM_BYTES_SFA_RAW // 128) - if p >= cutlass.Int32(LAG): - g = p - cutlass.Int32(LAG) - n_here = g_n - g * cutlass.Int32(4) - n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) - for bb in range(nb): - blk = blk0 + bb - for t2 in cutlass.range_constexpr(KT2): - _waitd(ab_empty.subview(stage), phase, tdbg, 2) - if prims.elect_sync(): - prims.mbarrier_arrive_expect_tx( - ab_full.subview(stage), - n_here - * cutlass.Int32( - NUM_BYTES_BOX1 - if SFA_CP - else NUM_BYTES_BOX1 + NUM_BYTES_SFA_RAW // 4 - ), - ) - for jj in cutlass.range_constexpr(4): - if cutlass.Int32(jj) < n_here: - e = s_el.load(idx=g * cutlass.Int32(4) + cutlass.Int32(jj)) - prims.cp_async_bulk_tensor_shared_cta_global( - sA.subview( - stage * cutlass.Int32(NUM_BYTES_A) - + cutlass.Int32(jj * NUM_BYTES_BOX1) - ), - tma_a2_desc.get_ptr(), - ( - cutlass.Int32(t2 * KT_BYTES), - blk * cutlass.Int32(R2), - e, - ), - ab_full.subview(stage), - l2_cache_hint=EVICT_FIRST, - ) - if cutlass.const_expr(not SFA_CP): - prims.cp_async_bulk_tensor_shared_cta_global( - sSFAraw.subview( - stage * cutlass.Int32(NUM_WORDS_SFA_RAW) - + cutlass.Int32(jj * NUM_WORDS_SFA_RAW // 4) - ), - tma_s2_desc.get_ptr(), - ( - cutlass.Int32(0), - ( - ( - e * cutlass.Int32(SF2_RB) - + (blk >> cutlass.Int32(2)) - ) - * cutlass.Int32(SF2_KA) - + cutlass.Int32(4 * t2) - ) - * cutlass.Int32(4), - ), - ab_full.subview(stage), - ) - if cutlass.const_expr(SFA_CP): - _sf2_cp(lane, stage, g, blk, t2, n_here, s_el, sSFA, s2w, ab_full) - stage, phase = _adv(stage, phase) - if lane == 0: - _stamp(tdbg, bx, 3) - - # ---- scale splicer (warp 6): every stage, atom kb word [r][q] = raw atom 4 q + kb word [r][m1] - if (warp == S_WARP) & cutlass.Boolean(not SFA_CP): # (SFA_CP: no splicer) - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - stage = cutlass.Int32(0) - phase = cutlass.Int32(0) - for p in range(ng + cutlass.Int32(LAG)): - i0 = cutlass.Int32(0) - i1 = cutlass.Int32(0) - for part in cutlass.range_constexpr(2): - cnt = cutlass.Int32(0) - m1 = cutlass.Int32(0) - nk = cutlass.Int32(0) - if cutlass.const_expr(part == 0): - if p < ng: - i0, i1 = _plan_range(s_plan, p) - nk = _nsegs(i0, i1) - else: - if p >= cutlass.Int32(LAG): - nk = nb - for k in range(nk): - if cutlass.const_expr(part == 0): - ug, j0, j1 = _seg(i0, i1, k) - m1 = ( - ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - ) & cutlass.Int32(3) - cnt = j1 - j0 - else: - m1 = (blk0 + k) & cutlass.Int32(3) - cnt = cutlass.Int32(KT2) - for _j in cutlass.range(cnt, unroll=1): - _waitd(ab_full.subview(stage), phase, tdbg, 3) - raw0 = ( - stage * cutlass.Int32(NUM_WORDS_SFA_RAW) - + lane * cutlass.Int32(4) - + m1 - ) - dst0 = stage * cutlass.Int32(NUM_WORDS_SFA) + lane * cutlass.Int32(4) - for kb in cutlass.range_constexpr(NUM_KB): - w0 = sSFAraw.load(idx=raw0 + cutlass.Int32((0 * 4 + kb) * 128)) - w1 = sSFAraw.load(idx=raw0 + cutlass.Int32((1 * 4 + kb) * 128)) - w2 = sSFAraw.load(idx=raw0 + cutlass.Int32((2 * 4 + kb) * 128)) - w3 = sSFAraw.load(idx=raw0 + cutlass.Int32((3 * 4 + kb) * 128)) - sSFA.store( - ( - cutlass.Int32(w0), - cutlass.Int32(w1), - cutlass.Int32(w2), - cutlass.Int32(w3), - ), - idx=dst0 + cutlass.Int32(kb * 128), - alignment=16, - ) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - cute.arch.sync_warp() - if prims.elect_sync(): - sd0 = prims.Tcgen05SmemDesc.build( - sSFA.subview(stage * cutlass.Int32(NUM_WORDS_SFA)), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - for kb in cutlass.range_constexpr(NUM_KB): - prims.tcgen05_cp( - s2t_shape, - _tptr( - base_row, - base_col - + SFA1_COL - + stage * cutlass.Int32(16) - + kb * 4, - ), - sd0 + (kb * 512 >> 4), - multicast=s2t_mc, - ) - prims.tcgen05_commit(sfa_ready.subview(stage)) - stage, phase = _adv(stage, phase) - - if cutlass.const_expr(B2IMG): - # ---- (B2IMG) group releases (warp 6): this CTA's finished units of group p, one release add per group - if warp == cutlass.Int32(S_WARP): - for p in range(ng): - i0, i1 = _plan_range(s_plan, p) - n_done = cutlass.Int32(0) - for k in range(_nsegs(i0, i1)): - ug, j0, j1 = _seg(i0, i1, k) - n_done = n_done + cutlass.select_( - (j0 != cutlass.Int32(0)) | (j1 == cutlass.Int32(S1)), - cutlass.Int32(1), - cutlass.Int32(0), - ) - _wait( - fc1_done.subview(p & cutlass.Int32(3)), - (p >> cutlass.Int32(2)) & cutlass.Int32(1), - ) - if cutlass.const_expr(B2_DLY > 0): - t0d = _gtime() - dl = cutlass.Int64( - ( - (bx * cutlass.Int32(37) + p * cutlass.Int32(11) + ep) - & cutlass.Int32(63) - ) - * cutlass.Int32(B2_DLY // 64) - ) - while _gtime() - t0d < dl: - pass - if (lane == 0) & (n_done > cutlass.Int32(0)): - _fence_proxy_async_global() - _red_release_add( - b2cnt.subview(p * cutlass.Int32(8)).data_ptr().toint(), n_done - ) - cute.arch.sync_warp() - # ---- (B2IMG) FC2 B loader (warp 5): group g's image once its counter is complete - if warp == cutlass.Int32(X_WARP): - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - cnt_target = (ep + cutlass.Int32(1)) * cutlass.Int32(CNT_PER_CALL) - for g in range(ng): - buf = g & cutlass.Int32(1) - if g >= cutlass.Int32(2): - _waitd( - b2_free.subview(buf), - ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), - tdbg, - 101, - ) - if lane == 0: - _stamp(tdbg, bx, 29) - caddr = b2cnt.subview(g * cutlass.Int32(8)).data_ptr().toint() - if lane == 0: - cv = _ld_acquire_u32(caddr) - while (cv - cnt_target) < cutlass.Int32(0): - _nanosleep(B2IMG_NS) - cv = _ld_acquire_u32(caddr) - _fence_proxy_async_global() - prims.mbarrier_arrive_expect_tx(img_bar.subview(buf), IMG_BYTES) - prims.cp_async_bulk_shared_cluster_global( - sB2.subview(buf * cutlass.Int32(IMG_B2)), - img32.subview(g * cutlass.Int32(IMG_BYTES // 4)), - img_bar.subview(buf), - IMG_B2, - ) - prims.cp_async_bulk_shared_cluster_global( - sSFB2.subview(buf * cutlass.Int32(IMG_SF // 4)), - img32.subview( - g * cutlass.Int32(IMG_BYTES // 4) + cutlass.Int32(IMG_B2 // 4) - ), - img_bar.subview(buf), - IMG_SF, - ) - cute.arch.sync_warp() - _waitd( - img_bar.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 108 - ) - if prims.elect_sync(): - sd0 = prims.Tcgen05SmemDesc.build( - sSFB2.subview(0), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - for a in cutlass.range_constexpr(KT2 * NUM_KB): - prims.tcgen05_cp( - s2t_shape, - _tptr( - base_row, - base_col - + SFB2_COL - + buf * cutlass.Int32(KT2 * NUM_KB * 4) - + a * 4, - ), - sd0 - + cutlass.Int64( - (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) - * cutlass.Int32(512 >> 4) - ), - multicast=s2t_mc, - ) - prims.tcgen05_commit(b2_ready.subview(buf)) - cute.arch.sync_warp() - if lane == 0: - _stamp(tdbg, bx, 25) - # ---- FC2 B builder (warps 5, 8 ..): group g's records (all 48 per (slot, token), the call's tag on - # every word) into B2 buffer g % 2 (row j M + t, k-tile t2 = record / 16) and its scale atoms, then - # tcgen05.cp to TMEM. - if ((warp == X_WARP) | (warp >= cutlass.Int32(8))) & cutlass.Boolean(not B2IMG): - xi = cutlass.select_( - warp == cutlass.Int32(X_WARP), cutlass.Int32(0), warp - cutlass.Int32(7) - ) - if cutlass.const_expr(G_CAP > TOP_K): - if (bx == 0) & (xi == cutlass.Int32(0)): - # rows of slots this call does not use: tag 0, never taken for a later call's - for i in range( - g_n * cutlass.Int32(M_MAX * UNITS_ROW) + lane, - cutlass.Int32(G_CAP * M_MAX * UNITS_ROW), - 32, - ): - for rr in cutlass.range_constexpr(HB_REPL): - hbuf32.store( - ( - cutlass.Int32(0), - cutlass.Int32(0), - cutlass.Int32(0), - cutlass.Int32(0), - ), - idx=i * cutlass.Int32(4) - + (cutlass.Int32(rr * HB_STRIDE) if rr else 0), - alignment=16, - ) - s2t_shape, s2t_mc = prims.S2TCopyMode.S2T_32x128b_WARPX4 - sSFB2_8 = cutlass.Array( - sSFB2.data_ptr(0), shape=(2 * KT2 * NUM_KB * 512,), dtype=cutlass.Int8, alignment=16 - ) - sB2_32 = cutlass.Array( - sB2.data_ptr(0), - shape=(2 * KT2 * NUM_BYTES_B2 // 4,), - dtype=cutlass.Int32, - alignment=16, - ) - nrec_cap = 4 * M_MAX * UNITS_ROW # records of one group - hb_off = cutlass.Int32(0) - if cutlass.const_expr(HB_REPL > 1): - # this CTA's copy of the records - hb_off = (bx % cutlass.Int32(HB_REPL)) * cutlass.Int32(HB_STRIDE) - NZ = (nrec_cap + 32 * XW - 1) // (32 * XW) # records per lane: lane + 32 (XW z + xi) - for g in range(ng): - buf = g & cutlass.Int32(1) - if g >= cutlass.Int32(2): - _waitd( - b2_free.subview(buf), - ((g >> cutlass.Int32(1)) - cutlass.Int32(1)) & cutlass.Int32(1), - tdbg, - 101, - ) - _waitd( - fc1_done.subview(g & cutlass.Int32(3)), - (g >> cutlass.Int32(2)) & cutlass.Int32(1), - tdbg, - 4, - ) # this CTA's own share of group g is out - if cutlass.const_expr(TIMING): - if (lane == 0) & (xi == cutlass.Int32(0)) & (g < cutlass.Int32(8)): - tdbg.store( - _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(32) + g - ) - n_here = g_n - g * cutlass.Int32(4) - n_here = cutlass.select_(n_here > cutlass.Int32(4), cutlass.Int32(4), n_here) - nrec = n_here * cutlass.Int32(M_MAX * UNITS_ROW) - rec0 = g * cutlass.Int32( - 4 * M_MAX * UNITS_ROW - ) # rows (4 g + j) M + t are contiguous - done = cutlass.Int32(0) - rspin = cutlass.Int32(0) - if (lane == 0) & (xi == cutlass.Int32(0)): - _stamp(tdbg, bx, 29) - npoll = cutlass.Int32(0) - got = [[cutlass.Int32(0) for _ in range(3)] for _ in range(NZ)] - while done == cutlass.Int32(0): - ok = cutlass.Boolean(True) - for z in cutlass.range_constexpr(NZ): - qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi - if qi < nrec: - v4 = prims.load_ext( - hbuf32.subview( - (rec0 + qi) * cutlass.Int32(4) + (hb_off if HB_REPL > 1 else 0) - ), - dtype=cutlass.Int32, - count=4, - order="relaxed", - scope="gpu", - ) - for wq in cutlass.range_constexpr(3): - got[z][wq] = cutlass.Int32(v4[wq]) - ok = ok & ( - ( - ((got[z][0] ^ tagv) | (got[z][1] ^ tagv) | (got[z][2] ^ tagv)) - & cutlass.Int32(-16777216) - ) - == cutlass.Int32(0) - ) - npoll = npoll + cutlass.Int32(1) - if prims.vote_sync(0xFFFFFFFF, ok, prims.VoteSync.BALLOT) == cutlass.Int32(-1): - done = cutlass.Int32(1) - else: - _nanosleep(B2_POLL_NS) - if cutlass.const_expr(HANG_DBG): - rspin = rspin + cutlass.Int32(1) - if rspin > cutlass.Int32(1 << 14): - done = cutlass.Int32(1) - tdbg.store( - cutlass.Int64(200), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(24) + warp, - ) - for z in cutlass.range_constexpr(NZ): - qi = lane + cutlass.Int32(32 * XW * z) + cutlass.Int32(32) * xi - if qi < nrec: - w0, w1, w2 = got[z][0], got[z][1], got[z][2] - rr = qi // cutlass.Int32( - UNITS_ROW - ) # row within the group: j M + t = B row n - c = qi - rr * cutlass.Int32(UNITS_ROW) - t2 = c >> cutlass.Int32(4) - cc = (c >> cutlass.Int32(1)) & cutlass.Int32(7) - lo = (w0 & cutlass.Int32(0xFFFFFF)) | (w1 << cutlass.Int32(24)) - hi = ((w1 >> cutlass.Int32(8)) & cutlass.Int32(0xFFFF)) | ( - (w2 & cutlass.Int32(0xFFFF)) << cutlass.Int32(16) - ) - sB2_32.store( - (lo, hi), - idx=( - (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_BYTES_B2) - + rr * cutlass.Int32(KT_BYTES) - + (cc ^ (rr & cutlass.Int32(7))) * cutlass.Int32(16) - + (c & cutlass.Int32(1)) * cutlass.Int32(8) - ) - // cutlass.Int32(4), - alignment=8, - ) - # scale of block c: atom (t2, kb = (c % 16) / 4), row rr, byte c % 4 - sSFB2_8.store( - cutlass.Int8((w2 >> cutlass.Int32(16)) & cutlass.Int32(0xFF)), - idx=( - ( - (buf * cutlass.Int32(KT2) + t2) * cutlass.Int32(NUM_KB) - + ((c & cutlass.Int32(15)) >> cutlass.Int32(2)) - ) - * cutlass.Int32(128) - + rr * cutlass.Int32(4) - ) - * cutlass.Int32(4) - + (c & cutlass.Int32(3)), - ) - prims.fence_proxy("async_shared", space=prims.SharedSpace.shared_cta) - if cutlass.const_expr(XW > 1): - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=B2_BAR, number_of_threads=32 * XW) - cute.arch.sync_warp() - if (xi == cutlass.Int32(0)) & prims.elect_sync(): - sd0 = prims.Tcgen05SmemDesc.build( - sSFB2.subview(0), - leading_byte_offset=16, - stride_byte_offset=128, - base_offset=0, - layout=0, - ) - for a in cutlass.range_constexpr(KT2 * NUM_KB): - prims.tcgen05_cp( - s2t_shape, - _tptr( - base_row, - base_col + SFB2_COL + buf * cutlass.Int32(KT2 * NUM_KB * 4) + a * 4, - ), - sd0 - + cutlass.Int64( - (buf * cutlass.Int32(KT2 * NUM_KB) + cutlass.Int32(a)) - * cutlass.Int32(512 >> 4) - ), - multicast=s2t_mc, - ) - prims.tcgen05_commit(b2_ready.subview(buf)) - if (lane == 0) & (xi == cutlass.Int32(0)): - _stamp(tdbg, bx, 25) - if cutlass.const_expr(TIMING): - if g < cutlass.Int32(8): - tdbg.store( - _gtime(), idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(11) + g - ) - if cutlass.const_expr(TIMING): - tdbg.store( - cutlass.Int64(npoll), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(30), - ) - - # ---- MMA (warp 7) - if warp == MMA_WARP: - idesc1 = prims.Tcgen05MxOmmaInstrDesc.build( - a_dtype=cutlass.Float4E2M1FN, - b_dtype=cutlass.Float4E2M1FN, - scale_format=0, - n_dim=N1, - m_dim=MMA_M, - sparsity_version=0, - ) - idesc2 = prims.Tcgen05MxOmmaInstrDesc.build( - a_dtype=cutlass.Float4E2M1FN, - b_dtype=cutlass.Float4E2M1FN, - scale_format=0, - n_dim=N2, - m_dim=MMA_M, - sparsity_version=0, - ) - desc_a0 = prims.Tcgen05SmemDesc.build( - sA.subview(0), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - desc_b10 = prims.Tcgen05SmemDesc.build( - sB1.subview(0), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - desc_b20 = prims.Tcgen05SmemDesc.build( - sB2.subview(0), - leading_byte_offset=16, - stride_byte_offset=1024, - base_offset=0, - layout=2, - ) - tcol = (base_row << 16) | base_col - _waitd(b1_ready, 0, tdbg, 5) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - stage = cutlass.Int32(0) - phase = cutlass.Int32(0) - k1 = cutlass.Int32(0) - k2 = cutlass.Int32(0) - acc_b2 = cutlass.Int64(0) - acc_sf = cutlass.Int64(0) - for p in range(ng + cutlass.Int32(LAG)): - if p < ng: - i0, i1 = _plan_range(s_plan, p) - for k in range(_nsegs(i0, i1)): - ug, j0, j1 = _seg(i0, i1, k) - ab = k1 & cutlass.Int32(1) - _waitd( - acc_empty.subview(ab), - ((k1 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), - tdbg, - 102, - ) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - acc_ptr = cutlass.inttoptr( - tcol + ACC1_COL + ab * cutlass.Int32(N1), 6, cutlass.Float32 - ) - for j in cutlass.range(j0, j1, unroll=1): - tw0 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) - if cutlass.const_expr(SFA_CP): - _waitd(ab_full.subview(stage), phase, tdbg, 6) - prims.fence_proxy( - "async_shared", space=prims.SharedSpace.shared_cta - ) - else: - _waitd(sfa_ready.subview(stage), phase, tdbg, 6) - if cutlass.const_expr(TIMING): - acc_sf = acc_sf + (_gtime() - tw0) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) - db = desc_b10 + cutlass.Int64(j * cutlass.Int32(NUM_BYTES_B1 >> 4)) - sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) - sfb_c = tcol + SFB1_COL + j * cutlass.Int32(NUM_KB * 4) - if prims.elect_sync(): - if cutlass.const_expr(SFA_CP): - _sfa_to_tmem(sSFA, stage, base_row, base_col) - for kb in cutlass.range_constexpr(NUM_KB): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF4NVF4, - prims.CTAGroup.CTA_1, - acc_ptr, - da + INC_K * kb, - db + INC_K * kb, - idesc1, - (j > j0) | cutlass.Boolean(kb > 0), - cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), - cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), - scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, - ) - prims.tcgen05_commit(ab_empty.subview(stage)) - stage, phase = _adv(stage, phase) - if prims.elect_sync(): - prims.tcgen05_commit(acc_full.subview(ab)) - k1 = k1 + cutlass.Int32(1) - if p >= cutlass.Int32(LAG): - g = p - cutlass.Int32(LAG) - buf = g & cutlass.Int32(1) - tw2 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) - _waitd( - b2_ready.subview(buf), (g >> cutlass.Int32(1)) & cutlass.Int32(1), tdbg, 103 - ) - if cutlass.const_expr(TIMING): - acc_b2 = acc_b2 + (_gtime() - tw2) - if (lane == 0) & (g == ng - cutlass.Int32(1)): - _stamp(tdbg, bx, 4) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - for bb in range(nb): - a2 = k2 & cutlass.Int32(1) - _waitd( - acc2_empty.subview(a2), - ((k2 >> cutlass.Int32(1)) & cutlass.Int32(1)) ^ cutlass.Int32(1), - tdbg, - 104, - ) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - acc2_ptr = cutlass.inttoptr( - tcol + ACC2_COL + a2 * cutlass.Int32(N2), 6, cutlass.Float32 - ) - for t2 in cutlass.range_constexpr(KT2): - tw1 = _gtime() if cutlass.const_expr(TIMING) else cutlass.Int64(0) - if cutlass.const_expr(SFA_CP): - _waitd(ab_full.subview(stage), phase, tdbg, 7) - prims.fence_proxy( - "async_shared", space=prims.SharedSpace.shared_cta - ) - else: - _waitd(sfa_ready.subview(stage), phase, tdbg, 7) - if cutlass.const_expr(TIMING): - acc_sf = acc_sf + (_gtime() - tw1) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - da = desc_a0 + cutlass.Int64(stage * cutlass.Int32(NUM_BYTES_A >> 4)) - db = desc_b20 + cutlass.Int64( - (buf * cutlass.Int32(KT2) + cutlass.Int32(t2)) - * cutlass.Int32(NUM_BYTES_B2 >> 4) - ) - sfa_c = tcol + SFA1_COL + stage * cutlass.Int32(16) - sfb_c = ( - tcol - + SFB2_COL - + buf * cutlass.Int32(KT2 * NUM_KB * 4) - + cutlass.Int32(t2 * NUM_KB * 4) - ) - if prims.elect_sync(): - if cutlass.const_expr(SFA_CP): - _sfa_to_tmem(sSFA, stage, base_row, base_col) - for kb in cutlass.range_constexpr(_nkb2(t2)): - prims.tcgen05_mma_block_scale( - prims.MMABlockScaleKind.MXF4NVF4, - prims.CTAGroup.CTA_1, - acc2_ptr, - da + INC_K * kb, - db + INC_K * kb, - idesc2, - cutlass.Boolean(t2 > 0 or kb > 0), - cutlass.inttoptr(sfa_c + kb * 4, 6, cutlass.Int32), - cutlass.inttoptr(sfb_c + kb * 4, 6, cutlass.Int32), - scale_vec_size=prims.Tcgen05MMABlockScale.BLOCK16, - ) - prims.tcgen05_commit(ab_empty.subview(stage)) - stage, phase = _adv(stage, phase) - if prims.elect_sync(): - prims.tcgen05_commit(acc2_full.subview(a2)) - k2 = k2 + cutlass.Int32(1) - if cutlass.const_expr(TIMING): - if lane == 0: - tdbg.store(acc_b2, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(26)) - tdbg.store(acc_sf, idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(27)) - - # ---- epilogue warps (0-3) - if warp < 4: - row_w = base_row + warp * 32 - for r in cutlass.range_constexpr((6 * G_CAP + EPI_THREADS - 1) // EPI_THREADS): - qi = tidx + cutlass.Int32(r * EPI_THREADS) - sl_e = qi // cutlass.Int32(6) - if (qi < cutlass.Int32(6 * G_CAP)) & (sl_e < g_n): - fld = qi - sl_e * cutlass.Int32(6) - s_esc.store( - cutlass.Float32( - escal.load(idx=fld * cutlass.Int32(E_LOCAL) + s_el.load(idx=sl_e)) - ), - idx=sl_e * cutlass.Int32(8) + fld, - ) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) - k1 = cutlass.Int32(0) - k2 = cutlass.Int32(0) - SH_ADD = SHARED - wys = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] - pars = [cutlass.Int32(0) for _ in range(NB2_CAP * M_MAX // 4 + 1)] - ard = (cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0), cutlass.Int64(0)) - for p in range(ng + cutlass.Int32(LAG)): - if cutlass.const_expr(AR_PRE): - if (p == ng + cutlass.Int32(LAG - 1)) & ar_on: - # one 8-byte load each (the vectorizer merges adjacent arp.load()s into 16-byte reads - # at 8-byte offsets: misaligned) - ard = tuple( - cutlass.Int64(_ld_relaxed_u64(arp.subview(i_).data_ptr().toint())) - for i_ in (1, 3, 4, 5) - ) - for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): - f_pair = warp + cutlass.Int32(bt * 4) - f_bb = f_pair // cutlass.Int32(M_MAX) - f_t = f_pair - f_bb * cutlass.Int32(M_MAX) - f_cl = cutlass.select_( - f_bb < nb, (blk0 + f_bb) * cutlass.Int32(R2), cutlass.Int32(0) - ) - pars[bt] = _ar_parity(arp, f_t, f_cl) - if cutlass.const_expr(SH_ADD): - if p == ng + cutlass.Int32(LAG - 1): - # this CTA's y_sh words, loaded while the last group's FC2 runs (re-polled in finalize) - for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): - f_pair = warp + cutlass.Int32(bt * 4) - f_bb = f_pair // cutlass.Int32(M_MAX) - f_t = f_pair - f_bb * cutlass.Int32(M_MAX) - f_ch = cutlass.select_( - f_bb < nb, - f_t * cutlass.Int32(H) - + (blk0 + f_bb) * cutlass.Int32(R2) - + cutlass.Int32(4) * (lane & cutlass.Int32(7)) - + (lane >> cutlass.Int32(3)), - cutlass.Int32(0), - ) - wys[bt] = _ld_relaxed_u32(dnx.subview(f_ch).data_ptr().toint()) - if p < ng: - i0, i1 = _plan_range(s_plan, p) - for k in range(_nsegs(i0, i1)): - ug, j0, j1 = _seg(i0, i1, k) - ab = k1 & cutlass.Int32(1) - slot = p * cutlass.Int32(4) + ug // cutlass.Int32(U1) - cblk = ug - (ug // cutlass.Int32(U1)) * cutlass.Int32(U1) - unit = slot * cutlass.Int32(U1) + cblk - _waitd( - acc_full.subview(ab), - (k1 >> cutlass.Int32(1)) & cutlass.Int32(1), - tdbg, - 105, - ) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - acc = prims.tcgen05_ld( - "32x32b", - _tptr( - row_w, base_col + ACC1_COL + ab * cutlass.Int32(N1), cutlass.Float32 - ), - num=N1, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - prims.mbarrier_arrive(acc_empty.subview(ab)) - pb = ab * cutlass.Int32(M_MAX * C1 * R1) - for t in cutlass.range_constexpr(M_MAX): - part = cutlass.Float32(0.0) - for w in cutlass.range_constexpr(4): - part = cutlass.select_( - warp == cutlass.Int32(w), cutlass.Float32(acc[QS * w + t]), part - ) - s_part.store(part, idx=pb + (t * C1 + warp) * R1 + lane) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) - for tr in cutlass.range_constexpr((M_MAX + 3) // 4): - t = warp + cutlass.Int32(4 * tr) - if t < cutlass.Int32(M_MAX): - xv = _add_rn( - _add_rn( - _add_rn( - s_part.load(idx=pb + (t * C1 + 0) * R1 + lane), - s_part.load(idx=pb + (t * C1 + 1) * R1 + lane), - ), - s_part.load(idx=pb + (t * C1 + 2) * R1 + lane), - ), - s_part.load(idx=pb + (t * C1 + 3) * R1 + lane), - ) - finish = cutlass.Int32(1) - if (j0 != cutlass.Int32(0)) | (j1 != cutlass.Int32(S1)): - paddr = ( - partials.subview( - (unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) - + lane - ) - .data_ptr() - .toint() - ) - if j0 == cutlass.Int32(0): - _st_relaxed_u32(paddr, xv.bitcast(cutlass.Int32)) - finish = cutlass.Int32(0) - else: - pbits = _ld_relaxed_u32(paddr) - spins = cutlass.Int32(0) - while pbits == cutlass.Int32(PARTIAL_EMPTY): - pbits = _ld_relaxed_u32(paddr) - spins = spins + cutlass.Int32(1) - _st_relaxed_u32(paddr, cutlass.Int32(PARTIAL_EMPTY)) - xv = _add_rn(pbits.bitcast(cutlass.Float32), xv) - cute.arch.sync_warp() - if cutlass.const_expr( - XVDUMP - ): # debug: FC1 sums into the (unused) partial slots - partials.store( - xv.bitcast(cutlass.Int32), - idx=(unit * cutlass.Int32(M_MAX) + t) * cutlass.Int32(R1) - + lane, - ) - if finish == cutlass.Int32(1): - gate_raw = cute.arch.shuffle_sync_bfly(xv, 8) - is_up = ( - (lane >> cutlass.Int32(3)) & cutlass.Int32(1) - ) == cutlass.Int32(0) - es = slot * cutlass.Int32(8) - qv, sfb = _swiglu_quant( - xv, - gate_raw, - is_up, - s_esc.load(idx=es), - s_esc.load(idx=es + cutlass.Int32(1)), - s_esc.load(idx=es + cutlass.Int32(2)), - s_esc.load(idx=es + cutlass.Int32(3)), - s_esc.load(idx=es + cutlass.Int32(4)), - ) - byte = _e2m1x2(cute.arch.shuffle_sync_down(qv, 16), qv) - bt = [ - cute.arch.shuffle_sync(byte, i) & cutlass.Int32(0xFF) - for i in range(8) - ] - rw0 = ( - bt[0] - | (bt[1] << cutlass.Int32(8)) - | (bt[2] << cutlass.Int32(16)) - | tagv - ) - rw1 = ( - bt[3] - | (bt[4] << cutlass.Int32(8)) - | (bt[5] << cutlass.Int32(16)) - | tagv - ) - rw2 = ( - bt[6] - | (bt[7] << cutlass.Int32(8)) - | ((sfb & cutlass.Int32(0xFF)) << cutlass.Int32(16)) - | tagv - ) - if cutlass.const_expr(B2IMG): - if lane == 0: - # B2 row n = j M + t, k-tile c / 16, 16-byte chunk (c / 2) % 8 - # (swizzled), half c % 2; scale: atom (t2, kb = (c % 16) / 4), - # row n at byte 16 n, byte c % 4 - n = (ug // cutlass.Int32(U1)) * cutlass.Int32(M_MAX) + t - t2i = cblk >> cutlass.Int32(4) - cci = (cblk >> cutlass.Int32(1)) & cutlass.Int32(7) - gbase = p * cutlass.Int32(IMG_BYTES) - lo = ( - bt[0] - | (bt[1] << cutlass.Int32(8)) - | (bt[2] << cutlass.Int32(16)) - | (bt[3] << cutlass.Int32(24)) - ) - hi = ( - bt[4] - | (bt[5] << cutlass.Int32(8)) - | (bt[6] << cutlass.Int32(16)) - | (bt[7] << cutlass.Int32(24)) - ) - img32.store( - (lo, hi), - idx=( - gbase - + t2i * cutlass.Int32(NUM_BYTES_B2) - + n * cutlass.Int32(KT_BYTES) - + (cci ^ (n & cutlass.Int32(7))) - * cutlass.Int32(16) - + (cblk & cutlass.Int32(1)) * cutlass.Int32(8) - ) - // cutlass.Int32(4), - alignment=8, - ) - img8 = cutlass.Array( - img32.data_ptr(0), - shape=(NG_CAP * IMG_BYTES,), - dtype=cutlass.Int8, - alignment=16, - ) - img8.store( - cutlass.Int8(sfb & cutlass.Int32(0xFF)), - idx=gbase - + cutlass.Int32(IMG_B2) - + ( - ( - t2i * cutlass.Int32(NUM_KB) - + ( - (cblk & cutlass.Int32(15)) - >> cutlass.Int32(2) - ) - ) - * cutlass.Int32(512) - ) - + n * cutlass.Int32(16) - + (cblk & cutlass.Int32(3)), - ) - else: - if cutlass.const_expr(HB_REPL == 1): - if lane == 0: - hbuf32.store( - (rw0, rw1, rw2, tagv), - idx=( - (slot * cutlass.Int32(M_MAX) + t) - * cutlass.Int32(UNITS_ROW) - + cblk - ) - * cutlass.Int32(4), - alignment=16, - ) - else: - if lane < HB_REPL: # copy `lane` - hbuf32.store( - (rw0, rw1, rw2, tagv), - idx=( - (slot * cutlass.Int32(M_MAX) + t) - * cutlass.Int32(UNITS_ROW) - + cblk - ) - * cutlass.Int32(4) - + lane * cutlass.Int32(HB_STRIDE), - alignment=16, - ) - k1 = k1 + cutlass.Int32(1) - prims.mbarrier_arrive( - fc1_done.subview(p & cutlass.Int32(3)) - ) # warp 5 may poll group p - if tidx == 0: - _stamp(tdbg, bx, 10) - if cutlass.const_expr(TIMING): - if p < cutlass.Int32(8): - tdbg.store( - _gtime(), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(40) + p, - ) - if p >= cutlass.Int32(LAG): - g = p - cutlass.Int32(LAG) - for bb in range(nb): - a2 = k2 & cutlass.Int32(1) - _waitd( - acc2_full.subview(a2), - (k2 >> cutlass.Int32(1)) & cutlass.Int32(1), - tdbg, - 106, - ) - if cutlass.const_expr(TIMING): - if (tidx == 0) & (g == ng - cutlass.Int32(1)): - tdbg.store( - _gtime(), - idx=bx * cutlass.Int32(TD_STRIDE) + cutlass.Int32(8) + bb, - ) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - yv = prims.tcgen05_ld( - "32x32b", - _tptr( - row_w, - base_col - + ACC2_COL - + a2 * cutlass.Int32(N2) - + warp * cutlass.Int32(M_MAX), - cutlass.Float32, - ), - num=LD_POW2, - ) - prims.tcgen05_wait(prims.Tcgen05Wait.LOAD) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - prims.mbarrier_arrive(acc2_empty.subview(a2)) - sl = g * cutlass.Int32(4) + warp - if sl < g_n: - a2s = s_esc.load(idx=sl * cutlass.Int32(8) + cutlass.Int32(5)) - for t in cutlass.range_constexpr(M_MAX): - kp = s_kpos.load(idx=sl * cutlass.Int32(M_MAX) + cutlass.Int32(t)) - if kp >= cutlass.Int32(0): - ys.store( - _bf16_round(_mul_rn(cutlass.Float32(yv[t]), a2s)), - idx=( - (bb * cutlass.Int32(M_MAX) + cutlass.Int32(t)) - * cutlass.Int32(TOP_K) - + kp - ) - * cutlass.Int32(R2) - + lane, - ) - k2 = k2 + cutlass.Int32(1) - prims.mbarrier_arrive(b2_free.subview(g & cutlass.Int32(1))) - if tidx == 0: - _stamp(tdbg, bx, 6) - cute.arch.sync_warp() - cute.arch.barrier(barrier_id=EPI_BAR, number_of_threads=EPI_THREADS) - if tidx == 0: - _stamp(tdbg, bx, 19) - # finalize: (block bb, token t) per warp round: the fma chain over t's top-4 in order - for bt in cutlass.range_constexpr((NB2_CAP * M_MAX + 3) // 4): - pair = warp + cutlass.Int32(bt * 4) - bb = pair // cutlass.Int32(M_MAX) - t = pair - bb * cutlass.Int32(M_MAX) - if bb < nb: - acc_o = cutlass.Float32(0.0) - for k in cutlass.range_constexpr(TOP_K): - acc_o = _fma_rn( - s_w.load(idx=t * cutlass.Int32(TOP_K) + cutlass.Int32(k)), - ys.load( - idx=( - (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) - + cutlass.Int32(k) - ) - * cutlass.Int32(R2) - + lane - ), - acc_o, - ) - if cutlass.const_expr( - YDUMP - ): # debug: y of (token, k) into the partials workspace - partials.store( - cutlass.Float32( - ys.load( - idx=( - (bb * cutlass.Int32(M_MAX) + t) * cutlass.Int32(TOP_K) - + cutlass.Int32(k) - ) - * cutlass.Int32(R2) - + lane - ) - ).bitcast(cutlass.Int32), - idx=(t * cutlass.Int32(TOP_K) + cutlass.Int32(k)) * cutlass.Int32(H) - + (blk0 + bb) * cutlass.Int32(R2) - + cutlass.Int32(4) * (lane & cutlass.Int32(7)) - + (lane >> cutlass.Int32(3)), - ) - ch = ( - t * cutlass.Int32(H) - + (blk0 + bb) * cutlass.Int32(R2) - + cutlass.Int32(4) * (lane & cutlass.Int32(7)) - + (lane >> cutlass.Int32(3)) - ) - if cutlass.const_expr(SH_ADD): - # + y_sh (this row's shared output, from its shared CTA), as shared_output.add_(routed_output) - wy = wys[bt] - while ((wy >> cutlass.Int32(16)) & cutlass.Int32(0xFFFF)) != tag16: - wy = _ld_relaxed_u32(dnx.subview(ch).data_ptr().toint()) - acc_o = _add_rn( - (wy << cutlass.Int32(16)).bitcast(cutlass.Float32), _bf16_round(acc_o) - ) - out.store(cutlass.BFloat16(acc_o), idx=ch) - if ar_on: - cl = ( - (blk0 + bb) * cutlass.Int32(R2) - + cutlass.Int32(4) * (lane & cutlass.Int32(7)) - + (lane >> cutlass.Int32(3)) - ) - if cutlass.const_expr(AR_PRE): - _ar_push_pre( - ard, - pars[bt], - _bf16_round(acc_o), - ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), - t, - cl, - ) - else: - _ar_push( - arp, - _bf16_round(acc_o), - ((lane >> cutlass.Int32(3)) & cutlass.Int32(1)) == cutlass.Int32(0), - t, - cl, - ) - - # ---- teardown - if tidx == 0: - _stamp(tdbg, bx, 7) - epochs.store(ep + cutlass.Int32(1), idx=bx) - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - cute.arch.sync_warp() - prims.barrier_cta_sync(0) - if ar_on | cutlass.Boolean(LATE_TRIGGER): - # the all-reduce polls the pushed slots without a grid wait: release it after every push of this CTA - cute.arch.griddepcontrol_launch_dependents() - if warp == MMA_WARP: - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - prims.tcgen05_dealloc( - cutlass.inttoptr(tmem_ptr_i32.load(), 6, cutlass.Float32), TMEM_COLS - ) - - -@cute.jit -def m3_moe( - a1_tensor: cute.Tensor, - sfa1_tensor: cute.Tensor, - a2_tensor: cute.Tensor, - sfa2_tensor: cute.Tensor, - logits: cute.Tensor, - bias: cute.Tensor, - x32: cute.Tensor, - xsf32: cute.Tensor, - escal: cute.Tensor, - plan: cute.Tensor, - partials: cute.Tensor, - hbuf32: cute.Tensor, - epochs: cute.Tensor, - out: cute.Tensor, - route_dbg: cute.Tensor, - tdbg: cute.Tensor, - xb32: cute.Tensor, - rw: cute.Tensor, - sfsc: cute.Tensor, - lgbuf: cute.Tensor, - guw: cute.Tensor, - gusf: cute.Tensor, - dnw: cute.Tensor, - dnsf: cute.Tensor, - hx: cute.Tensor, - dnx: cute.Tensor, - arp: cute.Tensor, - s1w: cute.Tensor, - s2w: cute.Tensor, - img32: cute.Tensor, - b2cnt: cute.Tensor, - stream: cuda_driver.CUstream, -): - tma_a1_desc = cuda.create_tensor_map_tiled( - global_address=a1_tensor.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[H // 2, TWO_I, E_LOCAL], - global_strides=[H // 2 // 16, TWO_I * (H // 2) // 16], - box_dims=(KT_BYTES, R1, 1), - swizzle=cuda.TensorMapSwizzle.s128b, - ) - tma_a2_desc = cuda.create_tensor_map_tiled( - global_address=a2_tensor.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[I_TP // 2, H, E_LOCAL], - global_strides=[I_TP // 2 // 16, H * (I_TP // 2) // 16], - box_dims=(KT_BYTES, R2, 1), - swizzle=cuda.TensorMapSwizzle.s128b, - ) - tma_s1_desc = cuda.create_tensor_map_tiled( - global_address=sfa1_tensor.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[128, E_LOCAL * TWO_I * (H // 16) // 128], - global_strides=[128 // 16], - box_dims=(128, NUM_BYTES_SFA_RAW // 128), - swizzle=cuda.TensorMapSwizzle.none, - ) - tma_s2_desc = cuda.create_tensor_map_tiled( - global_address=sfa2_tensor.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[128, E_LOCAL * H * (I_TP // 16) // 128], - global_strides=[128 // 16], - box_dims=(128, NUM_BYTES_SFA_RAW // 4 // 128), - swizzle=cuda.TensorMapSwizzle.none, - ) - tma_gu_desc = tma_a1_desc - tma_dn_desc = tma_a2_desc - if cutlass.const_expr(SHARED and not SH_PRETILE): - tma_gu_desc = cuda.create_tensor_map_tiled( - global_address=guw.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[H, 2 * SH_I], - global_strides=[H // 16], - box_dims=(KT_BYTES, 32), - swizzle=cuda.TensorMapSwizzle.s128b, - ) - tma_dn_desc = cuda.create_tensor_map_tiled( - global_address=dnw.iterator.toint(), - dtype=cutlass.Uint8, - global_dims=[SH_I, H], - global_strides=[SH_I // 16], - box_dims=(KT_BYTES, 128), - swizzle=cuda.TensorMapSwizzle.s128b, - ) - m3_moe_kernel( - tma_a1_desc, - tma_a2_desc, - tma_s1_desc, - tma_s2_desc, - tma_gu_desc, - tma_dn_desc, - logits, - bias, - x32, - xsf32, - escal, - plan, - partials, - hbuf32, - epochs, - out, - route_dbg, - tdbg, - xb32, - rw, - sfsc, - lgbuf, - guw, - gusf, - dnw, - dnsf, - hx, - dnx, - arp, - s1w, - s2w, - img32, - b2cnt, - ).launch( - grid=(NCTA + NSH, 1, 1), - block=(THREADS, 1, 1), - cluster=(1, 1, 1), - stream=stream, - use_pdl=True, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py index 9f176116d722..7c5f390ef65a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_moe/op.py @@ -53,6 +53,34 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.phase_plan import ( + MoePlanGeometry, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.phase_plan import ( + make_phase_plan_tp8 as _shared_make_phase_plan, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + M3MoeLayer as M3MoeLayer, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + _view as _view, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + ar_push_desc as ar_push_desc, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + fine_ranges, + pretile_shared, + shared_views, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.runtime import ( + is_supported as is_supported, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.state import ( + M3MoeStateBase, + MoeStateBindings, +) + HIDDEN = 6144 # capacity: fixed: the checkpoint's hidden size INTER = 384 # capacity: fixed: the routed experts' width per tp8 rank NUM_EXPERTS = 128 @@ -93,7 +121,9 @@ M3_MOE_LAG = 2 -_ENGINE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "m3_moe_kernel.py") +_ENGINE_PATH = os.path.abspath( + os.path.join(os.path.dirname(__file__), "../../../kernels/moe/m3_moe_kernel.py") +) def _kernel_module(config: dict, path: str): @@ -108,128 +138,13 @@ def _kernel_module(config: dict, path: str): ) spec = importlib.util.spec_from_file_location(name, path) mod = importlib.util.module_from_spec(spec) - mod.M3_CONFIG = dict(config) + mod.M3_CONFIG = dict(config, tp_size=8) sys.modules[name] = mod spec.loader.exec_module(mod) _modules[key] = mod return mod -def _view(t: torch.Tensor, align: int): - from cutlass.cute.runtime import from_dlpack - - return from_dlpack(t.detach(), assumed_align=align) - - -def is_supported(device: Optional[torch.device] = None) -> tuple: - if not torch.cuda.is_available(): - return False, "no CUDA" - major, _ = torch.cuda.get_device_capability(device) - if major != 10: - return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" - try: - import cutlass # noqa: F401 - except ImportError: - return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" - return True, "" - - -def ar_push_desc( - mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device -) -> torch.Tensor: - """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A - of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). - With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old - behaviour.""" - return torch.tensor( - [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device - ) - - -def _shared_views(shared_weights, who): - """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" - if shared_weights is None: - raise ValueError(f"{who} (shared): shared_weights are required") - gu_w, gu_sf, dn_w, dn_sf = shared_weights - want = {"gate_up": (gu_w, (2 * SHARED_INTER, HIDDEN)), "down": (dn_w, (HIDDEN, SHARED_INTER))} - for name, (t, shape) in want.items(): - if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): - raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") - if ( - gu_sf.numel() != 2 * SHARED_INTER * HIDDEN // 32 - or dn_sf.numel() != HIDDEN * SHARED_INTER // 32 - ): - raise ValueError( - f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" - ) - return ( - gu_w.view(torch.uint8).view(-1), - gu_sf.view(torch.uint8).reshape(-1), - dn_w.view(torch.uint8).view(-1), - dn_sf.view(torch.uint8).reshape(-1), - ) - - -def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24) -> torch.Tensor: - """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB - shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and - its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 - = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row - R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then - down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per - layer.""" - dev = gu8.device - H_, I_ = HIDDEN, SHARED_INTER - P = I_ // n_cta - split = 128 // (2 * P) - gt = H_ // 128 // split - dkt = I_ // 128 - gu = gu8.view(2 * I_, H_) - dn = dn8.view(H_, I_) - gkb, dkb = H_ // 32, I_ // 32 - - def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast - return sf[ - ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) - ] - - r = torch.arange(128, device=dev) - q = r >> 5 - ri = r & 31 - octant = r // P - chunk = torch.arange(8, device=dev) - swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) - out = torch.empty( - n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev - ) # capacity: fixed: a workspace of the kernel's fixed geometry - for s in range(n_cta): - # gate_up: MMA row R -> weight row, slice - nrow = (octant & 1) * I_ + P * s + r % P - sl = octant >> 1 - for g in range(gt): - kt = sl * gt + g # per row - tile = gu[ - nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) - ] # [128, 128] - img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) - out[s, g, :16384] = img.reshape(-1) - kb = torch.arange(4, device=dev) - codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] - atom = torch.empty(512, dtype=torch.uint8, device=dev) - atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) - out[s, g, 16384:] = atom - for gd in range(2 * dkt): - jt, kt = gd // dkt, gd % dkt - row0 = 256 * s + 128 * jt - tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] - img = tile.reshape(128, 8, 16)[r[:, None], swz] - out[s, gt + gd, :16384] = img.reshape(-1) - out[s, gt + gd, 16384:] = dnsf8[ - ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 - ] - return out.view(-1) - - # --------------------------------------------------------------------------------------------------------------------- # m3_moe (m3_moe_kernel.py): 1..8 tokens per call, group-phased weight stream. # --------------------------------------------------------------------------------------------------------------------- @@ -246,26 +161,6 @@ def _pretiled(gu8, gusf8, dn8, dnsf8) -> torch.Tensor: return hit[1] -def _fine_ranges(weights, total: int): - """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item - counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary - strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit - runs at least to that unit's end. Returns the item counts.""" - n = len(weights) - items = [0] * n - s = 0 - acc = 0.0 - for c in range(n): - acc += weights[c] - e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) - if s % _S1 and e < total: # started inside a unit: run to (at least) its end - e = max(e, min(total, (s // _S1 + 1) * _S1)) - items[c] = e - s - s = e - assert s == total - return items - - def make_phase_plan( num_ctas: int, m_max: int, @@ -281,76 +176,26 @@ def make_phase_plan( group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" - gcap = TOP_K * m_max - nb2 = HIDDEN // 32 - plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) - nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] - b0s = [sum(nbs[:c]) for c in range(num_ctas)] - blk_bytes = 32 * _ROW2_BYTES # one block, one slot - for g in range(TOP_K, gcap + 1): - ng = (g + 3) // 4 - for c in range(num_ctas): - plan[g, c, m_max] = b0s[c] | (nbs[c] << 16) - for p in range(ng): - n_p = min(4, g - 4 * p) - total = n_p * U1 * _S1 - fc2 = [0] * num_ctas - for c in range(num_ctas): - if p >= lag: - fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes - if p == ng - 1 and not ( - early0 and p == 0 and m_max >= 2 - ): # early0: group 0's ranges the same for every G - for q in range(max(0, ng - lag), ng): - fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) - items = [0] * num_ctas - if total >= _S1 * num_ctas: - - def counts_for(limit): - return [max(_S1, (limit - f) // _ITEM_BYTES) for f in fc2] - - lo, hi = 0, max(fc2) + total * _ITEM_BYTES - while lo < hi: - mid = (lo + hi) // 2 - if sum(counts_for(mid)) >= total: - hi = mid - else: - lo = mid + 1 - items = counts_for(lo) - extra = sum(items) - total - while extra > 0: - c = min( - (c for c in range(num_ctas) if items[c] > _S1), - key=lambda c: (-(fc2[c] + items[c] * _ITEM_BYTES), c), - ) - items[c] -= 1 - extra -= 1 - elif fine: - # water-fill: CTA c streams fc2[c] + w_c item bytes, equal where w_c > 0 - lo_f, hi_f = 0.0, float(max(fc2) + total * _ITEM_BYTES) - for _ in range(100): - mid = 0.5 * (lo_f + hi_f) - if sum(max(0.0, (mid - f) / _ITEM_BYTES) for f in fc2) >= total: - hi_f = mid - else: - lo_f = mid - wts = [max(0.0, (hi_f - f) / _ITEM_BYTES) for f in fc2] - sw = sum(wts) - items = _fine_ranges([w * total / sw for w in wts], total) - else: - chosen = sorted(sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // _S1]) - for c in chosen: - items[c] = _S1 - i0 = 0 - for c in range(num_ctas): - if items[c]: - plan[g, c, p] = i0 | ((i0 + items[c]) << 16) - i0 += items[c] - assert i0 == total - return plan, max(nbs) + return _shared_make_phase_plan( + num_ctas=num_ctas, + m_max=m_max, + lag=lag, + tail_w=tail_w, + early0=early0, + fine=fine, + geometry=MoePlanGeometry( + top_k=TOP_K, + hidden=HIDDEN, + u1=U1, + stages=_S1, + row2_bytes=_ROW2_BYTES, + item_bytes=_ITEM_BYTES, + ), + fine_ranges_fn=_fine_ranges, + ) -class M3MoeState: +class M3MoeState(M3MoeStateBase): """Workspace + compiled kernel of ``m3_moe`` (1..8 tokens per call) on one device, shared by the layers that run on it in one stream order. The call takes the bf16 hidden states and computes the router logits (bit-exact with m3_router_gemv_fp4) and the NVFP4 activation itself. ``shared``: the shared expert runs on SHARED_CTAS dedicated @@ -359,6 +204,15 @@ class M3MoeState: fused_router = True + BINDINGS = MoeStateBindings( + inter=INTER, + shared_inter=SHARED_INTER, + engine_path=_ENGINE_PATH, + compiled=_compiled, + pretiled=_pretiled, + cache_module_config=False, + ) + def __init__( self, device: torch.device, @@ -472,225 +326,14 @@ def __init__( self.dnx = torch.zeros(num_tokens * HIDDEN, dtype=torch.int32, **kw) self._fn = None - def layer( - self, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - e_score_correction_bias, - router_weight=None, - fc31_input_scale=None, - shared_weights=None, - ) -> "M3MoeLayer": - """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). - ``shared_weights`` (a ``shared`` state): the shared expert's (gate_up e4m3 [768, 6144], its swizzled UE8M0 - scales, down e4m3 [6144, 384], its swizzled UE8M0 scales).""" - layer = M3MoeLayer( - self, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - e_score_correction_bias, - router_weight, - fc31_input_scale, - ) - if self.shared: - layer.shared_weights = _shared_views(shared_weights, "m3_moe") - layer.shared_mxfp8 = tuple(shared_weights) - if self.mod.SH_PRETILE: - # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy - gu8, _, dn8, _ = layer.shared_weights - layer.shared_weights = ( - gu8, - _pretiled(*layer.shared_weights), - dn8, - self._dummy.view(torch.uint8), - ) - return layer - - def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: - m = self.m - if ( - hidden is None - or hidden.dtype != torch.bfloat16 - or tuple(hidden.shape) != (m, HIDDEN) - or not hidden.is_contiguous() - ): - raise ValueError( - f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" - ) - fz = ( - hidden.view(-1).view(torch.int32), - layer.router_weight.view(-1), - layer.fc31_input_scale.view(-1), - self.lgbuf, - ) - lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy - if out is None: - out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) - sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 - args = ( - layer.w3_w1_weight, - layer.w3_w1_weight_scale.view(-1), - layer.w2_weight, - layer.w2_weight_scale.view(-1), - lgv, - layer.bias, - xq, - xs, - layer.escal, - self.plan.view(-1), - self.partials, - self.hbuf32, - self.epochs, - out.view(-1), - self.route_dbg.view(-1), - self.tdbg.view(-1), - *fz, - *sh, - self.hx, - self.dnx, - self._ar_off if arp is None else arp, - layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), - layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), - self.img32, - self.cnt, - ) - if self._fn is None: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("m3_moe compiles on its first call: call it once before capture") - import cutlass.cute as cute - - aligns = [ - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 4, - 16, - 16, - 16, - 16, - 16, - 4, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - 16, - ] - sig = [_view(t, a) for t, a in zip(args, aligns)] - key = (_ENGINE_PATH,) + tuple(sorted(self.config.items())) - fn = _compiled.get(key) - if fn is None: - fn = _compiled[key] = cute.compile( - self.mod.m3_moe, - *sig, - cute.runtime.make_fake_stream(), - options="--enable-tvm-ffi", - ) - self._fn = fn - self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) - return out +def _shared_views(shared_weights, who): + return shared_views(shared_weights, who, shared_inter=SHARED_INTER) -class M3MoeLayer: - """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in - place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading - weights).""" - shared_weights = None # the kernel's view of the shared expert (a shared state) - shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given +def _pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24): + return pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta, shared_inter=SHARED_INTER) - def __init__( - self, - state, - w3_w1_weight, - w3_w1_weight_scale, - w2_weight, - w2_weight_scale, - fc31_alpha, - fc31_scale_c, - act_alpha, - act_beta, - act_clamp, - fc2_alpha, - bias, - router_weight=None, - fc31_input_scale=None, - ): - e = NUM_EXPERTS - expected = { - "w3_w1_weight": (w3_w1_weight, (e, 2 * INTER, HIDDEN // 2)), - "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * INTER, HIDDEN // 16)), - "w2_weight": (w2_weight, (e, HIDDEN, INTER // 2)), - "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, INTER // 16)), - } - for name, (t, shape) in expected.items(): - if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): - raise ValueError( - f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" - ) - self.state = state - self.w3_w1_weight = w3_w1_weight.view(torch.uint8) - self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) - self.w2_weight = w2_weight.view(torch.uint8) - self.w2_weight_scale = w2_weight_scale.view(torch.uint8) - self.escal = torch.stack( - [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] - ).float() - self.escal = self.escal.reshape(6, e).contiguous() - self.bias = bias.float().contiguous() - if ( - router_weight is None - or fc31_input_scale is None - or router_weight.dtype != torch.float32 - or tuple(router_weight.shape) != (e, HIDDEN) - or not router_weight.is_contiguous() - or fc31_input_scale.dtype != torch.float32 - or fc31_input_scale.numel() != 1 - ): - raise ValueError( - "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" - ) - self.router_weight = router_weight - self.fc31_input_scale = fc31_input_scale - def __call__( - self, - hidden: torch.Tensor, - out: Optional[torch.Tensor] = None, - arp: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output - into an all-reduce exchange.""" - return self.state.run(self, hidden, out, arp=arp) +def _fine_ranges(weights, total: int): + return fine_ranges(weights, total, stages=_S1) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py index 48a7827f59b3..a89f96fe0b10 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_prequant.py @@ -1,111 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Hand-off of an MXFP8-quantized activation from the kernel that produced it to the next MXFP8 GEMM. -A producer that already holds a BF16 row in registers (the multicast all-reduce + RMSNorm, the split-K attention -merge) also writes its MXFP8 form (E4M3 rows + UE8M0 scales in the 128x4-swizzled layout of -``trtllm::mxfp8_quantize(x, True)``) and publishes it here keyed by the BF16 tensor. The MXFP8 linear whose input -is that tensor takes it and skips its own quantize kernel. Every take clears the slot, so a hand-off is only ever -consumed by the first MXFP8 linear that runs after the producer. - -Module state shared across files (one instance each per process): - -* ``SLOT`` (PrequantSlot): published by ``m3_mc_allreduce`` (``quant=True``), the sparse-decode merge and the fused - decode attention (MXFP8 output); cleared by the sparse decode before it runs; taken by the MXFP8 linear dispatch - and by the one-launch QKV + attention + o_proj path (which also peeks ``pending``). -* ``AR_PUSH`` (ArPushSlot): armed by the fast forward around the attention call; the o_proj GEMV takes the target - and marks ``pushed``; the fast forward checks ``take_pushed`` before the attention boundary all-reduce. -* ``MOE_PUSH`` (MoePushSlot): armed by the fast forward before the MoE; the routed-expert engine takes ``engine`` - and the shared expert's down projection takes ``shared``, each marks its pushed output; the fast forward - checks ``take_pushed`` before the MoE boundary all-reduce. -""" - - -class PrequantSlot: - def __init__(self): - self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) - - def publish(self, x, xq, xsf): - self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) - - def clear(self): - self.pending = None - - def take(self, x): - """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" - pending, self.pending = self.pending, None - if ( - pending is None - or x.dim() != 2 - or x.data_ptr() != pending[0] - or x.shape[0] != pending[1] - or x.shape[1] != pending[2] - ): - return None - return pending[3], pending[4] +"""Target-local activation hand-off state.""" +from ...ops.gemm.prequant import ArPushSlot, MoePushSlot, PrequantSlot SLOT = PrequantSlot() - - -class ArPushSlot: - """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. - - The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output - it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" - - def __init__(self): - self.target = None - self.pushed = None # data_ptr of the pushed output - - def arm(self, target): - self.target = target - self.pushed = None - - def take_target(self): - target, self.target = self.target, None - return target - - def take_pushed(self, x): - pushed, self.pushed = self.pushed, None - return pushed is not None and x.data_ptr() == pushed - - AR_PUSH = ArPushSlot() - - -class MoePushSlot: - """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, - ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv - ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" - - def __init__(self): - self.engine = None - self.shared = None - self.engine_pushed = None # data_ptr of the pushed routed output - self.shared_pushed = None # data_ptr of the pushed shared output - - def arm(self, engine, shared): - self.engine, self.shared = engine, shared - self.engine_pushed = self.shared_pushed = None - - def disarm(self): - self.engine = self.shared = None - - def take_engine(self): - engine, self.engine = self.engine, None - return engine - - def take_shared(self): - shared, self.shared = self.shared, None - return shared - - def take_pushed(self, shared, routed): - """(routed pushed, shared pushed) for these outputs; clears the marks.""" - e = routed is not None and self.engine_pushed == routed.data_ptr() - s = shared is not None and self.shared_pushed == shared.data_ptr() - self.engine_pushed = self.shared_pushed = None - return e, s - - MOE_PUSH = MoePushSlot() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py index 30855eb7f05d..d92ac2eb018d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_router_gemv.py @@ -15,9 +15,15 @@ import torch -from ..ops.moe.m3_router_gemv import m3_router_gemv as _router_gemv_op -from ..ops.moe.m3_router_gemv_fp4 import m3_router_gemv_fp4 as _router_gemv_fp4_op -from ..ops.moe.m3_router_tc import m3_router_tc as _router_tc_op +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv import ( + m3_router_gemv as _router_gemv_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv_fp4 import ( + m3_router_gemv_fp4 as _router_gemv_fp4_op, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_tc import ( + m3_router_tc as _router_tc_op, +) # Opt-in dynamic shared memory left after the kernel's static reduction buffer (SM100: 227 KiB). _MAX_DYNAMIC_SMEM = 227 * 1024 - 1024 # capacity: fixed: the device's shared-memory limit diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py index 10bd966adf98..13758156ccc7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_select_long_dsl.py @@ -1,298 +1,47 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Top-16 block select for MiniMax-M3 decode rows of up to 9216 blocks, one row per 384-thread CTA (CuTe DSL). +"""Target-specific select policy and independent compilation cache.""" -Bit for bit the selection of this target's ``m3_select_long`` (``torch.ops.trtllm.minimax_m3_sm107_tp2_select_long``, -one 1024-thread CTA per row), from a CTA of 384 threads: +import sys +from typing import Dict, Tuple, cast -* Every block's score is forced (the first ``init_blocks`` and the last ``local_blocks`` of the row's valid blocks) and - mapped to an order key held in registers: thread t owns blocks t, t + 384, ... (4, 11 or 24 of them by the row's - valid count). Larger scores get smaller keys; NaN sorts above +inf. The key is the 1024-thread kernel's unsigned - key with its top bit flipped, so signed Int32 order here is that kernel's order. Every thread loads at a clamped - block and masks the result, so a thread's loads issue back to back. -* 384 is a multiple of 16, so the 16 interleaved groups ``block % 16`` are the lanes' ``lane % 16``. The weakest of - the 16 group maxima bounds the 16th-largest score, so the top 16 is among the keys at or above it. -* The candidates are compacted by a CTA prefix count, each as one int64 (key << 32 | block), and ranked exactly by - that value (one signed compare per pair); more than 1024 of them (keys tied at the bound) rank every candidate - against the whole row instead. -* Warp 0 emits the 16 picked block ids sorted ascending, -1 for an empty or -inf pick. - -The kernel waits on its grid dependency first; it releases the dependent sparse decode right after that wait -(``early_trigger``), once every score of its row is loaded (trigger phase 1) or at its exit (phase 0). Kernels are -compiled per (PDL launch, early trigger, trigger phase) with symbolic shapes and strides; ``prebuild`` compiles and -loads them before the executor's warmup, so a decode step never compiles one. -""" - -from typing import Dict, Tuple - -import cutlass -import cutlass.cute as cute +import cutlass as cutlass +import cutlass.cute as cute # noqa: F401 — used by the bound family implementation import torch -from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream, make_fake_tensor -from cutlass.utils.smem_allocator import SmemAllocator +from cutlass.cute.runtime import make_fake_compact_tensor as make_fake_compact_tensor +from cutlass.cute.runtime import make_fake_stream as make_fake_stream +from cutlass.cute.runtime import make_fake_tensor as make_fake_tensor -from tensorrt_llm._torch.kernels.blackwell.utils import ( - griddepcontrol_launch_dependents, - griddepcontrol_wait, +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + long_select_host as _select_shared, ) -# The 384-thread select reads its row's valid-block count before its grid wait instead of after it, which takes one -# dependent global load off the post-wait path; the selection is the same either way. Exact: the count -# (MsaMetadata.msa_n_valid_blocks) is written only before the step's first layer, behind a plain launch. -SELECT_N_VALID_PREWAIT = True - -SELECT_THREADS = 384 -SELECT_WARPS = SELECT_THREADS // 32 -# 1M tokens + draft tokens at 128 tokens per block. -MAX_BLOCKS = 9216 # capacity: fixed: the kernel's built blocks per row -MAX_CAND = 1024 # capacity: fixed: the kernel's candidate staging -TOPK = 16 # capacity: fixed: MSA's top-k -NO_KEY = 0x7FFFFFFF # capacity: fixed: a sentinel, not a size (an unused key slot (the reference's 0xffffffff)) -# The key of -inf (the reference's 0xff800000): never emitted. -NEG_INF_KEY = 0x7F800000 # capacity: fixed: a key encoding, not a size -INIT_SCORE = 1.0e30 -LOCAL_SCORE = 1.0e29 -# (keys per thread, rows up to this many valid blocks) -ITEM_CLASSES = ((4, 4 * SELECT_THREADS), (11, 11 * SELECT_THREADS), (24, MAX_BLOCKS)) - -# Shared memory of the phase, int32 words: candidates [MAX_CAND] int64 (key << 32 | block, so one signed compare -# orders them by (key, block)), picked keys / blocks [TOPK] each, per-warp group minima [SELECT_WARPS][16], group -# minima [16], per-warp candidate counts [SELECT_WARPS]. -_CAND = 0 -_PICK_KEY = _CAND + 2 * MAX_CAND -_PICK_BLOCK = _PICK_KEY + TOPK -_WARP_MIN = _PICK_BLOCK + TOPK -_GROUP_MIN = _WARP_MIN + SELECT_WARPS * 16 -_WARP_CNT = _GROUP_MIN + 16 -SMEM_WORDS = _WARP_CNT + SELECT_WARPS -SMEM_BYTES = 4 * SMEM_WORDS - - -def _key(score): - """Order key of a float score: smaller key = larger score, NaN first; signed order = the reference's unsigned.""" - bits = cutlass.Float32(score).bitcast(cutlass.Int32) - return cutlass.select_(bits < cutlass.Int32(0), bits & cutlass.Int32(0x7FFFFFFF), ~bits) - - -def _imin(a, b): - return cutlass.select_(b < a, b, a) - - -def _imax(a, b): - return cutlass.select_(b > a, b, a) - - -def _forced_key(scores, head, block, query, init_blocks, local_start): - raw = scores[head, block, query] - s = cutlass.select_( - block >= local_start, - cutlass.Float32(LOCAL_SCORE), - cutlass.select_(block < init_blocks, cutlass.Float32(INIT_SCORE), cutlass.Float32(raw)), - ) - return _key(s) - - -@cute.jit -def _select_row( - row, - scores: cute.Tensor, # float32 [num_kv_heads, num_blocks, total_q], any strides - n_valid: cute.Tensor, # int32 [>= total_q] - out: cute.Tensor, # int32 [total_q, num_kv_heads, 16] - init_blocks, - local_blocks, - num_kv_heads, - smem: cute.Tensor, # int32 [SMEM_WORDS], shared memory - bar_id: cutlass.Constexpr, - loaded_trigger: cutlass.Constexpr, - raw_prewait, - n_valid_prewait: cutlass.Constexpr, -): - """Select row ``row`` (= query * num_kv_heads + head) with this CTA's SELECT_THREADS threads. Uses named barrier - ``bar_id`` over SELECT_THREADS threads. ``loaded_trigger`` releases the dependent grid once every score of the - row is loaded (trigger phase 1). With ``n_valid_prewait`` the row's valid-block count is ``raw_prewait``, read by - the caller before its grid wait; otherwise it is read here.""" - tid, _, _ = cute.arch.thread_idx() - lane = tid % cutlass.Int32(32) - warp = tid // cutlass.Int32(32) - query = row // num_kv_heads - head = row % num_kv_heads - if cutlass.const_expr(n_valid_prewait): - raw = raw_prewait - else: - raw = n_valid[query] - valid = _imax(_imin(raw, cutlass.Int32(scores.shape[1])), cutlass.Int32(0)) - local_start = _imax(raw - local_blocks, cutlass.Int32(0)) - cand = cute.make_tensor( - cute.recast_ptr(smem.iterator + _CAND, dtype=cutlass.Int64), cute.make_layout(MAX_CAND) - ) - - if tid < TOPK: - smem[_PICK_KEY + tid] = cutlass.Int32(NO_KEY) - smem[_PICK_BLOCK + tid] = cutlass.Int32(0) - - if valid <= TOPK: - # Every valid block is picked (thread t holds block t). - if tid < valid: - smem[_PICK_KEY + tid] = _forced_key(scores, head, tid, query, init_blocks, local_start) - smem[_PICK_BLOCK + tid] = tid - if cutlass.const_expr(loaded_trigger): - griddepcontrol_launch_dependents() - for ci in cutlass.range_constexpr(len(ITEM_CLASSES)): - items = ITEM_CLASSES[ci][0] - lo = TOPK if ci == 0 else ITEM_CLASSES[ci - 1][1] - hi = ITEM_CLASSES[ci][1] - if (valid > lo) & (valid <= hi): - key = cute.make_rmem_tensor(items, cutlass.Int32) - local_min = cutlass.Int32(NO_KEY) - for s in cutlass.range_constexpr(items): - # Every thread loads (a clamped block past the row) so the loads issue back to back, unbranched. - block = tid + cutlass.Int32(s * SELECT_THREADS) - k = _forced_key( - scores, head, _imin(block, valid - 1), query, init_blocks, local_start - ) - key[s] = cutlass.select_(block < valid, k, cutlass.Int32(NO_KEY)) - local_min = _imin(local_min, key[s]) - local_min = _imin(local_min, cute.arch.shuffle_sync_bfly(local_min, 16)) - if lane < 16: - smem[_WARP_MIN + warp * cutlass.Int32(16) + lane] = local_min - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - if cutlass.const_expr(loaded_trigger): - griddepcontrol_launch_dependents() - if tid < 16: - g = cutlass.Int32(NO_KEY) - for w in cutlass.range_constexpr(SELECT_WARPS): - g = _imin(g, smem[_WARP_MIN + w * 16 + tid]) - smem[_GROUP_MIN + tid] = g - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - bound = cutlass.Int32( - -2147483648 - ) # the smallest key: signed keys, unlike the reference's unsigned ones - for g in cutlass.range_constexpr(16): - bound = _imax(bound, smem[_GROUP_MIN + g]) - - # Candidates (key <= bound; unused slots hold NO_KEY > any bound): a CTA prefix count gives the slots. - cnt = cutlass.Int32(0) - for s in cutlass.range_constexpr(items): - cnt += cutlass.select_(key[s] <= bound, cutlass.Int32(1), cutlass.Int32(0)) - incl = cnt - for d in cutlass.range_constexpr(5): - up = cute.arch.shuffle_sync( - incl, _imax(lane - cutlass.Int32(1 << d), cutlass.Int32(0)) - ) - incl += cutlass.select_(lane >= cutlass.Int32(1 << d), up, cutlass.Int32(0)) - if lane == 31: - smem[_WARP_CNT + warp] = incl - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - before = cutlass.Int32(0) - total = cutlass.Int32(0) - for w in cutlass.range_constexpr(SELECT_WARPS): - c = smem[_WARP_CNT + w] - before += cutlass.select_(cutlass.Int32(w) < warp, c, cutlass.Int32(0)) - total += c - slot = before + incl - cnt - for s in cutlass.range_constexpr(items): - if key[s] <= bound: - if slot < MAX_CAND: - cand[slot] = (cutlass.Int64(key[s]) << cutlass.Int64(32)) | cutlass.Int64( - tid + cutlass.Int32(s * SELECT_THREADS) - ) - slot += 1 - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - - if total <= MAX_CAND: - for i in cutlass.range(tid, total, SELECT_THREADS, unroll=1): - mine = cand[i] - rank = cutlass.Int32(0) - for j in cutlass.range(total, unroll=8): - rank += cutlass.select_(cand[j] < mine, cutlass.Int32(1), cutlass.Int32(0)) - if rank < TOPK: - smem[_PICK_KEY + rank] = cutlass.Int32(mine >> cutlass.Int64(32)) - smem[_PICK_BLOCK + rank] = cutlass.Int32(mine & cutlass.Int64(0xFFFFFFFF)) - else: - # More than MAX_CAND keys tied at the bound: rank each candidate against the whole row. - for s in cutlass.range_constexpr(items): - if key[s] <= bound: - block = tid + cutlass.Int32(s * SELECT_THREADS) - rank = cutlass.Int32(0) - for j in cutlass.range(valid, unroll=1): - kj = _forced_key(scores, head, j, query, init_blocks, local_start) - less = (kj < key[s]) | ((kj == key[s]) & (j < block)) - rank += cutlass.select_(less, cutlass.Int32(1), cutlass.Int32(0)) - if rank < TOPK: - smem[_PICK_KEY + rank] = key[s] - smem[_PICK_BLOCK + rank] = block - cute.arch.barrier(barrier_id=bar_id, number_of_threads=SELECT_THREADS) - - if warp == 0: - # The 16 picked block ids ascending (NO_KEY = empty or -inf, emitted as -1), bitonic over the warp. - v = cutlass.Int32(NO_KEY) - if lane < TOPK: - pk = smem[_PICK_KEY + lane] - if (pk != cutlass.Int32(NO_KEY)) & (pk != cutlass.Int32(NEG_INF_KEY)): - v = smem[_PICK_BLOCK + lane] - for size_log in cutlass.range_constexpr(1, 6): - for stride_log in cutlass.range_constexpr(size_log - 1, -1, -1): - other = cute.arch.shuffle_sync_bfly(v, 1 << stride_log) - ascending = (lane & cutlass.Int32(1 << size_log)) == cutlass.Int32(0) - lower = (lane & cutlass.Int32(1 << stride_log)) == cutlass.Int32(0) - v = cutlass.select_(lower == ascending, _imin(v, other), _imax(v, other)) - if lane < TOPK: - out[query, head, lane] = cutlass.select_( - v == cutlass.Int32(NO_KEY), cutlass.Int32(-1), v - ) - - -class _SelectLongDslKernel: - """One row per CTA: the grid wait, then ``_select_row``. ``n_valid_prewait`` reads the row's valid-block count - before the grid wait, so that dependent load leaves the post-wait critical path. That is exact because - ``n_valid`` is written only by ``on_update_kv_lens`` and the host copy in ``prepare()``, both before the step's - first layer and behind a plain (non-PDL) launch, so no grid this select can overlap writes it - (``attention/test_modeling_v2_m3_n_valid_contract.py``).""" - - def __init__( - self, use_pdl: bool, early_trigger: bool, trigger_phase: int, n_valid_prewait: bool - ): - self.use_pdl = use_pdl - self.early_trigger = early_trigger - self.loaded_trigger = (not early_trigger) and trigger_phase == 1 - self.n_valid_prewait = n_valid_prewait - - @cute.kernel - def kernel(self, scores, n_valid, out, init_blocks, local_blocks, num_kv_heads): - row, _, _ = cute.arch.block_idx() - smem = SmemAllocator() - words = smem.allocate_tensor(cutlass.Int32, cute.make_layout(SMEM_WORDS), byte_alignment=16) - raw_prewait = cutlass.Int32(0) - if cutlass.const_expr(self.n_valid_prewait): - raw_prewait = n_valid[row // num_kv_heads] - griddepcontrol_wait() - if cutlass.const_expr(self.early_trigger): - griddepcontrol_launch_dependents() - _select_row( - row, - scores, - n_valid, - out, - init_blocks, - local_blocks, - num_kv_heads, - words, - bar_id=1, - loaded_trigger=self.loaded_trigger, - raw_prewait=raw_prewait, - n_valid_prewait=self.n_valid_prewait, - ) - - @cute.jit - def __call__(self, scores, n_valid, out, init_blocks, local_blocks, stream): - num_kv_heads = scores.shape[0] - self.kernel(scores, n_valid, out, init_blocks, local_blocks, num_kv_heads).launch( - grid=(scores.shape[2] * num_kv_heads, 1, 1), - block=(SELECT_THREADS, 1, 1), - stream=stream, - use_pdl=self.use_pdl, - ) +from ...kernels.attention.m3_select_long_dsl import _CAND as _CAND +from ...kernels.attention.m3_select_long_dsl import _GROUP_MIN as _GROUP_MIN +from ...kernels.attention.m3_select_long_dsl import _PICK_BLOCK as _PICK_BLOCK +from ...kernels.attention.m3_select_long_dsl import _PICK_KEY as _PICK_KEY +from ...kernels.attention.m3_select_long_dsl import _WARP_CNT as _WARP_CNT +from ...kernels.attention.m3_select_long_dsl import _WARP_MIN as _WARP_MIN +from ...kernels.attention.m3_select_long_dsl import INIT_SCORE as INIT_SCORE +from ...kernels.attention.m3_select_long_dsl import ITEM_CLASSES as ITEM_CLASSES +from ...kernels.attention.m3_select_long_dsl import LOCAL_SCORE as LOCAL_SCORE +from ...kernels.attention.m3_select_long_dsl import MAX_BLOCKS as MAX_BLOCKS +from ...kernels.attention.m3_select_long_dsl import MAX_CAND as MAX_CAND +from ...kernels.attention.m3_select_long_dsl import NEG_INF_KEY as NEG_INF_KEY +from ...kernels.attention.m3_select_long_dsl import NO_KEY as NO_KEY +from ...kernels.attention.m3_select_long_dsl import SELECT_THREADS as SELECT_THREADS +from ...kernels.attention.m3_select_long_dsl import SELECT_WARPS as SELECT_WARPS +from ...kernels.attention.m3_select_long_dsl import SMEM_BYTES as SMEM_BYTES +from ...kernels.attention.m3_select_long_dsl import SMEM_WORDS as SMEM_WORDS +from ...kernels.attention.m3_select_long_dsl import TOPK as TOPK +from ...kernels.attention.m3_select_long_dsl import _forced_key as _forced_key +from ...kernels.attention.m3_select_long_dsl import _imax as _imax +from ...kernels.attention.m3_select_long_dsl import _imin as _imin +from ...kernels.attention.m3_select_long_dsl import _key as _key +from ...kernels.attention.m3_select_long_dsl import _select_row as _select_row +from ...kernels.attention.m3_select_long_dsl import _SelectLongDslKernel as _SelectLongDslKernel +SELECT_N_VALID_PREWAIT = True #: Compiled kernels by (PDL launch, early trigger, trigger phase, valid count read before the wait). kernel_cache: Dict[Tuple[bool, bool, int, bool], object] = {} @@ -312,28 +61,7 @@ def kernel_key( def _compile(key: Tuple[bool, bool, int, bool]): - if key not in kernel_cache: - heads, blocks, rows = cute.sym_int(), cute.sym_int(), cute.sym_int() - scores = make_fake_tensor( - cutlass.Float32, - (heads, blocks, rows), - stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64()), - ) - n_valid = make_fake_compact_tensor(cutlass.Int32, (cute.sym_int(),), stride_order=(0,)) - out = make_fake_tensor( - cutlass.Int32, (rows, heads, TOPK), stride=(cute.sym_int64(), cute.sym_int64(), 1) - ) - kernel_cache[key] = cute.compile( - _SelectLongDslKernel(*key), - scores, - n_valid, - out, - cutlass.Int32(0), - cutlass.Int32(0), - make_fake_stream(use_tvm_ffi_env_stream=True), - options="--enable-tvm-ffi", - ) - return kernel_cache[key] + return _select_shared._compile(_SELECT_TARGET, key) def select_long_dsl( @@ -347,21 +75,22 @@ def select_long_dsl( trigger_phase: int, ) -> None: """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" - _compile(kernel_key(use_pdl, early_trigger, trigger_phase))( - scores, n_valid_blocks, out, init_blocks, local_blocks + return _select_shared.select_long_dsl( + _SELECT_TARGET, + scores, + n_valid_blocks, + out, + init_blocks, + local_blocks, + use_pdl, + early_trigger, + trigger_phase, ) def prebuild(keys, device) -> list: """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" - scores = torch.zeros( - 1, 2048, 4, dtype=torch.float32, device=device - ) # capacity: fixed: a dummy prebuild row, not a size - n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) - out = torch.empty(1, 4, TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) - built = [] - for key in keys: - key = kernel_key(*key) - _compile(key)(scores, n_valid, out, 0, 1) - built.append(key) - return built + return _select_shared.prebuild(_SELECT_TARGET, keys, device) + + +_SELECT_TARGET = cast(_select_shared.LongSelectTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py index 8caaee389d0a..63463423a602 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_shard_accept.py @@ -21,6 +21,10 @@ import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.accept_workspace import ( + create_accept_workspace, +) + from . import m3_report #: Verify rows per call, and exchange rows per call half of the workspace. @@ -56,49 +60,7 @@ def _num_sms() -> int: def create_workspace(mapping, device, slot_rows: int) -> AcceptWorkspace: - """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the - local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call - outside CUDA-graph capture, in the same order on every rank.""" - words = 2 * slot_rows * max(mapping.tp_size, 1) - if mapping.tp_size > 1: - from tensorrt_llm._torch.distributed.ops import ( - _make_mnnvl_mcast_buffer, - _mnnvl_workspace_all_succeeded, - ) - from tensorrt_llm._utils import mpi_comm - - session = mpi_comm() - session_group = session.Get_group() - group = session_group.Incl(mapping.tp_group) - session_group.Free() - comm = session.Create_group(group) - group.Free() - if not _mnnvl_workspace_all_succeeded(comm, True): - comm.Free() - raise RuntimeError("exchange workspace creation failed on another rank") - handle = _make_mnnvl_mcast_buffer( - comm, max(words * 8, 4096), mapping, mapping.is_multi_node() - ) - uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) - mc = handle.get_mc_buffer((words,), torch.int64, 0) - uc.zero_() - torch.cuda.synchronize() - if not _mnnvl_workspace_all_succeeded(comm, True): - raise RuntimeError("exchange workspace creation failed on another rank") - mc_ptr, rank = mc.data_ptr(), mapping.tp_rank - else: - uc = torch.zeros(words, dtype=torch.int64, device=device) - mc_ptr, rank, handle, comm = uc.data_ptr(), 0, None, None - return AcceptWorkspace( - uc=uc, - mc_ptr=mc_ptr, - keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), - counters=torch.zeros(2, dtype=torch.int32, device=uc.device), - rank=rank, - world_size=max(mapping.tp_size, 1), - handle=handle, - comm=comm, - ) + return create_accept_workspace(mapping, device, slot_rows, AcceptWorkspace) def _workspace(mapping, device) -> Optional[AcceptWorkspace]: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py index f5a1b229c777..03cae9c855c9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_target_head.py @@ -12,13 +12,30 @@ from __future__ import annotations +import sys from dataclasses import dataclass -from typing import Optional +from typing import Optional, cast import torch -from ..ops.gemm.m3_target_head import CTA_WARPS, MAX_ROWS, encode_chunked, round_trip_ok -from ..ops.gemm.m3_target_head import m3_target_head as _target_head_op +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + target_head_state as _head_shared, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + CTA_WARPS as CTA_WARPS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + MAX_ROWS as MAX_ROWS, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + encode_chunked as encode_chunked, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + m3_target_head as _target_head_op, # noqa: F401 — used by the bound family implementation +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + round_trip_ok as round_trip_ok, +) @dataclass(eq=False) @@ -43,71 +60,24 @@ class TargetHead: def _num_sms() -> int: - global _NUM_SMS - if _NUM_SMS is None: - _NUM_SMS = torch.cuda.get_device_properties( - torch.cuda.current_device() - ).multi_processor_count - return _NUM_SMS + return _head_shared._num_sms(_HEAD_TARGET) def head_supported(lm_head) -> bool: """A bias-free 2-D BF16 head of whole 16-row tiles whose K the coding covers (a multiple of 64).""" - weight = getattr(lm_head, "weight", None) - if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: - return False - return ( - getattr(lm_head, "bias", None) is None - and weight.shape[0] % 16 == 0 - and weight.shape[1] % 64 == 0 - ) + return _head_shared.head_supported(_HEAD_TARGET, lm_head) def prepare_head(lm_head, grid: Optional[int] = None) -> TargetHead: """Code ``lm_head``'s local shard and check it decodes to the BF16 weight bit for bit (fails closed: a shard that does not is returned with ``coded=False`` and nothing kept). Built once, eagerly, before the executor's warmup (``M3Eagle3Worker.prebuild_target_head``).""" - weight = lm_head.weight - with torch.no_grad(): - w = weight.detach() - packed, base, exc_off, exc_pos, exc_w = encode_chunked(w) - if not round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): - return TargetHead(weight=weight, coded=False) - if grid is None: - grid = _num_sms() - n = weight.shape[0] - return TargetHead( - weight=weight, - coded=True, - w=packed, - base=base, - exc_off=exc_off, - exc_pos=exc_pos, - exc_w=exc_w.view(-1), - part=torch.zeros(grid * CTA_WARPS * 2 * 4 * 32, dtype=torch.float32, device=weight.device), - counter=torch.zeros((n + 15) // 16, dtype=torch.int32, device=weight.device), - grid=grid, - ) + return _head_shared.prepare_head(_HEAD_TARGET, lm_head, grid) def logits(head: TargetHead, x: torch.Tensor, use_pdl: bool = True) -> torch.Tensor: """This rank's BF16 logits shard [rows, n] of the rows x [rows <= MAX_ROWS, K] (BF16) from the coded shard.""" - assert head.coded and 1 <= x.shape[0] <= MAX_ROWS, (head.coded, x.shape) - if x.stride(1) != 1 or x.stride(0) % 8 or x.data_ptr() % 16: - x = x.contiguous() - n = head.weight.shape[0] - out = torch.empty((x.shape[0], n), dtype=torch.bfloat16, device=x.device) - return _target_head_op( - head.w, - x, - out, - head.part, - head.counter, - head.exc_off, - head.exc_pos, - head.exc_w, - head.base, - n, - head.grid, - use_pdl, - ) + return _head_shared.logits(_HEAD_TARGET, head, x, use_pdl) + + +_HEAD_TARGET = cast(_head_shared.HeadTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py index db0077117fe3..5d99aab162d7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/m3_tokmap_guard.py @@ -17,61 +17,20 @@ from __future__ import annotations -from typing import Callable, Dict, Optional, Tuple +import sys +from typing import Callable, Dict, Optional, Tuple, cast import torch -from tensorrt_llm.bindings.internal import thop -from tensorrt_llm.logger import logger +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention import ( + token_map as _shared, +) +from tensorrt_llm.bindings.internal import thop as thop +from tensorrt_llm.logger import logger as logger #: Argument names of ``thop.trtllm_gen_generation_preprocess`` in order (the nanobind binding). -_ARG_NAMES = ( - "qkv_input", - "workspace", - "sequence_lengths", - "spec_decoding_generation_lengths", - "spec_decoding_position_offsets", - "kv_cache_block_offsets", - "host_kv_cache_pool_pointers", - "host_kv_cache_pool_mapping", - "kv_scale_orig_quant", - "kv_scale_quant_orig", - "attention_output_orig_quant", - "rotary_inv_freq", - "rotary_cos_sin", - "mrope_position_deltas", - "layer_idx", - "seq_offset", - "num_heads", - "num_kv_heads", - "head_size", - "tokens_per_block", - "kv_cache_quant_mode", - "max_attention_window_size", - "cyclic_attention_window_size", - "num_tokens", - "batch_beam", - "input_seq_length", - "max_past_kv_length", - "rotary_embedding_dim", - "rotary_embedding_base", - "rotary_embedding_scale_type", - "rotary_embedding_scale", - "rotary_embedding_max_positions", - "position_embedding_type", - "bmm1_scale", - "bmm2_scale", - "fp8_context_fmha", - "predicted_tokens_per_seq", - "attention_chunk_size", - "multi_processor_count", - "total_num_blocks", - "kv_factor", - "need_build_kv_cache_metadata", - "cross_attention", - "skip_fmha_workspace", -) -_INT2_BYTES = 8 # capacity: fixed: an element size +_ARG_NAMES = _shared._ARG_NAMES +_INT2_BYTES = _shared._INT2_BYTES _original: Optional[Callable] = None _max_rows = 0 @@ -91,70 +50,20 @@ def _no_missed(rows: int, reason: str) -> None: _on_missed: Callable[[int, str], None] = _no_missed -def _indexed(device: torch.device) -> torch.device: - """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" - if device.type == "cuda" and device.index is None: - return torch.device("cuda", torch.cuda.current_device()) - return device +_indexed = _shared._indexed def _one_request_map(device: torch.device, rows: int) -> torch.Tensor: """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" - device = _indexed(device) - m = _maps.get(device) - if m is None or m.shape[0] < rows: - reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" - _on_missed(rows, reason) - raise RuntimeError(f"m3_tokmap_guard: {reason}") - return m[:rows] + return _shared._one_request_map(_TARGET, device, rows) def _map_offset(args: dict, rows: int) -> int: - x = args["qkv_input"] - skip = bool(args.get("skip_fmha_workspace", False)) - key = ( - x.dtype, - rows, - int(args["num_heads"]), - int(args["head_size"]), - int(args["rotary_embedding_dim"]), - int(args["num_kv_heads"]), - skip, - ) - off = _offsets.get(key) - if off is None: - layout = thop.get_trtllm_gen_generation_workspace_layout( - x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip - ) - off, size = int(layout["tokens_info_offset"]), int(layout["tokens_info_size"]) - if off < 0 or off % 4 != 0 or size < _INT2_BYTES * rows: - raise RuntimeError( - f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " - f"(offset {off}, {size} bytes): the layout changed under this guard" - ) - _offsets[key] = off - return off + return _shared._map_offset(_TARGET, args, rows) def _guarded(*args, **kwargs): - a = dict(zip(_ARG_NAMES, args)) - a.update(kwargs) - rows = int(a["num_tokens"]) - if ( - a.get("spec_decoding_generation_lengths") is not None - and int(a["predicted_tokens_per_seq"]) > 1 - and int(a["batch_beam"]) == 1 - and rows > 1 - and not bool(a.get("cross_attention", False)) - ): - ws = a["workspace"] - src = _one_request_map(ws.device, rows) - off = _map_offset(a, rows) - # view (never reshape: a copy would leave the op's own workspace unwritten); raises on a strided workspace. - dst = ws.view(-1).view(torch.uint8)[off : off + _INT2_BYTES * rows] - dst.view(torch.int32).view(rows, 2).copy_(src) - _on_taken(rows) - return _original(*args, **kwargs) + return _shared._guarded(_TARGET, *args, **kwargs) def install( @@ -166,20 +75,7 @@ def install( """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this call installed the wrapper.""" - global _original, _max_rows, _on_taken, _on_missed - device = _indexed(device) - if device not in _maps or _maps[device].shape[0] < max_rows: - t = torch.arange(max_rows, dtype=torch.int32, device=device) - _maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() - _max_rows = max(_max_rows, max_rows) - _on_taken = on_taken - _on_missed = on_missed - if _original is not None: - return False - _original = thop.trtllm_gen_generation_preprocess - thop.trtllm_gen_generation_preprocess = _guarded - logger.info( - "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " - f"(up to {max_rows} rows)" - ) - return True + return _shared.install(_TARGET, device, max_rows, on_taken, on_missed) + + +_TARGET = cast(_shared.TokenMapTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py index 113878a991f9..d14b0e64a43b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_moe_engine.py @@ -12,21 +12,34 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm.logger import logger -from tensorrt_llm.quantization.utils.fp4_utils import ( - get_reorder_rows_for_gated_act_gemm_row_indices, - get_shuffle_matrix_a_row_indices, -) -from ..m3_moe.op import ( - HIDDEN, - INTER, - SHARED_INTER, - TOP_K, - M3MoeLayer, - M3MoeState, - ar_push_desc, - is_supported, -) +from ....ops.moe.reference import _E2M1 as _E2M1 +from ....ops.moe.reference import _F32_MAX as _F32_MAX +from ....ops.moe.reference import _LOG2E as _LOG2E +from ....ops.moe.reference import _RCP6 as _RCP6 +from ....ops.moe.reference import ROUTED_SCALING as ROUTED_SCALING +from ....ops.moe.reference import SHARED_ALPHA as SHARED_ALPHA +from ....ops.moe.reference import SHARED_LIMIT as SHARED_LIMIT +from ....ops.moe.reference import _bf16 as _bf16 +from ....ops.moe.reference import _e2m1_decode as _e2m1_decode +from ....ops.moe.reference import _e2m1_round as _e2m1_round +from ....ops.moe.reference import _e4m3_round as _e4m3_round +from ....ops.moe.reference import _f32 as _f32 +from ....ops.moe.reference import _fc1_epilogue as _fc1_epilogue +from ....ops.moe.reference import _gemm as _gemm +from ....ops.moe.reference import _nvfp4_act as _nvfp4_act +from ....ops.moe.reference import _nvfp4_expert as _nvfp4_expert +from ....ops.moe.reference import _row_perm as _row_perm +from ....ops.moe.reference import _shared_expert as _shared_expert_reference +from ....ops.moe.reference import _ue8m0_up as _ue8m0_up +from ....ops.moe.reference import moe_reference as _moe_reference +from ....ops.moe.reference import moe_route as moe_route +from ....ops.moe.reference import mxfp8_dequant as mxfp8_dequant +from ....ops.moe.reference import mxfp8_quant_dequant as mxfp8_quant_dequant +from ....ops.moe.reference import unswizzle_sf as unswizzle_sf +from ..m3_moe.op import HIDDEN as HIDDEN +from ..m3_moe.op import INTER, SHARED_INTER, M3MoeLayer, M3MoeState, ar_push_desc, is_supported +from ..m3_moe.op import TOP_K as TOP_K __all__ = [ "M3MoeLayer", @@ -41,204 +54,16 @@ "unswizzle_sf", ] -# The model's constants the kernels are built with: routed_scaling_factor, the shared expert's SwiGLU-OAI -# (swiglu_alpha, swiglu_limit, the (up + 1) offset). -ROUTED_SCALING = 2.0 -SHARED_ALPHA = 1.702 -SHARED_LIMIT = 7.0 -_LOG2E = 1.4426950216293334961 -_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum -_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) -_F32_MAX = 3.40282346638528859812e38 - - -def _f32(x: float) -> torch.Tensor: - return torch.tensor(x, dtype=torch.float32) - - -def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: - """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" - lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) - lo = lut[(packed & 0xF).long()] - hi = lut[(packed >> 4).long()] - return torch.stack((lo, hi), dim=-1).flatten(-2) - - -def _e2m1_round(x: torch.Tensor) -> torch.Tensor: - """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" - a = x.abs() - q = torch.full_like(a, 6.0) - for hi, value, inclusive in ( - (5.0, 4.0, True), - (3.5, 3.0, False), - (2.5, 2.0, True), - (1.75, 1.5, False), - (1.25, 1.0, True), - (0.75, 0.5, False), - (0.25, 0.0, True), - ): - q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) - return torch.where(x < 0, -q, q) - - -def _e4m3_round(x: torch.Tensor) -> torch.Tensor: - """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" - return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() - - -def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: - """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" - bits = x.contiguous().view(torch.int32) - code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) - return code.clamp(0, 254) - - -def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: - """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" - flat = sf.reshape(-1).view(torch.uint8) - r = torch.arange(rows, device=flat.device).view(-1, 1) - b = torch.arange(blocks, device=flat.device).view(1, -1) - padded_blocks = (blocks + 3) // 4 * 4 - idx = ( - (r // 128) * (128 * padded_blocks) - + (b // 4) * 512 - + (r % 32) * 16 - + ((r % 128) // 32) * 4 - + (b % 4) - ) - return flat[idx] - - -def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: - """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" - rows, k = q.shape - codes = unswizzle_sf(sf, rows, k // 32).float() - return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( - 32, dim=1 - ) - - -def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: - """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to - fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" - rows, k = x.shape - blocks = x.float().view(rows, k // 32, 32) - code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) - scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) - return (_e4m3_round(blocks / scale) * scale).view(rows, k) - - -def _bf16(x: torch.Tensor) -> torch.Tensor: - return x.to(torch.bfloat16).float() - - -def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" - return (a.double() @ b.double().t()).float() - - -def _row_perm(rows: int, gated: bool) -> torch.Tensor: - """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle - of epilogue tile 128).""" - probe = torch.empty(rows, 1) - perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) - if gated: - perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] - return perm - - -def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: - """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row - order (the global scale is not applied: it is the per-expert alpha).""" - rows, half = w.shape - codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() - vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) - return vals[perm_inv] - - -def moe_route(x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor): - """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" - logits = _gemm(x, router_weight) - sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() - keys = sig + bias - ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] - sk = keys.gather(1, ids) - bias[ids] - total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) - w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) - return ids, _bf16(w) - - -def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: - """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" - rows, k = x.shape - blocks = x.view(rows, k // 16, 16) - vmax = blocks.abs().amax(dim=-1, keepdim=True) - sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) - # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element - # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). - osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) - return (_e2m1_round(blocks * osc) * sff).view(rows, k) - - -def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: - """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 - input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with - act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" - a, c, alpha, beta, lim = (esc[i] for i in range(5)) - gc = torch.where(gate > lim, lim, gate) - gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) - uc = torch.minimum(torch.maximum(up, -lim), lim) - c1 = a * (alpha * _f32(_LOG2E).to(up.device)) - sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) - h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) - blocks = h.view(-1, 16) - amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) - sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) - osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) - return (_e2m1_round((blocks * osc) * c) * sff).view(-1) - - -def _shared_expert(x: torch.Tensor, shared_mxfp8) -> torch.Tensor: - """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the - model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" - gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 - gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) - gate, up = gu[:, :SHARED_INTER], gu[:, SHARED_INTER:] - a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) - ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 - sg = 1.0 / ( - torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 - ) - h = _bf16(ub * (sg * a)) - return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) +def _shared_expert( + x: torch.Tensor, + shared_mxfp8: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], +) -> torch.Tensor: + return _shared_expert_reference(x, shared_mxfp8, shared_inter=SHARED_INTER) def moe_reference(hidden: torch.Tensor, layer: M3MoeLayer, shared: bool) -> torch.Tensor: - """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" - x = hidden.float() - m = x.shape[0] - ids, weights = moe_route(x, layer.router_weight, layer.bias) - xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) - inv31 = torch.argsort(_row_perm(2 * INTER, gated=True)).to(x.device) - inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) - y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) - for e in torch.unique(ids).tolist(): - w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) - w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) - esc = layer.escal[:, e] - for t, k in (ids == e).nonzero().tolist(): - acc1 = _gemm(xr[t : t + 1], w31)[0] - # The buffer holds [w3 (up); w1 (gate)]. - h = _fc1_epilogue(acc1[:INTER], acc1[INTER:], esc) - y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) - acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) - for k in range(TOP_K): - acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() - routed = acc.float() - if not shared: - return routed.to(torch.bfloat16) - return (_shared_expert(x, layer.shared_mxfp8) + _bf16(routed)).to(torch.bfloat16) + return _moe_reference(hidden, layer, shared, inter=INTER, shared_inter=SHARED_INTER) _ENGAGED: set = set() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py deleted file mode 100644 index 636dfcd0b614..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv.py +++ /dev/null @@ -1,101 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""fp32 router GEMV of the decode MoE: ``logits[M, E] = float(x) @ W^T`` at one token.""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) - -_E4M3_MAX = 448.0 -_E2M1_MAX = 6.0 -#: Midpoints between consecutive e2m1 magnitudes; ``_TIE_UP[i]``: round-to-nearest-even takes the upper code. -_E2M1_MIDPOINTS = (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0) -_TIE_UP = (False, True, False, True, False, True, False) -#: The kernel builds each block's output scale through two ``rcp.approx.ftz.f32`` reciprocals (~2^-23 relative): -#: a value this close to an e2m1 midpoint may round to either neighbour, everything further out must match. -TIE_WINDOW = 2.0**-20 - - -def nvfp4_quantize( - x: torch.Tensor, sf_scale: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """``fp4_quantize(x, sf_scale, 16, False, False)``: (packed e2m1 [M, K / 2], linear E4M3 scales [M * K / 16], - the pre-rounding scaled values [M, K] that ``near_tie`` reads). - - Per 16 columns: scale = e4m3(sf_scale * amax / 6), data = x * sf_scale / scale rounded to nearest-even e2m1. - """ - gs = sf_scale.float().reshape(()) - m, k = x.shape - blocks = x.float().view(m, k // 16, 16) - vecmax = blocks.abs().amax(-1) - sf = (gs * (vecmax / _E2M1_MAX)).clamp(max=_E4M3_MAX).to(torch.float8_e4m3fn) - out_scale = torch.where(vecmax != 0, gs / sf.float(), torch.zeros_like(vecmax)) - scaled = (blocks * out_scale.unsqueeze(-1)).view(m, k) - mid = torch.tensor(_E2M1_MIDPOINTS, device=x.device, dtype=torch.float32) - tie_up = torch.tensor(_TIE_UP, device=x.device, dtype=torch.bool) - a = scaled.abs().unsqueeze(-1) - code = ((a > mid).sum(-1) + ((a == mid) & tie_up).any(-1)).to(torch.uint8) - negative = (scaled < 0) | ((scaled == 0) & torch.signbit(scaled)) - code = code | (negative.to(torch.uint8) << 3) - packed = code[:, 0::2] | (code[:, 1::2] << 4) - return packed, sf.view(torch.uint8).reshape(-1), scaled - - -def near_tie(scaled: torch.Tensor) -> torch.Tensor: - """Elements whose e2m1 rounding the kernel's approximate reciprocals decide (see ``TIE_WINDOW``).""" - mid = torch.tensor(_E2M1_MIDPOINTS, device=scaled.device, dtype=torch.float32) - a = scaled.abs().unsqueeze(-1) - return ((a - mid).abs() <= TIE_WINDOW * a).any(-1) - - -def fp32_logits(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: - """``float(x) @ w^T`` as a true fp32 product (TF32 off).""" - with true_fp32_matmul(): - return x.float() @ w.t() - - -class M3RouterGemv(OpWrapper): - """fp32 router logits ``float(x[M, K]) @ w[E, K]^T`` with exact fp32 FMA, one CTA per expert. - - ``x`` is BF16 with 1..16 rows, unit column stride and 16-byte rows; ``w`` is the fp32 router weight; K a - multiple of 1024. ``rows`` experts per CTA; only 1 is built. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # The target takes the GEMV at exactly one token (2..48 go to the tensor-core router): 128 experts, hidden 6144. - CELLS: tuple[Cell, ...] = ( - Cell( - why="one decode token, 128 experts x hidden 6144", - spec=dict(m=1, e=128, k=6144), - ), - ) - - note = """ - The op accepts up to 16 rows; the target sends one, so wider calls are not cells. The logits differ from a true - fp32 product only in accumulation order (per-thread FMA chains over K / 1024 strided column groups, a warp tree - and an 8-warp sum). - """ - - def __call__(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: - return torch.ops.trtllm.minimax_m3_router_gemv(x, w, rows, pdl) - - def reference(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> torch.Tensor: - """True fp32 product; ``rows`` and ``pdl`` do not enter the arithmetic.""" - return fp32_logits(x, w) - - def is_valid(self, x: torch.Tensor, w: torch.Tensor, rows: int, pdl: bool) -> None: - # The activation rows are bulk-copied from x's base address; the op checks the row stride but not the base, - # and a misaligned bulk copy faults the context rather than raising. - assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" - - -m3_router_gemv = M3RouterGemv() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py deleted file mode 100644 index f776afce11da..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_tc.py +++ /dev/null @@ -1,106 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""TF32 tensor-core router of the decode MoE for 2..48 tokens, optionally with the NVFP4 activation.""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, -) - -from .m3_router_gemv import nvfp4_quantize - - -def tf32_rna(w: torch.Tensor) -> torch.Tensor: - """fp32 -> TF32 as ``cvt.rna.tf32.f32``: keep 10 mantissa bits, round to nearest with ties away from zero.""" - bits = w.contiguous().view(torch.int32) - rounded = (bits + 0x1000) & ~0x1FFF - finite = (bits & 0x7F800000) != 0x7F800000 - return torch.where(finite, rounded, bits).view(torch.float32) - - -class M3RouterTc(OpWrapper): - """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. - - ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' - partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result - does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns - ``[logits]`` or ``[logits, x_fp4, x_sf]``. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # The target routes 2..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. - CELLS: tuple[Cell, ...] = ( - Cell( - why="two tokens, logits only (the smallest call the target makes here)", - spec=dict(m=2, k=6144, fp4=False), - ), - Cell( - why="8 tokens with the NVFP4 activation: one 8-token MMA tile, every expert tile quantizes one row", - spec=dict(m=8, k=6144, fp4=True), - ), - Cell( - why="13 tokens: a partial second token tile", - spec=dict(m=13, k=6144, fp4=True), - ), - Cell( - why="the 48-token maximum, six token tiles, logits only", - spec=dict(m=48, k=6144, fp4=False), - ), - Cell( - why="the 48-token maximum with the NVFP4 activation", - spec=dict(m=48, k=6144, fp4=True), - ), - # Every other row count the router takes (modeling's ROUTER_TC), with and without the NVFP4 activation. - *( - Cell( - why=f"{m} rows: a router row count between the ones above", - spec=dict(m=m, k=6144, fp4=m % 2 == 1), - ) - for m in range(2, 49) - if m not in (2, 8, 13, 48) - ), - ) - - note = """ - The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both - sum the same 16 slice partials in the same order. Only the one the test device picks is driven. - """ - - def __call__( - self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool - ) -> list[torch.Tensor]: - return torch.ops.trtllm.minimax_m3_router_tc(x, w, sf_scale, pdl) - - def reference( - self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool - ) -> tuple[torch.Tensor, ...]: - """fp32 product of ``x`` and the TF32-rounded weight, plus the NVFP4 recipe with ``sf_scale``.""" - with true_fp32_matmul(): - logits = x.float() @ tf32_rna(w).t() - if sf_scale is None: - return (logits,) - packed, sf, _ = nvfp4_quantize(x, sf_scale) - return logits, packed, sf - - def is_valid( - self, x: torch.Tensor, w: torch.Tensor, sf_scale: Optional[torch.Tensor], pdl: bool - ) -> None: - assert x.data_ptr() % 16 == 0, "x must start 16-byte aligned (cp.async.bulk source)" - if sf_scale is not None: - value = float(sf_scale.reshape(())) - assert value > 0.0 and value != float("inf"), ( - f"sf_scale must be positive and finite, got {value}" - ) - - -m3_router_tc = M3RouterTc() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py index efe4879f9fda..f1cbef70313e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/mx_gemv/op.py @@ -21,12 +21,19 @@ from __future__ import annotations +import sys import threading -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Tuple, cast import torch -from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + mxfp8_core as _core, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm import ( + mxfp8_host as _host, +) +from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl as get_env_enable_pdl MAX_TOKENS = 48 # capacity: fixed: the decode kernels' row bound (decode steps) MMA_WIDTHS = ( @@ -39,16 +46,11 @@ def _mma_n(m: int) -> int: - return next(w for w in MMA_WIDTHS if m <= w) + return _host.gemv_mma_n(m, mma_widths=MMA_WIDTHS) def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: - return ( - ring * (128 * 128 + 512) - + tiles * (mma_n * 128 + 512) - + split * (128 // split) * (mma_n + 4) * 4 - + 1024 - ) + return _host.gemv_smem_bytes(ring, tiles, split, mma_n) _lock = threading.Lock() @@ -95,17 +97,12 @@ def _smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: def _dummy_timers(device) -> torch.Tensor: - """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" - t = _DUMMY.get(device) - if t is None: - t = _DUMMY[device] = torch.zeros(8, dtype=torch.int64, device=device) - return t + "Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on." + return _host.gemv_dummy_timers(device, dummy=_DUMMY) def _arg(t: torch.Tensor, align: int = 16): - from cutlass.cute.runtime import from_dlpack - - return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + return _host.gemv_arg(t, align) # (rule 1b, deterministic) The decode projections at up to STRIDED_MAX_ROWS rows split K over strided k-tiles per @@ -125,38 +122,32 @@ def _arg(t: torch.Tensor, align: int = 16): def get_config(n: int, k: int, m: Optional[int] = None) -> Optional[tuple]: - cfg = _CONFIGS_BY_M.get((n, k, m)) if m is not None else None - if cfg is None and m is not None and m > 8: - cfg = _CONFIGS_WIDE.get((n, k)) - cfg = cfg or _CONFIGS.get((n, k)) - # The split-K epilogue gives each of the split CTAs 128 // split rows of a 128-row output tile. - if cfg is None or n % 128 != 0 or k % 128 != 0 or (k // 128) % cfg[0] != 0 or 128 % cfg[0] != 0: - return None - if ( - STRIDED_K_TILES - and m is not None - and m <= STRIDED_MAX_ROWS - and (n, k) in _STRIDED - and len(cfg) == 2 - ): - cfg = cfg + (0,) - return cfg + return _host.gemv_get_config( + n, + k, + m, + configs_by_m=_CONFIGS_BY_M, + configs_wide=_CONFIGS_WIDE, + configs=_CONFIGS, + strided=STRIDED_K_TILES, + strided_max_rows=STRIDED_MAX_ROWS, + strided_shapes=_STRIDED, + ) def _declines(n: int, k: int, m: int, x_bf16: bool, cfg: tuple) -> bool: - """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" - split = cfg[0] - if not 0 < m <= MAX_TOKENS or 128 % split != 0 or (k // 128) % split != 0: - return True - tiles = k // 128 // split - mma_n = _mma_n(m) - # 17..32 rows at K < 1024 decline (an unsplit K below 1024 is slower than the FlashInfer GEMM there); no TP2 - # projection has K < 1024. - # The in-kernel bf16 -> MXFP8 quantization (x_bf16) costs one 32-element block per thread and round on 128 threads; - # beyond 8 rounds per CTA, mxfp8_quantize + this GEMV (or the fallback GEMM) is faster. - if (16 < m <= 32 and k < 1024) or (x_bf16 and tiles * 4 * mma_n > 8 * 128): - return True - return _smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > _SMEM_LIMIT + "Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds." + return _host.gemv_declines( + n, + k, + m, + x_bf16, + cfg, + max_tokens=MAX_TOKENS, + smem_limit=_SMEM_LIMIT, + select_mma_n=_mma_n, + smem_bytes=_smem_bytes, + ) def serves(n: int, k: int, m: int, x_bf16: bool) -> bool: @@ -181,137 +172,39 @@ def _resolve( ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the rank are runtime arguments.""" - cfg = cfg or get_config(n, k, m) - if cfg is None or _declines(n, k, m, x_bf16, cfg): - return None - split, ring = cfg[0], cfg[1] - contig = bool(cfg[2]) if len(cfg) > 2 else _CONTIG - nacc = int(cfg[3]) if len(cfg) > 3 else _NACC - a_tmem = bool(cfg[4]) if len(cfg) > 4 else _A_TMEM - st_async = bool(cfg[5]) if len(cfg) > 5 else _ST_ASYNC - if prod and not ( - n == PROD_N - and split == 8 - and m <= split - and st_async - and not x_bf16 - and not x_swiglu - and ar is None - ): - # The producer epilogue: the sparse QKV + index projection (one 128-row packed head per cluster of 8), at most - # one token per cluster rank, the st.async exchange, MXFP8 input, no all-reduce push. - return None - if nacc != 1: - # More than one TMEM accumulator is broken on SM107 (783556 / 783557 / 783579: nondeterministic or garbage - # outputs). No shipped table uses it; decline (the caller's fallback runs) until the kernel path is fixed. - return None - ring = min(ring, k // 128 // split) - tiles = k // 128 // split - mma_n = _mma_n(m) - while ring > 2 and _smem_bytes(ring, tiles, split, mma_n) > _SMEM_LIMIT: - ring -= 1 - if a_tmem: - from .m3_mx_gemv_kernel import tmem_cols - - if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: - a_tmem = False - # ar = (multicast workspace address, flags address, rank, world, bucket tokens, mode): push the output rows into - # the next m3_mc_allreduce call's exchange slots (mode 2 also defers the dependents' launch until after the push). - ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = 0, 0, 0, 0, 1, 8 - if ar is not None and split <= 4 and n % 2 == 0: - ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar - if ar_mode == 2: - trigger_early = False - if use_pdl is None: - use_pdl = get_env_enable_pdl() - # The kernel's ``timed`` flag (per-CTA milestone timers) is always off. - key = ( + return _core._resolve( + _CORE_TARGET, + m, n, k, - split, - ring, - trigger_early, - contig, - False, - nacc, - a_tmem, - st_async, x_bf16, - ar_mode, - ar_world, - ar_maxt, - mma_n, + trigger_early, + cfg, use_pdl, + ar, x_swiglu, prod, + has_swiglu=True, + has_producer=True, ) - return key, (ar_ws, ar_flags, ar_rank) def _compile( key: tuple, args: tuple, m: int, runtime_ar: tuple, stream, extra: tuple, prod_args: tuple ): """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" - import cutlass.cute as cute - - from . import m3_mx_gemv_kernel as kernel - - ( - n, - k, - split, - ring, - trigger_early, - contig, - timed, - nacc, - a_tmem, - st_async, - x_bf16, - ar_mode, - ar_world, - ar_maxt, - mma_n, - use_pdl, - x_swiglu, - prod, - ) = key - ar_ws, ar_flags, ar_rank = runtime_ar - with _lock: - fn = _compiled.get(key) - if fn is None: - fn = _compiled[key] = cute.compile( - kernel.m3_mx_gemv, - *args, - m, - n, - k, - split, - ring, - trigger_early, - contig, - timed, - nacc, - a_tmem, - st_async, - x_bf16, - ar_ws, - ar_flags, - ar_rank, - ar_mode, - ar_world, - ar_maxt, - mma_n, - use_pdl, - stream, - *extra, - False, - False, - x_swiglu, - *prod_args, - prod, - ) - return fn + return _core._compile( + _CORE_TARGET, + key, + args, + m, + runtime_ar, + stream, + extra, + prod_args, + has_swiglu=True, + has_producer=True, + ) def _views( @@ -324,22 +217,18 @@ def _views( x_swiglu: bool = False, ) -> tuple: # x_swiglu: x is [M, 2K] (gate | up); the GEMV's K is half its width. - k = x.shape[1] // 2 if x_swiglu else x.shape[1] - if x_bf16: - x_bytes = x.view(torch.uint8) - xsf = _placeholder_scales(k, x.device) - x_words = x.view(torch.int32) - else: - x_bytes = x.view(torch.uint8) - x_words = _placeholder_words(x.device) - return ( - _arg(weight.view(torch.uint8)), - _arg(weight_scale.view(torch.uint8).reshape(-1)), - _arg(x_bytes), - _arg(xsf.view(torch.uint8).reshape(-1)), - _arg(y.view(-1) if y.numel() else _placeholder_words(x.device).view(torch.bfloat16)), - _arg(x_words.reshape(-1)), - _arg(_dummy_timers(x.device)), + return _host.gemv_views( + x, + xsf, + weight, + weight_scale, + y, + x_bf16, + x_swiglu, + placeholder_scales=_placeholder_scales, + placeholder_words=_placeholder_words, + arg=_arg, + dummy_timers=_dummy_timers, ) @@ -356,46 +245,22 @@ def _run( x_swiglu: bool = False, prod_args: Optional[tuple] = None, ) -> Optional[torch.Tensor]: - global last_pushed, last_uncompiled - last_pushed = False - last_uncompiled = None - import cuda.bindings.driver as cuda_driver - - m, k = x.shape - if x_swiglu: - k //= 2 - n = weight.shape[0] - if weight.shape[1] != k: - return None - prod = prod_args is not None - resolved = _resolve(m, n, k, x_bf16, trigger_early, cfg, use_pdl, ar, x_swiglu, prod) - if resolved is None: - return None - key, runtime_ar = resolved - if not prod: - prod_args = PROD_OFF - fn = _compiled.get(key) - if fn is None: - if _frozen[0]: - # Compiles are closed (the model built every kernel it dispatches to before warmup): decline, the caller - # reports the miss. - last_uncompiled = key - return None - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "m3_mx_gemv must run once per shape outside CUDA-graph capture first" - ) - # prod: no packed row is stored (the producer epilogue consumes it); y is the placeholder. - y = torch.empty(0 if prod else m, n, dtype=torch.bfloat16, device=x.device) - args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) - extra = _extra_args(x.device) - stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) - if fn is None: - fn = _compile(key, args, m, runtime_ar, stream, extra, prod_args) - ar_ws, ar_flags, ar_rank = runtime_ar - fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra, *prod_args) - last_pushed = key[11] > 0 - return y + return _core._run( + _CORE_TARGET, + x, + xsf, + weight, + weight_scale, + x_bf16, + trigger_early, + cfg, + use_pdl, + ar, + x_swiglu, + prod_args, + has_swiglu=True, + has_producer=True, + ) def precompile( @@ -411,31 +276,19 @@ def precompile( """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" - import cuda.bindings.driver as cuda_driver - - n, k = weight.shape - push = None if ar is None else (0, 0, 0, *ar) - resolved = _resolve(m, n, k, x_bf16, True, None, use_pdl, push, x_swiglu, prod) - if resolved is None: - return None - key, runtime_ar = resolved - if key in _compiled: - return key - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") - dev = weight.device - x = torch.zeros( + return _core.precompile( + _CORE_TARGET, + weight, + weight_scale, m, - 2 * k if x_swiglu else k, - dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, - device=dev, + x_bf16, + use_pdl, + ar, + x_swiglu, + prod, + has_swiglu=True, + has_producer=True, ) - xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) - y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) - args = _views(x, xsf, weight, weight_scale, y, x_bf16, x_swiglu) - stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) - _compile(key, args, m, runtime_ar, stream, _extra_args(dev), PROD_OFF) - return key def freeze(frozen: bool = True) -> None: @@ -453,13 +306,8 @@ def is_compiled(key: tuple) -> bool: def _extra_args(device) -> tuple: - """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" - i32 = _placeholder_words(device) - f32 = i32.view(torch.float32) - u8 = i32.view(torch.uint8) - d = (_arg(i32), _arg(i32), _arg(i32), _arg(i32), _arg(f32), 0, _arg(i32), _arg(i32), 0) - r = (_arg(f32), _arg(i32), _arg(i32), _arg(u8), _arg(u8), _arg(f32), _arg(i32)) - return d + r + "Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off)." + return _host.gemv_extra_args(device, placeholder_words=_placeholder_words, arg=_arg) def m3_mx_gemv( @@ -599,47 +447,25 @@ def m3_mx_gemv_producer( def _placeholder_scales(k: int, device) -> torch.Tensor: - key = ("sf", k, device) - t = _PLACEHOLDERS.get(key) - if t is None: - t = _PLACEHOLDERS[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) - return t + return _host.gemv_placeholder_scales(k, device, placeholders=_PLACEHOLDERS) def _placeholder_words(device) -> torch.Tensor: - key = ("w", device) - t = _PLACEHOLDERS.get(key) - if t is None: - t = _PLACEHOLDERS[key] = torch.zeros(16, dtype=torch.int32, device=device) - return t + return _host.gemv_placeholder_words(device, placeholders=_PLACEHOLDERS) def ensure_compiled(weight: torch.Tensor, weight_scale: torch.Tensor, use_pdl: bool) -> None: - """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" - if ( - weight.dim() != 2 - or weight.dtype != torch.float8_e4m3fn - or get_config(*weight.shape) is None - ): - return - n, k = weight.shape - if (n, k, use_pdl) in _ensured: - return - _ensured.add((n, k, use_pdl)) - xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) - xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) - m3_mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) - for width in MMA_WIDTHS[1:]: - m3_mx_gemv( - torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), - xsf, - weight, - weight_scale, - use_pdl=use_pdl, - ) - m3_mx_gemv_bf16( - torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + "Compiles the kernel of this weight's shape outside capture (a one-token dummy call)." + return _host.gemv_ensure_compiled( weight, weight_scale, - use_pdl=use_pdl, + use_pdl, + get_config=get_config, + ensured=_ensured, + mma_widths=MMA_WIDTHS, + mx_gemv=m3_mx_gemv, + mx_gemv_bf16=m3_mx_gemv_bf16, ) + + +_CORE_TARGET = cast(_core.GemvTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py index 11b92b47fc50..d3b947b72add 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/norm/m3_embedding_firstnorm.py @@ -9,6 +9,9 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.decode_embedding import ( + embedding_firstnorm_reference as _embedding_reference, +) from ..embedding_firstnorm.op import DecodeEmbedding, m3_embed_norm @@ -78,13 +81,7 @@ def reference( use_pdl: bool = True, ) -> Tuple[torch.Tensor, torch.Tensor]: """New tensors (``raw`` / ``out`` are not written); `use_pdl` is ignored.""" - valid = (ids >= 0) & (ids < table.shape[0]) - rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( - table.dtype - ) - x = rows.double() - rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) - return (x * rstd * weight.double()).to(torch.bfloat16), rows + return _embedding_reference(ids, table, weight, eps) m3_embedding_firstnorm = M3EmbeddingFirstnorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py deleted file mode 100644 index be5aa46fe774..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_capture_norm.py +++ /dev/null @@ -1,80 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Eagle3 hidden-state capture with the drafter's fc_norm folded in (``m3_spec_tail.capture_norm``).""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .m3_chunk_norm import rmsnorm_rows - - -class M3CaptureNorm(OpWrapper): - """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture - buffer slice whose row stride may exceed H. Returns ``dst``.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the - # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. - CELLS: tuple[Cell, ...] = ( - Cell( - why="4 verify rows of one request, layer slot 0, residual added", - spec=dict(m=4, slot=0, residual=True), - ), - Cell( - why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) - ), - Cell( - why="no residual (a capture point without one): plain RMSNorm", - spec=dict(m=8, slot=1, residual=False), - ), - ) - - note = """ - fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional - one-ulp move against torch's statistic. - """ - - def __call__( - self, - h: torch.Tensor, - r: Optional[torch.Tensor], - w: torch.Tensor, - eps: float, - dst: torch.Tensor, - ) -> torch.Tensor: - torch.ops.trtllm.minimax_m3_capture_norm(h, r, w, eps, dst) - return dst - - def reference( - self, - h: torch.Tensor, - r: Optional[torch.Tensor], - w: torch.Tensor, - eps: float, - dst: torch.Tensor, - ) -> torch.Tensor: - """bf16 sum, then the normalized row; ``dst`` only says where the op writes.""" - x = (h.float() + r.float()).to(torch.bfloat16) if r is not None else h - return rmsnorm_rows(x, w, eps) - - def is_valid( - self, - h: torch.Tensor, - r: Optional[torch.Tensor], - w: torch.Tensor, - eps: float, - dst: torch.Tensor, - ) -> None: - # The norm weight is read as H dense bf16 values; a strided view (a column of a stacked weight) is read - # past its elements without an error. - assert w.is_contiguous(), "w must be dense [H]" - - -m3_capture_norm = M3CaptureNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py deleted file mode 100644 index 2f71f1b381b1..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_chunk_norm.py +++ /dev/null @@ -1,74 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The Eagle3 fc_norm over the captured layers in one launch: per-chunk RMSNorm of the concatenated hidden states -(``m3_spec_tail.chunk_norm``).""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). -MAX_HIDDEN = 8192 - - -def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" - xf = x.float() - r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) - return (xf * r * weight.float()).to(torch.bfloat16) - - -class M3ChunkNorm(OpWrapper): - """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may - have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers - # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. - CELLS: tuple[Cell, ...] = ( - Cell( - why="one request's 4 verify rows, dense [4, 3 x 6144]", - spec=dict(m=4, c=3, h=6144, extra_cols=0), - ), - Cell( - why="two requests' 8 verify rows (b2) from the capture buffer, dense [8, 3 x 6144]", - spec=dict(m=8, c=3, h=6144, extra_cols=0), - ), - Cell( - why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", - spec=dict(m=32, c=3, h=6144, extra_cols=6144), - ), - ) - - note = """ - fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then - against torch's statistic; the cells' band states it. - """ - - def __call__( - self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True - ): - torch.ops.trtllm.minimax_m3_chunk_norm(src, w, eps, dst, pdl) - return dst - - def reference( - self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True - ): - """Each chunk normalized on its own; ``dst`` only says where the op writes, ``pdl`` shapes the launch.""" - c, h = w.shape - x = src[:, : c * h].reshape(src.shape[0], c, h) - return rmsnorm_rows(x, w, eps).reshape(src.shape[0], c * h) - - def is_valid( - self, src: torch.Tensor, w: torch.Tensor, eps: float, dst: torch.Tensor, pdl: bool = True - ) -> None: - c, h = w.shape - assert w.is_contiguous(), "w must be dense [C, H] (chunk c read at c * H)" - assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" - - -m3_chunk_norm = M3ChunkNorm() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py deleted file mode 100644 index 4fe14d538a5f..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_input.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The Eagle3 draft pass's input in one launch: embedding lookup, both RMSNorms and the concat -(``m3_spec_tail.draft_input``).""" - -from __future__ import annotations - -from typing import Optional - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -#: Largest hidden size the kernel's register tile holds (4 x 256 threads x 8 bf16). -MAX_HIDDEN = 8192 - - -def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: - """bf16(x * rsqrt(mean(x^2) + eps) * w) per row, fp32 in between ((x * r) * w, as the kernels multiply).""" - xf = x.float() - r = torch.rsqrt(xf.pow(2).mean(dim=-1, keepdim=True) + eps) - return (xf * r * weight.float()).to(torch.bfloat16) - - -class M3DraftInput(OpWrapper): - """``out [M, 2H] = [rmsnorm(table[ids]) * w_embed, rmsnorm(hidden[rows]) * w_hidden]`` and ``residual [M, H] = - hidden[rows]`` (rows default to 0..M-1) in one launch; ids outside the table embed to zero (the masked lookup of - a vocab-local embedding). Returns ``(out, residual)``.""" - - ARCHS = frozenset({Arch.SM_107}) - - # Every pass of the fused draft loop (modeling's DRAFT_INPUT) at draft_len 3, drafter hidden 6144, eps 1e-6: the - # first pass of 1..8 requests (4..32 rows, the carried hidden read in place) and the later passes of 1..8 requests - # (one row each; the second pass gathers the last accepted rows from the first pass's 4 rows per request). - CELLS: tuple[Cell, ...] = ( - *( - Cell( - why=f"first pass of {n} request(s): {4 * n} rows, hidden read in place", - spec=dict(m=4 * n, h=6144, gather=False, invalid_id=False), - ) - for n in range(1, 9) - ), - *( - Cell( - why=f"second pass of {n} request(s): the last accepted rows gathered from {4 * n} rows", - spec=dict(m=n, h=6144, gather=True, invalid_id=False), - ) - for n in range(1, 9) - ), - Cell( - why="third pass of 3 requests: one row each, hidden read in place", - spec=dict(m=3, h=6144, gather=False, invalid_id=False), - ), - Cell( - why="one row with an out-of-table id: the embedding half is zero", - spec=dict(m=1, h=6144, gather=False, invalid_id=True), - ), - ) - - note = """ - The arithmetic is the in-tree RMSNorm's (fp32 statistic, (x * rsqrt) * w, one bf16 rounding); the statistic's - summation order differs from FlashInfer's, which moves an output by one bf16 ulp now and then (the test's band, - also held against the in-tree modules). The out-of-table id is the masked lookup of a vocab-local embedding; the - drafter's embedding is the whole table, so its ids are always in range. - """ - - def __call__( - self, - ids: torch.Tensor, - table: torch.Tensor, - hidden: torch.Tensor, - rows: Optional[torch.Tensor], - w_embed: torch.Tensor, - w_hidden: torch.Tensor, - eps_embed: float, - eps_hidden: float, - out: torch.Tensor, - residual: Optional[torch.Tensor], - pdl: bool = True, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - torch.ops.trtllm.minimax_m3_draft_input( - ids, table, hidden, rows, w_embed, w_hidden, eps_embed, eps_hidden, out, residual, pdl - ) - return out, residual - - def reference( - self, - ids: torch.Tensor, - table: torch.Tensor, - hidden: torch.Tensor, - rows: Optional[torch.Tensor], - w_embed: torch.Tensor, - w_hidden: torch.Tensor, - eps_embed: float, - eps_hidden: float, - out: torch.Tensor, - residual: Optional[torch.Tensor], - pdl: bool = True, - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """The two normalized halves (fp32 statistic, one bf16 rounding) and the gathered hidden rows. ``out`` / - ``residual`` only say where the op writes; ``pdl`` only shapes the launch.""" - valid = (ids >= 0) & (ids < table.shape[0]) - emb = rmsnorm_rows(table[torch.where(valid, ids, 0).long()], w_embed, eps_embed) - emb = torch.where(valid[:, None], emb, torch.zeros_like(emb)) - hid = hidden[rows.long()] if rows is not None else hidden[: ids.shape[0]] - ref = torch.cat([emb, rmsnorm_rows(hid, w_hidden, eps_hidden)], dim=-1) - return ref, (hid.clone() if residual is not None else None) - - def is_valid( - self, - ids: torch.Tensor, - table: torch.Tensor, - hidden: torch.Tensor, - rows: Optional[torch.Tensor], - w_embed: torch.Tensor, - w_hidden: torch.Tensor, - eps_embed: float, - eps_hidden: float, - out: torch.Tensor, - residual: Optional[torch.Tensor], - pdl: bool = True, - ) -> None: - h = table.shape[1] - # Rows are read as hidden + rows[r] * H and table + id * H: both must be dense H-wide rows. - assert table.is_contiguous() and hidden.is_contiguous(), ( - "table / hidden must be dense [*, H]" - ) - if rows is not None: - assert bool(((rows >= 0) & (rows < hidden.shape[0])).all()), "rows must index hidden" - else: - assert hidden.shape[0] >= ids.shape[0], "hidden must hold a row per id" - # Vector loads of 16 bytes: H must be a multiple of 8 within the register tile, and the weights aligned. - assert h % 8 == 0 and h <= MAX_HIDDEN, "H must be a multiple of 8, at most 8192" - assert residual is None or ( - residual.is_contiguous() and residual.shape == (ids.shape[0], h) - ), "residual must be contiguous [M, H]" - - -m3_draft_input = M3DraftInput() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py index b2487a213990..aaadebde1166 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_silu.py @@ -14,7 +14,7 @@ class M3DraftSilu(OpWrapper): """``silu(g) * u`` of a bf16 [rows, 2d] gate | up tensor, rounded once to bf16: - trtllm::minimax_m3_sm107_tp8_silu_and_mul (minimaxM3Sm107Tp8/siluAndMul.cu), bit for bit the Triton + trtllm::minimax_m3_sm107_tp8_silu_and_mul (minimaxM3Sm107/siluAndMul.cu), bit for bit the Triton ``torch.ops.trtllm.silu_and_mul`` at its defaults (no limit, alpha 1, beta 0), the activation the drafter's GatedMLP runs. Contiguous, 16-byte aligned rows and d a multiple of 8; ``__call__`` returns None otherwise (the caller runs the Triton op).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py index e63389ffc6eb..85ea602d94b5 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_draft_step.py @@ -7,19 +7,23 @@ from __future__ import annotations import functools -from typing import Optional +import sys +from typing import Optional, cast import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import ( - Arch, - Cell, - OpWrapper, - true_fp32_matmul, + true_fp32_matmul as true_fp32_matmul, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( + decode as decode_coded, # noqa: F401 — used by the bound family implementation +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec import ( + draft_head_reference as _draft_head_shared, ) -from ..gemm.m3_target_head import decode as decode_coded from .m3_shard_accept import exchange_tokens, local_argmax, make_keys #: Requests (head rows) per pass, and exchange rows per call half of the workspace. @@ -60,18 +64,12 @@ def bf16_fragments(weight: torch.Tensor) -> torch.Tensor: One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). """ - n, k = weight.shape - assert n % 16 == 0 and k % 16 == 0 and weight.dtype == torch.bfloat16, (n, k, weight.dtype) - q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() - return q.view(n // 16, k // 16, 256) + return _draft_head_shared.bf16_fragments(_DRAFT_HEAD_TARGET, weight) def from_fragments(wf: torch.Tensor) -> torch.Tensor: """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" - tiles, ks, _ = wf.shape - return ( - wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) - ) + return _draft_head_shared.from_fragments(_DRAFT_HEAD_TARGET, wf) def head_weight( @@ -84,10 +82,9 @@ def head_weight( ) -> torch.Tensor: """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the 12-bit coded stream with its exception lists decoded.""" - if exc_off is None: - return from_fragments(wf) - n = (exc_off.numel() - 1) * 16 - return decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + return _draft_head_shared.head_weight( + _DRAFT_HEAD_TARGET, wf, k, exc_off, exc_pos, exc_w, exp_base + ) def shard_logits( @@ -100,10 +97,9 @@ def shard_logits( ) -> torch.Tensor: """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, rounded once to BF16 (the BF16 GEMM's output), as fp32.""" - w = head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) - with true_fp32_matmul(): - logits = x.float() @ w.float().t() - return logits.to(torch.bfloat16).float() + return _draft_head_shared.shard_logits( + _DRAFT_HEAD_TARGET, x, wf, exc_off, exc_pos, exc_w, exp_base + ) class M3DraftStep(OpWrapper): @@ -438,3 +434,5 @@ def is_valid( m3_draft_step = M3DraftStep() + +_DRAFT_HEAD_TARGET = cast(_draft_head_shared.DraftHeadTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py index d86ee7e23625..ab2566951069 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_sampled_accept.py @@ -5,14 +5,18 @@ from __future__ import annotations -from typing import Optional +import sys +from typing import Optional, cast import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec import ( + sampling_reference as _sampling_shared, +) -from .m3_shard_accept import strict_accept +from .m3_shard_accept import strict_accept as strict_accept #: Verify rows per call. MAX_ROWS = 48 @@ -25,25 +29,7 @@ def nucleus_probs( ) -> torch.Tensor: """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" - inv = 1.0 / (temperatures.float() + TEMP_EPSILON) - x = logits.float() * inv[:, None] - w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() - z = w.sum(dim=1, keepdim=True) - order = torch.sort(w, dim=1, descending=True) - sw = order.values - # mass strictly above each sorted weight: the cumulative sum before the first entry of its tie group - csum = torch.cumsum(sw, dim=1) - sw - first = torch.ones_like(sw, dtype=torch.bool) - first[:, 1:] = sw[:, 1:] != sw[:, :-1] - group = torch.cumsum(first.long(), dim=1) - 1 - starts = torch.zeros_like(csum) - starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) - above = torch.gather(starts, 1, group) - keep_sorted = above <= top_ps.double()[:, None] * z - keep = torch.zeros_like(keep_sorted) - keep.scatter_(1, order.indices, keep_sorted) - p = torch.where(keep, w, torch.zeros_like(w)) - return (p / p.sum(dim=1, keepdim=True)).float() + return _sampling_shared.nucleus_probs(_REFERENCE_TARGET, logits, temperatures, top_ps) class M3SampledAccept(OpWrapper): @@ -213,13 +199,34 @@ def reference( """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" - rows = logits.shape[0] - probs = nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) - - def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: - return strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) - - return {"probs": probs, "accept": accept} + return _sampling_shared.sampled_reference( + self, + _REFERENCE_TARGET, + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + top_ks, + min_ps, + ) def is_valid( self, @@ -248,35 +255,36 @@ def is_valid( top_ks: Optional[torch.Tensor] = None, min_ps: Optional[torch.Tensor] = None, ) -> None: - rows = logits.shape[0] - # top-k and min-p are not applied: a row that has them gets a token outside its real sampling set. - if top_ks is not None: - k = top_ks[:rows] - assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( - "top-k rows must take the generic path" - ) - if min_ps is not None: - assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" - # Rows are request-major verify rows: tokens[b * (D + 1) + k]; a mismatched layout pairs drafts with the - # wrong rows. - assert rows == draft.shape[0] * (draft.shape[1] + 1), ( - "logits rows must be nreq * (draft_len + 1)" - ) - # Per-row parameters are read at the row index, unchecked. - for t in (temperatures, top_ps, seeds, offsets): - assert t.numel() >= rows and t.is_contiguous(), ( - "per-row parameters must cover every row" - ) - assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( - "fp32 sampling parameters" - ) - assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( - "int64 Philox seeds / offsets" - ) - # The finished-row counter elects the CTA that runs the acceptance. - assert int(counters[0].item()) == 0, ( - "counters[0] (finished rows) must be zero between calls" + return _sampling_shared.validate_sampling( + self, + _REFERENCE_TARGET, + logits, + vocab_offset, + temperatures, + top_ps, + seeds, + offsets, + ws_uc, + ws_mc_ptr, + counters, + tokens, + rank, + world, + draft, + batch_idx, + acc_tokens, + num_acc, + gather_ids, + force_total, + next_new, + force_vec, + extra_attempts, + use_pdl, + top_ks, + min_ps, ) m3_sampled_accept = M3SampledAccept() + +_REFERENCE_TARGET = cast(_sampling_shared.SamplingTarget, sys.modules[__name__]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py index 6df70b3f467b..b8b10a57e0a0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_shard_accept.py @@ -12,97 +12,18 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from ....ops.spec.acceptance_reference import _LOW32 as _LOW32 +from ....ops.spec.acceptance_reference import exchange_tokens as exchange_tokens +from ....ops.spec.acceptance_reference import key_token as key_token +from ....ops.spec.acceptance_reference import local_argmax as local_argmax +from ....ops.spec.acceptance_reference import make_keys as make_keys +from ....ops.spec.acceptance_reference import ordered_bits as ordered_bits +from ....ops.spec.acceptance_reference import strict_accept as strict_accept +from ....ops.spec.acceptance_reference import unsigned_max as unsigned_max + #: Verify rows per call, and exchange rows per call half of its workspace. SHARD_ROWS = 48 -_LOW32 = 0xFFFFFFFF - - -def ordered_bits(v: torch.Tensor) -> torch.Tensor: - """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" - b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 - b = torch.where(b == 0x80000000, torch.zeros_like(b), b) - return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) - - -def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: - """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" - return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) - - -def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: - """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" - sign = torch.iinfo(torch.int64).min - return (keys ^ sign).max(dim=dim).values ^ sign - - -def key_token(key: torch.Tensor) -> torch.Tensor: - return (_LOW32 - (key & _LOW32)).to(torch.int32) - - -def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: - """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" - best = logits.max(dim=1).values - idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max - return best, idx.to(torch.int64) + vocab_offset - - -def exchange_tokens( - own_keys: torch.Tensor, - ws_uc: torch.Tensor, - counters: torch.Tensor, - rank: int, - world: int, - slot_rows: int, -) -> torch.Tensor: - """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. - - Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A - single-device run stands in for the other ranks by writing their keys there before the call. - """ - m = own_keys.shape[0] - half = int(counters[1].item()) & 1 - slots = ( - ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] - .view(m, world) - .clone() - ) - slots[:, rank] = own_keys - return key_token(unsigned_max(slots, dim=1)) - - -def strict_accept( - tokens: torch.Tensor, - draft: torch.Tensor, - batch_idx: torch.Tensor, - force_total: int, - next_new: Optional[torch.Tensor], - force_vec: Optional[torch.Tensor], -) -> dict[str, torch.Tensor]: - """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). - - accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by - force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and - next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. - """ - nreq, d = draft.shape - k1 = d + 1 - tokens = tokens.to(torch.int32) - accepted = tokens.view(nreq, k1).clone() - match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) - num = 1 + match.cumprod(dim=1).sum(dim=1) - if force_vec is not None: - num = force_vec[:nreq].to(num.dtype) - elif force_total > 0: - num = torch.full_like(num, force_total) - gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 - out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} - if next_new is not None: - nn = next_new.clone() - nn[:, 0] = tokens[gather] - out["next_new"] = nn - return out - class M3ShardAccept(OpWrapper): """Greedy verification of ``nreq * (D + 1) <= 48`` verify rows from this rank's BF16 logits shard ``[rows, Vs]`` diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py index 56c5f54b1161..a0602ced8764 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_spec_offsets.py @@ -16,7 +16,7 @@ class M3SpecOffsets(OpWrapper): """``positions += position_offsets`` and ``kv_lens += kv_lens_offsets`` (int32, contiguous 1-D views, equal lengths - per pair) in one plain launch: trtllm::minimax_m3_sm107_tp8_spec_offsets (minimaxM3Sm107Tp8/specOffsets.cu), the + per pair) in one plain launch: trtllm::minimax_m3_sm107_tp8_spec_offsets (minimaxM3Sm107/specOffsets.cu), the same wrapping int32 adds the two ATen ops do. ``__call__`` returns False, having written nothing, for other layouts (the caller runs the ATen adds).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py deleted file mode 100644 index a0da8b337574..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_gather.py +++ /dev/null @@ -1,160 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The overlap scheduler's per-step input gathers from the previous step's device outputs, in one launch -(``M3SpecStepFold.gather`` -> ``m3_spec_tail.step_gather``).""" - -from __future__ import annotations - -import torch - -import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - - -class M3StepGather(OpWrapper): - """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all - int32:: - - ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) - draft[d0 + p*D + d] = next_draft[slots[p], d] - pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) - kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) - - Every other element of ``ids`` / ``draft`` is left as it was. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the - # max batch); the engine's flat buffers sized for the max batch. - CELLS: tuple[Cell, ...] = ( - Cell( - why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", - spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), - ), - Cell( - why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", - spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), - ), - Cell( - why="one request, the store's 2-D view of new tokens (beam dim squeezed)", - spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), - ), - Cell( - why="an unchanged batch of 3 previous requests (c3) in shuffled slots", - spec=dict(p=3, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), - ), - Cell( - why="an unchanged batch of 4 previous requests (c4) in shuffled slots", - spec=dict(p=4, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), - ), - ) - - note = """ - Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid - states them. - """ - - def __call__( - self, - new_tok, - next_draft, - lens, - slots, - ppi, - P, - T, - D, - ids, - ids0, - draft, - d0, - pos_off, - kv_off, - b0, - ): - torch.ops.trtllm.minimax_m3_step_gather( - new_tok, - next_draft, - lens, - slots, - ppi, - P, - T, - D, - ids, - ids0, - draft, - d0, - pos_off, - kv_off, - b0, - ) - - def reference( - self, - new_tok, - next_draft, - lens, - slots, - ppi, - P, - T, - D, - ids, - ids0, - draft, - d0, - pos_off, - kv_off, - b0, - ): - """The four destinations after the call, on copies (``ids`` / ``draft`` keep their other elements).""" - nt = new_tok.reshape(new_tok.shape[0], new_tok.shape[1]) - s = slots[:P].long() - ids_r, draft_r = ids.clone(), draft.clone() - ids_r[ids0 : ids0 + P * T] = nt[:T, s].t().reshape(-1) - draft_r[d0 : d0 + P * D] = next_draft[s, :D].reshape(-1) - pos_r = torch.zeros_like(pos_off) - pos_r[b0 * T : (b0 + P) * T] = lens[ppi[: P * T].long()] - kv_r = torch.zeros_like(kv_off) - kv_r[b0 : b0 + P] = lens[s] - T - return {"ids": ids_r, "draft": draft_r, "pos_off": pos_r, "kv_off": kv_r} - - def is_valid( - self, - new_tok, - next_draft, - lens, - slots, - ppi, - P, - T, - D, - ids, - ids0, - draft, - d0, - pos_off, - kv_off, - b0, - ): - # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or - # writes another buffer. - s = slots[:P] - assert bool( - ( - (s >= 0) & (s < next_draft.shape[0]) & (s < new_tok.shape[1]) & (s < lens.numel()) - ).all() - ), "slots must index the store" - assert bool(((ppi[: P * T] >= 0) & (ppi[: P * T] < lens.numel())).all()), ( - "ppi must index lens" - ) - assert new_tok.shape[0] >= T and next_draft.shape[1] >= D, "store narrower than T / D" - assert ids.numel() >= ids0 + P * T and draft.numel() >= d0 + P * D, "destinations too short" - assert pos_off.numel() >= (b0 + P) * T and kv_off.numel() >= b0 + P, ( - "offset buffers too short" - ) - - -m3_step_gather = M3StepGather() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py index a6c626ff2d2f..f3a7c2ec073d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_step_prologue.py @@ -14,7 +14,7 @@ class M3StepPrologue(OpWrapper): """``StepInputPack``'s early-copy prologue as one single-CTA C++ launch - (trtllm::minimax_m3_sm107_tp8_step_prologue, minimaxM3Sm107Tp8/stepPrologue.cu). + (trtllm::minimax_m3_sm107_tp8_step_prologue, minimaxM3Sm107/stepPrologue.cu). ``window`` is the pack's device window (``gdev``): an int64 header ``[host copies, device copies]``, then one ``(dst, src, bytes)`` descriptor per copy, host copies first (a host copy's ``src`` is a byte offset into the diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py index 5b979f361997..869637a1e2c0 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode.py @@ -45,15 +45,57 @@ from typing import Optional import torch -import triton import triton.language as tl +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _NVFP4_TUNED_CAPABILITIES as _NVFP4_TUNED_CAPABILITIES, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _TARGET_GRID as _TARGET_GRID, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + MergeMxfp8 as MergeMxfp8, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _check_nvfp4_inputs as _check_nvfp4_inputs, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_launch_options as _sm100f_nvfp4_launch_options, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_merge_launch_options as _sm100f_nvfp4_merge_launch_options, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_num_topk_chunks as _sm100f_nvfp4_num_topk_chunks, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_query_group_size as _sm100f_nvfp4_query_group_size, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + _sm100f_nvfp4_use_linear_softmax as _sm100f_nvfp4_use_linear_softmax, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.sparse_decode_policy import ( + resolve_num_topk_chunks as resolve_num_topk_chunks, +) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.msa_utils import ( check_decode_span_shape, ) from tensorrt_llm._torch.flashinfer_utils import get_env_enable_pdl from tensorrt_llm._torch.memory_buffer_utils import get_memory_buffers +from ...kernels.attention.triton_sparse_decode import _dequant_nvfp4_rows as _dequant_nvfp4_rows +from ...kernels.attention.triton_sparse_decode import _e2m1x2_to_f16x2 as _e2m1x2_to_f16x2 +from ...kernels.attention.triton_sparse_decode import ( + _gqa_sparse_decode_kernel as _gqa_sparse_decode_kernel, +) +from ...kernels.attention.triton_sparse_decode import ( + _gqa_sparse_decode_pair_kernel as _gqa_sparse_decode_pair_kernel, +) +from ...kernels.attention.triton_sparse_decode import ( + _merge_topk_attn_out_kernel as _merge_topk_attn_out_kernel, +) +from ...kernels.attention.triton_sparse_decode import _rcp_approx_ftz as _rcp_approx_ftz +from ...kernels.attention.triton_sparse_decode import _store_mxfp8_head as _store_mxfp8_head from . import m3_attn_lc8, m3_prequant, m3_report # One sparse block is exactly one KV page. @@ -68,7 +110,6 @@ } # Total CTAs the split-K partitioning aims for before it stops splitting. -_TARGET_GRID = 256 # capacity: fixed: a launch-grid target, not a size # NVFP4 quantizes in groups of 16 cache elements along head_dim, one E4M3 scale # byte each; see the fp4_quantize call in msa_scatter.fused_write_layer_caches_nvfp4. @@ -93,1142 +134,12 @@ PRODUCER_PATH = "sparse_decode_producer" -class MergeMxfp8: - """MXFP8 copy of the split-K merge output for the o_proj GEMM. - - From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and swizzled UE8M0 - scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips its own - quantization. One 128-row buffer per width serves every layer: each o_proj consumes it before the next layer's - merge rewrites it, in stream order. - """ - - MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) - - def __init__(self): - self._buffers = {} - - def buffers(self, width: int, device: torch.device): - key = (width, device) - if key not in self._buffers: - self._buffers[key] = ( - torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), - torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), - ) - return self._buffers[key] - - MERGE_MXFP8 = MergeMxfp8() -@triton.jit -def _e2m1x2_to_f16x2(packed): - """Convert one packed E2M1 byte into its low/high FP16 values.""" - return tl.inline_asm_elementwise( - asm="""{ - .reg .b8 src; - .reg .b32 pair; - mov.b32 {src, _, _, _}, $2; - cvt.rn.f16x2.e2m1x2 pair, src; - mov.b32 {$0, $1}, pair; - }""", - constraints="=h,=h,r", - args=(packed,), - dtype=(tl.float16, tl.float16), - is_pure=True, - pack=1, - ) - - -@triton.jit -def _dequant_nvfp4_rows( - packed, # [BLOCK_N, D // 2] uint8, two E2M1 nibbles per byte - scale_bytes, # [BLOCK_N, SCALE_COLS] uint8, one E4M3 scale per 16 elements - BLOCK_N: tl.constexpr, - D: tl.constexpr, - SCALE_COLS: tl.constexpr, - SF_VEC: tl.constexpr, -): - """Dequantize one packed NVFP4 [BLOCK_N, D] tile to fp16. - - The caller gathers the block scales in whichever layout its role stores - them, so this handles only the value math and the scale broadcast. The - per-tensor scale is applied by the caller. - """ - # Blackwell converts both E2M1 nibbles with one instruction. Joining the - # low/high results along a new innermost axis and flattening it rebuilds - # the original element order without scalar exponent/sign reconstruction. - value_lo, value_hi = _e2m1x2_to_f16x2(packed) - values = tl.reshape(tl.join(value_lo, value_hi), (BLOCK_N, D)) - - scales = scale_bytes.to(tl.float8e4nv, bitcast=True).to(tl.float16) - scales = tl.reshape( - tl.broadcast_to(scales[:, :, None], (BLOCK_N, SCALE_COLS, SF_VEC)), (BLOCK_N, D) - ) - return values * scales - - -@triton.heuristics( - { - "BLOCK_SIZE_H": lambda args: max(16, triton.next_power_of_2(args["gqa_group_size"])), - "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), - } -) -# Mixed ctx+gen steps launch on [num_ctx:] slices (unaligned seq_lens / block_table) with arbitrary -# total_q; specializing on those compiled a new variant mid-serving (~650 ms TTFT stall at c8). -@triton.jit( - do_not_specialize=["decode_query_len", "total_q"], - do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], -) -def _gqa_sparse_decode_kernel( - q_ptr, # [total_q, num_heads, head_dim] - k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] - v_ptr, # paged V: same layout as K - kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False - k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] - v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order - k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 - v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 - t_ptr, # topk_idx: [num_kv_heads, total_q, topk] - o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT - lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT - out_ptr, # final [total_q, num_heads, head_dim] for FUSED_MERGE, else unused - counter_ptr, # int32 [total_q, num_kv_heads] arrival counters for FUSED_MERGE, zero at rest - block_table_ptr, # [num_reqs, max_blocks] - seq_lens, # [num_reqs] - total_q, - gqa_group_size, - head_dim, - max_topk, - sm_scale, - decode_query_len, - stride_qn, - stride_qh, - stride_qd, - stride_k_blk, - stride_k_h, - stride_k_pos, - stride_k_d, - stride_v_blk, - stride_v_h, - stride_v_pos, - stride_v_d, - stride_ksc_blk, - stride_ksc_h, - stride_vsc_blk, - stride_vsc_h, - stride_th, - stride_tn, - stride_tk, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_bt_b, - num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request - num_pages, # KV pool pages: a page id outside [0, num_pages) is not read - stride_out_n, - stride_out_h, - stride_out_d, - num_kv_heads, - BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_H: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache - WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register - COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math - USE_PDL: tl.constexpr, - KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales - USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence - SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only - DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge - FUSED_MERGE: tl.constexpr, # the last chunk to finish merges the partials in place - EARLY_TRIGGER: tl.constexpr, # release PDL dependents (weight prefetch) at entry - IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) -): - sm_scale_log2e = sm_scale * 1.4426950409 - if KV_NVFP4: - # Both per-tensor dequant scales factor out of the inner loop: QK is - # linear in K, and V's scale commutes with the softmax weights. Reading - # them from device memory keeps the launch free of a host sync, so - # CUDA-graph capture stays unconstrained. - sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) - v_global_scale = tl.load(v_gs_ptr) - # Split-K over the topk dimension: pid(0) folds (query token, chunk). - pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) - pid_b = pid_bc % total_q - pid_c = pid_bc // total_q - req_id = pid_b // decode_query_len - q_offset = pid_b - req_id * decode_query_len - pid_h = pid_kh * gqa_group_size - chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_start = pid_c * chunk_size_topk - - if USE_PDL: - tl.extra.cuda.gdc_wait() - if EARLY_TRIGGER: - # Dependents still wait for this grid to complete; triggering here only - # lets them start their input-independent prologue (e.g. weight loads). - tl.extra.cuda.gdc_launch_dependents() - - seq_len = tl.load(seq_lens + req_id) - query_pos = seq_len - decode_query_len + q_offset - # Full-CG padding uses zero-length request rows. Clamp to an empty - # attention range instead of letting padded rows produce negative lengths. - kv_len = tl.maximum(query_pos + 1, 0) - - # Bound the walk by the token's own valid block count rather than a - # sentinel: the selector emits ascending ids and only pads with -1 past - # min(topk, cdiv(kv_len, blk)) entries, so those are never dereferenced. - num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk = tl.minimum(max_topk, num_blocks) - chunk_end = tl.minimum(chunk_start + chunk_size_topk, real_topk) - - off_n = tl.arange(0, BLOCK_SIZE_K) - off_d = tl.arange(0, BLOCK_SIZE_D) - off_h = tl.arange(0, BLOCK_SIZE_H) - if KV_NVFP4: - # Packed byte column and block-scale column. The NVFP4 path runs only at - # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. - off_dp = tl.arange(0, BLOCK_SIZE_D // 2) - off_s = tl.arange(0, SCALE_COLS) - d_mask = off_d < head_dim - h_mask = off_h < gqa_group_size - hd_mask = h_mask[:, None] & d_mask[None, :] - bt_row = block_table_ptr + req_id * stride_bt_b - - m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - if USE_LINEAR_SOFTMAX: - # Keep the linear softmax denominator and reuse the producer's alpha. - # This removes one exp2 plus one log2 from every visited KV page; only - # the final LSE written for split-K merging still needs a logarithm. - l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - else: - lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - q = tl.load( - q_ptr - + pid_b * stride_qn - + (pid_h + off_h[:, None]) * stride_qh - + off_d[None, :] * stride_qd, - mask=hd_mask, - other=0.0, - ) - # Widen before anything reads q.dtype: K, V and p are all cast to it below, - # so leaving q narrow would silently run the whole attention in FP8 and - # quantize the softmax probabilities. E4M3 -> BF16 is exact, so the result - # matches a caller that widened q itself before the call. - if WIDEN_Q: - q = q.to(COMPUTE_DTYPE) - kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 - - cur_idx_ptr = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk - for it in tl.range(chunk_start, chunk_end): - if IDENTITY: - # Slot `it` of the identity selection, bounded by this row's own kv extent above. - blk = it - else: - blk = tl.load(cur_idx_ptr) - cur_idx_ptr = cur_idx_ptr + stride_tk - # int64 page offsets: a large cache overflows int32 well before the - # per-page block offsets do. - # Fail closed: a top-k id outside the request's block-table row, or the page it names outside the pool, - # contributes nothing instead of a read through a wrong page (valid inputs are unaffected). - blk_ok = (blk >= 0) & (blk < num_bt_cols) - page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) - blk_ok = blk_ok & (page >= 0) & (page < num_pages) - page = tl.where(blk_ok, page, 0) - pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len - # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block - # would make exp2(-inf - -inf) = NaN). - if blk_ok: - if KV_NVFP4: - # K block scales are token-major linear within the page-head region; - # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. - k_deq = _dequant_nvfp4_rows( - tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[:, None] * stride_k_pos - + off_dp[None, :] * stride_k_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - k_sc_ptr - + page * stride_ksc_blk - + pid_kh * stride_ksc_h - + off_n[:, None] * SCALE_COLS - + off_s[None, :], - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ) - k = tl.trans(k_deq).to(q.dtype) - else: - k = tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[None, :] * stride_k_pos - + off_d[:, None] * stride_k_d, - mask=d_mask[:, None] & pos_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - k = (k * kv_scale).to(q.dtype) - qk = tl.where(pos_mask[None, :], 0.0, float("-inf")) - qk += tl.dot(q, k) * sm_scale_log2e - m_ij = tl.maximum(m_i, tl.max(qk, axis=1)) - p = tl.exp2(qk - m_ij[:, None]) - l_ij = tl.sum(p, axis=1) - if USE_LINEAR_SOFTMAX: - alpha = tl.exp2(m_i - m_ij) - acc_o = acc_o * alpha[:, None] - else: - acc_o = acc_o * tl.exp2(m_i - m_ij)[:, None] - if KV_NVFP4: - # V block scales use vLLM's 4x4 token-quad order rather than K's - # linear one: [token // 4, scale_col, token % 4]. - v = _dequant_nvfp4_rows( - tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_dp[None, :] * stride_v_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - v_sc_ptr - + page * stride_vsc_blk - + pid_kh * stride_vsc_h - + (off_n[:, None] // 4) * (4 * SCALE_COLS) - + 4 * off_s[None, :] - + (off_n[:, None] % 4), - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ).to(q.dtype) - else: - v = tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_d[None, :] * stride_v_d, - mask=pos_mask[:, None] & d_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - v = (v * kv_scale).to(q.dtype) - acc_o += tl.dot(p.to(v.dtype), v) - if USE_LINEAR_SOFTMAX: - l_i = l_i * alpha + l_ij - else: - lse_i = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) - m_i = m_ij - - # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. - if USE_LINEAR_SOFTMAX: - scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) - lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) - else: - scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) - lse_out = lse_i - if KV_NVFP4: - # Applying V's scale per partial is exact: the merge is a weighted - # average whose weights sum to one, so a constant survives it. - scale = scale * v_global_scale - acc_o = acc_o * scale[:, None] - o_base = o_ptr + pid_c * stride_o_c + pid_b * stride_o_b + pid_h * stride_o_h - tl.store( - o_base + off_h[:, None] * stride_o_h + off_d[None, :] * stride_o_d, - acc_o.to(o_ptr.dtype.element_ty), - mask=hd_mask, - ) - if not DIRECT_OUTPUT: - lse_base = lse_ptr + pid_c * stride_l_c + pid_b * stride_l_b + pid_h * stride_l_h - tl.store(lse_base + off_h * stride_l_h, lse_out, mask=h_mask) - - if FUSED_MERGE: - # Split-K serial reduction: every chunk publishes its partial, and the - # chunk that arrives last for this (token, KV head) merges all of them - # exactly as _merge_topk_attn_out_kernel does, then re-arms the counter. - tl.debug_barrier() - counter = counter_ptr + pid_b * num_kv_heads + pid_kh - arrived = tl.atomic_add(counter, 1, sem="acq_rel", scope="gpu") - if arrived == NUM_TOPK_CHUNKS - 1: - tl.debug_barrier() - o_rows = ( - o_ptr - + pid_b * stride_o_b - + (pid_h + off_h[:, None]) * stride_o_h - + off_d[None, :] * stride_o_d - ) - l_rows = lse_ptr + pid_b * stride_l_b + (pid_h + off_h) * stride_l_h - lse_max = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - lse_max = tl.maximum(lse_max, lse_c) - # All chunks empty (a CUDA-graph padding row): zero, never NaN. - lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) - denom = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - denom += tl.exp2(lse_c - lse_max) - denom = tl.where(denom > 0, denom, 1.0) - merged = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - for c in tl.static_range(NUM_TOPK_CHUNKS): - lse_c = tl.load( - l_rows + c * stride_l_c, mask=h_mask, other=float("-inf"), cache_modifier=".cg" - ) - w_c = tl.exp2(lse_c - lse_max) / denom - o_c = tl.load( - o_rows + c * stride_o_c, mask=hd_mask, other=0.0, cache_modifier=".cg" - ).to(tl.float32) - merged += o_c * w_c[:, None] - tl.store( - out_ptr - + pid_b * stride_out_n - + (pid_h + off_h[:, None]) * stride_out_h - + off_d[None, :] * stride_out_d, - merged.to(out_ptr.dtype.element_ty), - mask=hd_mask, - ) - tl.atomic_xchg(counter, 0, sem="relaxed", scope="gpu") - - # After the stores, never before: this releases either the merge grid or, - # for direct output, the next PDL-enabled consumer of the final tensor. - if USE_PDL and not EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.heuristics( - { - "BLOCK_SIZE_H": lambda args: max( - 16, - triton.next_power_of_2(args["gqa_group_size"] * args["QUERY_GROUP_SIZE"]), - ), - "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]), - } -) -@triton.jit( - do_not_specialize=["decode_query_len", "total_q"], - do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"], -) -def _gqa_sparse_decode_pair_kernel( - q_ptr, # [total_q, num_heads, head_dim] - k_ptr, # paged K: [num_pages, num_kv_heads, page_size, head_dim (or /2 packed)] - v_ptr, # paged V: same layout as K - kv_scale_ptr, # scalar dequant scale, or a dummy when USE_SCALE is False - k_sc_ptr, # paged K block scales: [num_pages, num_kv_heads, page_size, D/16] - v_sc_ptr, # paged V block scales: same shape, 4x4 token-quad order - k_gs_ptr, # NVFP4 per-tensor K dequant scale, one fp32 - v_gs_ptr, # NVFP4 per-tensor V dequant scale, one fp32 - t_ptr, # topk_idx: [num_kv_heads, total_q, topk] - o_ptr, # partial out, or final [total_q, num_heads, head_dim] when DIRECT_OUTPUT - lse_ptr, # partial lse (log2), unused when DIRECT_OUTPUT - block_table_ptr, # [num_reqs, max_blocks] - seq_lens, # [num_reqs] - total_q, - gqa_group_size, - head_dim, - max_topk, - sm_scale, - decode_query_len, - stride_qn, - stride_qh, - stride_qd, - stride_k_blk, - stride_k_h, - stride_k_pos, - stride_k_d, - stride_v_blk, - stride_v_h, - stride_v_pos, - stride_v_d, - stride_ksc_blk, - stride_ksc_h, - stride_vsc_blk, - stride_vsc_h, - stride_th, - stride_tn, - stride_tk, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_bt_b, - num_bt_cols, # block-table width: a top-k id at or past it is not a block of the request - num_pages, # KV pool pages: a page id outside [0, num_pages) is not read - BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128) - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_H: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_SCALE: tl.constexpr, # apply the scalar dequant scale to an fp8 cache - WIDEN_Q: tl.constexpr, # q arrives FP8 and is widened in-register - COMPUTE_DTYPE: tl.constexpr, # dtype q is widened to; drives the QK/PV math - USE_PDL: tl.constexpr, - KV_NVFP4: tl.constexpr, # cache holds packed E2M1 plus E4M3 block scales - USE_LINEAR_SOFTMAX: tl.constexpr, # measured NVFP4 online-softmax recurrence - QUERY_GROUP_SIZE: tl.constexpr, # consecutive Eagle queries sharing one block set - MAX_TOPK: tl.constexpr, - SCALE_COLS: tl.constexpr, # head_dim // NVFP4_SF_VEC_SIZE, NVFP4 only - DIRECT_OUTPUT: tl.constexpr, # one chunk needs neither FP32 scratch nor a merge - IDENTITY: tl.constexpr = False, # identity selection: slot i is block i (t_ptr is not read) -): - sm_scale_log2e = sm_scale * 1.4426950409 - if KV_NVFP4: - # Both per-tensor dequant scales factor out of the inner loop: QK is - # linear in K, and V's scale commutes with the softmax weights. Reading - # them from device memory keeps the launch free of a host sync, so - # CUDA-graph capture stays unconstrained. - sm_scale_log2e = sm_scale_log2e * tl.load(k_gs_ptr) - v_global_scale = tl.load(v_gs_ptr) - # Split-K over the topk dimension: pid(0) folds (query group, chunk). - pid_bc, pid_kh = tl.program_id(0), tl.program_id(1) - num_query_groups = total_q // QUERY_GROUP_SIZE - pid_g = pid_bc % num_query_groups - pid_c = pid_bc // num_query_groups - groups_per_req = decode_query_len // QUERY_GROUP_SIZE - req_id = pid_g // groups_per_req - q_group_offset = (pid_g - req_id * groups_per_req) * QUERY_GROUP_SIZE - pid_b = req_id * decode_query_len + q_group_offset - pid_h = pid_kh * gqa_group_size - chunk_size_topk = (max_topk + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_start = pid_c * chunk_size_topk - - if USE_PDL: - tl.extra.cuda.gdc_wait() - - seq_len = tl.load(seq_lens + req_id) - # Query rows are flattened as [query-in-group, GQA head]. - off_h = tl.arange(0, BLOCK_SIZE_H) - q_lane = off_h // gqa_group_size - h_lane = off_h - q_lane * gqa_group_size - h_mask = off_h < gqa_group_size * QUERY_GROUP_SIZE - q_token = pid_b + q_lane - query_pos = seq_len - decode_query_len + q_group_offset + q_lane - # Full-CG padding uses zero-length request rows. Clamp to an empty - # attention range instead of letting padded rows produce negative lengths. - kv_len = tl.maximum(query_pos + 1, 0) - kv_len_max = tl.max(tl.where(h_mask, kv_len, 0), axis=0) - - # Each query can have a different valid-block count at a page boundary. - # The selector puts those valid entries first and pads only the tail. - kv_len_0 = tl.maximum(seq_len - decode_query_len + q_group_offset + 1, 0) - num_blocks_0 = (kv_len_0 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk_0 = tl.minimum(max_topk, num_blocks_0) - chunk_end_0 = tl.minimum(chunk_start + chunk_size_topk, real_topk_0) - if QUERY_GROUP_SIZE == 2: - kv_len_1 = tl.maximum(seq_len - decode_query_len + q_group_offset + 2, 0) - num_blocks_1 = (kv_len_1 + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K - real_topk_1 = tl.minimum(max_topk, num_blocks_1) - chunk_end_1 = tl.minimum(chunk_start + chunk_size_topk, real_topk_1) - - off_n = tl.arange(0, BLOCK_SIZE_K) - off_d = tl.arange(0, BLOCK_SIZE_D) - if KV_NVFP4: - # Packed byte column and block-scale column. The NVFP4 path runs only at - # head_dim == BLOCK_SIZE_D (128), so neither needs a tail mask. - off_dp = tl.arange(0, BLOCK_SIZE_D // 2) - off_s = tl.arange(0, SCALE_COLS) - d_mask = off_d < head_dim - hd_mask = h_mask[:, None] & d_mask[None, :] - bt_row = block_table_ptr + req_id * stride_bt_b - - m_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - if USE_LINEAR_SOFTMAX: - # Keep the linear softmax denominator and reuse the producer's alpha. - # This removes one exp2 plus one log2 from every visited KV page; only - # the final LSE written for split-K merging still needs a logarithm. - l_i = tl.zeros((BLOCK_SIZE_H,), dtype=tl.float32) - else: - lse_i = tl.full((BLOCK_SIZE_H,), float("-inf"), dtype=tl.float32) - acc_o = tl.zeros((BLOCK_SIZE_H, BLOCK_SIZE_D), dtype=tl.float32) - q = tl.load( - q_ptr - + q_token[:, None] * stride_qn - + (pid_h + h_lane[:, None]) * stride_qh - + off_d[None, :] * stride_qd, - mask=hd_mask, - other=0.0, - ) - # Widen before anything reads q.dtype: K, V and p are all cast to it below, - # so leaving q narrow would silently run the whole attention in FP8 and - # quantize the softmax probabilities. E4M3 -> BF16 is exact, so this - # reproduces the values the caller's standalone widening kernel produced. - if WIDEN_Q: - q = q.to(COMPUTE_DTYPE) - kv_scale = tl.load(kv_scale_ptr) if USE_SCALE else 1.0 - - idx_ptr_0 = t_ptr + pid_kh * stride_th + pid_b * stride_tn + chunk_start * stride_tk - if QUERY_GROUP_SIZE == 2: - idx_ptr_1 = idx_ptr_0 + stride_tn - idx_0 = chunk_start - idx_1 = chunk_start - count_0 = chunk_end_0 - chunk_start - count_1 = chunk_end_1 - chunk_start - chunk_capacity: tl.constexpr = (MAX_TOPK + NUM_TOPK_CHUNKS - 1) // NUM_TOPK_CHUNKS - chunk_offsets = tl.arange(0, chunk_capacity) - sentinel = 2147483647 - if IDENTITY: - blocks_0 = tl.where(chunk_offsets < count_0, chunk_start + chunk_offsets, sentinel) - blocks_1 = tl.where(chunk_offsets < count_1, chunk_start + chunk_offsets, sentinel) - else: - blocks_0 = tl.load( - idx_ptr_0 + chunk_offsets * stride_tk, - mask=chunk_offsets < count_0, - other=sentinel, - ) - blocks_1 = tl.load( - idx_ptr_1 + chunk_offsets * stride_tk, - mask=chunk_offsets < count_1, - other=sentinel, - ) - chunk_equal = tl.sum(blocks_0 != blocks_1, axis=0) == 0 - # Equal chunks take the measured shared-page loop. A divergent chunk - # walks the sorted union, bounded by the sum of both row lengths. - num_candidates = tl.where(chunk_equal, count_0, count_0 + count_1) - else: - num_candidates = chunk_end_0 - chunk_start - - for cand in tl.range(0, num_candidates): - if QUERY_GROUP_SIZE == 2: - if chunk_equal: - has_block = True - if IDENTITY: - blk = chunk_start + cand - else: - blk = tl.load(idx_ptr_0) - idx_ptr_0 += stride_tk - query_active = h_mask - else: - valid_0 = idx_0 < chunk_end_0 - valid_1 = idx_1 < chunk_end_1 - if IDENTITY: - blk_0 = tl.where(valid_0, idx_0, sentinel) - blk_1 = tl.where(valid_1, idx_1, sentinel) - else: - blk_0 = tl.load(idx_ptr_0, mask=valid_0, other=sentinel) - blk_1 = tl.load(idx_ptr_1, mask=valid_1, other=sentinel) - use_0 = valid_0 & (blk_0 <= blk_1) - use_1 = valid_1 & (blk_1 <= blk_0) - has_block = use_0 | use_1 - blk = tl.minimum(blk_0, blk_1) - idx_0 += use_0.to(tl.int32) - idx_1 += use_1.to(tl.int32) - idx_ptr_0 += use_0.to(tl.int64) * stride_tk - idx_ptr_1 += use_1.to(tl.int64) * stride_tk - query_active = h_mask & (((q_lane == 0) & use_0) | ((q_lane == 1) & use_1)) - else: - has_block = True - if IDENTITY: - blk = chunk_start + cand - else: - blk = tl.load(idx_ptr_0) - idx_ptr_0 += stride_tk - query_active = h_mask - - if has_block: - # int64 page offsets: a large cache overflows int32 well before the - # per-page block offsets do. - # Fail closed as in _gqa_sparse_decode_kernel: an id outside the row, or a page outside the - # pool, reads nothing. - blk_ok = (blk >= 0) & (blk < num_bt_cols) - page = tl.load(bt_row + blk, mask=blk_ok, other=0).to(tl.int64) - blk_ok = blk_ok & (page >= 0) & (page < num_pages) - page = tl.where(blk_ok, page, 0) - pos_mask = blk * BLOCK_SIZE_K + off_n < kv_len_max - q_pos_mask = blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None] - # A slot that failed closed is skipped whole: no score, no rescale (an all-masked first block - # would make exp2(-inf - -inf) = NaN). - if blk_ok: - if KV_NVFP4: - # K scales are token-major linear within the page-head region; - # see _fused_nvfp4_paged_scatter_kernel in msa_scatter. - k_deq = _dequant_nvfp4_rows( - tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[:, None] * stride_k_pos - + off_dp[None, :] * stride_k_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - k_sc_ptr - + page * stride_ksc_blk - + pid_kh * stride_ksc_h - + off_n[:, None] * SCALE_COLS - + off_s[None, :], - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ) - k = tl.trans(k_deq).to(q.dtype) - else: - k = tl.load( - k_ptr - + page * stride_k_blk - + pid_kh * stride_k_h - + off_n[None, :] * stride_k_pos - + off_d[:, None] * stride_k_d, - mask=d_mask[:, None] & pos_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - k = (k * kv_scale).to(q.dtype) - active_pos = query_active[:, None] & q_pos_mask - qk = tl.where(active_pos, 0.0, float("-inf")) - qk += tl.dot(q, k) * sm_scale_log2e - m_ij = tl.where(query_active, tl.maximum(m_i, tl.max(qk, axis=1)), m_i) - p = tl.where(active_pos, tl.exp2(qk - m_ij[:, None]), 0.0) - l_ij = tl.sum(p, axis=1) - alpha = tl.where(query_active, tl.exp2(m_i - m_ij), 1.0) - acc_o = acc_o * alpha[:, None] - if KV_NVFP4: - # V scales use vLLM's 4x4 token-quad order rather than K's - # linear one: [token // 4, scale_col, token % 4]. - v = _dequant_nvfp4_rows( - tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_dp[None, :] * stride_v_d, - mask=pos_mask[:, None], - other=0, - ), - tl.load( - v_sc_ptr - + page * stride_vsc_blk - + pid_kh * stride_vsc_h - + (off_n[:, None] // 4) * (4 * SCALE_COLS) - + 4 * off_s[None, :] - + (off_n[:, None] % 4), - mask=pos_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=BLOCK_SIZE_D, - SCALE_COLS=SCALE_COLS, - SF_VEC=BLOCK_SIZE_D // SCALE_COLS, - ).to(q.dtype) - else: - v = tl.load( - v_ptr - + page * stride_v_blk - + pid_kh * stride_v_h - + off_n[:, None] * stride_v_pos - + off_d[None, :] * stride_v_d, - mask=pos_mask[:, None] & d_mask[None, :], - other=0.0, - ).to(q.dtype) - if USE_SCALE: - v = (v * kv_scale).to(q.dtype) - acc_o += tl.dot(p.to(v.dtype), v) - if USE_LINEAR_SOFTMAX: - l_i = tl.where(query_active, l_i * alpha + l_ij, l_i) - else: - lse_ij = m_ij + tl.log2(tl.exp2(lse_i - m_ij) + l_ij) - lse_i = tl.where(query_active, lse_ij, lse_i) - m_i = m_ij - - # An empty chunk of an active row must store zero, or the merge hits 0 * NaN. - if USE_LINEAR_SOFTMAX: - scale = tl.where(l_i > 0.0, 1.0 / l_i, 0.0) - lse_out = tl.where(l_i > 0.0, m_i + tl.log2(l_i), float("-inf")) - else: - scale = tl.where(lse_i > float("-inf"), tl.exp2(m_i - lse_i), 0.0) - lse_out = lse_i - if KV_NVFP4: - # Applying V's scale per partial is exact: the merge is a weighted - # average whose weights sum to one, so a constant survives it. - scale = scale * v_global_scale - acc_o = acc_o * scale[:, None] - tl.store( - o_ptr - + pid_c * stride_o_c - + q_token[:, None] * stride_o_b - + (pid_h + h_lane[:, None]) * stride_o_h - + off_d[None, :] * stride_o_d, - acc_o.to(o_ptr.dtype.element_ty), - mask=hd_mask, - ) - if not DIRECT_OUTPUT: - tl.store( - lse_ptr + pid_c * stride_l_c + q_token * stride_l_b + (pid_h + h_lane) * stride_l_h, - lse_out, - mask=h_mask, - ) - - # After the stores, never before: this releases either the merge grid or, - # for direct output, the next PDL-enabled consumer of the final tensor. - if USE_PDL: - tl.extra.cuda.gdc_launch_dependents() - - -@triton.jit -def _rcp_approx_ftz(x): - return tl.inline_asm_elementwise( - "rcp.approx.ftz.f32 $0, $1;", "=r,r", [x], dtype=tl.float32, is_pure=True, pack=1 - ) - - -@triton.jit -def _store_mxfp8_head(x, q8_ptr, sf_ptr, row, head, num_sf_k_tiles, HEAD_DIM: tl.constexpr): - """Quantize one row's [head * HEAD_DIM, (head + 1) * HEAD_DIM) slice to MXFP8 as - torch.ops.trtllm.mxfp8_quantize(x, True) does: UE8M0 scale = round-up(amax * rcp(448)), E4M3 RN-satfinite, - scales in the 128x4-swizzled layout.""" - NUM_SF: tl.constexpr = HEAD_DIM // 32 - xr = tl.reshape(x.to(tl.float32), (NUM_SF, 32)) - amax = tl.max(tl.abs(xr), axis=1) - sfv = amax * _rcp_approx_ftz(tl.full((NUM_SF,), 448.0, tl.float32)) - bits = sfv.to(tl.int32, bitcast=True) - e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) - e = tl.minimum(e, 254) - scale = ((254 - e) << 23).to(tl.float32, bitcast=True) - scale = tl.where(amax != 0, scale, 0.0) - q = (xr * scale[:, None]).to(tl.float8e4nv) - off_q = tl.reshape(tl.arange(0, HEAD_DIM), (NUM_SF, 32)) - tl.store(q8_ptr + row * (num_sf_k_tiles * 128) + head * HEAD_DIM + off_q, q) - k_sf = head * NUM_SF + tl.arange(0, NUM_SF) - off_sf = ( - (row // 128) * num_sf_k_tiles * 512 - + (k_sf // 4) * 512 - + (row % 32) * 16 - + ((row % 128) // 32) * 4 - + k_sf % 4 - ) - tl.store(sf_ptr + off_sf, e.to(tl.uint8)) - - -@triton.heuristics({"BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"])}) -@triton.jit -def _merge_topk_attn_out_kernel( - o_ptr, # partials: [NUM_TOPK_CHUNKS, total_q, num_heads, head_dim] - lse_ptr, # partials (log2): [NUM_TOPK_CHUNKS, total_q, num_heads] - out_ptr, # merged out: [total_q, num_heads, head_dim] - head_dim, - stride_o_c, - stride_o_b, - stride_o_h, - stride_o_d, - stride_l_c, - stride_l_b, - stride_l_h, - stride_out_n, - stride_out_h, - stride_out_d, - NUM_TOPK_CHUNKS: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - USE_PDL: tl.constexpr, - EARLY_TRIGGER: tl.constexpr = False, - q8_ptr=None, # MXFP8 copy of out: [total_q, num_heads * head_dim] E4M3 - sf_ptr=None, # its UE8M0 block scales, 128x4-swizzled - num_sf_k_tiles=0, # ceil(num_heads * head_dim / 128) - MXFP8_OUT: tl.constexpr = False, -): - pid_b, pid_h = tl.program_id(0), tl.program_id(1) - - if USE_PDL: - if EARLY_TRIGGER: - # Release the dependent (the o_proj GEMM) before waiting on the attention producer, so its weight - # prefetch overlaps the attention kernel as well as this merge; its griddepcontrol.wait still waits - # for this grid (and therefore the producer) to complete. Because this trigger precedes this grid's own - # wait, it says nothing about what upstream has finished: a dependent may read only step-invariant data - # (weights) before its griddepcontrol.wait, never this step's attention output, q, top-k table or lengths. - tl.extra.cuda.gdc_launch_dependents() - tl.extra.cuda.gdc_wait() - - off_c = tl.arange(0, NUM_TOPK_CHUNKS) - off_d = tl.arange(0, BLOCK_SIZE_D) - o = tl.load( - o_ptr - + pid_b * stride_o_b - + pid_h * stride_o_h - + off_c[:, None] * stride_o_c - + off_d[None, :] * stride_o_d, - mask=off_d[None, :] < head_dim, - other=0.0, - ).to(tl.float32) - # Empty chunks contribute -inf, hence weight 0. - lse = tl.load(lse_ptr + pid_b * stride_l_b + pid_h * stride_l_h + off_c * stride_l_c) - lse_max = tl.max(lse, axis=0) - # A row whose every chunk is empty (a CUDA-graph padding row) would give - # -inf - -inf here. Zero it instead: the padded output is discarded, but a - # NaN would survive into the residual stream and the all-reduce. - lse_max = tl.where(lse_max == float("-inf"), 0.0, lse_max) - weights = tl.exp2(lse - lse_max) - denom = tl.sum(weights, axis=0) - weights = weights / tl.where(denom > 0, denom, 1.0) - o_merged = tl.sum(o * weights[:, None], axis=0) - out_ptrs = out_ptr + pid_b * stride_out_n + pid_h * stride_out_h + off_d * stride_out_d - tl.store(out_ptrs, o_merged.to(out_ptr.dtype.element_ty), mask=off_d < head_dim) - if MXFP8_OUT: - _store_mxfp8_head( - o_merged.to(out_ptr.dtype.element_ty), - q8_ptr, - sf_ptr, - pid_b, - pid_h, - num_sf_k_tiles, - BLOCK_SIZE_D, - ) - - if USE_PDL and not EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - -def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: - """Split-K factor over the top-k blocks, as a power of two. - - Depends only on shapes that are fixed for a captured batch size, so the - launch geometry is frozen inside a CUDA graph. - """ - target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) - return 1 << (target.bit_length() - 1) - - # NVFP4 instructions are supported on the datacenter Blackwell and Rubin targets. _SM100F_CAPABILITIES = ((10, 0), (10, 3), (10, 7)) # Shape-specific tuning was measured on GB300 only. SM100 uses the defaults. -_NVFP4_TUNED_CAPABILITIES = ((10, 3),) - - -def _sm100f_nvfp4_num_topk_chunks( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> Optional[int]: - """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return None - local_batch = total_q // decode_query_len - return { - 2: 8, - 3: 8, - 4: 8, - 5: 16, - 6: 4, - 7: 4, - 8: 2, - 9: 2, - 10: 8, - 11: 4, - 12: 2, - 13: 2, - 14: 2, - }.get(local_batch) - - -def _sm100f_nvfp4_use_linear_softmax( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> bool: - """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads not in (1, 2, 4) - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return False - local_batch = total_q // decode_query_len - # The TP4/DEP32 single-request graph is 8.9% faster with the original - # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral - # or faster with the linear denominator. - return not (num_kv_heads == 1 and local_batch == 1) - - -def _sm100f_nvfp4_query_group_size( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> int: - """Group Eagle queries using the GB300-measured shared-page policy.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - or total_q % decode_query_len != 0 - ): - return 1 - # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches - # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or - # slower and retain the single-query producer. - return 2 if total_q // decode_query_len in (11, 12, 14) else 1 - - -def _sm100f_nvfp4_launch_options( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> dict[str, int]: - """Select native-NVFP4 launch geometry for the SM100 family. - - The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four - warps/two stages wins through local batch 28, while two warps/one stage - avoids the next residency cliff at local batches 29--32. Leave every - unmeasured architecture, model geometry, and larger batch on Triton's - established defaults. - """ - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability not in _NVFP4_TUNED_CAPABILITIES - or num_kv_heads != 4 - or gqa_group_size != 16 - or max_topk != 16 - or decode_query_len != 4 - ): - return {} - local_batch = total_q // decode_query_len - if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): - return {"num_warps": 4, "num_stages": 1} - if local_batch == 4: - return {"num_warps": 2, "num_stages": 1} - if 8 <= local_batch <= 28: - return {"num_warps": 4, "num_stages": 2} - if 29 <= local_batch <= 32: - return {"num_warps": 2, "num_stages": 1} - return {} - - -def _sm100f_nvfp4_merge_launch_options( - *, - total_q: int, - num_kv_heads: int, - gqa_group_size: int, - max_topk: int, - decode_query_len: int, - capability: Optional[tuple[int, int]] = None, -) -> dict[str, int]: - """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" - if capability is None: - capability = torch.cuda.get_device_capability() - if ( - capability in _NVFP4_TUNED_CAPABILITIES - and num_kv_heads == 4 - and gqa_group_size == 16 - and max_topk == 16 - and decode_query_len == 4 - and total_q % decode_query_len == 0 - and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) - ): - return {"num_warps": 1} - return {} - - -def _check_nvfp4_inputs( - k_paged: torch.Tensor, - v_paged: torch.Tensor, - nvfp4_args: tuple[Optional[torch.Tensor], ...], - *, - head_dim: int, - scale_cols: int, -) -> None: - """Validate the packed-NVFP4 cache against what the kernel assumes. - - The kernel indexes block scales by a flat within-page offset and reshapes a - [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows - and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them - rather than let a mismatch read neighboring pages. - """ - k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args - if any(arg is None for arg in nvfp4_args): - raise ValueError( - "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " - "v_block_scale, k_global_scale and v_global_scale." - ) - if head_dim != triton.next_power_of_2(head_dim): - raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") - if v_paged.shape[:3] != k_paged.shape[:3]: - raise ValueError("NVFP4 K and V caches must have matching page, head and token extents.") - for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): - if int(data.shape[3]) * 2 != head_dim: - raise ValueError( - f"NVFP4 {name} must pack two elements per byte: expected " - f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." - ) - for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): - if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): - raise ValueError( - f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " - f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." - ) - if scale.stride(3) != 1 or scale.stride(2) != scale_cols: - raise ValueError( - f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " - f"{tuple(scale.stride())}." - ) - for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): - if gs.dtype != torch.float32 or gs.numel() != 1: - raise ValueError( - f"NVFP4 {name} must be a single-element FP32 tensor; got " - f"{gs.numel()} element(s) of {gs.dtype}." - ) # Multi-query producer (triton_sparse_decode_mq) for identity-selection verify steps with at least this many query @@ -1598,7 +509,7 @@ def minimax_m3_sparse_attn_decode( else {} ) if use_mq: - from .triton_sparse_decode_mq import launch_mq_producer + from ...kernels.attention.triton_sparse_decode_mq import launch_mq_producer launch_mq_producer( q, diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py deleted file mode 100644 index 34794a95aadd..000000000000 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/triton_sparse_decode_mq.py +++ /dev/null @@ -1,239 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Multi-query block-sparse decode producer for MiniMax-M3 short-sequence verify steps. - -In short-sequence mode every request attends at most ``topk`` KV blocks and block selection is the identity, so -all ``decode_query_len`` query tokens of a request read the same blocks. The per-token producer -(``_gqa_sparse_decode_kernel``) then dequantizes every block once per query token. This producer runs one program -per (request, block) for all of the request's query tokens: rows are (token, head) pairs, each block is -dequantized once, and the causal limit is a per-row mask. - -It writes the same per-block partials (output scaled by 1/l and the V global scale, log2 LSE) as the per-token -producer with one block per chunk, so ``_merge_topk_attn_out_kernel`` merges them unchanged. -""" - -import triton -import triton.language as tl - -from .triton_sparse_decode import _dequant_nvfp4_rows - - -@triton.jit( - do_not_specialize=["num_reqs"], do_not_specialize_on_alignment=["block_table_ptr", "seq_lens"] -) -def _gqa_sparse_decode_mq_kernel( - q_ptr, # [num_reqs * QLEN, num_heads, head_dim], FP8 or FP16/BF16 - k_ptr, # packed NVFP4 K: [num_pages, 1, page_size, head_dim / 2] uint8 - v_ptr, # packed NVFP4 V: same layout - k_sc_ptr, # K block scales [num_pages, 1, page_size, head_dim / 16], token-linear - v_sc_ptr, # V block scales, 4x4 token-quad order - k_gs_ptr, - v_gs_ptr, - t_ptr, # topk table [1, total_q, topk]; the request's last token's row is read - o_ptr, # partials [topk, total_q, num_heads, head_dim] fp32 - lse_ptr, # partials [topk, total_q, num_heads] fp32 (log2) - block_table_ptr, - seq_lens, - num_reqs, - sm_scale, - stride_qn, - stride_qh, - stride_qd, - stride_k_blk, - stride_k_pos, - stride_v_blk, - stride_v_pos, - stride_ksc_blk, - stride_vsc_blk, - stride_tn, - stride_tk, - stride_o_c, - stride_o_b, - stride_o_h, - stride_l_c, - stride_l_b, - stride_l_h, - stride_bt_b, - QLEN: tl.constexpr, - GROUP: tl.constexpr, # query heads per KV head (16) - BLOCK_R: tl.constexpr, # next power of two of QLEN * GROUP - HEAD_DIM: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, - SCALE_COLS: tl.constexpr, - USE_PDL: tl.constexpr, - EARLY_TRIGGER: tl.constexpr, -): - pid_r = tl.program_id(0) - pid_c = tl.program_id(1) # top-k position == block index under identity selection - sm_scale_log2e = sm_scale * 1.4426950409 * tl.load(k_gs_ptr) - v_global_scale = tl.load(v_gs_ptr) - - off_r = tl.arange(0, BLOCK_R) - off_n = tl.arange(0, BLOCK_SIZE_K) - off_d = tl.arange(0, HEAD_DIM) - off_dp = tl.arange(0, HEAD_DIM // 2) - off_s = tl.arange(0, SCALE_COLS) - row_t = off_r // GROUP - row_h = off_r % GROUP - r_mask = off_r < QLEN * GROUP - tok = pid_r * QLEN + row_t - - if USE_PDL: - tl.extra.cuda.gdc_wait() - if EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - seq_len = tl.load(seq_lens + pid_r) - # Token t of the request sits at seq_len - QLEN + t and attends positions [0, kv_len_t). - kv_len = tl.maximum(seq_len - QLEN + row_t + 1, 0) - blk = tl.load(t_ptr + (pid_r * QLEN + QLEN - 1) * stride_tn + pid_c * stride_tk) - blk_ok = blk >= 0 - blk = tl.where(blk_ok, blk, 0) - page = tl.load(block_table_ptr + pid_r * stride_bt_b + blk, mask=blk_ok, other=0).to(tl.int64) - # Rows past the last token's position are never read: their cache bytes may be stale (0 * NaN in the PV dot). - n_valid = tl.where( - blk_ok, tl.minimum(tl.maximum(seq_len - blk * BLOCK_SIZE_K, 0), BLOCK_SIZE_K), 0 - ) - n_mask = off_n < n_valid - - q = tl.load( - q_ptr + tok[:, None] * stride_qn + row_h[:, None] * stride_qh + off_d[None, :] * stride_qd, - mask=r_mask[:, None], - other=0.0, - ).to(tl.float16) - k_deq = _dequant_nvfp4_rows( - tl.load( - k_ptr + page * stride_k_blk + off_n[:, None] * stride_k_pos + off_dp[None, :], - mask=n_mask[:, None], - other=0, - ), - tl.load( - k_sc_ptr + page * stride_ksc_blk + off_n[:, None] * SCALE_COLS + off_s[None, :], - mask=n_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=HEAD_DIM, - SCALE_COLS=SCALE_COLS, - SF_VEC=HEAD_DIM // SCALE_COLS, - ) - valid = ( - (blk * BLOCK_SIZE_K + off_n[None, :] < kv_len[:, None]) & n_mask[None, :] & r_mask[:, None] - ) - qk = tl.dot(q, tl.trans(k_deq).to(tl.float16)) * sm_scale_log2e - qk = tl.where(valid, qk, float("-inf")) - m = tl.max(qk, axis=1) - m_safe = tl.where(m == float("-inf"), 0.0, m) - p = tl.exp2(qk - m_safe[:, None]) - p_sum = tl.sum(p, axis=1) - v = _dequant_nvfp4_rows( - tl.load( - v_ptr + page * stride_v_blk + off_n[:, None] * stride_v_pos + off_dp[None, :], - mask=n_mask[:, None], - other=0, - ), - tl.load( - v_sc_ptr - + page * stride_vsc_blk - + (off_n[:, None] // 4) * (4 * SCALE_COLS) - + 4 * off_s[None, :] - + (off_n[:, None] % 4), - mask=n_mask[:, None], - other=0, - ), - BLOCK_N=BLOCK_SIZE_K, - D=HEAD_DIM, - SCALE_COLS=SCALE_COLS, - SF_VEC=HEAD_DIM // SCALE_COLS, - ).to(tl.float16) - acc = tl.dot(p.to(tl.float16), v) - # An empty partial stores zero with LSE -inf, so the merge gives it weight 0. - scale = tl.where(p_sum > 0.0, 1.0 / p_sum, 0.0) * v_global_scale - lse = tl.where(p_sum > 0.0, m + tl.log2(p_sum), float("-inf")) - acc = acc * scale[:, None] - tl.store( - o_ptr - + pid_c * stride_o_c - + tok[:, None] * stride_o_b - + row_h[:, None] * stride_o_h - + off_d[None, :], - acc, - mask=r_mask[:, None], - ) - tl.store(lse_ptr + pid_c * stride_l_c + tok * stride_l_b + row_h * stride_l_h, lse, mask=r_mask) - if USE_PDL and not EARLY_TRIGGER: - tl.extra.cuda.gdc_launch_dependents() - - -def launch_mq_producer( - q, - k_paged, - v_paged, - k_block_scale, - v_block_scale, - k_global_scale, - v_global_scale, - topk_idx, - block_table, - seq_lens, - o_arg, - lse_arg, - *, - sm_scale, - decode_query_len, - use_pdl, - early_trigger, - num_warps=8, - num_stages=1, -): - """Launch the multi-query producer; arguments as minimax_m3_sparse_attn_decode after its NVFP4 views.""" - total_q, num_heads, head_dim = q.shape - num_reqs = total_q // decode_query_len - max_topk = int(topk_idx.shape[-1]) - rows = decode_query_len * num_heads - pdl_launch = {"launch_pdl": True} if use_pdl else {} - _gqa_sparse_decode_mq_kernel[(num_reqs, max_topk)]( - q, - k_paged, - v_paged, - k_block_scale, - v_block_scale, - k_global_scale, - v_global_scale, - topk_idx, - o_arg, - lse_arg, - block_table, - seq_lens, - num_reqs, - sm_scale, - q.stride(0), - q.stride(1), - q.stride(2), - k_paged.stride(0), - k_paged.stride(2), - v_paged.stride(0), - v_paged.stride(2), - k_block_scale.stride(0), - v_block_scale.stride(0), - topk_idx.stride(1), - topk_idx.stride(2), - o_arg.stride(0), - o_arg.stride(1), - o_arg.stride(2), - lse_arg.stride(0), - lse_arg.stride(1), - lse_arg.stride(2), - block_table.stride(0), - QLEN=decode_query_len, - GROUP=num_heads, - BLOCK_R=triton.next_power_of_2(rows), - HEAD_DIM=head_dim, - BLOCK_SIZE_K=int(k_paged.shape[2]), - SCALE_COLS=head_dim // 16, - USE_PDL=use_pdl, - EARLY_TRIGGER=early_trigger, - num_warps=num_warps, - num_stages=num_stages, - **pdl_launch, - ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/modeling_support.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/modeling_support.py new file mode 100644 index 000000000000..1eb09a216756 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/modeling_support.py @@ -0,0 +1,287 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Common model support with explicit target-owned resources and backend bindings. + +Targets retain forward assembly, topology, weights, shape/tactic policy, mutable +runner caches, diagnostics and their concrete drafter/worker classes. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, List, Protocol + +import torch +from torch import nn + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + + from .ops.gemm.mxfp8_state import MxMode + + +class _CaptureTail(Protocol): + capture_norm: Callable[..., None] + + +class ModelingSupport(Protocol): + m3_spec_tail: _CaptureTail + _record: Callable[..., None] + m3_mxfp8_linear: Callable[..., torch.Tensor] + m3_fp8_indexer_qk_norm_rope: Callable[..., torch.Tensor] + TensorParallelMode: Any + M3Eagle3Draft: Callable[..., nn.Module] + M3Eagle3Worker: Callable[..., Any] + ModelConfig: type[ModelConfig] + _get_requested_draft_moe_backend: Callable[..., Any] + _set_draft_kv_cache_quant_algo: Callable[..., None] + should_use_separate_draft_kv_cache: Callable[..., bool] + + +class PinnedModelingSupport(ModelingSupport, Protocol): + """The existing pinned-backend policy of TP2 and TP8 only.""" + + _PINNED_RUNNERS: dict + FlashInferMXFP8GemmRunner: Callable[..., Any] + MXFP8_DECODE_GEMM_TACTIC: int + MXFP8QuantizeRunner: Callable[..., Any] + MXFP8_DECODE_QUANT_TACTIC: int + _pinned_runner: Callable[..., Any] + + +def _mxfp8_linear( + state: ModelingSupport, + /, + core: Any, + x: torch.Tensor, + w: torch.Tensor, + s: torch.Tensor, + mode: MxMode, +) -> torch.Tensor: + """One MXFP8 projection on the engine-tuned path: MXFP8 quantize + GEMM, backend per ``mode``.""" + state._record( + "m3_mxfp8_linear", + tuple(w.shape), + mode.tuned_quantize_gemm, + mode.flashinfer, + mode.native_tuning, + ) + return state.m3_mxfp8_linear( + x.contiguous(), + w, + s, + core.mx_alpha, + torch.bfloat16, + mode.tuned_quantize_gemm, + mode.autotune, + mode.flashinfer, + mode.native_tuning, + ) + + +def _pinned_runner(state: PinnedModelingSupport, /, cls: Any, dtype: Any) -> Any: + runner = state._PINNED_RUNNERS.get((cls, dtype)) + if runner is None: + runner = state._PINNED_RUNNERS[cls, dtype] = cls(dtype) + return runner + + +def _pinned_mxfp8_gemm( + state: PinnedModelingSupport, + /, + act: Any, + act_scale: Any, + weight: Any, + weight_scale: Any, + out_dtype: Any, +) -> Any: + """The decode-graph MXFP8 GEMM at the target's fixed tactic (``MXFP8_DECODE_GEMM_TACTIC``), no AutoTuner: the + same choice on every rank and every start.""" + state._record( + "pinned_mxfp8_gemm", tuple(weight.shape), act.shape[0], state.MXFP8_DECODE_GEMM_TACTIC + ) + runner = state._pinned_runner(state.FlashInferMXFP8GemmRunner, out_dtype) + return runner.forward( + [act, act_scale, weight, weight_scale], tactic=state.MXFP8_DECODE_GEMM_TACTIC + ) + + +def _pinned_mxfp8_quantize(state: PinnedModelingSupport, /, x: Any) -> Any: + """The decode-graph MXFP8 activation quantize at the target's fixed tactic (``MXFP8_DECODE_QUANT_TACTIC``).""" + state._record("pinned_mxfp8_quantize", x.shape[0], state.MXFP8_DECODE_QUANT_TACTIC) + return state._pinned_runner(state.MXFP8QuantizeRunner, x.dtype).forward( + [x], tactic=state.MXFP8_DECODE_QUANT_TACTIC + ) + + +def _split_chunk(state: ModelingSupport, /, rows: int, num_chunks: int) -> List[int]: + """The MoE scheduler's even split: the first ``rows % num_chunks`` chunks take one more row.""" + base, extra = divmod(rows, num_chunks) + return [base + 1] * extra + [base] * (num_chunks - extra) + + +def _nvfp4_layer(state: ModelingSupport, /, manager: Any, i: int) -> bool: + return bool(getattr(manager, "is_nvfp4_layer", lambda _layer_idx: False)(i)) + + +def _nvfp4_fused_write_layout( + state: ModelingSupport, /, data: torch.Tensor, scales: torch.Tensor, num_heads: int +) -> bool: + """Reject invalid HND pools; select the fused producer only for its packed layout.""" + if ( + data.dim() != 5 + or scales.dim() != 5 + or data.shape[0] <= 0 + or (data.shape[0] != scales.shape[0]) + or (tuple(data.shape[1:]) != (2, num_heads, 128, 64)) + or (tuple(scales.shape[1:]) != (2, num_heads, 128, 8)) + or (data.dtype not in (torch.int8, torch.uint8)) + or (scales.dtype != torch.uint8) + ): + raise ValueError("MiniMax-M3 NVFP4 data and scale caches must have matching HND geometry") + for cache in (data, scales): + if ( + cache.stride(4) != 1 + or cache.stride(3) != cache.shape[4] + or cache.stride(2) < 128 * cache.stride(3) + or (cache.stride(1) < num_heads * cache.stride(2)) + or (cache.stride(0) < 2 * cache.stride(1)) + ): + raise ValueError( + "MiniMax-M3 NVFP4 cache rows must be packed and heads/planes/pages must not overlap" + ) + if scales.stride(2) != 128 * 8: + raise ValueError("MiniMax-M3 NVFP4 MSA scale cache must have packed head blocks") + return data.stride(2) == 128 * 64 and data.stride(0) % 2 == 0 and (data.stride(1) % 2 == 0) + + +def _fp8_index_qk_norm_rope( + state: ModelingSupport, + /, + core: Any, + i: int, + idx_qk: torch.Tensor, + position_ids: Any, + msa_idx_k_cache: Any, + msa_out_cache_loc: Any, +) -> torch.Tensor: + """The index branch of an unfused sparse layer: index-Q in FP8 out, index-K into the FP8 index-K cache, normed and + roped in one kernel (the built-in ``_fused_fp8_index_qk_norm_rope``). ``msa_idx_k_cache`` / ``msa_out_cache_loc`` + are the calling target's ``step_args`` entries.""" + if position_ids is None or idx_qk.dtype != torch.bfloat16: + raise NotImplementedError( + "MiniMax-M3 fused FP8 indexer requires BF16 activations and position_ids." + ) + index_k_cache = msa_idx_k_cache(i) + if index_k_cache.dtype != torch.float8_e4m3fn: + raise ValueError( + f"MiniMax-M3 fused FP8 indexer requires an E4M3 index-K cache, got {index_k_cache.dtype}." + ) + rows = int(idx_qk.shape[0]) + state._record("m3_fp8_indexer_qk_norm_rope", rows) + return state.m3_fp8_indexer_qk_norm_rope( + idx_qk.contiguous(), + index_k_cache, + msa_out_cache_loc[:rows], + core.index_heads, + core.index_dim, + core.rotary_dim, + core.rms_eps, + core.w[f"l{i}_idx_q_norm"], + core.w[f"l{i}_idx_k_norm"], + core.rope_theta, + position_ids.reshape(-1).contiguous().to(torch.int32), + ).flatten(1) + + +def _eagle3_capture( + state: ModelingSupport, + /, + spec_metadata: Any, + layer_id: int, + hidden_states: Any, + residual: Any, + capture_norm: Any, +) -> None: + """The Eagle3 hidden-state capture of one target layer. + + Mirrors ``Eagle3OneModelSpecMetadata.maybe_capture_hidden_states``. With ``capture_norm`` = (the drafter's stacked + fc_norm weights [layers, H], eps), the capture writes rmsnorm(bf16(h + residual)) * fc_norm[i] in one kernel (the + in-tree capture's bf16 add, then the in-tree fc_norm arithmetic), which the fused first drafter inputs read as is; + otherwise it is the in-tree capture. + """ + if capture_norm is None: + spec_metadata.maybe_capture_hidden_states(layer_id, hidden_states, residual) + return + for i, captured_layer_id in enumerate(spec_metadata.layers_to_capture): + if captured_layer_id == layer_id: + assert spec_metadata.hidden_states is not None + num_tokens = spec_metadata.num_capture_tokens + if num_tokens > spec_metadata.hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds buffer capacity: " + f"num_tokens={num_tokens}, capacity={spec_metadata.hidden_states.shape[0]}" + ) + if num_tokens > hidden_states.shape[0]: + raise RuntimeError( + "EAGLE3 hidden-state capture token count exceeds available hidden states: " + f"num_tokens={num_tokens}, hidden_states={hidden_states.shape[0]}" + ) + hidden_size = spec_metadata.hidden_size + state.m3_spec_tail.capture_norm( + hidden_states[:num_tokens], + residual[:num_tokens] if residual is not None else None, + capture_norm[0][i], + capture_norm[1], + spec_metadata.hidden_states[:num_tokens, i * hidden_size : (i + 1) * hidden_size], + ) + break + + +def _sharded_head(state: ModelingSupport, /, lm_head: nn.Module) -> bool: + """A vocab-sharded TP head without padding (the fused kernels read the local shard).""" + return ( + getattr(lm_head, "tp_mode", None) == state.TensorParallelMode.COLUMN + and getattr(lm_head, "padding_size", 0) == 0 + and (getattr(lm_head, "tp_size", 1) > 1) + ) + + +def _build_speculation(state: ModelingSupport, /, shell: Any, model_config: ModelConfig) -> Any: + """The drafter config, the drafter and the worker of one-model Eagle3 (the in-tree one-engine construction for + the llama3 drafter architecture).""" + spec = model_config.spec_config + if ( + not spec.spec_dec_mode.use_one_engine() + or not spec.spec_dec_mode.is_eagle3_one_model() + or spec.eagle3_model_arch != "llama3" + or getattr(spec, "use_dynamic_tree", False) + ): + raise ValueError( + "this MiniMax-M3 target speculates with linear one-model Eagle3 (llama3 drafter) only" + ) + draft_config = state.ModelConfig.from_pretrained( + spec.speculative_model, + trust_remote_code=True, + attn_backend=model_config.attn_backend, + moe_backend=state._get_requested_draft_moe_backend(model_config, spec), + mapping=model_config.mapping, + spec_config=model_config.spec_config, + max_num_tokens=model_config.max_num_tokens, + moe_max_num_tokens=model_config.moe_max_num_tokens, + ) + state._set_draft_kv_cache_quant_algo(draft_config, model_config) + draft_config.extra_attrs = model_config.extra_attrs + use_separate = state.should_use_separate_draft_kv_cache(spec) + draft = state.M3Eagle3Draft( + shell.model, + draft_config, + model_config.pretrained_config.num_hidden_layers, + spec.num_capture_layers, + ) + worker = state.M3Eagle3Worker( + spec, model_config.mapping, use_separate_draft_kv_cache=use_separate + ) + worker.set_draft_model(draft) + shell.epilogue.append(draft) + shell.epilogue.append(worker) + return (draft_config, draft, worker, use_separate) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/activation/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/activation/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/activation/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/activation/m3_silu_and_mul.py similarity index 67% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/activation/m3_silu_and_mul.py index daea97df9d12..5a61f9684843 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/activation/m3_silu_and_mul.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/activation/m3_silu_and_mul.py @@ -11,6 +11,51 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +CELLS_BY_TP = { + 2: ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=1536, oai=True)), + Cell(why="shared expert, an 8-row decode verify", spec=dict(m=8, inter=1536, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=6144, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=6144, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=9216, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=9216, oai=False), + ), + ), + 4: ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=768, oai=True)), + Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=768, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=3072, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=3072, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=4608, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=4608, oai=False), + ), + ), + 8: ( + Cell(why="shared expert, one decode token", spec=dict(m=1, inter=384, oai=True)), + Cell(why="shared expert, a 48-row decode verify", spec=dict(m=48, inter=384, oai=True)), + Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=1536, oai=True)), + Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=1536, oai=True)), + Cell( + why="drafter MLP (plain SwiGLU), eight decode rows", + spec=dict(m=8, inter=2304, oai=False), + ), + Cell( + why="drafter MLP (plain SwiGLU), a prefill chunk", + spec=dict(m=1024, inter=2304, oai=False), + ), + ), +} + class M3SiluAndMul(OpWrapper): """``out = g * sigmoid(alpha * g) * (u + beta)`` over ``x = [g | u]`` (bf16 [M, 2I] -> [M, I]), with @@ -24,20 +69,8 @@ class M3SiluAndMul(OpWrapper): ARCHS = frozenset({Arch.SM_107}) - CELLS: tuple[Cell, ...] = ( - Cell(why="shared expert, one decode token", spec=dict(m=1, inter=1536, oai=True)), - Cell(why="shared expert, an 8-row decode verify", spec=dict(m=8, inter=1536, oai=True)), - Cell(why="dense MLP, a prefill chunk", spec=dict(m=2048, inter=6144, oai=True)), - Cell(why="dense MLP, one decode token", spec=dict(m=1, inter=6144, oai=True)), - Cell( - why="drafter MLP (plain SwiGLU), eight decode rows", - spec=dict(m=8, inter=9216, oai=False), - ), - Cell( - why="drafter MLP (plain SwiGLU), a prefill chunk", - spec=dict(m=1024, inter=9216, oai=False), - ), - ) + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ The kernel computes in fp32 and rounds once to bf16; `reference` does the same arithmetic in fp32, so the two diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/draft_attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/draft_attention.py new file mode 100644 index 000000000000..e66a316ecf71 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/draft_attention.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared Eagle3 attention dispatch and grouped-head reference.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + + +class DraftAttentionBase(OpWrapper): + def __call__(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + num_tokens = metadata.num_tokens + out = backend.forward( + qkv[:num_tokens, :], + None, + None, + metadata, + forward_args=AttentionForwardArgs(attention_mask=PredefinedAttentionMask.CAUSAL), + ) + return out[0] if isinstance(out, tuple) else out + + +class GroupedDraftAttentionBase(DraftAttentionBase): + def reference(self, backend, qkv: torch.Tensor, metadata) -> torch.Tensor: + """The call's [num_tokens, q heads x 128] output, in ``qkv``'s dtype. + + Query token ``j`` of request ``r`` sits at position ``cached_r + j`` (the metadata's cached lengths) and + attends positions up to it. Its Q is rotated here; every K / V row, the step's own included, is read from + the drafter layer's paged cache at the slots the metadata's KV cache manager holds, so this is evaluated + after the call (whose roped, quantized append it then reads back; the test checks that append itself). + Mirrors grouped KV heads (query head h reads KV head h // (heads / kv heads); two per rank at tp2) and + un-scaled NeoX RoPE, the drafter's configuration. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = (backend.num_heads, backend.head_dim) + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + rope = backend.rope_params + assert backend.position_embedding_type == int(PositionEmbeddingType.rope_gpt_neox) + assert rope.scale_type == RotaryScalingType.none and rope.alpha == 1.0 + page = int(manager.tokens_per_block) + sm_scale = dim ** (-0.5) / float(backend.q_scaling) + kv_scale = float(backend.kv_scale_quant_orig.reshape(-1)[0]) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = int(metadata.num_tokens) + assert sum(qlens) == total + device = qkv.device + rot = int(rope.dim) + inv_freq = float(rope.theta) ** ( + -torch.arange(0, rot, 2, dtype=torch.float64, device=device) / rot + ) + pool = manager.get_buffers(layer, "HND") + q = qkv[:total, : heads * dim].to(torch.float64).reshape(total, heads, dim) + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = (qlens[row], cached[row] + qlens[row]) + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][ + : (length + page - 1) // page + ], + dtype=torch.int64, + device=device, + ) + keys = ( + pool[pages, 0] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + values = ( + pool[pages, 1] + .to(torch.float64) + .transpose(0, 1) + .reshape(kv_heads, -1, dim)[:, :length] + * kv_scale + ) + keys = keys.repeat_interleave(group, dim=0) + values = values.repeat_interleave(group, dim=0) + qpos = cached[row] + torch.arange(n, device=device) + angles = qpos.to(torch.float64)[:, None] * inv_freq[None, :] + cos, sin = (angles.cos()[:, None, :], angles.sin()[:, None, :]) + rows = q[first : first + n] + x1, x2 = (rows[..., : rot // 2], rows[..., rot // 2 : rot]) + rows = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin, rows[..., rot:]], dim=-1) + allowed = torch.arange(length, device=device)[None, :] <= qpos[:, None] + logits = torch.einsum("thd,hkd->thk", rows, keys) * sm_scale + logits = logits.masked_fill(~allowed[:, None, :], float("-inf")) + out[first : first + n] = torch.einsum( + "thk,hkd->thd", torch.softmax(logits, dim=-1), values + ) + first += n + return out.reshape(total, heads * dim).to(qkv.dtype) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/drafter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/drafter.py new file mode 100644 index 000000000000..7377ef730cc9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/drafter.py @@ -0,0 +1,407 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Common Eagle3 generation preprocessing and full-context plan execution. + +Targets retain their geometry, fused-RoPE eligibility and native launch policy. +The bound full-context module retains its own plan dictionary and lock. +""" + +from __future__ import annotations + +from types import ModuleType +from typing import Any, Callable, Mapping, Optional, Protocol, Tuple, cast + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.flashinfer_trtllm_gen import ( + FlashInferTrtllmGenFmha, +) +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.fmha.utils import ( + get_attention_chunk_size, + get_bmm1_scale, +) +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.paged_cache import ( + fp8_e4m3_saturate, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor +from tensorrt_llm.bindings.internal import thop +from tensorrt_llm.quantization.mode import QuantMode + + +class DrafterFullCtxBase(FlashInferTrtllmGenFmha): + """Shared execution; target subclasses provide their policy and module bindings.""" + + FULLCTX: ModuleType + GEOMETRY_ERROR: str + + def _shape(self, q: torch.Tensor, metadata) -> Optional[Tuple[int, int]]: + nreq = int(metadata.num_generations) + if nreq <= 0 or int(q.shape[0]) % nreq: + return None + return int(q.shape[0]) // nreq, nreq + + def _is_supported(self, q, k, v, metadata, forward_args, *, phase=None) -> bool: + if phase == FmhaPhase.CONTEXT or metadata.num_contexts > 0: + return False + if phase is None and forward_args.attention_input_type == AttentionInputType.context_only: + return False + attn = self.attn + if ( + attn.num_heads != self.FULLCTX.HEADS + or attn.num_kv_heads != self.FULLCTX.KV_HEADS + or attn.head_dim != self.FULLCTX.HEAD_DIM + or not QuantMode(attn.quant_mode).has_fp8_kv_cache() + ): + self._on_uncovered(self.GEOMETRY_ERROR) + return False + # The arithmetic matches the stock decode only at unit KV scales (bmm1 and bmm2 carry none here); the layer's + # own scales are checked to be 1.0 at install, so any other per-step scale falls through to the stock path. + if ( + forward_args.kv_scale_orig_quant is not attn.kv_scale_orig_quant + or forward_args.kv_scale_quant_orig is not attn.kv_scale_quant_orig + ): + self._on_uncovered("per-step KV cache scales") + return False + if int(metadata.max_seq_len) > self._max_kv_len: + self._on_uncovered("the cache length bound is beyond the plans' K/V capacity") + return False + shape = self._shape(q, metadata) + if shape is None or shape not in self._recipes: + self._on_uncovered(f"{shape} is not a covered (query rows, requests) shape") + return False + # Compiled before warmup (``install``); the forward never compiles one. + if self.FULLCTX.plan_for(q.device, *shape, self._own_reducer) is None: + self._on_missed(shape[0] * shape[1], "the PrimTS decode plan is not prebuilt") + return False + return super()._is_supported(q, k, v, metadata, forward_args, phase=FmhaPhase.GENERATION) + + @staticmethod + def _multi_token(params: FmhaParams) -> bool: + return ( + params.spec_decoding_generation_lengths is not None + and int(params.attn.predicted_tokens_per_seq) > 1 + ) + + def _run_fused_rope(self, params: FmhaParams, attention_input: torch.Tensor) -> bool: + """The pass with the one-launch preprocess; False (nothing run) where it does not apply.""" + if not self._fused_rope_applies(params, attention_input): + return False + attn = params.attn + meta = params.meta + qlen, nreq = int(params.input_seq_length), int(params.batch_size) + plan = self.FULLCTX.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + if plan is None: + return False + # The preprocess's view of the layer's cache (buildFlashinferTrtllmGenPagedKvCacheBuffers): the pool's primary + # pointer plus the layer's offset inside its pool, sub-pages of Hkv x tokens_per_block x 128 E4M3 bytes; the + # block table is the pool's [requests, 2, blocks] K / V sub-page indices. + pool_index, layer_in_pool = ( + int(x) for x in meta.host_kv_cache_pool_mapping[attn.local_layer_idx, :2] + ) + tpb = int(params.tokens_per_block) + sub_page_bytes = attn.num_kv_heads * tpb * attn.head_dim + base = ( + int(meta.host_kv_cache_pool_pointers[pool_index, 0]) + + layer_in_pool * int(params.kv_factor) * sub_page_bytes + ) + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(base, torch.float8_e4m3fn, [pages, attn.num_kv_heads, tpb, attn.head_dim]) + ) + tables = meta.kv_cache_block_offsets[ + pool_index, params.seq_offset : params.seq_offset + nreq + ] + multi = self._multi_token(params) + if not torch.cuda.is_current_stream_capturing(): + # The decode reads V page p as sub-page p + 1 of K's (kv_pages[1:] with the K table), as the preprocess + # path checks; the op itself writes V at the table's V sub-page. + tpb_held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(tables.shape[-1], device=tables.device) + in_use = cols[None, :] < tpb_held[:, None] + offsets = (tables[:, 1] - tables[:, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + q = torch.empty( + int(params.num_tokens), + attn.num_heads, + attn.head_dim, + dtype=torch.float8_e4m3fn, + device=attention_input.device, + ) + self._launch_rope( + attention_input, + attn.rotary_cos_sin, + params.sequence_lengths, + params.spec_decoding_position_offsets if multi else None, + kv_pages, + tables, + q, + qlen, + attn.num_heads, + attn.num_kv_heads, + int(attn.rope_params.dim), + tpb, + self._fused_rope_mode, + self._pdl, + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + tables[:, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + if self._on_rope_taken is not None: + self._on_rope_taken(params.num_tokens, "spec" if multi else "plain") + self._on_taken(params.num_tokens) + return True + + def run_generation(self, params: FmhaParams) -> None: + attention_input = params.qkv_input if params.qkv_input is not None else params.query_input + if attention_input is None: + raise RuntimeError("the drafter decode requires QKV or query input.") + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + if self._run_fused_rope(params, attention_input): + return + # The BF16 post-RoPE query, rounded to E4M3 below. At the unit KV scales this layer runs with, that is the FP8 + # query the stock decode gets (the preprocess's FP8 variant is e4m3(q * kv_scale_orig_quant)). + fp8_context_fmha = False + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + _fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + _window_left, + is_multi_token_gen, + ) = thop.trtllm_gen_generation_preprocess( + attention_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + get_bmm1_scale(attn), + 1.0, + fp8_context_fmha, + attn.predicted_tokens_per_seq, + get_attention_chunk_size(attn), + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, + params.is_cross, + ) + qlen, nreq = params.input_seq_length, params.batch_size + plan = self.FULLCTX.plan_for(attention_input.device, qlen, nreq, self._own_reducer) + # A first draft pass arrives as qlen rows per request either as a plain multi-row generation step (the + # drafter's own, input_seq_length = qlen) or with spec-decoding lengths (multi-token generation); a 1-row + # pass is never multi-token. + if plan is None or (is_multi_token_gen and qlen == 1): + raise RuntimeError( + f"the drafter decode was selected for {qlen} query rows x {nreq} requests without a uniform plan" + ) + rows = params.num_tokens * attn.num_heads * attn.head_dim + q = fp8_e4m3_saturate(q_processed.flatten()[:rows]) + q = q.view(params.num_tokens, attn.num_heads, attn.head_dim) + # K pages in row 0 of the table, V in row 1; this cache keeps V one sub-page after K, so V page p is page p + 1 + # of the pool (the two views have the same page count, as the kernel requires). + if not torch.cuda.is_current_stream_capturing(): + # The plan runs every request at qlen rows; a multi-token pass must give each request exactly that many. + lengths = params.spec_decoding_generation_lengths + if is_multi_token_gen and not bool((lengths[:nreq] == qlen).all()): + raise RuntimeError( + f"the drafter pass rows per request are not uniform: {lengths[:nreq].tolist()}" + ) + # Only the blocks a request holds are laid out; the table's tail entries are unset. + tpb = int(params.tokens_per_block) + held = (params.sequence_lengths[:nreq].to(torch.long) + tpb - 1) // tpb + cols = torch.arange(block_tables.shape[-1], device=block_tables.device) + in_use = cols[None, :] < held[:, None] + offsets = (block_tables[:nreq, 1] - block_tables[:nreq, 0])[in_use].unique() + if offsets.numel() != 1 or int(offsets) != 1: + raise RuntimeError( + f"the drafter cache does not keep V one sub-page after K: {offsets.tolist()}" + ) + # The preprocess sizes kv_pool by blocks; its table indexes the K and V sub-pages of every block. + # The view stops two sub-pages past the last block's K, as get_kv_subpage_pool's does. + pages = (int(params.total_num_blocks) - 1) * int(params.kv_factor) + 2 + kv_pages = convert_to_torch_tensor( + TensorWrapper(kv_pool.data_ptr(), kv_pool.dtype, [pages, *kv_pool.shape[1:]]) + ) + plan.run( + q, + kv_pages[:-1], + kv_pages[1:], + block_tables[:nreq, 0], + params.sequence_lengths[:nreq], + get_bmm1_scale(attn), + params.output.view(params.num_tokens, attn.num_heads, attn.head_dim), + ) + self._on_taken(params.num_tokens) + + +def install( + draft_model: torch.nn.Module, + recipes, + max_kv_len: int, + own_reducer: bool, + pdl: bool, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode: Optional[int] = None, + on_rope_taken=None, + *, + prebuild: Callable[..., None], + fmha_class: type[DrafterFullCtxBase], +) -> int: + """Prebuild the plans of ``recipes`` and put ``M3DrafterFullCtxFmha`` first in the FMHA list of the TrtllmAttention + backend of every attention module of ``draft_model``. Call before warmup. Returns the number of attention layers it + was installed on.""" + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + count = 0 + device = next(draft_model.parameters()).device + for module in draft_model.modules(): + backend = getattr(module, "attn", None) + if not isinstance(backend, TrtllmAttention) or not hasattr(backend, "_fmha_manager"): + continue + if not bool((backend.kv_cache_scaling_factor == 1.0).all()): + on_uncovered("a scaled FP8 KV cache") + continue + prebuild(device, recipes, max_kv_len, own_reducer, pdl) + libs = backend._fmha_manager.fmha_libs + if not any((isinstance(lib, fmha_class) for lib in libs)): + libs.insert( + 0, + fmha_class( + backend, + recipes, + max_kv_len, + own_reducer, + on_taken, + on_uncovered, + on_missed, + fused_rope_mode=fused_rope_mode, + pdl=pdl, + on_rope_taken=on_rope_taken, + ), + ) + count += 1 + return count + + +class DrafterBindings(Protocol): + PositionEmbeddingType: Any + fwd_has_no_mrope: Callable[..., bool] + + +class MultiRequestDrafterBindings(DrafterBindings, Protocol): + ROPE_ONE_REQUEST_MULTI: bool + + +def initialize_drafter_state( + self: Any, + state: DrafterBindings, + /, + recipes: Mapping[Tuple[int, int], Mapping], + max_kv_len: int, + own_reducer: bool, + on_taken: Any, + on_uncovered: Any, + on_missed: Any, + fused_rope_mode: Optional[int] = None, + pdl: bool = False, + on_rope_taken: Any = None, + *, + has_one_request_multi: bool, +) -> None: + self._recipes = recipes + self._max_kv_len = max_kv_len + self._own_reducer = own_reducer + self._fused_rope_mode = fused_rope_mode + self._pdl = pdl + self._on_rope_taken = on_rope_taken + if has_one_request_multi: + self._one_request_multi = cast(MultiRequestDrafterBindings, state).ROPE_ONE_REQUEST_MULTI + self._on_taken = on_taken + self._on_uncovered = on_uncovered + self._on_missed = on_missed + + +def drafter_rope_applies( + self: Any, + state: DrafterBindings, + /, + params: FmhaParams, + attention_input: torch.Tensor, + *, + has_one_request_multi: bool, +) -> bool: + attn = params.attn + width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim + rotary_dim = int(attn.rope_params.dim) + return ( + self._fused_rope_mode is not None + and params.qkv_input is not None + and (attention_input.dim() == 2) + and (attention_input.dtype == torch.bfloat16) + and attention_input.is_contiguous() + and (int(attention_input.shape[1]) == width) + and (not params.is_cross) + and (int(attn.position_embedding_type) == int(state.PositionEmbeddingType.rope_gpt_neox)) + and (attn.rotary_cos_sin is not None) + and (attn.rotary_cos_sin.dtype == torch.float32) + and state.fwd_has_no_mrope(params.fwd) + and (0 < rotary_dim <= attn.head_dim) + and (rotary_dim % 16 == 0) + and (int(params.num_tokens) == int(params.batch_size) * int(params.input_seq_length)) + and ( + has_one_request_multi + and self._one_request_multi + or not (self._multi_token(params) and int(params.batch_size) == 1) + ) + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/fullctx.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/fullctx.py new file mode 100644 index 000000000000..a9e76caeb959 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/fullctx.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared PrimTS plan construction; geometry, reducer policy and caches stay target-owned.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Mapping, Tuple + +import torch + + +@dataclass(frozen=True) +class FullCtxGeometry: + heads: int + kv_heads: int + head_dim: int + page: int + policy_kv_len: int + heads_per_kv: int + + +class FullCtxPlanBase: + """Construct one plan with its own workspace and target compile geometry.""" + + GEOMETRY: FullCtxGeometry + + def __init__( + self, + device: torch.device, + qlen: int, + nreq: int, + max_kv_len: int, + recipe: Mapping, + own_reducer: bool, + pdl: bool, + ): + from tensorrt_llm._torch.attention.backends.prims_ts import decode as tsdec + from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( + get_max_active_clusters_for_cluster_size, + make_decode_config, + ) + + self._tsdec = tsdec + self.qlen = qlen + self.nreq = nreq + self.pdl = pdl + self.recipe = dict(recipe) + splits = int(self.recipe["splits"]) + fp8 = tsdec._cutlass_dtype("float8_e4m3fn") + # PDL contract of the main kernel (prims_ts fmha_decode_kernel.py): every CTA acquires its producer + # (griddepcontrol.wait) before any task reads q, lengths or tables, and triggers its dependents + # (launch_dependents) only after its own tail: an active CTA after its last store and TMEM dealloc, an + # inactive split or zero-work CTA right after its wait. The trigger lets the merge launch before the + # slowest CTA ends, so a dependent may read only step-invariant data (weights) before its own + # griddepcontrol.wait, never this step's partial O/LSE, output, q, lengths or tables; the merge + # (minimax_m3_sm107_tp2_fullctx_reduce) waits as its first instruction. + cfg = make_decode_config( + headdim=self.GEOMETRY.head_dim, + args={"use_pdl": pdl, **self.recipe.get("config", {})}, + seq_len_q=qlen, + seq_len_kv=min(max_kv_len, self.GEOMETRY.policy_kv_len), + batch_size=nreq, + num_heads_q=self.GEOMETRY.heads, + num_heads_kv=self.GEOMETRY.kv_heads, + q_dtype=fp8, + k_dtype=fp8, + v_dtype=fp8, + o_dtype=tsdec._cutlass_dtype("float16"), + qkv_layout="pagedKv", + num_tokens_per_page=self.GEOMETRY.page, + storage_tokens_per_page=self.GEOMETRY.page, + split_kv_mode="gmem_reduction_with_separate_kernel", + splits_kv=splits, + max_splits_kv=splits, + min_loop_iters_per_split=1, + mask_type="causal" if qlen > 1 else "dense", + auto_tuner=True, + split_kv=True, + ) + check_layout(cfg, qlen, splits, heads_per_kv=self.GEOMETRY.heads_per_kv) + with torch.cuda.device(device): + max_active = get_max_active_clusters_for_cluster_size(1) + spec = tsdec._decode_launch_spec_from_config( + cfg, + batch_size=nreq, + num_qo_heads=self.GEOMETRY.heads, + num_kv_heads=self.GEOMETRY.kv_heads, + head_dim=self.GEOMETRY.head_dim, + seq_len_q=qlen, + max_active_clusters=int(max_active), + ) + compile_spec = tsdec._make_decode_compile_spec( + spec, + device_index=device_index(device), + num_qo_heads=self.GEOMETRY.heads, + num_kv_heads=self.GEOMETRY.kv_heads, + head_dim=self.GEOMETRY.head_dim, + page_size=self.GEOMETRY.page, + max_kv_len=max_kv_len, + seq_len_q=qlen, + q_dtype_key="float8_e4m3fn", + k_dtype_key="float8_e4m3fn", + v_dtype_key="float8_e4m3fn", + output_dtype_key="float16", + use_packed_q=False, + kv_prefix_mode="dynamic", + kv_lengths_mode="dynamic", + ) + self._main, self._primts_reducer = tsdec._get_compiled_decode(compile_spec) + layout = tsdec._make_decode_workspace_layout( + spec.scratch_shapes, + torch.float16, + use_separate_reduction_kernel=True, + use_split_kv=True, + ) + self._workspace_buffer = torch.empty(layout.total_bytes, device=device, dtype=torch.int8) + self.workspace = tsdec._bind_decode_workspace(self._workspace_buffer, layout) + self.workspace.split_kv_counter.zero_() + self.workspace.cu_seqlens_q.zero_() + self.workspace.attention_sinks.zero_() + rows = self.GEOMETRY.heads_per_kv * qlen + if tuple(self.workspace.partial_o.shape) != ( + nreq, + self.GEOMETRY.kv_heads, + splits, + rows, + self.GEOMETRY.head_dim, + ): + raise RuntimeError( + f"PrimTS split scratch {tuple(self.workspace.partial_o.shape)} is not the " + f"[requests, KV heads, splits, {rows} rows, {self.GEOMETRY.head_dim}] layout the BF16 merge reads" + ) + self.own_reducer = own_reducer + self.q_tokens_per_cta = max(int(cfg.q_tokens_per_cta), 1) + self.tile_kv = int(cfg.tile_size_kv) + self.insts_kv = int(cfg.num_insts_kv) + self.causal = qlen > 1 + shape = ( + (nreq, qlen, self.GEOMETRY.heads, self.GEOMETRY.head_dim) + if qlen > 1 + else (nreq, self.GEOMETRY.heads, self.GEOMETRY.head_dim) + ) + self.staging = torch.empty(shape, dtype=torch.float16, device=device) + + +def check_layout(cfg, qlen: int, splits: int, *, heads_per_kv: int) -> None: + """The scratch layout and split ranges the BF16 merge (``fullCtxReduce.cu`` ``activeSplits``) assumes. + + Rows are [token, local head]; a row's split ranges are those of its grouped Q CTA, whose first token is + ``(token // q_tokens_per_cta) * q_tokens_per_cta``; tiles are grouped ``num_insts_kv`` per loop step; no sliding + window, attention sinks or packed Q. + """ + problems = [] + if not (cfg.use_split_kv and cfg.use_separate_reduction_kernel): + problems.append("not split-KV with the separate GMEM reducer") + if int(cfg.splits_kv) != splits or int(cfg.max_splits_kv) != splits: + problems.append(f"splits {cfg.splits_kv}/{cfg.max_splits_kv}, recipe {splits}") + if cfg.use_cluster_smem_reduction or cfg.use_persistent_scheduler: + problems.append("cluster-SMEM reduction or a persistent scheduler") + if cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal or cfg.use_attention_sinks: + problems.append("packed Q, a sliding window or attention sinks") + if int(cfg.heads_q_per_kv) not in (0, heads_per_kv): + problems.append(f"{cfg.heads_q_per_kv} query heads per KV head") + if qlen > 1: + if not (cfg.groups_tokens_heads_q and cfg.uses_nontrivial_grouped_q_layout): + problems.append("multi-token Q without the grouped [token, head] tile layout") + elif int(cfg.q_tokens_per_cta) * heads_per_kv != int(cfg.tile_size_q): + problems.append( + f"{cfg.q_tokens_per_cta} tokens per CTA in a {cfg.tile_size_q}-row tile" + ) + if problems: + raise RuntimeError("PrimTS recipe outside the BF16 merge's layout: " + "; ".join(problems)) + + +def kv_page_views(pool: torch.Tensor, subpages_per_slot: int) -> Tuple[torch.Tensor, torch.Tensor]: + """K and V page views of a flat sub-page pool (K of slot s at sub-page s * f, V one later), indexed by slot.""" + return pool[0::subpages_per_slot], pool[1::subpages_per_slot] + + +def device_index(device: torch.device) -> int: + return device.index if device.index is not None else torch.cuda.current_device() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/long_select_host.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/long_select_host.py new file mode 100644 index 000000000000..63dbb7414b95 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/long_select_host.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared long-select compilation and launch; targets own policy and kernel caches.""" + +from __future__ import annotations + +from typing import Any, Callable, Iterable, Protocol, Tuple + +import torch + + +class LongSelectTarget(Protocol): + kernel_cache: dict[Tuple[bool, bool, int, bool], Any] + cute: Any + cutlass: Any + make_fake_tensor: Callable[..., Any] + make_fake_compact_tensor: Callable[..., Any] + make_fake_stream: Callable[..., Any] + TOPK: int + _SelectLongDslKernel: Callable[..., Any] + kernel_key: Callable[..., Tuple[bool, bool, int, bool]] + _compile: Callable[..., Any] + + +def _compile(state: LongSelectTarget, /, key: Tuple[bool, bool, int, bool]) -> Any: + if key not in state.kernel_cache: + heads, blocks, rows = (state.cute.sym_int(), state.cute.sym_int(), state.cute.sym_int()) + scores = state.make_fake_tensor( + state.cutlass.Float32, + (heads, blocks, rows), + stride=(state.cute.sym_int64(), state.cute.sym_int64(), state.cute.sym_int64()), + ) + n_valid = state.make_fake_compact_tensor( + state.cutlass.Int32, (state.cute.sym_int(),), stride_order=(0,) + ) + out = state.make_fake_tensor( + state.cutlass.Int32, + (rows, heads, state.TOPK), + stride=(state.cute.sym_int64(), state.cute.sym_int64(), 1), + ) + state.kernel_cache[key] = state.cute.compile( + state._SelectLongDslKernel(*key), + scores, + n_valid, + out, + state.cutlass.Int32(0), + state.cutlass.Int32(0), + state.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + return state.kernel_cache[key] + + +def select_long_dsl( + state: LongSelectTarget, + /, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool, + early_trigger: bool, + trigger_phase: int, +) -> None: + """Launch the select on the current stream (compiling its key first if ``prebuild`` did not).""" + state._compile(state.kernel_key(use_pdl, early_trigger, trigger_phase))( + scores, n_valid_blocks, out, init_blocks, local_blocks + ) + + +def prebuild( + state: LongSelectTarget, /, keys: Iterable[Tuple[bool, bool, int]], device: torch.device | str +) -> list[Tuple[bool, bool, int, bool]]: + """Compile every key and launch each once on a small dummy row (loading its module). Returns the keys.""" + scores = torch.zeros(1, 2048, 4, dtype=torch.float32, device=device) + n_valid = torch.full((4,), 2048, dtype=torch.int32, device=device) + out = torch.empty(1, 4, state.TOPK, dtype=torch.int32, device=device).permute(1, 0, 2) + built = [] + for key in keys: + key = state.kernel_key(*key) + state._compile(key)(scores, n_valid, out, 0, 1) + built.append(key) + return built diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_dense_decode.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_dense_decode.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_dense_decode.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_fp8_indexer_qk_norm_rope.py similarity index 85% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_fp8_indexer_qk_norm_rope.py index 2fecad0e9afa..108ff80e095c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_fp8_indexer_qk_norm_rope.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_fp8_indexer_qk_norm_rope.py @@ -13,6 +13,24 @@ fused_qk_norm_rope, ) +CELLS_BY_TP = { + 2: ( + Cell(why="one decode token, two index heads per rank", spec=dict(m=1, heads=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, heads=2)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=2)), + ), + 4: ( + Cell(why="one decode token, one index head per rank", spec=dict(m=1, heads=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, heads=1)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=1)), + ), + 8: ( + Cell(why="one decode token, two index heads per rank", spec=dict(m=1, heads=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, heads=2)), + Cell(why="a prefill chunk", spec=dict(m=2048, heads=2)), + ), +} + class M3Fp8IndexerQkNormRope(OpWrapper): """``qk`` [T, (heads + 1) * 128] BF16 (the index projection's output): returns index-Q [T, heads, 128] E4M3 and @@ -25,11 +43,8 @@ class M3Fp8IndexerQkNormRope(OpWrapper): ARCHS = frozenset({Arch.SM_107}) - CELLS: tuple[Cell, ...] = ( - Cell(why="one decode token, two index heads per rank", spec=dict(m=1, heads=2)), - Cell(why="an 8-row decode verify", spec=dict(m=8, heads=2)), - Cell(why="a prefill chunk", spec=dict(m=2048, heads=2)), - ) + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ `reference` is the in-place fused norm + RoPE in BF16 followed by the E4M3 cast (the kernel rounds through BF16 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_index_score.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_index_score.py new file mode 100644 index 000000000000..f778e7701da4 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_index_score.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/nvidia/ops/index_decode_score.py +"""Shared index-score operation; targets supply scheduling/cache policy explicitly.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper, true_fp32_matmul + +from ...kernels.attention.index_score_runtime import IndexScoreRunner + +PAGE = 128 +HEAD_DIM = 128 + + +def compiled_scorer(key: tuple, runner: type[IndexScoreRunner]): + if key not in runner.kernel_cache: + runner._compile(*key) + return runner.kernel_cache[key] + + +class M3IndexScoreBase(OpWrapper): + def __call__( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + *, + runner: type[IndexScoreRunner], + ) -> torch.Tensor: + _, num_heads, head_dim = idx_q.shape + key = (idx_q.dtype, num_heads, max_decode_query_len, head_dim, split_k, early) + compiled_scorer(key, runner)(idx_q, index_k, block_table, score, seq_lens) + return score + + def reference( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + *, + runner: type[IndexScoreRunner], + ) -> torch.Tensor: + """The fp32 [heads, total_q, blocks] scores (``-inf`` where nothing is written). ``score``, ``split_k`` and + ``early`` do not change them.""" + total_q, heads, _ = idx_q.shape + qlen = max_decode_query_len + out = torch.full( + (heads, total_q, score.shape[2]), + -float("inf"), + dtype=torch.float32, + device=idx_q.device, + ) + with true_fp32_matmul(): + for b, seq in enumerate(seq_lens.tolist()): + nblocks = (seq + PAGE - 1) // PAGE + if nblocks == 0: + continue + t0 = b * qlen + q = idx_q[t0 : t0 + qlen].float() + k = index_k[block_table[b, :nblocks].long()].reshape(nblocks * PAGE, -1).float() + s = torch.einsum("qhd,kd->hqk", q, k) + q_pos = seq - qlen + torch.arange(qlen, device=idx_q.device) + k_pos = torch.arange(k.shape[0], device=idx_q.device) + s.masked_fill_(k_pos[None, :] > q_pos[:, None], -float("inf")) + out[:, t0 : t0 + qlen, :nblocks] = s.reshape(heads, qlen, nblocks, PAGE).amax(dim=3) + return out + + def is_valid( + self, + idx_q: torch.Tensor, + index_k: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + score: torch.Tensor, + max_decode_query_len: int, + split_k: int, + early: bool, + *, + runner: type[IndexScoreRunner], + ) -> None: + # The block-table row and the score row are read and written up to ceil(seq_len / 128) with no bound. + nblocks = (seq_lens + PAGE - 1) // PAGE + assert bool((nblocks <= block_table.shape[1]).all()), ( + "a request runs past its block-table row" + ) + assert bool((nblocks <= score.shape[2]).all()), "a request runs past the score buffer" + assert idx_q.shape[0] == seq_lens.shape[0] * max_decode_query_len, ( + "idx_q must hold max_decode_query_len rows per request" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_qk_norm_rope_fp8.py similarity index 84% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_qk_norm_rope_fp8.py index f92a3e6efaa1..09f3444f0cb2 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/attention/m3_qk_norm_rope_fp8.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_qk_norm_rope_fp8.py @@ -12,6 +12,24 @@ fused_qk_norm_rope, ) +CELLS_BY_TP = { + 2: ( + Cell(why="one decode token, tp2 heads (32 q, 2 kv)", spec=dict(m=1, q=32, kv=2)), + Cell(why="an 8-row decode verify", spec=dict(m=8, q=32, kv=2)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=32, kv=2)), + ), + 4: ( + Cell(why="one decode token, tp4 heads (16 q, 1 kv)", spec=dict(m=1, q=16, kv=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=16, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=16, kv=1)), + ), + 8: ( + Cell(why="one decode token, tp8 heads (8 q, 1 kv)", spec=dict(m=1, q=8, kv=1)), + Cell(why="a 48-row decode verify", spec=dict(m=48, q=8, kv=1)), + Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=8, kv=1)), + ), +} + class M3QkNormRopeFp8(OpWrapper): """``fused_qk_norm_rope`` written out of place in FP8 E4M3: Q and K heads normed and roped in fp32 and rounded @@ -24,11 +42,8 @@ class M3QkNormRopeFp8(OpWrapper): ARCHS = frozenset({Arch.SM_107}) - CELLS: tuple[Cell, ...] = ( - Cell(why="one decode token, tp4 heads (16 q, 1 kv)", spec=dict(m=1, q=16, kv=1)), - Cell(why="a 48-row decode verify", spec=dict(m=48, q=16, kv=1)), - Cell(why="a prefill chunk on an FP8 cache", spec=dict(m=2048, q=16, kv=1)), - ) + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ `reference` is the in-place kernel followed by a cast to E4M3. V matches it bit for bit; Q and K round once from diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_qkv_insert.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_qkv_insert.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_qkv_insert.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_select_blocks.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/m3_select_blocks.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/m3_select_blocks.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa.py new file mode 100644 index 000000000000..3f04e9d011b3 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa.py @@ -0,0 +1,198 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared MSA dispatch and grouped-head reference; target cache readers remain bound locally.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + SparseBackendForwardArgs, +) + + +class MsaAttentionBase(OpWrapper): + def __call__( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + if k is not None: + backend.write_layer_caches( + k, v, None, metadata, kv_scale_orig_quant=kv_scale_orig_quant + ) + if sparse: + blocks = backend.run_indexer(idx_q, None, metadata, idx_k_prewritten=True) + args = AttentionForwardArgs( + output=output, sparse_backend_args=SparseBackendForwardArgs(topk_indices=blocks) + ) + else: + args = AttentionForwardArgs(output=output) + args.kv_scale_orig_quant = kv_scale_orig_quant + args.kv_scale_quant_orig = kv_scale_quant_orig + backend.forward(q, None, None, metadata, forward_args=args) + return output + + +class GroupedMsaAttentionBase(MsaAttentionBase): + def reference( + self, + backend, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + idx_q: Optional[torch.Tensor], + metadata, + output: torch.Tensor, + sparse: bool, + kv_scale_orig_quant: Optional[torch.Tensor], + kv_scale_quant_orig: Optional[torch.Tensor], + ) -> torch.Tensor: + """``output`` as the call is supposed to write it, [T, heads * 128] in ``output``'s dtype. + + Every attended row -- the step's own included -- is read from the layer's paged cache, at the slots the + metadata's KV cache manager holds for each request, so this is evaluated after the call (whose write of + ``k`` / ``v`` it then reads back). ``k``, ``v`` and ``kv_scale_orig_quant`` are therefore not read; the + test checks the write itself. Query token ``j`` of request ``r`` sits at position ``cached_r + j`` and + attends positions up to it; on a sparse layer only those in its top-k blocks. Mirrors grouped KV heads (query + head h reads KV head h // (heads / kv heads); two per rank at tp2), each KV head selecting its own top-k + blocks from its own index heads against the rank's single index-K head. + """ + manager = metadata.kv_cache_manager + layer = backend.layer_idx + heads, dim = (backend.num_heads, backend.head_dim) + kv_heads = int(backend.num_kv_heads) + assert heads % kv_heads == 0 + group = heads // kv_heads + page = int(manager.tokens_per_block) + sm_scale = dim ** (-0.5) / float(backend.q_scaling) + num_seqs = int(metadata.num_seqs) + qlens = [int(n) for n in metadata.seq_lens[:num_seqs].tolist()] + cached = [int(c) for c in metadata.kv_cache_params.num_cached_tokens_per_seq[:num_seqs]] + total = sum(qlens) + assert q.shape[0] == total and output.shape[0] == total + pool = manager.get_buffers(layer, "HND") + nvfp4 = manager.is_nvfp4_layer(layer) + if nvfp4: + scale_pool = manager.get_block_scale_buffers(layer, "HND") + k_global = float(kv_scale_quant_orig.reshape(-1)[1]) + v_global = float(kv_scale_quant_orig.reshape(-1)[2]) + q_rows = q.reshape(total, heads, dim) + if q_rows.dtype != torch.float8_e4m3fn and pool.dtype == torch.float8_e4m3fn: + q_rows = q_rows.to(torch.float8_e4m3fn) + q64 = q_rows.to(torch.float64) + if sparse: + config = backend.m3_config + idx_cache = metadata.msa_idx_k_cache(layer) + idx64 = idx_q.reshape(total, config.num_index_heads, config.sparse_index_dim).to( + torch.float64 + ) + assert config.num_index_heads % kv_heads == 0 + idx_group = config.num_index_heads // kv_heads + out = torch.zeros(total, heads, dim, dtype=torch.float64, device=q.device) + first = 0 + for row, request_id in enumerate(list(metadata.request_ids)[:num_seqs]): + n, length = (qlens[row], cached[row] + qlens[row]) + blocks = (length + page - 1) // page + pages = torch.tensor( + manager.get_batch_cache_indices([request_id], layer_idx=layer)[0][:blocks], + dtype=torch.int64, + device=q.device, + ) + qpos = cached[row] + torch.arange(n, device=q.device) + causal = torch.arange(length, device=q.device)[None, :] <= qpos[:, None] + for g in range(kv_heads): + if nvfp4: + keys = self._dequant_pages( + pool[:, 0].view(torch.uint8), + scale_pool[:, 0], + pages, + quad_scales=False, + kv_head=g, + ) + values = self._dequant_pages( + pool[:, 1].view(torch.uint8), + scale_pool[:, 1], + pages, + quad_scales=True, + kv_head=g, + ) + keys, values = (keys * k_global, values * v_global) + else: + keys = pool[pages, 0, g].to(torch.float64) + values = pool[pages, 1, g].to(torch.float64) + keys = keys.reshape(-1, dim)[:length] + values = values.reshape(-1, dim)[:length] + allowed = causal + if sparse: + idx_keys = idx_cache[pages, 0].to(torch.float64).reshape(-1, dim)[:length] + chosen = self._select_blocks( + idx64[first : first + n, g * idx_group : (g + 1) * idx_group], + idx_keys, + causal, + qpos, + page, + config, + ) + allowed = causal & chosen.repeat_interleave(page, dim=1)[:, :length] + h0, h1 = (g * group, (g + 1) * group) + for lo in range(0, n, 256): + hi = min(n, lo + 256) + logits = ( + torch.einsum("thd,kd->thk", q64[first + lo : first + hi, h0:h1], keys) + * sm_scale + ) + logits = logits.masked_fill(~allowed[lo:hi, None, :], float("-inf")) + out[first + lo : first + hi, h0:h1] = torch.softmax(logits, dim=-1) @ values + first += n + return out.reshape(total, heads * dim).to(output.dtype) + + +def _selected_blocks( + idx_q: torch.Tensor, + idx_keys: torch.Tensor, + allowed: torch.Tensor, + qpos: torch.Tensor, + page: int, + config, +) -> torch.Tensor: + """The blocks each query token of one request attends, [n, blocks] bool. + + A block's score is the maximum, over the KV head's index heads and the block's causally visible keys, of + ``idx_q . idx_k``. Of the token's ``ceil((pos + 1) / page)`` valid blocks the first ``init_blocks`` and the last + ``local_blocks`` are forced in, and the ``topk`` highest-scoring are kept, ties to the lower block (the + backend's ``select_blocks_from_maxscore`` rule). The positive score scale the scorers apply does not change the + ranking, so none is applied. + """ + n, length = allowed.shape + blocks = (length + page - 1) // page + scores = torch.einsum("thd,kd->thk", idx_q, idx_keys).amax(dim=1) + scores = scores.masked_fill(~allowed, float("-inf")) + padded = torch.full((n, blocks * page), float("-inf"), dtype=scores.dtype, device=scores.device) + padded[:, :length] = scores + block_scores = padded.view(n, blocks, page).amax(dim=-1) + ids = torch.arange(blocks, device=scores.device)[None, :] + n_valid = torch.div(qpos + page, page, rounding_mode="floor")[:, None] + block_scores = block_scores.masked_fill(ids >= n_valid, float("-inf")) + forced = (ids < torch.clamp(n_valid, max=config.init_blocks)) | ( + ids >= n_valid - config.local_blocks + ) & (ids < n_valid) + block_scores = block_scores.masked_fill(forced, float("inf")) + order = torch.sort(block_scores, dim=1, descending=True, stable=True).indices[:, : config.topk] + keep = torch.arange(order.shape[1], device=scores.device)[None, :] < n_valid + chosen = torch.zeros(n, blocks, dtype=torch.bool, device=scores.device) + rows = torch.arange(n, device=scores.device)[:, None].expand_as(order) + chosen[rows[keep], order[keep]] = True + return chosen diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa_kv_update.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/attention/msa_kv_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa_kv_update.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa_kv_update_hook.py similarity index 95% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa_kv_update_hook.py index 8c2a67356d3f..befa47f69acc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/msa_kv_update.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/msa_kv_update_hook.py @@ -20,7 +20,9 @@ import torch -from ..ops.attention.msa_kv_update import msa_kv_update as _kv_update_op +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update import ( + msa_kv_update as _kv_update_op, +) def msa_kv_update( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/producer_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/producer_reference.py new file mode 100644 index 000000000000..f8268ecf3022 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/producer_reference.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared unchanged producer references and test-only input validation.""" + +from __future__ import annotations + +from typing import Any, Callable, Protocol + +import torch + + +class ProducerTarget(Protocol): + HEAD_DIM: int + # Multi-KV geometry fields are read only by the TP2 reference variant. + HEADS_PACKED: int + HEADS_Q: int + HEADS_KV: int + HEADS_INDEX: int + FP4_BLOCK: int + E2M1_VALUES: tuple[float, ...] + _norm_rope: Callable[..., torch.Tensor] + PAGE: int + + +def _norm_rope( + state: ProducerTarget, + /, + x: torch.Tensor, + weight: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + eps: float, +): + """Gemma RMSNorm (``x * rsqrt(mean(x^2) + eps) * (1 + w)``) then NeoX RoPE on the first 64 dims, in fp32. + x [T, 128] fp32; cos / sin [T, 32].""" + y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) * (1.0 + weight.float()) + lo, hi, rest = (y[:, :32], y[:, 32:64], y[:, 64:]) + return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin, rest], dim=-1) + + +def decode_nvfp4_token( + state: ProducerTarget, /, data: torch.Tensor, scales: torch.Tensor, quant_scale: float +) -> torch.Tensor: + """fp32 [128] of one token's packed E2M1 row (64 bytes, low nibble first) and its 8 E4M3 block scales.""" + table = torch.tensor(state.E2M1_VALUES, dtype=torch.float32, device=data.device) + nib = torch.stack([data & 15, data >> 4], dim=-1).reshape(-1).long() + vals = table[nib & 7] * torch.where(nib >= 8, -1.0, 1.0) + s = scales.view(torch.float8_e4m3fn).float() / quant_scale + return vals * s.repeat_interleave(state.FP4_BLOCK) + + +def scale_index(state: ProducerTarget, /, within_page: int, plane: int) -> torch.Tensor: + """Positions of a token's 8 block scales in its page's [128 * 8] scale plane: row-major for K, the 4-token + swizzle for V.""" + cols = torch.arange(8) + if plane == 0: + return within_page * 8 + cols + return within_page // 4 * 32 + cols * 4 + within_page % 4 + + +def producer_reference_single_kv( + self: Any, + state: ProducerTarget, + /, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, +) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, heads_q, 128], "index_q" [T, 128], and per token "k", "v", + "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl`` and ``early`` do not change + them.""" + t = packed.shape[0] + x = packed.float().view(t, heads_q + 4, state.HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = (rope[:, 0], rope[:, 1]) + q = torch.stack( + [state._norm_rope(x[:, h], q_weight, cos, sin, eps) for h in range(heads_q)], dim=1 + ) + return { + "q": q, + "k": state._norm_rope(x[:, heads_q], k_weight, cos, sin, eps), + "v": x[:, heads_q + 1], + "index_q": state._norm_rope(x[:, heads_q + 2], index_q_weight, cos, sin, eps), + "index_k": state._norm_rope(x[:, heads_q + 3], index_k_weight, cos, sin, eps), + } + + +def producer_reference_multi_kv( + self: Any, + state: ProducerTarget, + /, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, +) -> dict[str, torch.Tensor]: + """fp32 values before their quantization: "q" [T, 32, 128], "index_q" [T, 2, 128], per token and KV head "k" + and "v" [T, 2, 128], and "index_k" [T, 128] (the values written to the token's slot). The caches, ``pdl``, + ``early`` and ``prewait_step_inputs`` do not change them.""" + t = packed.shape[0] + x = packed.float().view(t, state.HEADS_PACKED, state.HEAD_DIM) + rope = rotary_cos_sin[position_ids[:t].long()] + cos, sin = (rope[:, 0], rope[:, 1]) + + def heads(first, count, weight): + return torch.stack( + [state._norm_rope(x[:, h], weight, cos, sin, eps) for h in range(first, first + count)], + dim=1, + ) + + k0 = state.HEADS_Q + v0 = k0 + state.HEADS_KV + iq0 = v0 + state.HEADS_KV + return { + "q": heads(0, state.HEADS_Q, q_weight), + "k": heads(k0, state.HEADS_KV, k_weight), + "v": x[:, v0 : v0 + state.HEADS_KV], + "index_q": heads(iq0, state.HEADS_INDEX, index_q_weight), + "index_k": state._norm_rope(x[:, iq0 + state.HEADS_INDEX], index_k_weight, cos, sin, eps), + } + + +def validate_producer( + self: Any, + state: ProducerTarget, + /, + packed: torch.Tensor, + kv_data: torch.Tensor, + kv_scale: torch.Tensor, + index_k: torch.Tensor, + out_cache_loc: torch.Tensor, + kv_quant_scale: torch.Tensor, + eps: float, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + index_q_weight: torch.Tensor, + index_k_weight: torch.Tensor, + rotary_cos_sin: torch.Tensor, + position_ids: torch.Tensor, + heads_q: int, + pdl: bool = True, + early: bool = True, + prewait_step_inputs: bool = False, +) -> None: + t = packed.shape[0] + pos = position_ids[:t] + assert bool(((pos >= 0) & (pos < rotary_cos_sin.shape[0])).all()), ( + "position_ids must index rotary_cos_sin" + ) + slots = out_cache_loc[:t] + live = slots[(slots >= 0) & (slots < kv_data.shape[0] * state.PAGE)] + assert live.unique().numel() == live.numel(), "out_cache_loc must not repeat a slot" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/select_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/select_reference.py new file mode 100644 index 000000000000..8d89d281eaa0 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/select_reference.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unchanged integer-ranking reference, using each target's forcing constants.""" + +from __future__ import annotations + +from typing import Any, Protocol + +import torch + + +class SelectReferenceTarget(Protocol): + TOPK: int + INIT_SCORE: float + LOCAL_SCORE: float + + +def select_reference( + self: Any, + state: SelectReferenceTarget, + /, + scores: torch.Tensor, + n_valid_blocks: torch.Tensor, + out: torch.Tensor, + init_blocks: int, + local_blocks: int, + use_pdl: bool = False, + early_trigger: bool = False, + trigger_phase: int = 0, +) -> torch.Tensor: + """The [total_q, num_kv_heads, 16] int32 table the op is supposed to write. + + Exact integer ranking, no floating point beyond the forcing: each + score's bit pattern maps to an unsigned key that is smaller for a larger + score (NaN above +inf), packed with the block id so one sort ranks by + (score descending, block ascending). ``out``, ``use_pdl``, ``early_trigger`` and + ``trigger_phase`` only say where and how the op writes; they do not enter the result. + """ + num_kv_heads, num_blocks, total_q = scores.shape + device = scores.device + blocks = torch.arange(num_blocks, device=device, dtype=torch.int64) + raw = n_valid_blocks[:total_q].to(torch.int64) + valid = raw.clamp(0, num_blocks) + local_start = (raw - local_blocks).clamp_min(0) + result = torch.full((total_q, num_kv_heads, state.TOPK), -1, dtype=torch.int32, device=device) + for head in range(num_kv_heads): + forced = scores[head].float().t().clone() + forced = torch.where(blocks[None, :] < init_blocks, state.INIT_SCORE, forced) + forced = torch.where(blocks[None, :] >= local_start[:, None], state.LOCAL_SCORE, forced) + bits = forced.view(torch.int32).to(torch.int64) & 4294967295 + key = torch.where(bits >= 2147483648, bits, ~bits & 2147483647) + in_row = blocks[None, :] < valid[:, None] + packed = torch.where(in_row, key << 16 | blocks[None, :], torch.iinfo(torch.int64).max) + k = min(state.TOPK, num_blocks) + picked = torch.topk(packed, k, dim=1, largest=False).values + picked_key = picked >> 16 + keep = (picked != torch.iinfo(torch.int64).max) & (picked_key != 4286578688) + ids = torch.where(keep, picked & 65535, torch.iinfo(torch.int64).max) + ids = ids.sort(dim=1).values + ids = torch.where(ids == torch.iinfo(torch.int64).max, -1, ids) + result[:, head, :k] = ids.to(torch.int32) + return result diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/sparse_decode_policy.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/sparse_decode_policy.py new file mode 100644 index 000000000000..8907d928b80a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/sparse_decode_policy.py @@ -0,0 +1,267 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-License-Identifier: Apache-2.0 +# Vendored from vLLM (Apache-2.0): +# https://github.com/vllm-project/vllm/blob/6f91edf96d3f3272945809c04702380053bff4de/vllm/models/minimax_m3/common/ops/sparse_attn.py +"""Shared sparse-decode policies and buffer type; target instances remain isolated.""" + +from typing import Optional + +import torch +import triton + +_TARGET_GRID = 256 + + +class MergeMxfp8: + """MXFP8 copy of the split-K merge output for the o_proj GEMM. + + From ``MERGE_MXFP8_MIN_TOKENS`` to ``MAX_ROWS`` rows the merge kernel writes the E4M3 rows and + swizzled UE8M0 scales next to the BF16 output and publishes them to ``m3_prequant.SLOT``, so the MXFP8 o_proj skips + its own quantization. One buffer of ``MAX_ROWS`` rows (one 128-row scale tile) per width serves every layer: each + o_proj consumes it before the next layer's merge rewrites it, in stream order. + """ + + MAX_ROWS = 128 # capacity: fixed: the merge hand-off's rows (decode steps) + + def __init__(self): + self._buffers = {} + + def buffers(self, width: int, device: torch.device): + key = (width, device) + if key not in self._buffers: + self._buffers[key] = ( + torch.empty((self.MAX_ROWS, width), dtype=torch.float8_e4m3fn, device=device), + torch.zeros((self.MAX_ROWS * (width // 32),), dtype=torch.uint8, device=device), + ) + return self._buffers[key] + + @staticmethod + def handoff(mxfp8, rows: int): + """The hand-off of a ``rows``-row step: its E4M3 rows and the 128-row scale tiles that cover them. The o_proj + GEMM checks the swizzled scales against the step's padded rows, so the step gets the tiles that cover its rows, + never the whole buffer.""" + q8, sf = mxfp8 + width = q8.shape[-1] + return q8[:rows], sf[: (rows + 127) // 128 * 128 * (width // 32)] + + +def resolve_num_topk_chunks(total_q: int, num_kv_heads: int, max_topk: int) -> int: + """Split-K factor over the top-k blocks, as a power of two. + + Depends only on shapes that are fixed for a captured batch size, so the + launch geometry is frozen inside a CUDA graph. + """ + target = max(1, min(max_topk, _TARGET_GRID // max(1, total_q * num_kv_heads))) + return 1 << (target.bit_length() - 1) + + +_NVFP4_TUNED_CAPABILITIES = ((10, 3),) + + +def _sm100f_nvfp4_num_topk_chunks( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> Optional[int]: + """Select GB300-measured split-K factors for M3's SM100-family decode shape.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return None + local_batch = total_q // decode_query_len + return { + 2: 8, + 3: 8, + 4: 8, + 5: 16, + 6: 4, + 7: 4, + 8: 2, + 9: 2, + 10: 8, + 11: 4, + 12: 2, + 13: 2, + 14: 2, + }.get(local_batch) + + +def _sm100f_nvfp4_use_linear_softmax( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> bool: + """Use the GB300-measured lower-register recurrence on validated M3 shapes.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads not in (1, 2, 4) + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return False + local_batch = total_q // decode_query_len + # The TP4/DEP32 single-request graph is 8.9% faster with the original + # log-domain recurrence; every other measured DEP8/16/32 bucket is neutral + # or faster with the linear denominator. + return not (num_kv_heads == 1 and local_batch == 1) + + +def _sm100f_nvfp4_query_group_size( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> int: + """Group Eagle queries using the GB300-measured shared-page policy.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + or total_q % decode_query_len != 0 + ): + return 1 + # The sorted-union kernel is 6.5%, 1.6%, and 24.2% faster at local batches + # 11, 12, and 14 with split-K 8. Other measured buckets are neutral or + # slower and retain the single-query producer. + return 2 if total_q // decode_query_len in (11, 12, 14) else 1 + + +def _sm100f_nvfp4_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select native-NVFP4 launch geometry for the SM100 family. + + The FP16 QK/PV path is register-sensitive. GB300 dense replay shows four + warps/two stages wins through local batch 28, while two warps/one stage + avoids the next residency cliff at local batches 29--32. Leave every + unmeasured architecture, model geometry, and larger batch on Triton's + established defaults. + """ + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability not in _NVFP4_TUNED_CAPABILITIES + or num_kv_heads != 4 + or gqa_group_size != 16 + or max_topk != 16 + or decode_query_len != 4 + ): + return {} + local_batch = total_q // decode_query_len + if local_batch == 1 or local_batch in (2, 3, 5, 6, 7, 10, 11, 12): + return {"num_warps": 4, "num_stages": 1} + if local_batch == 4: + return {"num_warps": 2, "num_stages": 1} + if 8 <= local_batch <= 28: + return {"num_warps": 4, "num_stages": 2} + if 29 <= local_batch <= 32: + return {"num_warps": 2, "num_stages": 1} + return {} + + +def _sm100f_nvfp4_merge_launch_options( + *, + total_q: int, + num_kv_heads: int, + gqa_group_size: int, + max_topk: int, + decode_query_len: int, + capability: Optional[tuple[int, int]] = None, +) -> dict[str, int]: + """Select GB300-measured one-warp merges for M3's SM100-family graphs.""" + if capability is None: + capability = torch.cuda.get_device_capability() + if ( + capability in _NVFP4_TUNED_CAPABILITIES + and num_kv_heads == 4 + and gqa_group_size == 16 + and max_topk == 16 + and decode_query_len == 4 + and total_q % decode_query_len == 0 + and total_q // decode_query_len in (2, 4, 5, 6, 7, 10, 11) + ): + return {"num_warps": 1} + return {} + + +def _check_nvfp4_inputs( + k_paged: torch.Tensor, + v_paged: torch.Tensor, + nvfp4_args: tuple[Optional[torch.Tensor], ...], + *, + head_dim: int, + scale_cols: int, +) -> None: + """Validate the packed-NVFP4 cache against what the kernel assumes. + + The kernel indexes block scales by a flat within-page offset and reshapes a + [., SCALE_COLS] scale tile up to head_dim, so it needs contiguous scale rows + and head_dim exactly BLOCK_SIZE_D. Both hold for MiniMax-M3; check them + rather than let a mismatch read neighboring pages. + """ + k_block_scale, v_block_scale, k_global_scale, v_global_scale = nvfp4_args + if any(arg is None for arg in nvfp4_args): + raise ValueError( + "MiniMax-M3 NVFP4 sparse decode needs all four of k_block_scale, " + "v_block_scale, k_global_scale and v_global_scale." + ) + if head_dim != triton.next_power_of_2(head_dim): + raise ValueError(f"NVFP4 sparse decode requires a power-of-two head_dim; got {head_dim}.") + if v_paged.shape[:3] != k_paged.shape[:3]: + raise ValueError("NVFP4 K and V caches must have matching page, head and token extents.") + for name, data in (("k_paged", k_paged), ("v_paged", v_paged)): + if int(data.shape[3]) * 2 != head_dim: + raise ValueError( + f"NVFP4 {name} must pack two elements per byte: expected " + f"head_dim/2 ({head_dim // 2}) columns, got {int(data.shape[3])}." + ) + for name, scale in (("k_block_scale", k_block_scale), ("v_block_scale", v_block_scale)): + if tuple(scale.shape) != tuple(k_paged.shape[:3]) + (scale_cols,): + raise ValueError( + f"NVFP4 {name} must be [num_pages, num_kv_heads, page_size, " + f"{scale_cols}] to match the cache; got {tuple(scale.shape)}." + ) + if scale.stride(3) != 1 or scale.stride(2) != scale_cols: + raise ValueError( + f"NVFP4 {name} needs contiguous per-page scale bytes; got strides " + f"{tuple(scale.stride())}." + ) + for name, gs in (("k_global_scale", k_global_scale), ("v_global_scale", v_global_scale)): + if gs.dtype != torch.float32 or gs.numel() != 1: + raise ValueError( + f"NVFP4 {name} must be a single-element FP32 tensor; got " + f"{gs.numel()} element(s) of {gs.dtype}." + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/sparse_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/sparse_reference.py new file mode 100644 index 000000000000..f1e997ea776b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/sparse_reference.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared Torch mirrors of NVFP4 sparse attention and its MXFP8 handoff.""" + +from __future__ import annotations + +import torch + +HEAD_DIM = 128 +PAGE = 128 +TOPK = 16 +SF_VEC = 16 +MX_BLOCK = 32 + +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + +def dequant_nvfp4_pages( + packed: torch.Tensor, + block_scale: torch.Tensor, + pages: torch.Tensor, + quad_scales: bool, + kv_head: int = 0, +) -> torch.Tensor: + """Dequantize KV head ``kv_head`` of the given pages of a packed NVFP4 cache to float64 [n, PAGE, HEAD_DIM]. + + ``packed`` is the [pages, kv_heads, PAGE, HEAD_DIM // 2] byte view (low + nibble = even element), ``block_scale`` the [pages, kv_heads, PAGE, + HEAD_DIM // 16] E4M3 bytes ``fused_write_layer_caches_nvfp4`` writes: token + ``t``'s column ``c`` at byte ``t * 8 + c`` of the page's region for K, and at + ``(t // 4) * 32 + c * 4 + t % 4`` for V (``quad_scales``). The per-tensor + global scale is not applied. Every product is exact in float64. + """ + table = torch.tensor(_E2M1, dtype=torch.float64, device=packed.device) + rows = packed[pages, kv_head].to(torch.int64) # [n, PAGE, HEAD_DIM // 2] + nibbles = torch.stack([rows & 0x0F, rows >> 4], dim=-1).reshape(rows.shape[0], PAGE, HEAD_DIM) + cols = HEAD_DIM // SF_VEC + region = block_scale[pages, kv_head].reshape(rows.shape[0], PAGE * cols) + t = torch.arange(PAGE, device=packed.device)[:, None] + c = torch.arange(cols, device=packed.device)[None, :] + index = (t // 4) * 32 + c * 4 + t % 4 if quad_scales else t * cols + c + scales = region[:, index.reshape(-1)].reshape(rows.shape[0], PAGE, cols) + scales = scales.view(torch.float8_e4m3fn).to(torch.float64) + return table[nibbles] * scales.repeat_interleave(SF_VEC, dim=-1) + + +def sparse_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, + k_global_scale: torch.Tensor, + v_global_scale: torch.Tensor, + topk: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + sm_scale: float, + kv_head: int = 0, +) -> torch.Tensor: + """Block-sparse decode attention over KV head ``kv_head`` of a packed NVFP4 cache, float64 [T, heads, HEAD_DIM]. + + Query token ``tok`` belongs to request ``tok // decode_query_len`` at offset + ``tok % decode_query_len`` and attends the first ``min(16, ceil(kv_len / + PAGE))`` entries of ``topk[tok]`` that are not -1, each block truncated at + ``kv_len = max(seq_len - decode_query_len + offset + 1, 0)``. A token that + attends nothing (a CUDA-graph padding row) is all zeros. ``q`` is E4M3 (or + its byte view); softmax and both products are float64. + """ + q64 = ( + q.view(torch.float8_e4m3fn).to(torch.float64) + if q.element_size() == 1 + else q.to(torch.float64) + ) + total_q, num_heads, _ = q64.shape + kgs = float(k_global_scale.reshape(-1)[0]) + vgs = float(v_global_scale.reshape(-1)[0]) + out = torch.zeros(total_q, num_heads, HEAD_DIM, dtype=torch.float64, device=q.device) + seq = seq_lens.tolist() + positions = torch.arange(PAGE, device=q.device) + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + real_topk = min(TOPK, (kv_len + PAGE - 1) // PAGE) + blocks = topk[tok, :real_topk] + blocks = blocks[blocks >= 0].to(torch.int64) + if blocks.numel() == 0: + continue + pages = block_table[req, blocks].to(torch.int64) + keys = ( + dequant_nvfp4_pages( + k, k_block_scale, pages, quad_scales=False, kv_head=kv_head + ).reshape(-1, HEAD_DIM) + * kgs + ) + values = ( + dequant_nvfp4_pages(v, v_block_scale, pages, quad_scales=True, kv_head=kv_head).reshape( + -1, HEAD_DIM + ) + * vgs + ) + valid = ((blocks[:, None] * PAGE + positions[None, :]) < kv_len).reshape(-1) + logits = (q64[tok] @ keys.T) * sm_scale + logits = logits.masked_fill(~valid[None, :], float("-inf")) + values = values.masked_fill(~valid[:, None], 0) + out[tok] = torch.softmax(logits, dim=-1) @ values + return out + + +def mxfp8_handoff_reference( + out: torch.Tensor, num_rows_alloc: int, *, width: int +) -> tuple[torch.Tensor, torch.Tensor]: + """The MXFP8 o_proj hand-off the merge writes for BF16 rows ``out`` [T, width]: ``(q8 [T, width] uint8, + sf [num_rows_alloc * width // 32] uint8)``. + + Per 32-element block: ``e`` is the biased exponent of ``amax * rcp(448)`` + rounded up when its mantissa is non-zero (capped at 254), the elements are + scaled by ``2^(127 - e)`` and rounded to E4M3 (saturating), and ``e`` lands + in the 128x4-swizzled scale layout the MXFP8 GEMM reads. Entries of ``sf`` + for rows past ``T`` are zero. Computed from the op's own BF16 rows, so a + comparison against it is exact. + """ + total_q = out.shape[0] + x = out.float().reshape(total_q, width // MX_BLOCK, MX_BLOCK) + amax = x.abs().amax(dim=-1) + sfv = amax * torch.tensor(1.0 / 448.0, dtype=torch.float32, device=out.device) + bits = sfv.view(torch.int32) + e = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + e = e.clamp(max=254) + qs = torch.where(amax != 0, ((254 - e) << 23).view(torch.float32), torch.zeros_like(amax)) + xq = (x * qs[..., None]).clamp(-448.0, 448.0) + q8 = xq.to(torch.float8_e4m3fn).view(torch.uint8).reshape(total_q, width) + num_sf_k_tiles = width // 128 + sf = torch.zeros(num_rows_alloc * (width // MX_BLOCK), dtype=torch.uint8, device=out.device) + tok = torch.arange(total_q, device=out.device)[:, None] + k_sf = torch.arange(width // MX_BLOCK, device=out.device)[None, :] + offset = ( + (tok // 128) * num_sf_k_tiles * 512 + + (k_sf // 4) * 512 + + (tok % 32) * 16 + + ((tok % 128) // 32) * 4 + + k_sf % 4 + ) + sf[offset.reshape(-1)] = e.to(torch.uint8).reshape(-1) + return q8, sf + + +def validate_sparse_selection( + q: torch.Tensor, + topk_idx: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + decode_query_len: int, + identity_blocks: bool, + page: int, +) -> None: + """The inputs this call takes and answers wrongly. + + Every producer walks a token's first ``min(topk, ceil(kv_len / 128))`` + table entries and indexes the block table with them unchecked, so they + must be real blocks of the request. ``identity_blocks`` is a promise the + multi-query producer relies on: every token of a request reads the same + blocks, all of them. Reads ``topk_idx`` and ``seq_lens`` back to the + host: a test-only cost. + """ + total_q = q.shape[0] + seq = seq_lens.tolist() + max_topk = topk_idx.shape[-1] + for tok in range(total_q): + req, offset = divmod(tok, decode_query_len) + kv_len = max(seq[req] - decode_query_len + offset + 1, 0) + nblocks = (kv_len + page - 1) // page + ids = topk_idx[:, tok, : min(max_topk, nblocks)] + assert bool(((ids >= 0) & (ids < min(nblocks, block_table.shape[1]))).all()), ( + f"token {tok}: the first min(topk, block count) selected ids must be blocks of the request; " + "the block table is indexed with them unchecked" + ) + if identity_blocks: + assert nblocks <= max_topk and bool( + (ids == torch.arange(nblocks, device=ids.device, dtype=ids.dtype)).all() + ), ( + f"token {tok}: identity_blocks promises the table holds every block of the request in order" + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/token_map.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/token_map.py new file mode 100644 index 000000000000..393e422dab7b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/attention/token_map.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Token-map preprocessing using each target's live state and callback bindings.""" + +from __future__ import annotations + +from typing import Any, Callable, Optional, Protocol, cast + +import torch + + +class _PreprocessOps(Protocol): + get_trtllm_gen_generation_workspace_layout: Callable[..., dict[str, int]] + trtllm_gen_generation_preprocess: Callable[..., object] + + +class _LogSink(Protocol): + def info(self, message: str) -> None: ... + + +class TokenMapTarget(Protocol): + """State and callback fields owned by one SM107 target module.""" + + _original: Optional[Callable[..., object]] + _max_rows: int + _maps: dict[torch.device, torch.Tensor] + _offsets: dict[tuple, int] + _on_taken: Callable[[int], None] + _on_missed: Callable[[int, str], None] + _ARG_NAMES: tuple[str, ...] + _INT2_BYTES: int + _indexed: Callable[[torch.device], torch.device] + _one_request_map: Callable[[torch.device, int], torch.Tensor] + _map_offset: Callable[[dict, int], int] + _guarded: Callable[..., object] + thop: _PreprocessOps + logger: _LogSink + + +_ARG_NAMES = ( + "qkv_input", + "workspace", + "sequence_lengths", + "spec_decoding_generation_lengths", + "spec_decoding_position_offsets", + "kv_cache_block_offsets", + "host_kv_cache_pool_pointers", + "host_kv_cache_pool_mapping", + "kv_scale_orig_quant", + "kv_scale_quant_orig", + "attention_output_orig_quant", + "rotary_inv_freq", + "rotary_cos_sin", + "mrope_position_deltas", + "layer_idx", + "seq_offset", + "num_heads", + "num_kv_heads", + "head_size", + "tokens_per_block", + "kv_cache_quant_mode", + "max_attention_window_size", + "cyclic_attention_window_size", + "num_tokens", + "batch_beam", + "input_seq_length", + "max_past_kv_length", + "rotary_embedding_dim", + "rotary_embedding_base", + "rotary_embedding_scale_type", + "rotary_embedding_scale", + "rotary_embedding_max_positions", + "position_embedding_type", + "bmm1_scale", + "bmm2_scale", + "fp8_context_fmha", + "predicted_tokens_per_seq", + "attention_chunk_size", + "multi_processor_count", + "total_num_blocks", + "kv_factor", + "need_build_kv_cache_metadata", + "cross_attention", + "skip_fmha_workspace", +) + +_INT2_BYTES = 8 + + +def _indexed(device: torch.device) -> torch.device: + """``device`` with its index (an unindexed ``cuda`` means the current device), the key ``_maps`` is built under.""" + if device.type == "cuda" and device.index is None: + return torch.device("cuda", torch.cuda.current_device()) + return device + + +def _one_request_map(state: TokenMapTarget, /, device: torch.device, rows: int) -> torch.Tensor: + """The prebuilt (0, t) map's first ``rows`` rows on ``device``.""" + device = state._indexed(device) + m = state._maps.get(device) + if m is None or m.shape[0] < rows: + reason = f"no {rows}-row token map was prebuilt on {device} (install() builds it before the warmup)" + state._on_missed(rows, reason) + raise RuntimeError(f"m3_tokmap_guard: {reason}") + return m[:rows] + + +def _map_offset(state: TokenMapTarget, /, args: dict, rows: int) -> int: + x = args["qkv_input"] + skip = bool(args.get("skip_fmha_workspace", False)) + key = ( + x.dtype, + rows, + int(args["num_heads"]), + int(args["head_size"]), + int(args["rotary_embedding_dim"]), + int(args["num_kv_heads"]), + skip, + ) + off = state._offsets.get(key) + if off is None: + layout = state.thop.get_trtllm_gen_generation_workspace_layout( + x.dtype, 1, rows, key[2], key[3], key[4], key[5], 0, False, skip + ) + off, size = (int(layout["tokens_info_offset"]), int(layout["tokens_info_size"])) + if off < 0 or off % 4 != 0 or size < state._INT2_BYTES * rows: + raise RuntimeError( + f"m3_tokmap_guard: the op's generation layout has no {rows}-row token map " + f"(offset {off}, {size} bytes): the layout changed under this guard" + ) + state._offsets[key] = off + return off + + +def _guarded(state: TokenMapTarget, /, *args: Any, **kwargs: Any) -> object: + a = dict(zip(state._ARG_NAMES, args)) + a.update(kwargs) + rows = int(a["num_tokens"]) + if ( + a.get("spec_decoding_generation_lengths") is not None + and int(a["predicted_tokens_per_seq"]) > 1 + and (int(a["batch_beam"]) == 1) + and (rows > 1) + and (not bool(a.get("cross_attention", False))) + ): + ws = a["workspace"] + src = state._one_request_map(ws.device, rows) + off = state._map_offset(a, rows) + dst = ws.view(-1).view(torch.uint8)[off : off + state._INT2_BYTES * rows] + dst.view(torch.int32).view(rows, 2).copy_(src) + state._on_taken(rows) + return cast(Callable[..., object], state._original)(*args, **kwargs) + + +def install( + state: TokenMapTarget, + /, + device: torch.device, + max_rows: int, + on_taken: Callable[[int], None], + on_missed: Callable[[int, str], None], +) -> bool: + """Wrap ``thop.trtllm_gen_generation_preprocess`` (once per process) and prebuild the (0, t) map for up to + ``max_rows`` rows on ``device`` (a one-request call past it is ``on_missed`` and raises). Returns whether this + call installed the wrapper.""" + device = state._indexed(device) + if device not in state._maps or state._maps[device].shape[0] < max_rows: + t = torch.arange(max_rows, dtype=torch.int32, device=device) + state._maps[device] = torch.stack([torch.zeros_like(t), t], dim=1).contiguous() + state._max_rows = max(state._max_rows, max_rows) + state._on_taken = on_taken + state._on_missed = on_missed + if state._original is not None: + return False + state._original = state.thop.trtllm_gen_generation_preprocess + state.thop.trtllm_gen_generation_preprocess = state._guarded + state.logger.info( + "[M3 prebuild] m3_tokmap_guard installed: one-request multi-token generations get their (0, t) token map " + f"(up to {max_rows} rows)" + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/m3_allreduce.py similarity index 77% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/m3_allreduce.py index d67b040e8bfd..d0ef0fd93c38 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/comm/m3_allreduce.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/m3_allreduce.py @@ -11,6 +11,30 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch.distributed import AllReduce, AllReduceParams +CELLS_BY_TP = { + 2: ( + Cell(why="tp2, a prefill chunk", spec=dict(world=2, m=2048, hidden=6144)), + Cell(why="tp2, an 8-row generation step", spec=dict(world=2, m=8, hidden=6144)), + Cell( + why="tp2, plain sum of eight drafter rows", + spec=dict(world=2, m=8, hidden=6144, plain=True), + ), + ), + 4: ( + Cell(why="tp4, a prefill chunk", spec=dict(world=4, m=2048, hidden=6144)), + Cell(why="tp4, a 64-row generation step", spec=dict(world=4, m=64, hidden=6144)), + Cell( + why="tp4, plain sum of eight drafter rows", + spec=dict(world=4, m=8, hidden=6144, plain=True), + ), + ), + 8: ( + Cell(why="a prefill chunk", spec=dict(m=2048, hidden=6144)), + Cell(why="a 64-row generation step", spec=dict(m=64, hidden=6144)), + Cell(why="plain sum of eight drafter rows", spec=dict(m=8, hidden=6144, plain=True)), + ), +} + class M3AllReduce(OpWrapper): """``comm(x, all_reduce_params=params)``: the runtime's ``AllReduce`` over the TP group. ``params`` None: the @@ -25,14 +49,8 @@ class M3AllReduce(OpWrapper): ARCHS = frozenset({Arch.SM_107}) - CELLS: tuple[Cell, ...] = ( - Cell(why="tp4, a prefill chunk", spec=dict(world=4, m=2048, hidden=6144)), - Cell(why="tp4, a 64-row generation step", spec=dict(world=4, m=64, hidden=6144)), - Cell( - why="tp4, plain sum of eight drafter rows", - spec=dict(world=4, m=8, hidden=6144, plain=True), - ), - ) + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ The reduction order across ranks is the strategy's; `reference` sums the ranks' rows in fp32 in rank order, then diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/mc_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/mc_workspace.py new file mode 100644 index 000000000000..9f2983035daa --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/mc_workspace.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Multicast workspace storage shared by the sm_107 TP topologies.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +EMPTY_WORD = -2147483648 + + +@dataclass(eq=False) +class McAllReduceWorkspace: + """Two halves of [max_tokens][world][H] BF16 (int32 words) per rank behind one multicast mapping, plus one call + counter per CTA per cluster size. Create collectively and eagerly on every TP rank.""" + + uc: torch.Tensor + mc: torch.Tensor + flags: dict + rank: int + world_size: int + hidden: int + handle: Any + comm: Any + max_tokens: int = 8 + + @classmethod + def create(cls, mapping, hidden: int, max_tokens: int = 8) -> "McAllReduceWorkspace": + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + words = 4 * max_tokens * mapping.tp_size * (hidden // 2) # two exchange sets of two halves + ok = not torch.cuda.is_current_stream_capturing() + if not _mnnvl_workspace_all_succeeded(comm, ok): + comm.Free() + raise RuntimeError( + "McAllReduceWorkspace.create must run outside CUDA-graph capture on every rank" + ) + handle = _make_mnnvl_mcast_buffer(comm, words * 4, mapping, mapping.is_multi_node()) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int32, 0) + mc = handle.get_mc_buffer((words,), torch.int32, 0) + uc.fill_(EMPTY_WORD) + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("McAllReduceWorkspace.create failed on another rank") + return cls( + uc=uc, + mc=mc, + flags={}, + rank=mapping.tp_rank, + world_size=mapping.tp_size, + hidden=hidden, + handle=handle, + comm=comm, + max_tokens=max_tokens, + ) + + def flags_for(self, cluster: int): + """This cluster size's per-CTA call counters, or None when they do not exist yet during a capture. + get_workspace creates the cluster-8 counters with the workspace, so for cluster 8 (the only size built) + that None never happens and cannot differ between ranks.""" + flags = self.flags.get(cluster) + if flags is None: + if torch.cuda.is_current_stream_capturing(): + return None + flags = torch.zeros(self.max_tokens * cluster, dtype=torch.int32, device=self.uc.device) + self.flags[cluster] = flags + return flags diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/reference.py new file mode 100644 index 000000000000..300f79625855 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/comm/reference.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared Torch mirrors of the MiniMax-M3 multicast boundaries.""" + +from __future__ import annotations + +from typing import Optional, Sequence + +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import OpWrapper + +_E4M3_MAX = 448.0 + + +def mxfp8_rows(normed: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """MXFP8 of BF16 rows as ``trtllm::mxfp8_quantize(normed, True)``: E4M3 [M, H] and 128x4-swizzled UE8M0 bytes. + + One scale per 32 columns, the smallest power of two at or above ``amax / 448`` (biased exponent, round up), + data = row / scale rounded to nearest E4M3. The scale buffer is padded to 128 rows, padding bytes left as the + op leaves them (the comparison masks them out). + """ + m, h = normed.shape + blocks = normed.float().view(m, h // 32, 32) + amax = blocks.abs().amax(-1) + bits = (amax / _E4M3_MAX).view(torch.int32) + byte = (((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) > 0).to(torch.int32)).clamp(0, 254) + scale = torch.exp2(byte.float() - 127.0).unsqueeze(-1) + data = torch.where(amax.unsqueeze(-1) != 0, blocks / scale, torch.zeros_like(blocks)) + data = data.view(m, h).to(torch.float8_e4m3fn) + return data, swizzle_sf(byte.to(torch.uint8)) + + +def swizzle_sf(byte: torch.Tensor) -> torch.Tensor: + """Scatter [rows, H / 32] scale bytes into the 128x4-swizzled buffer the op writes (zero padding).""" + rows, cols = byte.shape + padded_rows = (rows + 127) // 128 * 128 + r = torch.arange(rows, device=byte.device).view(-1, 1) + c = torch.arange(cols, device=byte.device).view(1, -1) + index = ( + (c % 4) + (c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (r // 128) * (128 * cols) + ) + out = torch.zeros(padded_rows * cols, dtype=torch.uint8, device=byte.device) + out[index.reshape(-1)] = byte.reshape(-1) + return out + + +def sf_written(rows: int, cols: int, device) -> torch.Tensor: + """Mask of the swizzled scale bytes the op writes for ``rows`` rows (the padding rows are left unwritten).""" + return swizzle_sf(torch.ones(rows, cols, dtype=torch.uint8, device=device)).bool() + + +def boundary( + partials: Sequence[torch.Tensor], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, +) -> tuple[torch.Tensor, ...]: + """The boundary arithmetic once each rank's BF16 partial is known (shared with the finalize entry). + + ``drafter``: the Eagle3 drafter boundary (opts bit 16), whose stock is a plain BF16 all-reduce, then FlashInfer + fused_add_rmsnorm: the RMS statistic and the normalization take the unrounded FP32 sum bf16(rank sum) + residual; + only the stored residual is rounded to BF16. + + The stock rounding points of this target (a BF16 all-reduce, then the fused residual add + RMSNorm): ranks are + summed in FP32 in rank order and rounded to BF16, the residual added and rounded to BF16 again; the RMS statistic + and the normalization are done in float64 and rounded to BF16 once. + """ + total = torch.zeros(partials[0].shape, dtype=torch.float32, device=residual.device) + for p in partials: + total = total + p.float() + pre = total.to(torch.bfloat16).float() + residual.float() + updated = pre.to(torch.bfloat16) + u = pre.double() if drafter else updated.double() + inv = torch.rsqrt(u.pow(2).mean(-1, keepdim=True) + eps) + normed = (u * inv * norm_w.double()).to(torch.bfloat16) + if not quant: + return normed, updated + return (normed, updated, *mxfp8_rows(normed)) + + +def finalize_partial( + fc2: torch.Tensor, perm: torch.Tensor, scales: torch.Tensor, shared: Optional[torch.Tensor] +) -> torch.Tensor: + """One rank's BF16 partial: ``bf16(bf16(sum_k scale_k * fc2[perm_k]) + shared)``, the routed sum in FP32 in k + order and rounded to BF16, then the shared expert added as stock adds it (one more BF16 rounding). + + The kernel's per-k step is a fused multiply-add; it is reproduced exactly by the float64 product (exact for two + FP32 operands) plus the FP32 accumulator, rounded back to FP32. Dropped entries (perm -1) are skipped. + """ + m, topk = perm.shape + acc = torch.zeros(m, fc2.shape[1], dtype=torch.float32, device=fc2.device) + sc = scales.reshape(m, topk).float() + for k in range(topk): + rows = perm[:, k].long() + keep = (rows >= 0).unsqueeze(-1) + gathered = fc2[rows.clamp_min(0)].double() + step = (gathered * sc[:, k : k + 1].double() + acc.double()).float() + acc = torch.where(keep, step, acc) + if shared is not None: + acc = acc.to(torch.bfloat16).float() + shared.float() + return acc.to(torch.bfloat16) + + +class McAllreduceNormReference(OpWrapper): + """Reference arithmetic; targets supply their native call and certification cells.""" + + def reference( + self, + per_rank_x: Sequence[torch.Tensor], + per_rank_x2: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + drafter: bool = False, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute (``drafter``: the opts-16 drafter boundary). + + The documented collective exception: it takes every rank's ``x`` (and ``x2``) in rank order, not the calling + rank's. Each rank's partial is ``bf16(x + x2)`` (``x`` alone without ``x2``), whatever pushed it. + """ + partials = [] + for r, x in enumerate(per_rank_x): + if per_rank_x2 is None: + partials.append(x) + else: + partials.append((x.float() + per_rank_x2[r].float()).to(torch.bfloat16)) + return boundary(partials, residual, norm_w, eps, quant, drafter) + + +class McFinalizeAllreduceNormReference(OpWrapper): + """Reference arithmetic; targets supply their native call and certification cells.""" + + def reference( + self, + per_rank_fc2: Sequence[torch.Tensor], + per_rank_perm: Sequence[torch.Tensor], + per_rank_scales: Sequence[torch.Tensor], + per_rank_shared: Optional[Sequence[torch.Tensor]], + residual: torch.Tensor, + norm_w: torch.Tensor, + eps: float, + quant: bool, + ) -> tuple[torch.Tensor, ...]: + """The boundary the group is supposed to compute, from every rank's MoE outputs in rank order (the + collective exception, as in ``m3_mc_allreduce_norm.reference``).""" + partials = [ + finalize_partial( + fc2, + per_rank_perm[r], + per_rank_scales[r], + None if per_rank_shared is None else per_rank_shared[r], + ) + for r, fc2 in enumerate(per_rank_fc2) + ] + return boundary(partials, residual, norm_w, eps, quant) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/decode_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/decode_embedding.py new file mode 100644 index 000000000000..35fc7dfd331d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/decode_embedding.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Common embedding host implementation, bound to per-target state and callbacks.""" + +from __future__ import annotations + +from contextlib import AbstractContextManager +from typing import Any, Callable, Protocol, Tuple + +import torch + + +class EmbeddingTarget(Protocol): + HIDDEN: int + MAX_ROWS: int + _lock: AbstractContextManager[bool] + _compiled: dict[tuple, Callable[..., None]] + _compile: Callable[..., Callable[..., None]] + _check: Callable[..., None] + arg: Callable[[torch.Tensor], Any] + m3_embed_norm: Callable[..., Tuple[torch.Tensor, torch.Tensor]] + _m3_embed_norm_jit: Any + cuda_driver: Any + cute: Any + from_dlpack: Callable[..., Any] + + +def arg(state: EmbeddingTarget, /, tensor: torch.Tensor) -> Any: + return state.from_dlpack( + tensor.detach(), + assumed_align=tensor.element_size() + if tensor.ndim == 1 and tensor.dtype in (torch.int32, torch.int64) + else 16, + ) + + +def _compile( + state: EmbeddingTarget, + /, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool, +) -> Callable[..., None]: + key = (ids.shape[0], ids.dtype, tuple(table.shape), table.device, bool(use_pdl)) + fn = state._compiled.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_embed_norm compiles on its first call per shape: call it once before capture" + ) + with state._lock: + fn = state._compiled.get(key) + if fn is None: + stream = state.cuda_driver.CUstream( + torch.cuda.current_stream(table.device).cuda_stream + ) + fn = state._compiled[key] = state.cute.compile( + state._m3_embed_norm_jit, + state.arg(ids), + state.arg(table), + state.arg(weight), + state.arg(raw), + state.arg(out), + int(table.shape[0]), + float(eps), + ids.shape[0], + state.HIDDEN, + bool(use_pdl), + stream, + ) + return fn + + +def m3_embed_norm( + state: EmbeddingTarget, + /, + ids: torch.Tensor, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + raw: torch.Tensor, + out: torch.Tensor, + use_pdl: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Write ``table[ids]`` into ``raw`` (a zero row for an id outside ``[0, V)``) and its RMSNorm with ``weight`` / + ``eps`` into ``out``; returns ``(out, raw)``.""" + state._check(ids, table, weight, raw, out) + fn = state._compile(ids, table, weight, eps, raw, out, use_pdl) + stream = state.cuda_driver.CUstream(torch.cuda.current_stream(ids.device).cuda_stream) + fn( + state.arg(ids), + state.arg(table), + state.arg(weight), + state.arg(raw), + state.arg(out), + int(table.shape[0]), + float(eps), + stream, + ) + return (out, raw) + + +def init_decode_embedding( + self: Any, + state: EmbeddingTarget, + /, + table: torch.Tensor, + weight: torch.Tensor, + eps: float, + rows: int, + id_dtype: torch.dtype, + use_pdl: bool = True, +) -> None: + assert not torch.cuda.is_current_stream_capturing() + assert ( + table.dtype == torch.bfloat16 and table.is_contiguous() and (table.shape[1] == state.HIDDEN) + ) + assert weight.dtype == torch.bfloat16 and weight.shape == (state.HIDDEN,) + assert 1 <= rows <= state.MAX_ROWS and id_dtype in (torch.int32, torch.int64) + self.table, self.weight, self.rows, self.id_dtype, self.use_pdl = ( + table, + weight, + rows, + id_dtype, + use_pdl, + ) + self.eps = float(eps) + self.raw = torch.empty(rows, state.HIDDEN, device=table.device, dtype=torch.bfloat16) + self.output = torch.empty_like(self.raw) + ids = torch.empty(rows, device=table.device, dtype=id_dtype) + state._compile(ids, table, weight, self.eps, self.raw, self.output, use_pdl) + + +def call_decode_embedding( + self: Any, state: EmbeddingTarget, /, ids: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + assert ids.shape == (self.rows,) and ids.dtype == self.id_dtype and ids.is_contiguous() + return state.m3_embed_norm( + ids, self.table, self.weight, self.eps, self.raw, self.output, self.use_pdl + ) + + +def embedding_firstnorm_reference( + ids: torch.Tensor, table: torch.Tensor, weight: torch.Tensor, eps: float +) -> Tuple[torch.Tensor, torch.Tensor]: + """The original float64 embedding and first-layer RMSNorm reference math.""" + valid = (ids >= 0) & (ids < table.shape[0]) + rows = table[torch.where(valid, ids, torch.zeros_like(ids)).long()] * valid.unsqueeze(1).to( + table.dtype + ) + x = rows.double() + rstd = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) + return ((x * rstd * weight.double()).to(torch.bfloat16), rows) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/m3_embedding.py similarity index 65% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/m3_embedding.py index 3060c17acdff..5c12084a1cfc 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/embedding/m3_embedding.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/embedding/m3_embedding.py @@ -9,17 +9,20 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._torch.modules.embedding import pre_comm_embedding_ops - -class M3Embedding(OpWrapper): - """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). - - The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode - embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( +CELLS_BY_TP = { + 2: ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="an 8-row decode verify, int64 ids", + spec=dict(rows=8, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ), + 4: ( Cell( why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) ), @@ -30,7 +33,33 @@ class M3Embedding(OpWrapper): Cell( why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) ), - ) + ), + 8: ( + Cell( + why="a prefill chunk, int32 ids", spec=dict(rows=2048, id_dtype=torch.int32, pad=False) + ), + Cell( + why="an 8-row decode verify, int64 ids", + spec=dict(rows=8, id_dtype=torch.int64, pad=False), + ), + Cell( + why="draft rows padded with -1 ids", spec=dict(rows=32, id_dtype=torch.int32, pad=True) + ), + ), +} + + +class M3Embedding(OpWrapper): + """``table[ids]`` with a zero row for an id outside ``[0, V)`` (the speculative sampler pads with -1). + + The table is the full (replicated) bf16 embedding [V, 6144]; ids int32 / int64 of any count. The fused decode + embedding (``m3_embedding_firstnorm``) takes decode steps of up to 8 rows; this entry takes every other step. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ An exact gather: `reference` must match bit for bit. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/locality.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/locality.py new file mode 100644 index 000000000000..de0501d76100 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/locality.py @@ -0,0 +1,431 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Locality domains ("uGPU") of a multi-die GPU (Rubin VR200: 2 domains of 100 SMs + 12 remainder SMs) for the +streaming decode kernels of this target: the SM -> domain map, per-domain weight placement, and the CTA -> SM +assignment a persistent kernel uses to read only its own die's memory. + +A kernel launched with one CTA per SM reads ``%smid`` once and takes its logical block index from a host-built +permutation (:meth:`LocalityMap.logical_map`), so every per-CTA structure (work plan row, epochs, roles) follows its +die. Weights that are split by output rows into per-domain halves (:meth:`LocalityMap.split_rows`) live in one virtual +range whose halves are physically in the two domains. Placement only: the arithmetic is unchanged, so results are +bitwise those of the unplaced weights. + +The SM -> domain map is attested by the driver: the SMs are split by locality domain +(``cuDevSmResourceSplit`` with ``CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID``), and a ``%smid`` probe runs on a green +context of each split. Allocations use driver VMM with ``CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN`` and are checked +with ``CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL``; any failure raises (fail closed).""" + +from __future__ import annotations + +from typing import List, Sequence, Tuple + +import torch + +_PROBE_SRC = r""" +extern "C" __global__ void smid_probe(int* out) { + if (threadIdx.x == 0) { + unsigned s; + asm volatile("mov.u32 %0, %%smid;" : "=r"(s)); + out[blockIdx.x] = (int)s; + } +} +""" + + +def _drv(): + from cuda.bindings import driver + + return driver + + +def _ck(res, what: str): + """cuda-python returns (CUresult, *values): raise on an error, return the values.""" + err = res[0] if isinstance(res, tuple) else res + if int(err) != 0: + raise RuntimeError(f"locality: {what} failed: {err}") + if isinstance(res, tuple): + return res[1] if len(res) == 2 else res[1:] + return None + + +class _Mapping: + """One virtual range backed by per-domain physical allocations, exposed to torch through + ``__cuda_array_interface__``. Lives as long as a tensor that views it (the tensor holds this object).""" + + def __init__(self, device: int, pieces: Sequence[Tuple[int, int]]): + """``pieces``: (bytes, domain) in address order; each a multiple of the allocation granularity.""" + d = _drv() + self.device = device + self.size = sum(n for n, _ in pieces) + self.va = _ck(d.cuMemAddressReserve(self.size, 0, 0, 0), "cuMemAddressReserve") + self.handles = [] + off = 0 + try: + for nbytes, dom in pieces: + prop = _alloc_prop(device, dom) + h = _ck(d.cuMemCreate(nbytes, prop, 0), f"cuMemCreate(domain {dom})") + self.handles.append((h, nbytes)) + _ck(d.cuMemMap(int(self.va) + off, nbytes, 0, h, 0), "cuMemMap") + off += nbytes + acc = d.CUmemAccessDesc() + acc.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + acc.location.id = device + acc.flags = d.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + _ck(d.cuMemSetAccess(self.va, self.size, [acc], 1), "cuMemSetAccess") + except Exception: + self.close() + raise + self.__cuda_array_interface__ = { + "shape": (self.size,), + "typestr": "|u1", + "data": (int(self.va), False), + "version": 3, + "strides": None, + "stream": None, + } + + def close(self): + d = _drv() + off = 0 + for h, n in self.handles: + d.cuMemUnmap(int(self.va) + off, n) + d.cuMemRelease(h) + off += n + self.handles = [] + if self.va is not None: + d.cuMemAddressFree(self.va, self.size) + self.va = None + + def __del__(self): + try: + self.close() + except Exception: + pass + + +_LIBCUDA = None + + +def _libcuda(): + global _LIBCUDA + if _LIBCUDA is None: + import ctypes + + _LIBCUDA = ctypes.CDLL("libcuda.so.1") + return _LIBCUDA + + +def _alloc_prop(device: int, domain: int): + d = _drv() + prop = d.CUmemAllocationProp() + prop.type = d.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.requestedHandleTypes = d.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE + prop.location.type = d.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN + prop.location.localized.deviceId = device + prop.location.localized.localityDomainId = domain + return prop + + +class LocalityMap: + """The locality domains of one device. Build with :meth:`get` before CUDA-graph capture.""" + + def __init__(self, device: int): + d = _drv() + _ck(d.cuInit(0), "cuInit") + dev = _ck(d.cuDeviceGet(device), "cuDeviceGet") + attr = d.CUdevice_attribute + self.device = device + self.n_domains = int( + _ck( + d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_COUNT, dev), + "domain count", + ) + ) + self.sms_per_domain = int( + _ck( + d.cuDeviceGetAttribute( + attr.CU_DEVICE_ATTRIBUTE_LOCALITY_DOMAIN_MULTIPROCESSOR_COUNT, dev + ), + "SMs per domain", + ) + ) + self.num_sms = int( + _ck(d.cuDeviceGetAttribute(attr.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev), "SMs") + ) + if self.n_domains != 2: + raise RuntimeError( + f"locality: {self.n_domains} locality domains (this layer handles 2)" + ) + self._domain_sms = self._probe(d, dev) + dom = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for s in sms: + if dom[s] != -1: + raise RuntimeError(f"locality: SM {s} probed in two domains") + dom[s] = k + self.sm_domain: List[int] = dom + self.remainder: List[int] = [s for s in range(self.num_sms) if dom[s] < 0] + rt = _alloc_prop(device, 0) + self.granularity = int( + _ck( + d.cuMemGetAllocationGranularity( + rt, d.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM + ), + "granularity", + ) + ) + + def sms(self, domain: int) -> List[int]: + return list(self._domain_sms[domain]) + + @property + def sm_table(self) -> torch.Tensor: + """int32 [num_sms] on the device: (domain << 16) | rank within the domain's sorted SMs; -1 = remainder SM. + For kernels that pick work by %smid themselves.""" + t = getattr(self, "_sm_table", None) + if t is None: + v = [-1] * self.num_sms + for k, sms in enumerate(self._domain_sms): + for i, s in enumerate(sms): + v[s] = (k << 16) | i + t = self._sm_table = torch.tensor(v, dtype=torch.int32, device=f"cuda:{self.device}") + return t + + def _probe(self, d, dev) -> List[List[int]]: + """Split the SMs by locality domain (strict, no backfill) and record %smid of a probe kernel launched on a + green context of each part: the driver's own SM -> domain assignment.""" + from cuda.bindings import nvrtc + + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + # the probe module, compiled for this device + major = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev + ), + "cc", + ) + minor = _ck( + d.cuDeviceGetAttribute( + d.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev + ), + "cc", + ) + prog = _ck( + nvrtc.nvrtcCreateProgram(_PROBE_SRC.encode(), b"smid_probe.cu", 0, [], []), + "nvrtcCreateProgram", + ) + opts = [f"--gpu-architecture=sm_{major}{minor}".encode()] + err = nvrtc.nvrtcCompileProgram(prog, len(opts), opts)[0] + if int(err) != 0: + raise RuntimeError(f"locality: probe compile failed: {err}") + n = _ck(nvrtc.nvrtcGetCUBINSize(prog), "cubin size") + cubin = b" " * n + _ck(nvrtc.nvrtcGetCUBIN(prog, cubin), "cubin") + out: List[List[int]] = [] + nb = 4 * self.sms_per_domain # more blocks than SMs: every SM of the part shows up + torch.cuda.set_device(self.device) + buf = torch.full((nb,), -1, dtype=torch.int32, device=f"cuda:{self.device}") + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate(desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM), + "green ctx", + ) + st = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + ctx = _ck(d.cuCtxFromGreenCtx(g), "ctx from green ctx") + _ck(d.cuCtxPushCurrent(ctx), "push ctx") + try: + mod = _ck(d.cuModuleLoadData(cubin), "probe module") + fn = _ck(d.cuModuleGetFunction(mod, b"smid_probe"), "probe fn") + buf.fill_(-1) + torch.cuda.synchronize(self.device) + import ctypes + + arg = ctypes.c_void_p(buf.data_ptr()) + _ck( + d.cuLaunchKernel( + fn, nb, 1, 1, 32, 1, 1, 0, st, ((arg,), (ctypes.c_void_p,)), 0 + ), + "probe launch", + ) + _ck(d.cuStreamSynchronize(st), "probe sync") + sms = sorted(set(int(x) for x in buf.tolist() if x >= 0)) + d.cuModuleUnload(mod) + finally: + d.cuCtxPopCurrent() + d.cuStreamDestroy(st) + d.cuGreenCtxDestroy(g) + if len(sms) != self.sms_per_domain: + raise RuntimeError( + f"locality: domain {k} probed {len(sms)} SMs, expected {self.sms_per_domain}" + ) + out.append(sms) + return out + + def logical_map(self, per_domain: int, extra: int) -> Tuple[torch.Tensor, List[int]]: + """The SM -> logical block index permutation for a kernel of ``2 * per_domain + extra`` CTAs, one per SM: + logical blocks [d * per_domain, (d + 1) * per_domain) run on domain d's first ``per_domain`` SMs; logical + blocks 2 * per_domain .. run on the rest (the leftover domain SMs, alternating domains, then the remainder + SMs). Returns (int32 [num_sms] on the device, -1 = no block; the domain of each extra block, -1 remainder).""" + if 2 * per_domain + extra != self.num_sms: + raise ValueError( + f"locality: {2 * per_domain} + {extra} CTAs must cover the {self.num_sms} SMs" + ) + if per_domain > self.sms_per_domain: + raise ValueError(f"locality: {per_domain} CTAs per domain > {self.sms_per_domain} SMs") + lmap = [-1] * self.num_sms + for k in range(2): + for i, s in enumerate(self._domain_sms[k][:per_domain]): + lmap[s] = k * per_domain + i + left = [self._domain_sms[k][per_domain:] for k in range(2)] + rest = [x for pair in zip(*left) for x in pair] + self.remainder + ext_dom = [] + for i, s in enumerate(rest): + lmap[s] = 2 * per_domain + i + ext_dom.append(self.sm_domain[s]) + assert sorted(lmap) == list(range(self.num_sms)) + return torch.tensor(lmap, dtype=torch.int32, device=f"cuda:{self.device}"), ext_dom + + def domain_streams(self) -> List[torch.cuda.ExternalStream]: + """One non-blocking stream per domain on a green context holding exactly that domain's SMs (kept alive with + this map): work launched on stream d runs only on domain d's SMs. Fork / join with events from the caller's + stream; graph capture follows the forks. Built once, before capture.""" + st = getattr(self, "_streams", None) + if st is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.domain_streams: build before CUDA-graph capture") + d = _drv() + dev = _ck(d.cuDeviceGet(self.device), "cuDeviceGet") + full = _ck( + d.cuDeviceGetDevResource(dev, d.CUdevResourceType.CU_DEV_RESOURCE_TYPE_SM), + "dev resource", + ) + params = [] + for k in range(self.n_domains): + p = d.CU_DEV_SM_RESOURCE_GROUP_PARAMS() + p.smCount = self.sms_per_domain + p.flags = d.CUdevSmResourceGroup_flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID + p.localityDomainId = k + params.append(p) + res, _rem = _ck( + d.cuDevSmResourceSplit(self.n_domains, full, 0, params), "cuDevSmResourceSplit" + ) + self._green = [] + st = [] + for k in range(self.n_domains): + desc = _ck(d.cuDevResourceGenerateDesc([res[k]], 1), "resource desc") + g = _ck( + d.cuGreenCtxCreate( + desc, dev, d.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM + ), + "green", + ) + s_ = _ck( + d.cuGreenCtxStreamCreate(g, d.CUstream_flags.CU_STREAM_NON_BLOCKING, 0), + "green stream", + ) + self._green.append((g, s_)) + st.append(torch.cuda.ExternalStream(int(s_), device=f"cuda:{self.device}")) + self._streams = st + return st + + def domain_of(self, ptr: int) -> int: + """The locality domain ordinal of a device address (-1: not placed in one domain). Through libcuda + directly: cuda-python's attribute helper does not know this attribute.""" + import ctypes + + d = _drv() + lib = _libcuda() + out = ctypes.c_int(-9) + attr = int(d.CUpointer_attribute.CU_POINTER_ATTRIBUTE_LOCALITY_DOMAIN_ORDINAL) + err = lib.cuPointerGetAttribute(ctypes.byref(out), ctypes.c_int(attr), ctypes.c_uint64(ptr)) + if err != 0: + raise RuntimeError( + f"locality: cuPointerGetAttribute(LOCALITY_DOMAIN_ORDINAL) failed: {err}" + ) + return int(out.value) + + def place_split(self, t: torch.Tensor, split_row: int) -> torch.Tensor: + """Split-mapped placement: a copy of ``t`` (contiguous, rows along dim 0) with the SAME shape and layout, in one + virtual range whose rows [0, split_row) are physically in domain 0 and rows [split_row, N) in domain 1. The + split must fall on an allocation granule (rows [0, split_row) a multiple of 2 MiB on VR200); the tail is padded + to a granule in domain 1. Stock consumers read it as a plain tensor; no extra memory once ``t`` is dropped.""" + if not t.is_contiguous(): + raise ValueError("locality.place_split: needs a contiguous tensor") + row_bytes = t[0].numel() * t.element_size() if t.dim() > 1 else t.element_size() + head = split_row * row_bytes + total = t.numel() * t.element_size() + if not 0 < head < total or head % self.granularity: + raise ValueError( + f"locality.place_split: rows [0, {split_row}) = {head} B must be a granule multiple " + f"({self.granularity} B) inside the {total} B tensor" + ) + tail = -(-(total - head) // self.granularity) * self.granularity + m = _Mapping(self.device, [(head, 0), (tail, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + out = flat[:total].view(t.dtype).view(t.shape) + out.copy_(t) + out._locality_mapping = m + for q, k in ( + (out.data_ptr(), 0), + (out.data_ptr() + head - 1, 0), + (out.data_ptr() + head, 1), + (out.data_ptr() + total - 1, 1), + ): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.place_split: byte {q:#x} is not in domain {k}") + return out + + def replace_param_(self, param: torch.nn.Parameter, split_row: int) -> torch.nn.Parameter: + """Re-place a parameter in place with :meth:`place_split` (swaps ``param.data``; the old block returns to the + caching allocator, so call ``torch.cuda.empty_cache()`` between large re-placements).""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("locality.replace_param_: not during CUDA-graph capture") + param.data = self.place_split(param.data, split_row) + return param + + def split_rows(self, t: torch.Tensor, dim: int = 1) -> torch.Tensor: + """Place ``t`` (contiguous, 1-byte elements, rows along ``dim`` split into two equal halves) as + [2, *t.shape[:dim], t.shape[dim] // 2, *t.shape[dim + 1:]]: half k (rows [k h, (k + 1) h) of every leading + index) in domain k's memory, one virtual range. Each half's byte size must be a multiple of the allocation + granularity (2 MiB on VR200). The copy is exact; the caller drops ``t`` afterwards.""" + if t.element_size() != 1 or not t.is_contiguous(): + raise ValueError("locality.split_rows: needs a contiguous 1-byte tensor") + n = t.shape[dim] + if n % 2: + raise ValueError(f"locality.split_rows: {n} rows don't split in halves") + half_bytes = t.numel() // 2 + if half_bytes % self.granularity: + raise ValueError( + f"locality.split_rows: a half is {half_bytes} B, not a multiple of the {self.granularity} B granularity" + ) + m = _Mapping(self.device, [(half_bytes, 0), (half_bytes, 1)]) + flat = torch.as_tensor(m, device=f"cuda:{self.device}") + shape = (2,) + tuple(t.shape[:dim]) + (n // 2,) + tuple(t.shape[dim + 1 :]) + out = flat.view(torch.uint8).view(shape) + src = t.view(torch.uint8) + for k in range(2): + out[k].copy_(src.narrow(dim, k * (n // 2), n // 2)) + out._locality_mapping = m # the tensor keeps its mapping alive + for k in range(2): + p = out[k].data_ptr() + for q in (p, p + half_bytes - 1): + if self.domain_of(q) != k: + raise RuntimeError(f"locality.split_rows: byte {q:#x} is not in domain {k}") + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_bf16_linear.py similarity index 61% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_bf16_linear.py index 3ed3cc97a229..cc29a3e9a1c7 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_bf16_linear.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_bf16_linear.py @@ -9,6 +9,30 @@ from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +CELLS_BY_TP = { + 2: ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=4608, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=4096)), + Cell(why="drafter gate_up, one row", spec=dict(m=1, n=18432, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=9216)), + ), + 4: ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=2304, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=2048)), + Cell(why="drafter gate_up, one row", spec=dict(m=1, n=9216, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=4608)), + ), + 8: ( + Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), + Cell(why="drafter QKV, nine draft rows", spec=dict(m=9, n=1280, k=12288)), + Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=1024)), + Cell(why="drafter gate_up, 16 rows", spec=dict(m=16, n=4608, k=6144)), + Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=2304)), + ), +} + class M3Bf16Linear(OpWrapper): """``F.linear(x, w)`` for bf16 ``x`` [M, K] and a bf16 weight [N, K] (cuBLAS, fp32 accumulation), bf16 out. @@ -19,13 +43,8 @@ class M3Bf16Linear(OpWrapper): ARCHS = frozenset({Arch.SM_107}) - CELLS: tuple[Cell, ...] = ( - Cell(why="drafter fc, 32 verify rows", spec=dict(m=32, n=6144, k=18432)), - Cell(why="drafter QKV, eight draft rows", spec=dict(m=8, n=4608, k=12288)), - Cell(why="drafter o_proj, a prefill chunk", spec=dict(m=1024, n=6144, k=4096)), - Cell(why="drafter gate_up, one row", spec=dict(m=1, n=18432, k=6144)), - Cell(why="drafter down, eight rows", spec=dict(m=8, n=6144, k=9216)), - ) + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ cuBLAS's split and accumulation order: `reference` multiplies in fp32 (one bf16 rounding), within the test's band. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_dense_swiglu.py similarity index 66% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_dense_swiglu.py index a0013e3ded3b..2685cc55c224 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_dense_swiglu.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_dense_swiglu.py @@ -15,9 +15,28 @@ LIMIT, ALPHA, BETA = 7.0, 1.702, 1.0 +CELLS_BY_TP = { + 2: ( + Cell(why="the TP2 dense MLP at 1 row", spec=dict(rows=1, d=6144)), + Cell(why="the TP2 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=6144)), + Cell(why="the TP2 dense MLP at 8 rows", spec=dict(rows=8, d=6144)), + ), + 4: ( + Cell(why="the TP4 dense MLP at 1 row", spec=dict(rows=1, d=3072)), + Cell(why="the TP4 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=3072)), + Cell(why="the TP4 dense MLP at 8 rows", spec=dict(rows=8, d=3072)), + ), + 8: ( + Cell(why="the TP8 dense MLP at 1 row", spec=dict(rows=1, d=1536)), + Cell(why="the TP8 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=1536)), + Cell(why="the TP8 dense MLP at 8 rows", spec=dict(rows=8, d=1536)), + ), +} + + class M3DenseSwiglu(OpWrapper): """``min(g, 7) * sigmoid(1.702 * min(g, 7)) * (clamp(u, -7, 7) + 1)`` of a bf16 [rows, 2d] gate | up tensor, - rounded once to bf16: trtllm::minimax_m3_sm107_tp8_swiglu_oai (minimaxM3Sm107Tp8/swigluOai.cu), bit for bit the + rounded once to bf16: the per-TP trtllm::swiglu_oai entry (minimaxM3Sm107/swigluOai.cu), bit for bit the Triton ``torch.ops.trtllm.silu_and_mul`` at swiglu_limit 7, swiglu_alpha 1.702, swiglu_beta 1. With ``use_pdl`` the grid waits on its predecessor and releases its dependents before it computes, so the down GEMV stages its weights during the activation (the GEMV reads the activation only after its own grid wait). Contiguous rows with an even @@ -26,12 +45,9 @@ class M3DenseSwiglu(OpWrapper): ARCHS = frozenset({Arch.SM_107}) - # The TP8 dense MLP shard: gate | up 3072 wide (d = 1536), at the decode rows. - CELLS: tuple[Cell, ...] = ( - Cell(why="the TP8 dense MLP at 1 row", spec=dict(rows=1, d=1536)), - Cell(why="the TP8 dense MLP at 4 rows (1 request x 4 tokens)", spec=dict(rows=4, d=1536)), - Cell(why="the TP8 dense MLP at 8 rows", spec=dict(rows=8, d=1536)), - ) + # Preserve each target's certified dense MLP shard dimensions and decode rows. + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ Elementwise, one rounding: bit for bit the Triton kernel (its instruction sequence: the clamps, (u + 1), the @@ -45,6 +61,7 @@ def __call__( alpha: Optional[float] = ALPHA, beta: Optional[float] = BETA, use_pdl: bool = True, + tp_size: int = 2, ) -> Optional[torch.Tensor]: if ( (limit, alpha, beta) != (LIMIT, ALPHA, BETA) @@ -57,15 +74,23 @@ def __call__( ): return None out = torch.empty((x.shape[0], x.shape[1] // 2), dtype=x.dtype, device=x.device) - torch.ops.trtllm.minimax_m3_sm107_tp8_swiglu_oai(x, out, bool(use_pdl)) + if tp_size == 2: + torch.ops.trtllm.minimax_m3_sm107_tp2_swiglu_oai(x, out, bool(use_pdl)) + elif tp_size == 4: + torch.ops.trtllm.minimax_m3_sm107_tp4_swiglu_oai(x, out, bool(use_pdl)) + else: + assert tp_size == 8, tp_size + torch.ops.trtllm.minimax_m3_sm107_tp8_swiglu_oai(x, out, bool(use_pdl)) return out - def reference(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> torch.Tensor: + def reference( + self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True, tp_size=2 + ) -> torch.Tensor: return torch.ops.trtllm.silu_and_mul( x, swiglu_limit=limit, swiglu_alpha=alpha, swiglu_beta=beta ) - def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True) -> None: + def is_valid(self, x, limit=LIMIT, alpha=ALPHA, beta=BETA, use_pdl=True, tp_size=2) -> None: assert x.dtype == torch.bfloat16 and x.dim() == 2, "x: bf16 [rows, 2d]" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_fp32_linear.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_fp32_linear.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_fp32_linear.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_mxfp8_gemm.py similarity index 74% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_mxfp8_gemm.py index 2bfc6ee495fb..4b5bebff2310 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/gemm/m3_mxfp8_gemm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_mxfp8_gemm.py @@ -7,9 +7,27 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_linear import ( + mxfp8_reference, +) from tensorrt_llm._torch.custom_ops.torch_custom_ops import mxfp8_gemm_autotuned -from .m3_mxfp8_linear import mxfp8_reference +CELLS_BY_TP = { + 2: ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=1536)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=1536)), + ), + 4: ( + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=768)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=768)), + ), + 8: ( + Cell(why="sparse-layer QKV, 4 rows", spec=dict(m=4, n=1536, k=6144)), + Cell(why="o_proj, 32 rows", spec=dict(m=32, n=6144, k=1024)), + Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=384)), + Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=384)), + ), +} class M3Mxfp8Gemm(OpWrapper): @@ -18,15 +36,11 @@ class M3Mxfp8Gemm(OpWrapper): and tactic cache as the engine-tuned quantize + GEMM); ``tune`` profiles in the engine's tuning pass.""" ARCHS = frozenset({Arch.SM_107}) + CELLS_BY_TP = CELLS_BY_TP - # The decode dispatch hands it the boundary / merge MXFP8 activation when the GEMV declines rows it serves: the - # sparse-layer QKV and the o_proj at the tp8 shards. - CELLS: tuple[Cell, ...] = ( - Cell(why="sparse-layer QKV, 4 rows", spec=dict(m=4, n=1536, k=6144)), - Cell(why="o_proj, 32 rows", spec=dict(m=32, n=6144, k=1024)), - Cell(why="shared expert down projection, 17 rows", spec=dict(m=17, n=6144, k=384)), - Cell(why="shared expert down projection, 32 rows", spec=dict(m=32, n=6144, k=384)), - ) + # The decode dispatch hands it the boundary / merge MXFP8 activation on rows the GEMV does not serve: the + # unsplit shared down projection at 17..32 rows. + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ Exactly the MXFP8 product (fp32 accumulation, one bf16 rounding): `reference` dequantizes both operands and diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_mxfp8_linear.py similarity index 77% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_mxfp8_linear.py index d7a70d725189..2020e8381e08 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/ops/gemm/m3_mxfp8_linear.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_mxfp8_linear.py @@ -41,23 +41,27 @@ def mxfp8_reference(act, act_scale, weight, weight_scale) -> torch.Tensor: return _dequant(act, act_scale) @ _dequant(weight, weight_scale).t() -class M3Mxfp8Linear(OpWrapper): - """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 - out, on the backend the engine selected: - - * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize - and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); - * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or - the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by - ``alpha`` (fp32 [1], 1.0 for these checkpoints). - - Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it - outside the decode-step kernels. - """ - - ARCHS = frozenset({Arch.SM_107}) - - CELLS: tuple[Cell, ...] = ( +CELLS_BY_TP = { + 2: ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=4992, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=4096, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=6144, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=3072, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=4096, path="tuned")), + ), + 4: ( Cell( why="sparse-layer QKV, a prefill chunk, native GEMM", spec=dict(m=2048, n=2560, k=6144, path="native"), @@ -75,7 +79,47 @@ class M3Mxfp8Linear(OpWrapper): spec=dict(m=64, n=1536, k=6144, path="flashinfer"), ), Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=2048, path="tuned")), - ) + ), + 8: ( + Cell( + why="sparse-layer QKV, a prefill chunk, native GEMM", + spec=dict(m=2048, n=1536, k=6144, path="native"), + ), + Cell( + why="o_proj, a prefill chunk, native GEMM", + spec=dict(m=2048, n=6144, k=1024, path="native"), + ), + Cell( + why="dense down projection, eager decode, native GEMM", + spec=dict(m=4, n=6144, k=1536, path="native"), + ), + Cell( + why="shared gate_up, FlashInfer GEMM", + spec=dict(m=64, n=768, k=6144, path="flashinfer"), + ), + Cell(why="o_proj, tuned quantize + GEMM", spec=dict(m=64, n=6144, k=1024, path="tuned")), + ), +} + + +class M3Mxfp8Linear(OpWrapper): + """``x @ weight^T`` for bf16 ``x`` [M, K] and an MXFP8 weight [N, K] (e4m3 + 128x4-swizzled UE8M0 scales), bf16 + out, on the backend the engine selected: + + * ``tuned_quantize_gemm`` (decode CUDA graphs once the engine tuned per-bucket backends): the AutoTuner's quantize + and GEMM per token bucket (``tune`` profiles them in the engine's tuning pass); + * otherwise the native ``mxfp8_quantize`` (swizzled scales) and then FlashInfer's MXFP8 GEMM (``flashinfer``) or + the CUTLASS block-scaled GEMM, through the autotuner while the engine tunes it (``native_tuning``), scaled by + ``alpha`` (fp32 [1], 1.0 for these checkpoints). + + Every MXFP8 projection of the target (the attention QKV and o_proj, the dense MLPs and the shared expert) runs it + outside the decode-step kernels. + """ + + ARCHS = frozenset({Arch.SM_107}) + + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ Exactly the MXFP8 product of the MXFP8 round trip of x (fp32 accumulation, one bf16 rounding): `reference` diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_target_head.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/gemm/m3_target_head.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/m3_target_head.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_core.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_core.py new file mode 100644 index 000000000000..4ccf0e106005 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_core.py @@ -0,0 +1,384 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""GEMV resolution, compilation and launch shared by the SM107 thin targets. + +Feature bindings preserve each target's original compile-key arity and exact +native argument tail. Tables, compilation locks/caches and diagnostic state stay +in the target. Producer fields are read only for a producer-capable target. +""" + +from __future__ import annotations + +from contextlib import AbstractContextManager +from typing import Any, Callable, Optional, Protocol, Tuple + +import torch + + +class GemvTarget(Protocol): + get_config: Callable[..., Optional[tuple]] + _declines: Callable[..., bool] + _CONTIG: bool + _NACC: int + _A_TMEM: bool + _ST_ASYNC: bool + _mma_n: Callable[[int], int] + _smem_bytes: Callable[..., int] + _SMEM_LIMIT: int + get_env_enable_pdl: Callable[[], bool] + PROD_N: int + PROD_OFF: tuple + _lock: AbstractContextManager[bool] + _compiled: dict[tuple, Any] + _frozen: list[bool] + last_pushed: bool + last_uncompiled: Optional[tuple] + _resolve: Callable[..., Optional[tuple]] + _views: Callable[..., tuple] + _extra_args: Callable[..., tuple] + _compile: Callable[..., Any] + + +def _resolve( + state: GemvTarget, + /, + m: int, + n: int, + k: int, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple], + x_swiglu: bool = False, + prod: bool = False, + *, + has_swiglu: bool, + has_producer: bool, +) -> Optional[tuple]: + """The compile key of an [M, K] x [N, K]^T call and its runtime push words, or None when the call is declined. + + ``ar`` is the push tuple of ``_run``; only its (world, bucket tokens, mode) enter the key, the addresses and the + rank are runtime arguments.""" + cfg = cfg or state.get_config(n, k, m) + if cfg is None or state._declines(n, k, m, x_bf16, cfg): + return None + split, ring = (cfg[0], cfg[1]) + contig = bool(cfg[2]) if len(cfg) > 2 else state._CONTIG + nacc = int(cfg[3]) if len(cfg) > 3 else state._NACC + a_tmem = bool(cfg[4]) if len(cfg) > 4 else state._A_TMEM + st_async = bool(cfg[5]) if len(cfg) > 5 else state._ST_ASYNC + if has_producer and ( + prod + and ( + not ( + n == state.PROD_N + and split == 8 + and (m <= split) + and st_async + and (not x_bf16) + and (not x_swiglu) + and (ar is None) + ) + ) + ): + return None + if nacc != 1: + return None + ring = min(ring, k // 128 // split) + tiles = k // 128 // split + mma_n = state._mma_n(m) + while ring > 2 and state._smem_bytes(ring, tiles, split, mma_n) > state._SMEM_LIMIT: + ring -= 1 + if a_tmem: + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_mx_gemv_kernel import ( + tmem_cols, + ) + + if tmem_cols(ring, nacc, True, tiles, mma_n) > 512: + a_tmem = False + ar_ws, ar_flags, ar_rank, ar_mode, ar_world, ar_maxt = (0, 0, 0, 0, 1, 8) + if ar is not None and split <= 4 and (n % 2 == 0): + ar_ws, ar_flags, ar_rank, ar_world, ar_maxt, ar_mode = ar + if ar_mode == 2: + trigger_early = False + if use_pdl is None: + use_pdl = state.get_env_enable_pdl() + key = ( + n, + k, + split, + ring, + trigger_early, + contig, + False, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + *((x_swiglu,) if has_swiglu else ()), + *((prod,) if has_producer else ()), + ) + return (key, (ar_ws, ar_flags, ar_rank)) + + +def _compile( + state: GemvTarget, + /, + key: tuple, + args: tuple, + m: int, + runtime_ar: tuple, + stream, + extra: tuple, + prod_args: tuple = (), + *, + has_swiglu: bool, + has_producer: bool, +): + """cute.compile the kernel of ``key`` (see ``_resolve``) on these argument views; cached.""" + import cutlass.cute as cute + + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm import ( + m3_mx_gemv_kernel as kernel, + ) + + if has_producer: + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + prod, + ) = key + elif has_swiglu: + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + x_swiglu, + ) = key + else: + ( + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + ) = key + ar_ws, ar_flags, ar_rank = runtime_ar + with state._lock: + fn = state._compiled.get(key) + if fn is None: + fn = state._compiled[key] = cute.compile( + kernel.m3_mx_gemv_producer if has_producer else kernel.m3_mx_gemv, + *args, + m, + n, + k, + split, + ring, + trigger_early, + contig, + timed, + nacc, + a_tmem, + st_async, + x_bf16, + ar_ws, + ar_flags, + ar_rank, + ar_mode, + ar_world, + ar_maxt, + mma_n, + use_pdl, + stream, + *extra, + False, + False, + *((x_swiglu,) if has_swiglu else ()), + *(prod_args if has_producer else ()), + *((prod,) if has_producer else ()), + ) + return fn + + +def _run( + state: GemvTarget, + /, + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + x_bf16: bool, + trigger_early: bool, + cfg: Optional[tuple], + use_pdl: Optional[bool], + ar: Optional[tuple] = None, + x_swiglu: bool = False, + prod_args: Optional[tuple] = None, + *, + has_swiglu: bool, + has_producer: bool, +) -> Optional[torch.Tensor]: + state.last_pushed = False + state.last_uncompiled = None + import cuda.bindings.driver as cuda_driver + + m, k = x.shape + if has_swiglu and x_swiglu: + k //= 2 + n = weight.shape[0] + if weight.shape[1] != k: + return None + if has_producer: + prod = prod_args is not None + resolved = state._resolve( + m, + n, + k, + x_bf16, + trigger_early, + cfg, + use_pdl, + ar, + *((x_swiglu,) if has_swiglu else ()), + *((prod,) if has_producer else ()), + ) + if resolved is None: + return None + key, runtime_ar = resolved + if has_producer and (not prod): + prod_args = state.PROD_OFF + fn = state._compiled.get(key) + if fn is None: + if state._frozen[0]: + state.last_uncompiled = key + return None + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "m3_mx_gemv must run once per shape outside CUDA-graph capture first" + ) + y = torch.empty(0 if has_producer and prod else m, n, dtype=torch.bfloat16, device=x.device) + args = state._views( + x, xsf, weight, weight_scale, y, x_bf16, *((x_swiglu,) if has_swiglu else ()) + ) + extra = state._extra_args(x.device) + stream = cuda_driver.CUstream(torch.cuda.current_stream(x.device).cuda_stream) + if fn is None: + fn = state._compile( + key, args, m, runtime_ar, stream, extra, *((prod_args,) if has_producer else ()) + ) + ar_ws, ar_flags, ar_rank = runtime_ar + fn(*args, m, ar_ws, ar_flags, ar_rank, stream, *extra, *(prod_args if has_producer else ())) + state.last_pushed = key[11] > 0 + return y + + +def precompile( + state: GemvTarget, + /, + weight: torch.Tensor, + weight_scale: torch.Tensor, + m: int, + x_bf16: bool, + use_pdl: Optional[bool] = None, + ar: Optional[Tuple[int, int, int]] = None, + x_swiglu: bool = False, + prod: bool = False, + *, + has_swiglu: bool, + has_producer: bool, +) -> Optional[tuple]: + """Compile, without launching, the kernel an M-row call against ``weight`` takes; returns its key, or None when + such a call is declined. ``ar`` = (world, bucket tokens, mode) of a pushing call (its addresses and rank are + runtime arguments, so none is needed here). Runs outside capture, even while compiles are frozen.""" + import cuda.bindings.driver as cuda_driver + + n, k = weight.shape + push = None if ar is None else (0, 0, 0, *ar) + resolved = state._resolve( + m, + n, + k, + x_bf16, + True, + None, + use_pdl, + push, + *((x_swiglu,) if has_swiglu else ()), + *((prod,) if has_producer else ()), + ) + if resolved is None: + return None + key, runtime_ar = resolved + if key in state._compiled: + return key + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_mx_gemv.precompile allocates and compiles: call it outside capture") + dev = weight.device + x = torch.zeros( + m, + 2 * k if has_swiglu and x_swiglu else k, + dtype=torch.bfloat16 if x_bf16 else torch.float8_e4m3fn, + device=dev, + ) + xsf = None if x_bf16 else torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=dev) + y = torch.empty(m, n, dtype=torch.bfloat16, device=dev) + args = state._views( + x, xsf, weight, weight_scale, y, x_bf16, *((x_swiglu,) if has_swiglu else ()) + ) + stream = cuda_driver.CUstream(torch.cuda.current_stream(dev).cuda_stream) + state._compile( + key, + args, + m, + runtime_ar, + stream, + state._extra_args(dev), + *((state.PROD_OFF,) if has_producer else ()), + ) + return key diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_host.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_host.py new file mode 100644 index 000000000000..2450531d5de9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_host.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared MXFP8 host helpers; each target supplies its own tuning tables and caches.""" + +from __future__ import annotations + +from typing import Callable, Collection, Optional + +import torch + + +def gemv_mma_n(m: int, *, mma_widths: tuple[int, ...]) -> int: + """Choose the smallest compiled MMA width covering the activation rows.""" + return next((w for w in mma_widths if m <= w)) + + +def gemv_smem_bytes(ring: int, tiles: int, split: int, mma_n: int) -> int: + """Compute the shared-memory bytes for the supplied launch geometry.""" + return ( + ring * (128 * 128 + 512) + + tiles * (mma_n * 128 + 512) + + split * (128 // split) * (mma_n + 4) * 4 + + 1024 + ) + + +def gemv_dummy_timers( + device: torch.device, *, dummy: dict[torch.device, torch.Tensor] +) -> torch.Tensor: + """Placeholder for the kernel's timer buffer, which it only writes when compiled with timers on.""" + t = dummy.get(device) + if t is None: + t = dummy[device] = torch.zeros(8, dtype=torch.int64, device=device) + return t + + +def gemv_arg(t: torch.Tensor, align: int = 16) -> object: + """Expose a tensor with the launch alignment and dynamic leading layout.""" + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align).mark_layout_dynamic(leading_dim=t.dim() - 1) + + +def gemv_get_config( + n: int, + k: int, + m: Optional[int] = None, + *, + configs_by_m: dict[tuple[int, int, int], tuple], + configs_wide: dict[tuple[int, int], tuple], + configs: dict[tuple[int, int], tuple], + strided: bool, + strided_max_rows: int, + strided_shapes: Collection[tuple[int, int]], +) -> Optional[tuple]: + """Select a configuration from the target tables in their original order.""" + cfg = configs_by_m.get((n, k, m)) if m is not None else None + if cfg is None and m is not None and (m > 8): + cfg = configs_wide.get((n, k)) + cfg = cfg or configs.get((n, k)) + if ( + cfg is None + or n % 128 != 0 + or k % 128 != 0 + or (k // 128 % cfg[0] != 0) + or (128 % cfg[0] != 0) + ): + return None + if ( + strided + and m is not None + and (m <= strided_max_rows) + and ((n, k) in strided_shapes) + and (len(cfg) == 2) + ): + cfg = cfg + (0,) + return cfg + + +def gemv_declines( + n: int, + k: int, + m: int, + x_bf16: bool, + cfg: tuple, + *, + max_tokens: int, + smem_limit: int, + select_mma_n: Callable[[int], int], + smem_bytes: Callable[[int, int, int, int], int], +) -> bool: + """Whether a (split, ring, ...) config declines an [M, K] x [N, K]^T call on shape grounds.""" + split = cfg[0] + if not 0 < m <= max_tokens or 128 % split != 0 or k // 128 % split != 0: + return True + tiles = k // 128 // split + mma_n = select_mma_n(m) + if 16 < m <= 32 and k < 1024 or (x_bf16 and tiles * 4 * mma_n > 8 * 128): + return True + return smem_bytes(min(cfg[1], tiles, 2), tiles, split, mma_n) > smem_limit + + +def gemv_views( + x: torch.Tensor, + xsf: Optional[torch.Tensor], + weight: torch.Tensor, + weight_scale: torch.Tensor, + y: torch.Tensor, + x_bf16: bool, + x_swiglu: bool = False, + *, + placeholder_scales: Callable[[int, torch.device], torch.Tensor], + placeholder_words: Callable[[torch.device], torch.Tensor], + arg: Callable[..., object], + dummy_timers: Callable[[torch.device], torch.Tensor], +) -> tuple: + """Build argument views using the target placeholder and timer caches.""" + k = x.shape[1] // 2 if x_swiglu else x.shape[1] + if x_bf16: + x_bytes = x.view(torch.uint8) + xsf = placeholder_scales(k, x.device) + x_words = x.view(torch.int32) + else: + x_bytes = x.view(torch.uint8) + x_words = placeholder_words(x.device) + return ( + arg(weight.view(torch.uint8)), + arg(weight_scale.view(torch.uint8).reshape(-1)), + arg(x_bytes), + arg(xsf.view(torch.uint8).reshape(-1)), + arg(y.view(-1) if y.numel() else placeholder_words(x.device).view(torch.bfloat16)), + arg(x_words.reshape(-1)), + arg(dummy_timers(x.device)), + ) + + +def gemv_extra_args( + device: torch.device, + *, + placeholder_words: Callable[[torch.device], torch.Tensor], + arg: Callable[..., object], +) -> tuple: + """Placeholder arguments of the kernel's draft-input prologue and RoPE / KV epilogue (compiled off).""" + i32 = placeholder_words(device) + f32 = i32.view(torch.float32) + u8 = i32.view(torch.uint8) + d = (arg(i32), arg(i32), arg(i32), arg(i32), arg(f32), 0, arg(i32), arg(i32), 0) + r = (arg(f32), arg(i32), arg(i32), arg(u8), arg(u8), arg(f32), arg(i32)) + return d + r + + +def gemv_placeholder_scales( + k: int, device: torch.device, *, placeholders: dict[tuple, torch.Tensor] +) -> torch.Tensor: + """Get the target cache's zeroed scale placeholder for this device and width.""" + key = ("sf", k, device) + t = placeholders.get(key) + if t is None: + t = placeholders[key] = torch.zeros(k // 128 * 512, dtype=torch.uint8, device=device) + return t + + +def gemv_placeholder_words( + device: torch.device, *, placeholders: dict[tuple, torch.Tensor] +) -> torch.Tensor: + """Get the target cache's zeroed word placeholder for this device.""" + key = ("w", device) + t = placeholders.get(key) + if t is None: + t = placeholders[key] = torch.zeros(16, dtype=torch.int32, device=device) + return t + + +def gemv_ensure_compiled( + weight: torch.Tensor, + weight_scale: torch.Tensor, + use_pdl: bool, + *, + get_config: Callable[..., Optional[tuple]], + ensured: set[tuple], + mma_widths: tuple[int, ...], + mx_gemv: Callable[..., Optional[torch.Tensor]], + mx_gemv_bf16: Callable[..., Optional[torch.Tensor]], +) -> None: + """Compiles the kernel of this weight's shape outside capture (a one-token dummy call).""" + if ( + weight.dim() != 2 + or weight.dtype != torch.float8_e4m3fn + or get_config(*weight.shape) is None + ): + return + n, k = weight.shape + if (n, k, use_pdl) in ensured: + return + ensured.add((n, k, use_pdl)) + xq = torch.zeros(1, k, dtype=torch.float8_e4m3fn, device=weight.device) + xsf = torch.full((128 * (k // 32),), 127, dtype=torch.uint8, device=weight.device) + mx_gemv(xq, xsf, weight, weight_scale, use_pdl=use_pdl) + for width in mma_widths[1:]: + mx_gemv( + torch.zeros(width, k, dtype=torch.float8_e4m3fn, device=weight.device), + xsf, + weight, + weight_scale, + use_pdl=use_pdl, + ) + mx_gemv_bf16( + torch.zeros(1, k, dtype=torch.bfloat16, device=weight.device), + weight, + weight_scale, + use_pdl=use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_state.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_state.py new file mode 100644 index 000000000000..a70aa0540216 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/mxfp8_state.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 backend dispatch definitions; each model constructs its own mutable state.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from torch import nn + +from tensorrt_llm._torch.modules.linear import ( + MXFP8LinearMethod, + flashinfer_mxfp8_autotune_active, + flashinfer_mxfp8_decode_graph_capture_active, +) + + +class _Mxfp8GemmState(nn.Module): + """The MXFP8 GEMM backend state the engine's warmup drives. + + The engine finds MXFP8 projections by a module's ``quant_method`` and, for every one, opts it into native GEMM + tuning and (on decode CUDA graphs) the auto-dispatched FlashInfer GEMM, then marks each tuned. Every projection of + this target dispatches identically, so one method object carries that state for all of them; the forward reads it + at each call (``_mx_mode``) and passes the resolved backend to the GEMM entries. + """ + + def __init__(self) -> None: + super().__init__() + self.quant_method = MXFP8LinearMethod() + + +@dataclass(frozen=True) +class MxMode: + """The MXFP8 GEMM dispatch of one call, resolved from ``_Mxfp8GemmState`` and the engine's contexts.""" + + decode_graph: bool # inside a decode CUDA-graph warmup or capture + autotune: bool # inside the engine's FlashInfer MXFP8 tuning pass + tuned_quantize_gemm: ( + bool # per-bucket tuned quantize + GEMM (decode graphs, tune_decode_graph_backends) + ) + flashinfer: bool # FlashInfer's MXFP8 GEMM on the quantized activation + native_tuning: ( + bool # the native GEMM through the autotuner (warmup, before its tactics are registered) + ) + + +def _mx_mode(state: _Mxfp8GemmState) -> MxMode: + method = state.quant_method + decode_graph = flashinfer_mxfp8_decode_graph_capture_active() + autotune = flashinfer_mxfp8_autotune_active() + return MxMode( + decode_graph=decode_graph, + autotune=autotune, + tuned_quantize_gemm=method.tune_decode_graph_backends and decode_graph, + flashinfer=method.backend == "flashinfer" + or ( + method.backend == "auto" + and (autotune or (method._flashinfer_autotuned and decode_graph)) + ), + native_tuning=method.needs_native_autotune, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/prequant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/prequant.py new file mode 100644 index 000000000000..62ab8f7c9e95 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/prequant.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Activation hand-off slot types; targets own independent slot instances.""" + + +class PrequantSlot: + def __init__(self): + self.pending = None # (data_ptr, rows, width, e4m3 [rows, width], scales) + + def publish(self, x, xq, xsf): + self.pending = (x.data_ptr(), x.shape[0], x.shape[-1], xq, xsf) + + def clear(self): + self.pending = None + + def take(self, x): + """(e4m3 [x.shape[0], width], scales) if x is the published tensor, else None.""" + pending, self.pending = self.pending, None + if ( + pending is None + or x.dim() != 2 + or x.data_ptr() != pending[0] + or x.shape[0] != pending[1] + or x.shape[1] != pending[2] + ): + return None + return pending[3], pending[4] + + +class ArPushSlot: + """The attention boundary's all-reduce exchange for the o_proj GEMV to push into. + + The fast forward arms ``target`` (m3_mx_gemv ``ar`` tuple) around the attention call; the o_proj marks the output + it pushed in ``pushed``; the boundary then runs the all-reduce with ``prepushed`` = the ``ar`` tuple's mode.""" + + def __init__(self): + self.target = None + self.pushed = None # data_ptr of the pushed output + + def arm(self, target): + self.target = target + self.pushed = None + + def take_target(self): + target, self.target = self.target, None + return target + + def take_pushed(self, x): + pushed, self.pushed = self.pushed, None + return pushed is not None and x.data_ptr() == pushed + + +class MoePushSlot: + """The MoE boundary's all-reduce exchange for its two producers to push into: the routed-expert engine (set A, + ``engine`` = its ``arp`` descriptor) and the shared expert's down projection (set B, ``shared`` = m3_mx_gemv + ``ar`` tuple). Each marks the output it pushed; the boundary then runs the all-reduce with ``prepushed`` = 3.""" + + def __init__(self): + self.engine = None + self.shared = None + self.engine_pushed = None # data_ptr of the pushed routed output + self.shared_pushed = None # data_ptr of the pushed shared output + + def arm(self, engine, shared): + self.engine, self.shared = engine, shared + self.engine_pushed = self.shared_pushed = None + + def disarm(self): + self.engine = self.shared = None + + def take_engine(self): + engine, self.engine = self.engine, None + return engine + + def take_shared(self): + shared, self.shared = self.shared, None + return shared + + def take_pushed(self, shared, routed): + """(routed pushed, shared pushed) for these outputs; clears the marks.""" + e = routed is not None and self.engine_pushed == routed.data_ptr() + s = shared is not None and self.shared_pushed == shared.data_ptr() + self.engine_pushed = self.shared_pushed = None + return e, s diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/target_head_state.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/target_head_state.py new file mode 100644 index 000000000000..a6977b8c8b44 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/gemm/target_head_state.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Common head host implementation, bound to per-target state and callbacks.""" + +from __future__ import annotations + +from typing import Callable, Optional, Protocol + +import torch + + +class HeadState(Protocol): + weight: torch.Tensor + coded: bool + w: Optional[torch.Tensor] + base: int + exc_off: Optional[torch.Tensor] + exc_pos: Optional[torch.Tensor] + exc_w: Optional[torch.Tensor] + part: Optional[torch.Tensor] + counter: Optional[torch.Tensor] + grid: int + + +class HeadTarget(Protocol): + _NUM_SMS: Optional[int] + _num_sms: Callable[[], int] + TargetHead: Callable[..., HeadState] + CTA_WARPS: int + MAX_ROWS: int + encode_chunked: Callable[..., tuple] + round_trip_ok: Callable[..., bool] + _target_head_op: Callable[..., torch.Tensor] + + +def _num_sms(state: HeadTarget, /) -> int: + if state._NUM_SMS is None: + state._NUM_SMS = torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + return state._NUM_SMS + + +def head_supported(state: HeadTarget, /, lm_head: object) -> bool: + """A bias-free 2-D BF16 head of whole 16-row tiles whose K the coding covers (a multiple of 64).""" + weight = getattr(lm_head, "weight", None) + if weight is None or weight.dim() != 2 or weight.dtype != torch.bfloat16: + return False + return ( + getattr(lm_head, "bias", None) is None + and weight.shape[0] % 16 == 0 + and (weight.shape[1] % 64 == 0) + ) + + +def prepare_head( + state: HeadTarget, /, lm_head: torch.nn.Module, grid: Optional[int] = None +) -> HeadState: + """Code ``lm_head``'s local shard and check it decodes to the BF16 weight bit for bit (fails closed: a shard that + does not is returned with ``coded=False`` and nothing kept). Built once, eagerly, before the executor's warmup + (``M3Eagle3Worker.prebuild_target_head``).""" + weight = lm_head.weight + with torch.no_grad(): + w = weight.detach() + packed, base, exc_off, exc_pos, exc_w = state.encode_chunked(w) + if not state.round_trip_ok(packed, base, exc_off, exc_pos, exc_w, w): + return state.TargetHead(weight=weight, coded=False) + if grid is None: + grid = state._num_sms() + n = weight.shape[0] + return state.TargetHead( + weight=weight, + coded=True, + w=packed, + base=base, + exc_off=exc_off, + exc_pos=exc_pos, + exc_w=exc_w.view(-1), + part=torch.zeros( + grid * state.CTA_WARPS * 2 * 4 * 32, dtype=torch.float32, device=weight.device + ), + counter=torch.zeros((n + 15) // 16, dtype=torch.int32, device=weight.device), + grid=grid, + ) + + +def logits( + state: HeadTarget, /, head: HeadState, x: torch.Tensor, use_pdl: bool = True +) -> torch.Tensor: + """This rank's BF16 logits shard [rows, n] of the rows x [rows <= MAX_ROWS, K] (BF16) from the coded shard.""" + assert head.coded and 1 <= x.shape[0] <= state.MAX_ROWS, (head.coded, x.shape) + if x.stride(1) != 1 or x.stride(0) % 8 or x.data_ptr() % 16: + x = x.contiguous() + n = head.weight.shape[0] + out = torch.empty((x.shape[0], n), dtype=torch.bfloat16, device=x.device) + return state._target_head_op( + head.w, + x, + out, + head.part, + head.counter, + head.exc_off, + head.exc_pos, + head.exc_w, + head.base, + n, + head.grid, + use_pdl, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/m3_router_gemv.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_gemv.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/m3_router_gemv.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/m3_router_gemv_fp4.py similarity index 95% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/m3_router_gemv_fp4.py index 0306210adaab..462fe6a19ce3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/moe/m3_router_gemv_fp4.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/m3_router_gemv_fp4.py @@ -8,8 +8,10 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from .m3_router_gemv import fp32_logits, nvfp4_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv import ( + fp32_logits, + nvfp4_quantize, +) class M3RouterGemvFp4(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/m3_router_tc.py similarity index 67% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/m3_router_tc.py index 9e21113d8483..17a044e6b1d1 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/moe/m3_router_tc.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/m3_router_tc.py @@ -15,8 +15,9 @@ OpWrapper, true_fp32_matmul, ) - -from .m3_router_gemv import nvfp4_quantize +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv import ( + nvfp4_quantize, +) def tf32_rna(w: torch.Tensor) -> torch.Tensor: @@ -27,19 +28,8 @@ def tf32_rna(w: torch.Tensor) -> torch.Tensor: return torch.where(finite, rounded, bits).view(torch.float32) -class M3RouterTc(OpWrapper): - """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. - - ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' - partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result - does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns - ``[logits]`` or ``[logits, x_fp4, x_sf]``. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # The target routes 1..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. - CELLS: tuple[Cell, ...] = ( +CELLS_BY_TP = { + 2: ( Cell( why="one token, logits only (the wide target's router_path(1) with no fp32 GEMV rows)", spec=dict(m=1, k=6144, fp4=False), @@ -77,7 +67,70 @@ class M3RouterTc(OpWrapper): for m in range(2, 49) if m not in (2, 8, 13, 48) ), - ) + ), + 4: ( + Cell( + why="9 rows: a 3-request verify past the MoE engine", spec=dict(m=9, k=6144, fp4=False) + ), + Cell(why="10 rows", spec=dict(m=10, k=6144, fp4=False)), + Cell(why="11 rows", spec=dict(m=11, k=6144, fp4=False)), + Cell(why="12 rows: 3 requests x 4 tokens", spec=dict(m=12, k=6144, fp4=False)), + Cell(why="13 rows: a partial second token tile", spec=dict(m=13, k=6144, fp4=False)), + Cell(why="14 rows", spec=dict(m=14, k=6144, fp4=False)), + Cell(why="15 rows", spec=dict(m=15, k=6144, fp4=False)), + Cell( + why="16 rows: 4 requests x 4 tokens, two full token tiles", + spec=dict(m=16, k=6144, fp4=False), + ), + ), + 8: ( + Cell( + why="two tokens, logits only (the smallest call the target makes here)", + spec=dict(m=2, k=6144, fp4=False), + ), + Cell( + why="8 tokens with the NVFP4 activation: one 8-token MMA tile, every expert tile quantizes one row", + spec=dict(m=8, k=6144, fp4=True), + ), + Cell( + why="13 tokens: a partial second token tile", + spec=dict(m=13, k=6144, fp4=True), + ), + Cell( + why="the 48-token maximum, six token tiles, logits only", + spec=dict(m=48, k=6144, fp4=False), + ), + Cell( + why="the 48-token maximum with the NVFP4 activation", + spec=dict(m=48, k=6144, fp4=True), + ), + # Every other row count the router takes (modeling's ROUTER_TC), with and without the NVFP4 activation. + *( + Cell( + why=f"{m} rows: a router row count between the ones above", + spec=dict(m=m, k=6144, fp4=m % 2 == 1), + ) + for m in range(2, 49) + if m not in (2, 8, 13, 48) + ), + ), +} + + +class M3RouterTc(OpWrapper): + """Router logits on TF32 tensor cores for 128 experts and 1..48 tokens: ``x @ tf32(w)^T``, fp32 accumulate. + + ``x`` is BF16 (exact in TF32), ``w`` the fp32 router weight, rounded to TF32 in the kernel. The 16 K slices' + partials are summed in slice order whatever cluster size the device co-schedules (16 or 8 CTAs), so the result + does not depend on it. With ``sf_scale`` also the NVFP4 activation as ``m3_router_gemv_fp4``: returns + ``[logits]`` or ``[logits, x_fp4, x_sf]``. + """ + + ARCHS = frozenset({Arch.SM_107}) + CELLS_BY_TP = CELLS_BY_TP + + # The target routes 1..48 tokens here: plain logits on the unfused path, logits + NVFP4 on the fused-router path. + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ The kernel uses 16-CTA clusters when the device co-schedules eight of them and 8-CTA clusters otherwise; both diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/phase_plan.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/phase_plan.py new file mode 100644 index 000000000000..31257c1fbecd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/phase_plan.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared routed-expert phase planners, with target geometry supplied explicitly.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +import torch + + +@dataclass(frozen=True) +class MoePlanGeometry: + top_k: int + hidden: int + u1: int + stages: int + row2_bytes: int + item_bytes: int + + +def make_phase_plan( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = True, + heavy_fc2_only: bool = False, + *, + geometry: MoePlanGeometry, + fine_ranges_fn: Callable[[list[float], int], list[int]], +) -> tuple[torch.Tensor, int]: + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = geometry.top_k * m_max + nb2 = geometry.hidden // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + min_nb = min(nbs) + no_fc1 = [heavy_fc2_only and nbs[c] > min_nb for c in range(num_ctas)] + BIG_W = 1 << 40 + blk_bytes = 32 * geometry.row2_bytes + for g in range(geometry.top_k, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | nbs[c] << 16 + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * geometry.u1 * geometry.stages + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and (not (early0 and p == 0 and (m_max >= 2))): + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + for c in range(num_ctas): + if no_fc1[c]: + fc2[c] += BIG_W + items = [0] * num_ctas + if total >= geometry.stages * num_ctas: + + def counts_for(limit): + return [ + 0 if no_fc1[c] else max(geometry.stages, (limit - f) // geometry.item_bytes) + for c, f in enumerate(fc2) + ] + + lo, hi = ( + 0, + max((f for c, f in enumerate(fc2) if not no_fc1[c])) + + total * geometry.item_bytes, + ) + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > geometry.stages), + key=lambda c: (-(fc2[c] + items[c] * geometry.item_bytes), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + lo_f, hi_f = (0.0, float(max(fc2) + total * geometry.item_bytes)) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum((max(0.0, (mid - f) / geometry.item_bytes) for f in fc2)) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / geometry.item_bytes) for f in fc2] + sw = sum(wts) + items = fine_ranges_fn([w * total / sw for w in wts], total) + else: + chosen = sorted( + sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // geometry.stages] + ) + for c in chosen: + items[c] = geometry.stages + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | i0 + items[c] << 16 + i0 += items[c] + assert i0 == total + return (plan, max(nbs)) + + +def make_phase_plan_tp8( + num_ctas: int, + m_max: int, + lag: int = 2, + tail_w: float = 1.0, + early0: bool = False, + fine: bool = True, + *, + geometry: MoePlanGeometry, + fine_ranges_fn: Callable[[list[float], int], list[int]], +) -> tuple[torch.Tensor, int]: + """The work split of ``m3_moe`` for each distinct-expert count G (4 .. 4 m_max): FC2 32-row output blocks in + contiguous per-CTA ranges (the same for every group), and per group p (slots 4 p ..) a range of its unit-major FC1 + stages per CTA, balancing the bytes each CTA streams in phase p (FC1 of group p plus FC2 of group p - lag, and the + trailing FC2s with the last group). Ranges are empty or >= 6 stages (a unit is split over at most two CTAs); a + group too small for every CTA gets whole units on the CTAs with the least FC2 work, or (``fine``) byte-balanced + ranges of any length with at most one CTA boundary inside a unit (:func:`_fine_ranges`). Returns int32 + [4 m_max + 1, num_ctas, m_max + 1] (per group i0 | i1 << 16, then b0 | nb << 16) and the largest block count.""" + gcap = geometry.top_k * m_max + nb2 = geometry.hidden // 32 + plan = torch.zeros(gcap + 1, num_ctas, m_max + 1, dtype=torch.int32) + nbs = [nb2 // num_ctas + (1 if c < nb2 % num_ctas else 0) for c in range(num_ctas)] + b0s = [sum(nbs[:c]) for c in range(num_ctas)] + blk_bytes = 32 * geometry.row2_bytes + for g in range(geometry.top_k, gcap + 1): + ng = (g + 3) // 4 + for c in range(num_ctas): + plan[g, c, m_max] = b0s[c] | nbs[c] << 16 + for p in range(ng): + n_p = min(4, g - 4 * p) + total = n_p * geometry.u1 * geometry.stages + fc2 = [0] * num_ctas + for c in range(num_ctas): + if p >= lag: + fc2[c] += nbs[c] * min(4, g - 4 * (p - lag)) * blk_bytes + if p == ng - 1 and (not (early0 and p == 0 and (m_max >= 2))): + for q in range(max(0, ng - lag), ng): + fc2[c] += int(tail_w * nbs[c] * min(4, g - 4 * q) * blk_bytes) + items = [0] * num_ctas + if total >= geometry.stages * num_ctas: + + def counts_for(limit): + return [max(geometry.stages, (limit - f) // geometry.item_bytes) for f in fc2] + + lo, hi = (0, max(fc2) + total * geometry.item_bytes) + while lo < hi: + mid = (lo + hi) // 2 + if sum(counts_for(mid)) >= total: + hi = mid + else: + lo = mid + 1 + items = counts_for(lo) + extra = sum(items) - total + while extra > 0: + c = min( + (c for c in range(num_ctas) if items[c] > geometry.stages), + key=lambda c: (-(fc2[c] + items[c] * geometry.item_bytes), c), + ) + items[c] -= 1 + extra -= 1 + elif fine: + lo_f, hi_f = (0.0, float(max(fc2) + total * geometry.item_bytes)) + for _ in range(100): + mid = 0.5 * (lo_f + hi_f) + if sum((max(0.0, (mid - f) / geometry.item_bytes) for f in fc2)) >= total: + hi_f = mid + else: + lo_f = mid + wts = [max(0.0, (hi_f - f) / geometry.item_bytes) for f in fc2] + sw = sum(wts) + items = fine_ranges_fn([w * total / sw for w in wts], total) + else: + chosen = sorted( + sorted(range(num_ctas), key=lambda c: (fc2[c], c))[: total // geometry.stages] + ) + for c in chosen: + items[c] = geometry.stages + i0 = 0 + for c in range(num_ctas): + if items[c]: + plan[g, c, p] = i0 | i0 + items[c] << 16 + i0 += items[c] + assert i0 == total + return (plan, max(nbs)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/reference.py new file mode 100644 index 000000000000..30ce9fa1fd54 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/reference.py @@ -0,0 +1,233 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Torch reference arithmetic shared by the SM107 MiniMax-M3 MoE targets.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from tensorrt_llm.quantization.utils.fp4_utils import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + get_shuffle_matrix_a_row_indices, +) + +if TYPE_CHECKING: + from .runtime import M3MoeLayer + +HIDDEN = 6144 +TOP_K = 4 + +ROUTED_SCALING = 2.0 +SHARED_ALPHA = 1.702 +SHARED_LIMIT = 7.0 + +_LOG2E = 1.4426950216293334961 +_RCP6 = 0.16666667163372039795 # fp32(1 / 6), the kernels' reciprocal of the E2M1 maximum +_E2M1 = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) +_F32_MAX = 3.40282346638528859812e38 + + +def _f32(x: float) -> torch.Tensor: + return torch.tensor(x, dtype=torch.float32) + + +def _e2m1_decode(packed: torch.Tensor) -> torch.Tensor: + """fp32 values of packed E2M1 bytes [..., K / 2] (low nibble first) -> [..., K].""" + lut = torch.tensor(_E2M1, dtype=torch.float32, device=packed.device) + lo = lut[(packed & 0xF).long()] + hi = lut[(packed >> 4).long()] + return torch.stack((lo, hi), dim=-1).flatten(-2) + + +def _e2m1_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to the nearest E2M1 value, ties to even, saturating at 6 (cvt.rn.satfinite.e2m1x2).""" + a = x.abs() + q = torch.full_like(a, 6.0) + for hi, value, inclusive in ( + (5.0, 4.0, True), + (3.5, 3.0, False), + (2.5, 2.0, True), + (1.75, 1.5, False), + (1.25, 1.0, True), + (0.75, 0.5, False), + (0.25, 0.0, True), + ): + q = torch.where(a <= hi if inclusive else a < hi, torch.full_like(a, value), q) + return torch.where(x < 0, -q, q) + + +def _e4m3_round(x: torch.Tensor) -> torch.Tensor: + """x (fp32) to E4M3, nearest even, saturating (cvt.rn.satfinite.e4m3x2), as fp32.""" + return x.clamp(-448.0, 448.0).to(torch.float8_e4m3fn).float() + + +def _ue8m0_up(x: torch.Tensor) -> torch.Tensor: + """UE8M0 code of x >= 0 rounded toward +inf, saturating (cvt.rp.satfinite.ue8m0x2), int32.""" + bits = x.contiguous().view(torch.int32) + code = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(torch.int32) + return code.clamp(0, 254) + + +def unswizzle_sf(sf: torch.Tensor, rows: int, blocks: int) -> torch.Tensor: + """The [rows, blocks] scale bytes of a 128x4-swizzled scale buffer (int64 indices, any byte dtype).""" + flat = sf.reshape(-1).view(torch.uint8) + r = torch.arange(rows, device=flat.device).view(-1, 1) + b = torch.arange(blocks, device=flat.device).view(1, -1) + padded_blocks = (blocks + 3) // 4 * 4 + idx = ( + (r // 128) * (128 * padded_blocks) + + (b // 4) * 512 + + (r % 32) * 16 + + ((r % 128) // 32) * 4 + + (b % 4) + ) + return flat[idx] + + +def mxfp8_dequant(q: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """fp32 values of e4m3 ``q`` [R, K] with its 128x4-swizzled UE8M0 scales (one per 32 elements).""" + rows, k = q.shape + codes = unswizzle_sf(sf, rows, k // 32).float() + return q.view(torch.float8_e4m3fn).float() * torch.exp2(codes - 127.0).repeat_interleave( + 32, dim=1 + ) + + +def mxfp8_quant_dequant(x: torch.Tensor) -> torch.Tensor: + """x [R, K] through MXFP8 (one UE8M0 scale per 32, the smallest power of two >= amax / 448; e4m3 data), back to + fp32 -- mxfp8_quantize's recipe, which every MXFP8 activation here reproduces.""" + rows, k = x.shape + blocks = x.float().view(rows, k // 32, 32) + code = _ue8m0_up(blocks.abs().amax(dim=-1) * _f32(1.0 / 448.0).to(x.device)) + scale = torch.exp2(code.float() - 127.0).unsqueeze(-1) + return (_e4m3_round(blocks / scale) * scale).view(rows, k) + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16).float() + + +def _gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """a [M, K] @ b [N, K]^T, accumulated in fp64 and rounded once to fp32.""" + return (a.double() @ b.double().t()).float() + + +def _row_perm(rows: int, gated: bool) -> torch.Tensor: + """The TRTLLM-Gen row order: stored row j holds logical row perm[j] (the gated interleave, then the 32-row shuffle + of epilogue tile 128).""" + probe = torch.empty(rows, 1) + perm = get_shuffle_matrix_a_row_indices(probe, epilogue_tile_m=128) + if gated: + perm = get_reorder_rows_for_gated_act_gemm_row_indices(probe)[perm] + return perm + + +def _nvfp4_expert(w: torch.Tensor, sf: torch.Tensor, perm_inv: torch.Tensor) -> torch.Tensor: + """One expert's stored NVFP4 matrix (values [R, K / 2], swizzled e4m3 scales) as raw fp32 [R, K] in logical row + order (the global scale is not applied: it is the per-expert alpha).""" + rows, half = w.shape + codes = unswizzle_sf(sf, rows, 2 * half // 16).view(torch.float8_e4m3fn).float() + vals = _e2m1_decode(w) * codes.repeat_interleave(16, dim=1) + return vals[perm_inv] + + +def moe_route( + x: torch.Tensor, router_weight: torch.Tensor, bias: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + """MiniMax2 routing on the fp32 router GEMV: (top-4 expert ids [M, 4] int64, bf16-rounded weights [M, 4] fp32).""" + logits = _gemm(x, router_weight) + sig = (0.5 * torch.tanh(0.5 * logits).double() + 0.5).float() + keys = sig + bias + ids = torch.sort(keys, dim=-1, descending=True, stable=True).indices[:, :TOP_K] + sk = keys.gather(1, ids) - bias[ids] + total = (sk[:, 0] + sk[:, 2]) + (sk[:, 1] + sk[:, 3]) + w = (sk * _f32(ROUTED_SCALING).to(x.device)) / (total + _f32(1e-20).to(x.device)).unsqueeze(1) + return ids, _bf16(w) + + +def _nvfp4_act(x: torch.Tensor, sc: torch.Tensor) -> torch.Tensor: + """The activation's NVFP4 quantization (fp4_quantize(x, sc, 16)), as raw fp32 values (q x block scale).""" + rows, k = x.shape + blocks = x.view(rows, k // 16, 16) + vmax = blocks.abs().amax(dim=-1, keepdim=True) + sff = _e4m3_round(sc * (vmax * _f32(_RCP6).to(x.device))) + # The kernels' rcp.approx(sf x rcp.approx(sc)) lands where the correctly rounded sc / sf does on every element + # measured (the 1 / (sf x (1 / sc)) spelling differs on ~0.06% of them). + osc = torch.where(vmax != 0, sc / sff, torch.zeros_like(vmax)) + return (_e2m1_round(blocks * osc) * sff).view(rows, k) + + +def _fc1_epilogue(up: torch.Tensor, gate: torch.Tensor, esc: torch.Tensor) -> torch.Tensor: + """TRTLLM-Gen's SwiGLU-OAI FC1 epilogue on the raw FC1 sums (fp32 [I]) and its NVFP4 requantization: the raw FC2 + input values. ``esc`` = (fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha) of the expert, with + act_beta and act_clamp already divided by fc31_alpha (the loader's convention).""" + a, c, alpha, beta, lim = (esc[i] for i in range(5)) + gc = torch.where(gate > lim, lim, gate) + gc = torch.where(gate < -_F32_MAX, torch.full_like(gc, -_F32_MAX), gc) + uc = torch.minimum(torch.maximum(up, -lim), lim) + c1 = a * (alpha * _f32(_LOG2E).to(up.device)) + sgm = 1.0 / (torch.exp2(-(gc * c1)) + 1.0) + h = (uc.double() * a.double() + (a * beta).double()).float() * (gc * sgm) + blocks = h.view(-1, 16) + amax = (blocks.abs() * c).amax(dim=-1, keepdim=True) + sff = _e4m3_round(amax * _f32(_RCP6).to(up.device)) + osc = torch.where(sff != 0, 1.0 / sff, torch.zeros_like(sff)) + return (_e2m1_round((blocks * osc) * c) * sff).view(-1) + + +def _shared_expert( + x: torch.Tensor, + shared_mxfp8: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + *, + shared_inter: int, +) -> torch.Tensor: + """The shared expert as the target runs it outside the engine: MXFP8 gate_up (bf16 out), SwiGLU-OAI with the + model's alpha / limit / (up + 1) in fp32 (bf16 out), MXFP8 down (bf16 out).""" + gu_w, gu_sf, dn_w, dn_sf = shared_mxfp8 + gu = _bf16(_gemm(mxfp8_quant_dequant(x), mxfp8_dequant(gu_w, gu_sf))) + gate, up = gu[:, :shared_inter], gu[:, shared_inter:] + a = torch.minimum(gate, torch.full_like(gate, SHARED_LIMIT)) + ub = up.clamp(-SHARED_LIMIT, SHARED_LIMIT) + 1.0 + sg = 1.0 / ( + torch.exp2((a * _f32(-SHARED_ALPHA).to(x.device)) * _f32(_LOG2E).to(x.device)) + 1.0 + ) + h = _bf16(ub * (sg * a)) + return _bf16(_gemm(mxfp8_quant_dequant(h), mxfp8_dequant(dn_w, dn_sf))) + + +def moe_reference( + hidden: torch.Tensor, + layer: M3MoeLayer, + shared: bool, + *, + inter: int, + shared_inter: int, +) -> torch.Tensor: + """The MoE layer of the decode engines (both token ranges compute this), in plain torch: bf16 [M, 6144].""" + x = hidden.float() + m = x.shape[0] + ids, weights = moe_route(x, layer.router_weight, layer.bias) + xr = _nvfp4_act(x, layer.fc31_input_scale.reshape(()).float()) + inv31 = torch.argsort(_row_perm(2 * inter, gated=True)).to(x.device) + inv2 = torch.argsort(_row_perm(HIDDEN, gated=False)).to(x.device) + y = torch.zeros(m, TOP_K, HIDDEN, dtype=torch.float32, device=x.device) + for e in torch.unique(ids).tolist(): + w31 = _nvfp4_expert(layer.w3_w1_weight[e], layer.w3_w1_weight_scale[e], inv31) + w2 = _nvfp4_expert(layer.w2_weight[e], layer.w2_weight_scale[e], inv2) + esc = layer.escal[:, e] + for t, k in (ids == e).nonzero().tolist(): + acc1 = _gemm(xr[t : t + 1], w31)[0] + # The buffer holds [w3 (up); w1 (gate)]. + h = _fc1_epilogue(acc1[:inter], acc1[inter:], esc) + y[t, k] = _bf16(_gemm(h.unsqueeze(0), w2)[0] * esc[5]) + acc = torch.zeros(m, HIDDEN, dtype=torch.float64, device=x.device) + for k in range(TOP_K): + acc = (acc + weights[:, k : k + 1].double() * y[:, k].double()).float().double() + routed = acc.float() + if not shared: + return routed.to(torch.bfloat16) + return (_shared_expert(x, layer.shared_mxfp8, shared_inter=shared_inter) + _bf16(routed)).to( + torch.bfloat16 + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/runtime.py new file mode 100644 index 000000000000..4d92c1b7790d --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/runtime.py @@ -0,0 +1,232 @@ +# 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. +"""Shared MoE operand preparation; targets retain policy and cache state.""" + +from typing import Optional + +import torch + +HIDDEN = 6144 +NUM_EXPERTS = 128 + + +def _view(t: torch.Tensor, align: int): + from cutlass.cute.runtime import from_dlpack + + return from_dlpack(t.detach(), assumed_align=align) + + +def is_supported(device: Optional[torch.device] = None) -> tuple: + if not torch.cuda.is_available(): + return False, "no CUDA" + major, _ = torch.cuda.get_device_capability(device) + if major != 10: + return False, "needs the sm_100 family (tcgen05 block-scaled NVFP4 MMA)" + try: + import cutlass # noqa: F401 + except ImportError: + return False, "CuTe DSL (nvidia-cutlass-dsl) is not installed" + return True, "" + + +def ar_push_desc( + mc_ptr: int, flags_ptr: int, rank: int, world: int, max_tokens: int, device +) -> torch.Tensor: + """The ``arp`` descriptor that makes a call also push its bf16 output into a multicast all-reduce exchange (set A + of m3_mc_allreduce's workspace at ``mc_ptr``, slot half = the all-reduce CTA's call parity in ``flags_ptr``). + With it the kernel releases its PDL dependents only after every push; the default (enable 0) keeps the old + behaviour.""" + return torch.tensor( + [1, mc_ptr, flags_ptr, rank, world, max_tokens, 0, 0], dtype=torch.int64, device=device + ) + + +def shared_views(shared_weights, who, *, shared_inter: int): + """(gate_up, its scales, down, its scales) of the shared expert as flat uint8 views, checked.""" + if shared_weights is None: + raise ValueError(f"{who} (shared): shared_weights are required") + gu_w, gu_sf, dn_w, dn_sf = shared_weights + want = {"gate_up": (gu_w, (2 * shared_inter, HIDDEN)), "down": (dn_w, (HIDDEN, shared_inter))} + for name, (t, shape) in want.items(): + if t.dtype != torch.float8_e4m3fn or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError(f"{who} (shared): {name} weight must be contiguous e4m3 {shape}") + if ( + gu_sf.numel() != 2 * shared_inter * HIDDEN // 32 + or dn_sf.numel() != HIDDEN * shared_inter // 32 + ): + raise ValueError( + f"{who} (shared): the UE8M0 scales must be the swizzled [N, K / 32] buffers" + ) + return ( + gu_w.view(torch.uint8).view(-1), + gu_sf.view(torch.uint8).reshape(-1), + dn_w.view(torch.uint8).view(-1), + dn_sf.view(torch.uint8).reshape(-1), + ) + + +def pretile_shared(gu8, gusf8, dn8, dnsf8, n_cta: int = 24, *, shared_inter: int) -> torch.Tensor: + """The shared expert's weights as m3_moe's shared CTAs stream them (SH_PRETILE): per CTA s, per stage g, the 16 KB + shared-memory image (128-byte swizzle: 16-byte chunk j of row r at j ^ (r % 8)) of the stage's 128-row MMA tile and + its 512-byte instruction scale atom (row R's 4 codes at word 4 (R % 32) + R / 32). CTA s owns P = SHARED_INTER / 24 + = 16 positions in SPLIT = 128 / (2 P) = 4 K slices of GT = 48 / SPLIT = 12 k-tiles. Stages 0..GT-1: MMA row + R = P O + i (O = 2 q + u: slice q, u 0 gate / 1 up) = weight row u SHARED_INTER + P s + i at k-tile GT q + g; then + down tile j (rows 256 s + 128 j ..) at k-tile kt, j-major (SHARED_INTER / 128 k-tiles each). Built once per + layer.""" + dev = gu8.device + H_, I_ = HIDDEN, shared_inter + P = I_ // n_cta + split = 128 // (2 * P) + gt = H_ // 128 // split + dkt = I_ // 128 + gu = gu8.view(2 * I_, H_) + dn = dn8.view(H_, I_) + gkb, dkb = H_ // 32, I_ // 32 + + def sf_bytes(sf, n, b, kbn): # code of (row n, K block b), broadcast + return sf[ + ((n >> 7) * (kbn // 4) + (b >> 2)) * 512 + (n & 31) * 16 + ((n >> 5) & 3) * 4 + (b & 3) + ] + + r = torch.arange(128, device=dev) + q = r >> 5 + ri = r & 31 + octant = r // P + chunk = torch.arange(8, device=dev) + swz = chunk[None, :] ^ (r[:, None] & 7) # source chunk of image chunk j: j ^ (r % 8) + out = torch.empty( + n_cta, gt + 2 * dkt, 16384 + 512, dtype=torch.uint8, device=dev + ) # capacity: fixed: a workspace of the kernel's fixed geometry + for s in range(n_cta): + # gate_up: MMA row R -> weight row, slice + nrow = (octant & 1) * I_ + P * s + r % P + sl = octant >> 1 + for g in range(gt): + kt = sl * gt + g # per row + tile = gu[ + nrow[:, None], (kt[:, None] * 128 + torch.arange(128, device=dev)[None, :]) + ] # [128, 128] + img = tile.view(128, 8, 16)[r[:, None], swz] # image chunk j = tile chunk j ^ (r % 8) + out[s, g, :16384] = img.reshape(-1) + kb = torch.arange(4, device=dev) + codes = sf_bytes(gusf8, nrow[:, None], kt[:, None] * 4 + kb[None, :], gkb) # [128, 4] + atom = torch.empty(512, dtype=torch.uint8, device=dev) + atom[(ri[:, None] * 16 + q[:, None] * 4 + kb[None, :]).reshape(-1)] = codes.reshape(-1) + out[s, g, 16384:] = atom + for gd in range(2 * dkt): + jt, kt = gd // dkt, gd % dkt + row0 = 256 * s + 128 * jt + tile = dn[row0 : row0 + 128, kt * 128 : (kt + 1) * 128] + img = tile.reshape(128, 8, 16)[r[:, None], swz] + out[s, gt + gd, :16384] = img.reshape(-1) + out[s, gt + gd, 16384:] = dnsf8[ + ((row0 >> 7) * (dkb // 4) + kt) * 512 : ((row0 >> 7) * (dkb // 4) + kt) * 512 + 512 + ] + return out.view(-1) + + +def fine_ranges(weights, total: int, *, stages: int): + """Contiguous FC1 item ranges [s_c, e_c) over the CTAs, sized as close to ``weights`` (ideal real-valued item + counts summing to ``total``) as the split-unit protocol allows: every unit (S1 items) has at most one CTA boundary + strictly inside it (its head and tail parts, summed by the tail's holder), so a range that starts inside a unit + runs at least to that unit's end. Returns the item counts.""" + n = len(weights) + items = [0] * n + s = 0 + acc = 0.0 + for c in range(n): + acc += weights[c] + e = total if c == n - 1 else min(total, max(s, int(acc + 0.5))) + if s % stages and e < total: # started inside a unit: run to (at least) its end + e = max(e, min(total, (s // stages + 1) * stages)) + items[c] = e - s + s = e + assert s == total + return items + + +class M3MoeLayer: + """One MoE layer on a :class:`M3MoeState`: ``layer(hidden)``. The expert buffers are read in + place, never copied; the per-expert scalars are packed once (layer constants: rebuild the layer after reloading + weights).""" + + shared_weights = None # the kernel's view of the shared expert (a shared state) + shared_mxfp8 = None # the shared expert's (gate_up, its scales, down, its scales) as given + + def __init__( + self, + state, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + bias, + router_weight=None, + fc31_input_scale=None, + *, + inter: int, + ): + e = NUM_EXPERTS + expected = { + "w3_w1_weight": (w3_w1_weight, (e, 2 * inter, HIDDEN // 2)), + "w3_w1_weight_scale": (w3_w1_weight_scale, (e, 2 * inter, HIDDEN // 16)), + "w2_weight": (w2_weight, (e, HIDDEN, inter // 2)), + "w2_weight_scale": (w2_weight_scale, (e, HIDDEN, inter // 16)), + } + for name, (t, shape) in expected.items(): + if t.element_size() != 1 or tuple(t.shape) != shape or not t.is_contiguous(): + raise ValueError( + f"m3_moe: {name} is {t.dtype} {tuple(t.shape)}, expected contiguous 1-byte {shape}" + ) + self.state = state + self.w3_w1_weight = w3_w1_weight.view(torch.uint8) + self.w3_w1_weight_scale = w3_w1_weight_scale.view(torch.uint8) + self.w2_weight = w2_weight.view(torch.uint8) + self.w2_weight_scale = w2_weight_scale.view(torch.uint8) + self.escal = torch.stack( + [fc31_alpha, fc31_scale_c, act_alpha, act_beta, act_clamp, fc2_alpha] + ).float() + self.escal = self.escal.reshape(6, e).contiguous() + self.bias = bias.float().contiguous() + if ( + router_weight is None + or fc31_input_scale is None + or router_weight.dtype != torch.float32 + or tuple(router_weight.shape) != (e, HIDDEN) + or not router_weight.is_contiguous() + or fc31_input_scale.dtype != torch.float32 + or fc31_input_scale.numel() != 1 + ): + raise ValueError( + "m3_moe (fused router): needs router_weight f32 [128, 6144] and fc31_input_scale f32 [1]" + ) + self.router_weight = router_weight + self.fc31_input_scale = fc31_input_scale + + def __call__( + self, + hidden: torch.Tensor, + out: Optional[torch.Tensor] = None, + arp: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """``hidden``: contiguous bf16 [M, 6144]. ``arp`` (int64 [8], see :func:`ar_push_desc`): also push the output + into an all-reduce exchange.""" + return self.state.run(self, hidden, out, arp=arp) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/state.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/state.py new file mode 100644 index 000000000000..b3cead770713 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/moe/state.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Common MoE layer assembly and launches with target-owned compile caches.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable, ClassVar + +import torch + +from .runtime import HIDDEN, M3MoeLayer, _view, shared_views + + +@dataclass(frozen=True) +class MoeStateBindings: + """Geometry, packing and existing compile-cache policy selected by a target.""" + + inter: int + shared_inter: int + engine_path: str + compiled: dict[tuple, object] + pretiled: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor] + cache_module_config: bool + + +class M3MoeStateBase: + """The targets allocate their own state and retain their own shape/tactic choices.""" + + BINDINGS: ClassVar[MoeStateBindings] + + def layer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight=None, + fc31_input_scale=None, + shared_weights=None, + ) -> "M3MoeLayer": + """One MoE layer's operands (the TRTLLM-Gen NVFP4 expert buffers, read in place, never copied). + ``shared_weights`` (a ``shared`` state): gate_up e4m3 [2 * shared_inter, 6144], + its swizzled UE8M0 scales, down e4m3 [6144, shared_inter], and its swizzled scales. + """ + layer = M3MoeLayer( + self, + w3_w1_weight, + w3_w1_weight_scale, + w2_weight, + w2_weight_scale, + fc31_alpha, + fc31_scale_c, + act_alpha, + act_beta, + act_clamp, + fc2_alpha, + e_score_correction_bias, + router_weight, + fc31_input_scale, + inter=self.BINDINGS.inter, + ) + if self.shared: + layer.shared_weights = shared_views( + shared_weights, "m3_moe", shared_inter=self.BINDINGS.shared_inter + ) + layer.shared_mxfp8 = tuple(shared_weights) + if self.mod.SH_PRETILE: + # (gate_up, pre-tiled stages, down, dummy): the kernel streams the pre-tiled copy + gu8, _, dn8, _ = layer.shared_weights + layer.shared_weights = ( + gu8, + self.BINDINGS.pretiled(*layer.shared_weights), + dn8, + self._dummy.view(torch.uint8), + ) + return layer + + def run(self, layer, hidden, out=None, arp=None) -> torch.Tensor: + m = self.m + if ( + hidden is None + or hidden.dtype != torch.bfloat16 + or tuple(hidden.shape) != (m, HIDDEN) + or not hidden.is_contiguous() + ): + raise ValueError( + f"m3_moe (fused router): expects contiguous bf16 hidden states [{m}, {HIDDEN}]" + ) + fz = ( + hidden.view(-1).view(torch.int32), + layer.router_weight.view(-1), + layer.fc31_input_scale.view(-1), + self.lgbuf, + ) + lgv, xq, xs = self._dummy.view(torch.float32), self._dummy, self._dummy + if out is None: + out = torch.empty(m, HIDDEN, dtype=torch.bfloat16, device=self.device) + sh = layer.shared_weights if self.shared else (self._dummy.view(torch.uint8),) * 4 + args = ( + layer.w3_w1_weight, + layer.w3_w1_weight_scale.view(-1), + layer.w2_weight, + layer.w2_weight_scale.view(-1), + lgv, + layer.bias, + xq, + xs, + layer.escal, + self.plan.view(-1), + self.partials, + self.hbuf32, + self.epochs, + out.view(-1), + self.route_dbg.view(-1), + self.tdbg.view(-1), + *fz, + *sh, + self.hx, + self.dnx, + self._ar_off if arp is None else arp, + layer.w3_w1_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + layer.w2_weight_scale.view(torch.uint8).reshape(-1).view(torch.int32), + self.img32, + self.cnt, + ) + if self._fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("m3_moe compiles on its first call: call it once before capture") + import cutlass.cute as cute + + aligns = [ + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 4, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + 16, + ] + sig = [_view(t, a) for t, a in zip(args, aligns)] + # Preserve each target's cache-key policy: the module config for TP2/4, self.config for TP8. + # Each binding retains its target-owned cache dictionary. + key = (self.BINDINGS.engine_path,) + tuple( + sorted( + ( + self.mod.M3_CONFIG if self.BINDINGS.cache_module_config else self.config + ).items() + ) + ) + fn = self.BINDINGS.compiled.get(key) + if fn is None: + fn = self.BINDINGS.compiled[key] = cute.compile( + self.mod.m3_moe, + *sig, + cute.runtime.make_fake_stream(), + options="--enable-tvm-ffi", + ) + self._fn = fn + self._fn(*args, torch.cuda.current_stream(self.device).cuda_stream) + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/accept_workspace.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/accept_workspace.py new file mode 100644 index 000000000000..7936763c67ed --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/accept_workspace.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Collective acceptance workspace allocation; each target retains its type and cache.""" + +from __future__ import annotations + +from typing import Callable, TypeVar + +import torch + +from tensorrt_llm.mapping import Mapping + +_WorkspaceT = TypeVar("_WorkspaceT") + + +def create_accept_workspace( + mapping: Mapping, + device: torch.device | str, + slot_rows: int, + workspace_factory: Callable[..., _WorkspaceT], +) -> _WorkspaceT: + """Exchange slots ([2 halves][slot_rows rows][world] uint64) behind one multicast mapping over the TP group, the + local key accumulators [slot_rows] and the (finished CTAs, call) counters. Collective over the TP group; call + outside CUDA-graph capture, in the same order on every rank.""" + words = 2 * slot_rows * max(mapping.tp_size, 1) + if mapping.tp_size > 1: + from tensorrt_llm._torch.distributed.ops import ( + _make_mnnvl_mcast_buffer, + _mnnvl_workspace_all_succeeded, + ) + from tensorrt_llm._utils import mpi_comm + + session = mpi_comm() + session_group = session.Get_group() + group = session_group.Incl(mapping.tp_group) + session_group.Free() + comm = session.Create_group(group) + group.Free() + if not _mnnvl_workspace_all_succeeded(comm, True): + comm.Free() + raise RuntimeError("exchange workspace creation failed on another rank") + handle = _make_mnnvl_mcast_buffer( + comm, max(words * 8, 4096), mapping, mapping.is_multi_node() + ) + uc = handle.get_uc_buffer(mapping.tp_rank, (words,), torch.int64, 0) + mc = handle.get_mc_buffer((words,), torch.int64, 0) + uc.zero_() + torch.cuda.synchronize() + if not _mnnvl_workspace_all_succeeded(comm, True): + raise RuntimeError("exchange workspace creation failed on another rank") + mc_ptr, rank = (mc.data_ptr(), mapping.tp_rank) + else: + uc = torch.zeros(words, dtype=torch.int64, device=device) + mc_ptr, rank, handle, comm = (uc.data_ptr(), 0, None, None) + return workspace_factory( + uc=uc, + mc_ptr=mc_ptr, + keys=torch.zeros(slot_rows, dtype=torch.int64, device=uc.device), + counters=torch.zeros(2, dtype=torch.int32, device=uc.device), + rank=rank, + world_size=max(mapping.tp_size, 1), + handle=handle, + comm=comm, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/acceptance_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/acceptance_reference.py new file mode 100644 index 000000000000..7e42720613fd --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/acceptance_reference.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared Torch mirrors of the SM107 shard-pick and strict-acceptance epilogues.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +_LOW32 = 0xFFFFFFFF + + +def ordered_bits(v: torch.Tensor) -> torch.Tensor: + """The kernel's order-preserving uint32 image of fp32 values, as int64 (-0.0 maps to +0.0's image: they tie).""" + b = v.float().contiguous().view(torch.int32).to(torch.int64) & _LOW32 + b = torch.where(b == 0x80000000, torch.zeros_like(b), b) + return torch.where(b >= 0x80000000, _LOW32 - b, b | 0x80000000) + + +def make_keys(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """Exchange keys of (fp32 value, global token index) pairs: larger value wins, ties take the lower index.""" + return (ordered_bits(values) << 32) | (_LOW32 - index.to(torch.int64)) + + +def unsigned_max(keys: torch.Tensor, dim: int) -> torch.Tensor: + """Max of int64-stored uint64 keys (the kernel compares them unsigned).""" + sign = torch.iinfo(torch.int64).min + return (keys ^ sign).max(dim=dim).values ^ sign + + +def key_token(key: torch.Tensor) -> torch.Tensor: + return (_LOW32 - (key & _LOW32)).to(torch.int32) + + +def local_argmax(logits: torch.Tensor, vocab_offset: int) -> tuple[torch.Tensor, torch.Tensor]: + """(max value, lowest global index reaching it) per row of a [M, Vs] fp32 shard.""" + best = logits.max(dim=1).values + idx = (logits == best[:, None]).float().argmax(dim=1) # first (lowest) index at the max + return best, idx.to(torch.int64) + vocab_offset + + +def exchange_tokens( + own_keys: torch.Tensor, + ws_uc: torch.Tensor, + counters: torch.Tensor, + rank: int, + world: int, + slot_rows: int, +) -> torch.Tensor: + """Global tokens of the cross-rank pick: this rank's keys against the other ranks' keys already in its slots. + + Slot (half * slot_rows + t) * world + r of ``ws_uc`` holds rank r's key of row t; half = call counter & 1. A + single-device run stands in for the other ranks by writing their keys there before the call. + """ + m = own_keys.shape[0] + half = int(counters[1].item()) & 1 + slots = ( + ws_uc.view(torch.int64)[half * slot_rows * world : (half * slot_rows + m) * world] + .view(m, world) + .clone() + ) + slots[:, rank] = own_keys + return key_token(unsigned_max(slots, dim=1)) + + +def strict_accept( + tokens: torch.Tensor, + draft: torch.Tensor, + batch_idx: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], +) -> dict[str, torch.Tensor]: + """The strict acceptance epilogue from the verify-row tokens [nreq * (D + 1)] (request-major). + + accepted[b] = tokens of request b's rows; num_accepted[b] = 1 + the matching prefix of draft[b] (overridden by + force_vec[b], else by force_total > 0); gather_ids[b] = batch_idx[b] * (D + 1) + num_accepted[b] - 1, and + next_new[b, 0] (when given) = the token at that row. Returns the outputs plus ``next_new`` as written. + """ + nreq, d = draft.shape + k1 = d + 1 + tokens = tokens.to(torch.int32) + accepted = tokens.view(nreq, k1).clone() + match = (draft.to(torch.int32) == accepted[:, :d]).to(torch.int32) + num = 1 + match.cumprod(dim=1).sum(dim=1) + if force_vec is not None: + num = force_vec[:nreq].to(num.dtype) + elif force_total > 0: + num = torch.full_like(num, force_total) + gather = batch_idx.to(torch.int64) * k1 + num.to(torch.int64) - 1 + out = {"accepted": accepted, "num_accepted": num.to(torch.int32), "gather_ids": gather} + if next_new is not None: + nn = next_new.clone() + nn[:, 0] = tokens[gather] + out["next_new"] = nn + return out diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/draft_head_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/draft_head_reference.py new file mode 100644 index 000000000000..3541092bd77a --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/draft_head_reference.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared lossless BF16 fragment layout and draft-head reference arithmetic.""" + +from __future__ import annotations + +from contextlib import AbstractContextManager +from typing import Callable, Optional, Protocol + +import torch + + +class DraftHeadTarget(Protocol): + from_fragments: Callable[..., torch.Tensor] + decode_coded: Callable[..., torch.Tensor] + head_weight: Callable[..., torch.Tensor] + true_fp32_matmul: Callable[..., AbstractContextManager[None]] + + +def bf16_fragments(state: DraftHeadTarget, /, weight: torch.Tensor) -> torch.Tensor: + """BF16 [N, K] -> the same values in A-fragment order [N / 16, K / 16, 256]. + + One lane's 16 bytes of a 16-row tile's k16 step are the m16n8k16 A fragment: element (row h * 8 + g, k ks * 16 + + hk * 8 + jj * 2 + e) goes to [tile][ks][g][jj][hk][h][e] (lane = g * 4 + jj). + """ + n, k = weight.shape + assert n % 16 == 0 and k % 16 == 0 and (weight.dtype == torch.bfloat16), (n, k, weight.dtype) + q = weight.view(n // 16, 2, 8, k // 16, 2, 4, 2).permute(0, 3, 2, 5, 4, 1, 6).contiguous() + return q.view(n // 16, k // 16, 256) + + +def from_fragments(state: DraftHeadTarget, /, wf: torch.Tensor) -> torch.Tensor: + """The inverse of ``bf16_fragments``: [N / 16, K / 16, 256] -> [N, K].""" + tiles, ks, _ = wf.shape + return ( + wf.view(tiles, ks, 8, 4, 2, 2, 2).permute(0, 5, 2, 1, 4, 3, 6).reshape(tiles * 16, ks * 16) + ) + + +def head_weight( + state: DraftHeadTarget, + /, + wf: torch.Tensor, + k: int, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The BF16 head shard [N, K] the kernel multiplies by: the fragment-ordered copy, or (``exc_off`` given) the + 12-bit coded stream with its exception lists decoded.""" + if exc_off is None: + return state.from_fragments(wf) + n = (exc_off.numel() - 1) * 16 + return state.decode_coded(wf, exp_base, exc_off, exc_pos, exc_w, n, k) + + +def shard_logits( + state: DraftHeadTarget, + /, + x: torch.Tensor, + wf: torch.Tensor, + exc_off: Optional[torch.Tensor] = None, + exc_pos: Optional[torch.Tensor] = None, + exc_w: Optional[torch.Tensor] = None, + exp_base: int = 0, +) -> torch.Tensor: + """The in-tree head's BF16 logits of the rows x [M, K] against the shard (``head_weight``): the fp32 product, + rounded once to BF16 (the BF16 GEMM's output), as fp32.""" + w = state.head_weight(wf, x.shape[1], exc_off, exc_pos, exc_w, exp_base) + with state.true_fp32_matmul(): + logits = x.float() @ w.float().t() + return logits.to(torch.bfloat16).float() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/fold_tail.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/fold_tail.py new file mode 100644 index 000000000000..3890741abdb9 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/fold_tail.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared speculative tail guards and launch; targets bind their original native entry.""" + +from typing import Any, Callable + +import torch +import triton + + +def fold_tail( + ctrl: torch.Tensor, + slots: torch.Tensor, + rows: int, + src_new: torch.Tensor, + dst_new: torch.Tensor, + src_next: torch.Tensor, + dst_next: torch.Tensor, + src_draft: torch.Tensor, + dst_draft: torch.Tensor, + src_lens: torch.Tensor, + dst_lens: torch.Tensor, + ids: torch.Tensor, + draft_ids: torch.Tensor, + pos_off: torch.Tensor, + kv_off: torch.Tensor, + tokens_per_step: int, + cpp: bool = False, + *, + native_launch: Callable[..., None], + triton_kernel: Any, +) -> bool: + """Launch over ``rows`` programs (the first ``ctrl[0]`` write); False, having launched nothing, when the inputs + are outside what the kernel covers. Store views are 2-D [slots, width] (any strides), sources [rows, width].""" + two = (src_new, dst_new, src_next, dst_next, src_draft, dst_draft) + one = (src_lens, dst_lens, ids, draft_ids, pos_off, kv_off) + if any((t.dim() != 2 for t in two)) or any((t.dim() != 1 or t.stride(0) != 1 for t in one)): + return False + if ( + any((t.dtype != torch.int32 or not t.is_cuda for t in two + one + (ctrl,))) + or slots.dtype != torch.long + ): + return False + t_step = tokens_per_step + d_step = t_step - 1 + if ( + rows <= 0 + or slots.numel() < rows + or any((t.shape[0] < rows for t in (src_new, src_next, src_draft, src_lens))) + or (ids.numel() < rows * t_step) + or (pos_off.numel() < rows * t_step) + or (draft_ids.numel() < rows * d_step) + or (kv_off.numel() < rows) + ): + return False + widths = (src_new.shape[1], dst_new.shape[1], src_next.shape[1], dst_next.shape[1]) + widths += (src_draft.shape[1], dst_draft.shape[1], t_step) + if cpp and max(widths) <= 32: + native_launch( + ctrl, + slots, + rows, + src_new, + dst_new, + src_next, + dst_next, + src_draft, + dst_draft, + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + t_step, + ) + return True + block = triton.next_power_of_2(max(widths)) + triton_kernel[rows,]( + ctrl, + slots, + src_new, + src_new.stride(0), + src_new.stride(1), + src_new.shape[1], + dst_new, + dst_new.stride(0), + dst_new.stride(1), + dst_new.shape[1], + src_next, + src_next.stride(0), + src_next.stride(1), + src_next.shape[1], + dst_next, + dst_next.stride(0), + dst_next.stride(1), + dst_next.shape[1], + src_draft, + src_draft.stride(0), + src_draft.stride(1), + src_draft.shape[1], + dst_draft, + dst_draft.stride(0), + dst_draft.stride(1), + dst_draft.shape[1], + src_lens, + dst_lens, + ids, + draft_ids, + pos_off, + kv_off, + t_step, + d_step, + block, + num_warps=1, + ) + return True diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_capture_norm.py similarity index 71% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_capture_norm.py index a1993924a112..77b15bcf64f4 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_capture_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_capture_norm.py @@ -10,19 +10,12 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_chunk_norm import ( + rmsnorm_rows, +) -from .m3_chunk_norm import rmsnorm_rows - - -class M3CaptureNorm(OpWrapper): - """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture - buffer slice whose row stride may exceed H. Returns ``dst``.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the - # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. - CELLS: tuple[Cell, ...] = ( +CELLS_BY_TP = { + 2: ( *( Cell( why=f"{4 * n} verify rows of {n} requests (b4 / b8), layer slot {n % 3}, residual added", @@ -41,7 +34,46 @@ class M3CaptureNorm(OpWrapper): why="no residual (a capture point without one): plain RMSNorm", spec=dict(m=8, slot=1, residual=False), ), - ) + ), + 4: ( + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ), + 8: ( + Cell( + why="4 verify rows of one request, layer slot 0, residual added", + spec=dict(m=4, slot=0, residual=True), + ), + Cell( + why="cap 8 (32 rows) into the last layer slot", spec=dict(m=32, slot=2, residual=True) + ), + Cell( + why="no residual (a capture point without one): plain RMSNorm", + spec=dict(m=8, slot=1, residual=False), + ), + ), +} + + +class M3CaptureNorm(OpWrapper): + """``dst [M, H] = rmsnorm(bf16(h + r)) * w`` (``r`` optional), the rounding chunk_norm applies, into a capture + buffer slice whose row stride may exceed H. Returns ``dst``.""" + + ARCHS = frozenset({Arch.SM_107}) + CELLS_BY_TP = CELLS_BY_TP + + # The target's captured layers (hidden 6144) into the 3-layer capture buffer (row stride 3 x 6144), over the + # verify rows of draft_len 3 (4 .. 32 rows); the residual is present for every captured decoder layer. + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ fp32 statistic (block sum, rsqrtf), one bf16 rounding after the bf16 sum h + r; the band states the occasional diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_chunk_norm.py similarity index 76% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_chunk_norm.py index 453edec27764..b7b859c31b5e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_chunk_norm.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_chunk_norm.py @@ -21,15 +21,8 @@ def rmsnorm_rows(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Ten return (xf * r * weight.float()).to(torch.bfloat16) -class M3ChunkNorm(OpWrapper): - """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may - have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" - - ARCHS = frozenset({Arch.SM_107}) - - # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers - # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. - CELLS: tuple[Cell, ...] = ( +CELLS_BY_TP = { + 2: ( *( Cell( why=f"{n} requests' {4 * n} verify rows (b4 / b8), dense [{4 * n}, 3 x 6144]", @@ -49,7 +42,44 @@ class M3ChunkNorm(OpWrapper): why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", spec=dict(m=32, c=3, h=6144, extra_cols=6144), ), - ) + ), + 4: ( + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ), + 8: ( + Cell( + why="one request's 4 verify rows, dense [4, 3 x 6144]", + spec=dict(m=4, c=3, h=6144, extra_cols=0), + ), + Cell( + why="two requests' 8 verify rows (b2) from the capture buffer, dense [8, 3 x 6144]", + spec=dict(m=8, c=3, h=6144, extra_cols=0), + ), + Cell( + why="cap 8 (32 rows) read from the capture buffer with a fourth (final-hidden) slot in the row stride", + spec=dict(m=32, c=3, h=6144, extra_cols=6144), + ), + ), +} + + +class M3ChunkNorm(OpWrapper): + """``dst[r, c*H:(c+1)*H] = rmsnorm(src[r, c*H:(c+1)*H]) * w[c]`` for the C chunks of ``w [C, H]``; ``src`` may + have a row stride wider than C * H (the capture buffer). Writes ``dst [M, C*H]`` and returns it.""" + + ARCHS = frozenset({Arch.SM_107}) + + # The Eagle3 capture buffer of the target (3 captured layers x hidden 6144, buffer rows wider than the layers + # when the final hidden is also captured), over the verify rows of draft_len 3: 4 / 8 / 32 rows. + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ fp32 statistic per chunk (block sum, rsqrtf), one bf16 rounding: an output moves by one bf16 ulp now and then diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_draft_input.py similarity index 100% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_draft_input.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_draft_input.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_force_vec.py similarity index 96% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_force_vec.py index a192e89c5c76..192474b610ad 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_force_vec.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_force_vec.py @@ -7,8 +7,10 @@ import torch from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -from ..m3_force_vec import force_vec_kernel, force_vec_req_kernel +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.spec.m3_force_vec import ( + force_vec_kernel, + force_vec_req_kernel, +) class M3ForceVec(OpWrapper): diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_slot_scatter.py similarity index 87% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_slot_scatter.py index ff07c54de1f0..1b9822373410 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/ops/spec/m3_slot_scatter.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_slot_scatter.py @@ -61,12 +61,20 @@ def reference(self, slots, sa, da, sb, db, sl, dl, sd, dd): return out def is_valid(self, slots, sa, da, sb, db, sl, dl, sd, dd): - assert bool(((slots >= 0) & (slots < da.shape[0]) & (slots < dl.numel())).all()), ( - "slots must index the store" - ) + assert bool( + ( + (slots >= 0) + & (slots < da.shape[0]) + & (slots < db.shape[0]) + & (slots < dd.shape[0]) + & (slots < dl.numel()) + ).all() + ), "slots must index the store" # Two rows into one slot race (which one lands is unordered). assert slots.unique().numel() == slots.numel(), "slots must be distinct" - assert sa.shape[0] >= slots.numel() and sl.numel() >= slots.numel(), "a source row per slot" + assert min(sa.shape[0], sb.shape[0], sd.shape[0], sl.numel()) >= slots.numel(), ( + "a source row per slot" + ) m3_slot_scatter = M3SlotScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_step_gather.py similarity index 71% rename from tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py rename to tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_step_gather.py index a0da8b337574..d33551c53e33 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/ops/spec/m3_step_gather.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/m3_step_gather.py @@ -10,24 +10,8 @@ import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper - -class M3StepGather(OpWrapper): - """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all - int32:: - - ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) - draft[d0 + p*D + d] = next_draft[slots[p], d] - pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) - kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) - - Every other element of ``ids`` / ``draft`` is left as it was. - """ - - ARCHS = frozenset({Arch.SM_107}) - - # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the - # max batch); the engine's flat buffers sized for the max batch. - CELLS: tuple[Cell, ...] = ( +CELLS_BY_TP = { + 2: ( Cell( why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), @@ -48,7 +32,64 @@ class M3StepGather(OpWrapper): why="an unchanged batch of 4 previous requests (c4) in shuffled slots", spec=dict(p=4, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), ), - ) + ), + 4: ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + ), + 8: ( + Cell( + why="an unchanged batch of 8 previous requests in shuffled slots, nothing before them (b0 = 0)", + spec=dict(p=8, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="2 previous requests after 3 new ones: offsets ids0 / d0 / b0 nonzero, zero-filled offsets outside", + spec=dict(p=2, t=4, d=3, slots=8, b0=3, ids0=12, d0=9, extra_rows=3), + ), + Cell( + why="one request, the store's 2-D view of new tokens (beam dim squeezed)", + spec=dict(p=1, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0, two_d=True), + ), + Cell( + why="an unchanged batch of 3 previous requests (c3) in shuffled slots", + spec=dict(p=3, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + Cell( + why="an unchanged batch of 4 previous requests (c4) in shuffled slots", + spec=dict(p=4, t=4, d=3, slots=8, b0=0, ids0=0, d0=0, extra_rows=0), + ), + ), +} + + +class M3StepGather(OpWrapper): + """For the ``P`` previous requests (sequence slots ``slots``, ``T`` tokens and ``D`` draft tokens each), all + int32:: + + ids[ids0 + p*T + t] = new_tok[t, slots[p]] (new_tok [T', slots, (1)]) + draft[d0 + p*D + d] = next_draft[slots[p], d] + pos_off[i] = lens[ppi[i - b0*T]] for i in [b0*T, (b0+P)*T), else 0 (every i < pos_off.numel()) + kv_off[i] = lens[slots[i - b0]] - T for i in [b0, b0+P), else 0 (every i < kv_off.numel()) + + Every other element of ``ids`` / ``draft`` is left as it was. + """ + + ARCHS = frozenset({Arch.SM_107}) + + # draft_len 3: T = 4, D = 3; the store's [T, slots, 1] new tokens and [slots, D] next drafts (8 slots, the + # max batch); the engine's flat buffers sized for the max batch. + CELLS_BY_TP = CELLS_BY_TP + CELLS: tuple[Cell, ...] = tuple(cell for cells in CELLS_BY_TP.values() for cell in cells) note = """ Pure index copies: the reference is exact. The op does not bound slots, ppi or the destination offsets; is_valid @@ -141,6 +182,8 @@ def is_valid( ): # Reads through slots / ppi and writes at the offsets are unchecked: any of them out of range reads or # writes another buffer. + assert min(P, T, D, ids0, d0, b0) >= 0, "counts and destination offsets must be nonnegative" + assert slots.numel() >= P and ppi.numel() >= P * T, "index buffers too short" s = slots[:P] assert bool( ( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/sampling_reference.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/sampling_reference.py new file mode 100644 index 000000000000..d61de84f2db8 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/ops/spec/sampling_reference.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared unchanged sampling references and test-only input validation.""" + +from __future__ import annotations + +from typing import Any, Callable, Optional, Protocol + +import torch + + +class SamplingTarget(Protocol): + TEMP_EPSILON: float + nucleus_probs: Callable[..., torch.Tensor] + strict_accept: Callable[..., dict[str, torch.Tensor]] + + +def nucleus_probs( + state: SamplingTarget, /, logits: torch.Tensor, temperatures: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + """Per-row sampling distribution [rows, Vs] (fp32): w = exp(x / (T + 1e-6) - max), kept iff the mass strictly + above its weight is at most top_p * sum(w) (ties included), renormalized over the kept tokens.""" + inv = 1.0 / (temperatures.float() + state.TEMP_EPSILON) + x = logits.float() * inv[:, None] + w = torch.exp(x - x.max(dim=1, keepdim=True).values).double() + z = w.sum(dim=1, keepdim=True) + order = torch.sort(w, dim=1, descending=True) + sw = order.values + csum = torch.cumsum(sw, dim=1) - sw + first = torch.ones_like(sw, dtype=torch.bool) + first[:, 1:] = sw[:, 1:] != sw[:, :-1] + group = torch.cumsum(first.long(), dim=1) - 1 + starts = torch.zeros_like(csum) + starts.scatter_reduce_(1, group, csum, reduce="amin", include_self=False) + above = torch.gather(starts, 1, group) + keep_sorted = above <= top_ps.double()[:, None] * z + keep = torch.zeros_like(keep_sorted) + keep.scatter_(1, order.indices, keep_sorted) + p = torch.where(keep, w, torch.zeros_like(w)) + return (p / p.sum(dim=1, keepdim=True)).float() + + +def sampled_reference( + self: Any, + state: SamplingTarget, + /, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, +) -> dict: + """{"probs": the per-row nucleus distribution over this shard [rows, Vs] (global id = vocab_offset + + column), "accept": tokens -> the strict acceptance outputs}. The draw itself is random; ``seeds`` / + ``offsets`` select it and do not change the distribution. World 1: this shard is the whole vocabulary.""" + rows = logits.shape[0] + probs = state.nucleus_probs(logits, temperatures[:rows], top_ps[:rows]) + + def accept(drawn: torch.Tensor) -> dict[str, torch.Tensor]: + return state.strict_accept(drawn, draft, batch_idx, force_total, next_new, force_vec) + + return {"probs": probs, "accept": accept} + + +def validate_sampling( + self: Any, + state: SamplingTarget, + /, + logits: torch.Tensor, + vocab_offset: int, + temperatures: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws_uc: torch.Tensor, + ws_mc_ptr: int, + counters: torch.Tensor, + tokens: torch.Tensor, + rank: int, + world: int, + draft: torch.Tensor, + batch_idx: torch.Tensor, + acc_tokens: torch.Tensor, + num_acc: torch.Tensor, + gather_ids: torch.Tensor, + force_total: int, + next_new: Optional[torch.Tensor], + force_vec: Optional[torch.Tensor], + extra_attempts: Optional[torch.Tensor], + use_pdl: bool = True, + top_ks: Optional[torch.Tensor] = None, + min_ps: Optional[torch.Tensor] = None, +) -> None: + rows = logits.shape[0] + if top_ks is not None: + k = top_ks[:rows] + assert bool(((k <= 0) | (k >= logits.shape[1] * world)).all()), ( + "top-k rows must take the generic path" + ) + if min_ps is not None: + assert not bool((min_ps[:rows] > 0).any()), "min-p rows must take the generic path" + assert rows == draft.shape[0] * (draft.shape[1] + 1), ( + "logits rows must be nreq * (draft_len + 1)" + ) + for t in (temperatures, top_ps, seeds, offsets): + assert t.numel() >= rows and t.is_contiguous(), "per-row parameters must cover every row" + assert temperatures.dtype == torch.float32 and top_ps.dtype == torch.float32, ( + "fp32 sampling parameters" + ) + assert seeds.dtype == torch.int64 and offsets.dtype == torch.int64, ( + "int64 Philox seeds / offsets" + ) + assert int(counters[0].item()) == 0, "counters[0] (finished rows) must be zero between calls" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/reporting.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/reporting.py new file mode 100644 index 000000000000..e45ca82bc93b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3/reporting.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Common decode reporting policy with explicit target-owned state and callbacks.""" + +from __future__ import annotations + +from contextlib import AbstractContextManager +from typing import Callable, Optional, Protocol + + +class _Logger(Protocol): + info: Callable[[str], None] + warning: Callable[[str], None] + + +class ReportingTarget(Protocol): + """Each target keeps its own flags, lock, logged keys, exception and callbacks.""" + + _require: list[bool] + _lock: AbstractContextManager[bool] + _logged: set + _log_once: Callable[[Callable[[str], None], str, str, str], None] + missed: Callable[[str, Optional[int], str], None] + logger: _Logger + mpi_rank: Callable[[], int] + FastPathUnavailable: type[RuntimeError] + + +def set_require(state: ReportingTarget, /, require: bool) -> None: + """Whether a miss raises (``modeling_v2='require'``); set by the target at construction.""" + state._require[0] = bool(require) + + +def require_mode(state: ReportingTarget, /) -> bool: + return state._require[0] + + +def _log_once( + state: ReportingTarget, /, level: Callable[[str], None], path: str, reason: str, message: str +) -> None: + key = (path, reason) + with state._lock: + if key in state._logged: + return + state._logged.add(key) + level(message) + + +def missed(state: ReportingTarget, /, path: str, rows: Optional[int], reason: str) -> None: + """``path`` covers this step but cannot run: raise under require, else log once; the caller falls back.""" + at = "" if rows is None else f" at {rows} rows" + if state._require[0]: + raise state.FastPathUnavailable( + f"modeling_v2='require': the MiniMax-M3 decode path {path!r} covers this step{at} but cannot run: {reason}" + ) + state._log_once( + state.logger.warning, + path, + reason, + f"[M3 dispatch] {path}{at}: {reason}; falling back (logged once per reason)", + ) + + +def capture_missed(state: ReportingTarget, /, path: str, rows: Optional[int], what: str) -> None: + """``what`` is built on eager steps but a CUDA-graph capture found it missing (see ``missed``).""" + state.missed(path, rows, f"{what} is not built and a CUDA graph is being captured") + + +def taken(state: ReportingTarget, /, path: str, rows: Optional[int], what: str = "") -> None: + """``path`` ran on this rank at ``rows``: one rank-tagged info line per (path, rows, what), the engagement line + every GEN rank prints for each fast path it takes (``what``: an optional detail, e.g. the projection).""" + at = "" if rows is None else f" at {rows} rows" + detail = f" ({what})" if what else "" + state._log_once( + state.logger.info, + path, + f"taken: {rows}" + (f": {what}" if what else ""), + f"[M3 dispatch] rank {state.mpi_rank()}: {path} taken{at}{detail}", + ) + + +def uncovered(state: ReportingTarget, /, path: str, reason: str) -> None: + """The step is outside ``path``'s covered shapes: log once per (path, reason); the generic path runs.""" + state._log_once( + state.logger.info, path, reason, f"[M3 dispatch] {path} not taken: {reason} (logged once)" + ) + + +def reset_logged(state: ReportingTarget, /) -> None: + """Forget which (path, reason) pairs were logged.""" + with state._lock: + state._logged.clear() diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 1fb8f281ecad..334e73062b0a 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -1353,9 +1353,9 @@ def __init__(self, # AllReduceRunner.forward can use during the autotuner warm-up phase. extra_attrs = get_model_extra_attrs() # Declared by the model: never run or tune NCCL_SYMMETRIC (plain NCCL instead). - self._exclude_nccl_symmetric = bool( - extra_attrs - and extra_attrs.get("allreduce_exclude_nccl_symmetric", False)) + self._exclude_nccl_symmetric = bool(extra_attrs and extra_attrs.get( + "allreduce_exclude_nccl_symmetric", False)) or (os.environ.get( + "TRTLLM_AR_TUNE_NCCL_SYMMETRIC", "1") != "1") # Declared by the model: the MNNVL path runs its residual RMSNorm fusions on the CuTe DSL kernels. fused_norm = bool( extra_attrs diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index eb180f9b11ee..fcc3e89b335a 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -2934,11 +2934,12 @@ def preprocess_inputs( # RuntimeHooks.kv_lens_update_once, a step that corrects the KV lengths below (and then calls it again) calls # it only there: the second call rewrites everything the first one writes, and nothing reads the metadata # in between. - if attn_meta is not None and not ( - self.runtime_hooks.kv_lens_update_once - and self._kv_lens_corrected_in_preprocess(inputs, enable_spec_decode) - ): - attn_meta.on_update_kv_lens() + if attn_meta is not None: + if not ( + self.runtime_hooks.kv_lens_update_once + and self._kv_lens_corrected_in_preprocess(inputs, enable_spec_decode) + ): + attn_meta.on_update_kv_lens() if enable_spec_decode and not self._config.disable_overlap_scheduler: # When enabling overlap scheduler, the kv cache for draft tokens will diff --git a/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py index 7c3ff30b3c30..cbe8b9986b01 100644 --- a/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py +++ b/tests/integration/defs/accuracy/modeling_v2/test_minimax_m3.py @@ -26,6 +26,8 @@ """ import json +from collections.abc import Iterator +from contextlib import contextmanager import pytest @@ -44,6 +46,7 @@ from ...conftest import llm_models_root from ..accuracy_core import GSM8K, LlmapiAccuracyTestHarness +from ._env import scoped_env_overrides # The targets assert their own SM at construction: certification is per GPU architecture. skip_not_sm103 = pytest.mark.skipif( @@ -77,33 +80,38 @@ class _MiniMaxM3Target(LlmapiAccuracyTestHarness): EAGLE_PATH = f"{llm_models_root()}/MiniMax-M3-EAGLE3-GQA" TP = 4 - def _llm(self, eagle3: bool) -> LLM: + @contextmanager + def _llm(self, eagle3: bool) -> Iterator[LLM]: spec = ( Eagle3DecodingConfig(max_draft_len=3, speculative_model=self.EAGLE_PATH) if eagle3 else None ) - return LLM( - self.MODEL_PATH, - tensor_parallel_size=self.TP, - moe_expert_parallel_size=1, - kv_cache_config=KvCacheConfig( - free_gpu_memory_fraction=0.6, enable_block_reuse=False, dtype="nvfp4" - ), - sparse_attention_config=MiniMaxM3SparseAttentionConfig( - implementation="msa", indexer_kv_dtype="fp8", fuse_qkv_index_projection=True - ), - moe_config=MoeConfig(backend="TRTLLM"), - max_seq_len=4096, - max_batch_size=64, - cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=64), - speculative_config=spec, - enable_iter_perf_stats=eagle3, - max_stats_len=5000, - trust_remote_code=True, - modeling_v2="require", - env_overrides={MODELING_V2_VALIDATE_ENV: "1"}, - ) + with ( + scoped_env_overrides({MODELING_V2_VALIDATE_ENV: "1"}) as rank_env, + LLM( + self.MODEL_PATH, + tensor_parallel_size=self.TP, + moe_expert_parallel_size=1, + kv_cache_config=KvCacheConfig( + free_gpu_memory_fraction=0.6, enable_block_reuse=False, dtype="nvfp4" + ), + sparse_attention_config=MiniMaxM3SparseAttentionConfig( + implementation="msa", indexer_kv_dtype="fp8", fuse_qkv_index_projection=True + ), + moe_config=MoeConfig(backend="TRTLLM"), + max_seq_len=4096, + max_batch_size=64, + cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=64), + speculative_config=spec, + enable_iter_perf_stats=eagle3, + max_stats_len=5000, + trust_remote_code=True, + modeling_v2="require", + env_overrides=rank_env, + ) as llm, + ): + yield llm def _gsm8k(self, eagle3: bool) -> None: with self._llm(eagle3) as llm: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 20cd5290bf09..f47aa596797a 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -13,6 +13,9 @@ l0_cpu: orchestrator: mpi tests: - unittest/_torch/modeling_v2/models/minimax_m3/test_shared_kernel_contracts.py + - unittest/_torch/modeling_v2/models/minimax_m3/test_collective_launcher_cleanup.py + - unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py + - unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_spec_tail_contracts.py # MiniMax-M3 Rubin target unit tests. - unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/test_modeling_v2_m3_drafter_fullctx_cells.py # MiniMax-M3 target unit tests. diff --git a/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py b/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py index 319bb4abd8b6..9eb8dc252388 100644 --- a/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py +++ b/tests/unittest/_torch/distributed/test_allreduce_tune_nccl_symmetric.py @@ -10,10 +10,10 @@ from tensorrt_llm.functional import AllReduceStrategy -def _runner(): - runner = object.__new__(torch_custom_ops.AllReduceRunner) - runner.tp_size = 2 - return runner +def _runner(exclude: bool = False): + return torch_custom_ops.AllReduceRunner( + 2, [0, 1], torch.bfloat16, 0, 1e-6, False, exclude_nccl_symmetric=exclude + ) @pytest.mark.parametrize("rows", [8, 4096, 16384]) @@ -32,3 +32,88 @@ def test_nccl_symmetric_candidate_follows_knob( assert [t for t in stock if t != AllReduceStrategy.NCCL_SYMMETRIC.value] == [ t for t in tactics if t != AllReduceStrategy.NCCL_SYMMETRIC.value ] + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("exclude", [False, True]) +@pytest.mark.parametrize( + "cached", + [ + AllReduceStrategy.NCCL_SYMMETRIC.value, + AllReduceStrategy.NCCL.value, + AllReduceStrategy.ONESHOT.value, + ], +) +def test_cached_tactic_obeys_symmetric_policy(monkeypatch, enabled, exclude, cached): + monkeypatch.setattr(torch_custom_ops, "_AR_TUNE_NCCL_SYMMETRIC", enabled) + seen = [] + + def raw( + input_tensor, + residual, + norm_weight, + scale, + bias, + workspace, + group, + tactic, + op, + eps, + trigger_completion_at_end, + ): + seen.append(tactic) + return input_tensor + + monkeypatch.setattr(torch.ops.trtllm, "allreduce", raw) + x = torch.empty(2, 16, dtype=torch.bfloat16) + assert _runner(exclude).forward([x, None, None, None, None, None], tactic=cached) is x + forbidden = cached == AllReduceStrategy.NCCL_SYMMETRIC.value and (exclude or not enabled) + assert seen == [AllReduceStrategy.NCCL.value if forbidden else cached] + + +def test_environment_policy_separates_cache_identity(monkeypatch): + monkeypatch.setattr(torch_custom_ops, "_AR_TUNE_NCCL_SYMMETRIC", True) + enabled = _runner().unique_id() + explicit_exclusion = _runner(True).unique_id() + monkeypatch.setattr(torch_custom_ops, "_AR_TUNE_NCCL_SYMMETRIC", False) + disabled = _runner().unique_id() + assert disabled != enabled + assert disabled == explicit_exclusion + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("exclude", [False, True]) +def test_native_registration_rejects_cached_symmetric_tactic(monkeypatch, enabled, exclude): + monkeypatch.setattr(torch_custom_ops, "_AR_TUNE_NCCL_SYMMETRIC", enabled) + runner = _runner(exclude) + tuner = torch_custom_ops.AutoTuner.get() + cache = tuner.profiling_cache + monkeypatch.setattr(cache, "cache", {}) + inputs = [torch.empty(2, 16, dtype=torch.bfloat16), None, None, None, None, None] + shapes = tuple(tuner._get_input_sizes(inputs)) + buckets = runner.tuning_config.dynamic_tensor_specs[0].gen_tuning_buckets + for bucket in buckets: + bucket_shapes = (torch.Size((bucket, *shapes[0][1:])), *shapes[1:]) + key = cache.get_cache_key( + "trtllm::tunable_allreduce::allreduce", + runner, + bucket_shapes, + runner.tuning_config, + apply_map_to_tuning_buckets=False, + ) + cache.cache[key] = (0, AllReduceStrategy.NCCL_SYMMETRIC.value, 0.0) + registrations = [] + + def register( + input_tensor, residual, norm_weight, scale, bias, workspace, group, op, bucket, tactic + ): + registrations.append((bucket, tactic)) + + monkeypatch.setattr(torch.ops.trtllm, "register_allreduce_tactic", register) + runner._register_native_tactics(inputs) + expected = ( + AllReduceStrategy.NCCL_SYMMETRIC.value + if enabled and not exclude + else AllReduceStrategy.NCCL.value + ) + assert registrations == [(bucket, expected) for bucket in buckets] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/_m3_mxfp8_operands.py similarity index 100% rename from tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/_m3_mxfp8_operands.py rename to tests/unittest/_torch/modeling_v2/models/minimax_m3/_m3_mxfp8_operands.py diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/_m3_test_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/_m3_test_inputs.py new file mode 100644 index 000000000000..74a4994bdd16 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/_m3_test_inputs.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared deterministic inputs and unchanged numerical oracles for M3 tests.""" + +import zlib + +import torch + + +def seed_of(why: str) -> int: + return zlib.crc32(why.encode()) % 2**31 + + +def gen(seed: int) -> torch.Generator: + return torch.Generator(device="cuda").manual_seed(seed) + + +def num_sms() -> int: + return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + + +def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: + return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) + + +def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: + """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. + + Kernel and reference compute the same fp32 expression and round once; they differ only where a different + summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. + """ + assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 + o, r = out.float(), ref.float() + ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) + ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) + diff = (o - r).abs() + assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" + frac = (diff > 0).float().mean().item() + assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" + + +def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: + """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from + there on, so the acceptance counts vary across requests.""" + t = tokens.view(nreq, d + 1)[:, :d].clone() + for b in range(nreq): + n = (b + 2) % (d + 1) + t[b, n:] += 1 + return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_distributed_sampled_reference.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_distributed_sampled_reference.py new file mode 100644 index 000000000000..40ac165739d2 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_distributed_sampled_reference.py @@ -0,0 +1,255 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Real TP2/4/8 sampled verification against a full-vocabulary CPU oracle.""" + +import gc +import hashlib +import importlib +import json +import math +import os + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("distributed sampling requires CUDA", allow_module_level=True) +torch.cuda.set_device( + int(os.environ.get("SLURM_LOCALID", os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", "0"))) +) + +# Select the local GPU before TensorRT-LLM initializes its CUDA libraries. +from tensorrt_llm._utils import mpi_comm, mpi_rank, mpi_world_size # noqa: E402 +from tensorrt_llm.mapping import Mapping # noqa: E402 + +_WORLD = mpi_world_size() +_COMM = mpi_comm() +if _WORLD not in (2, 4, 8): + if os.environ.get("M3_REQUIRE_DISTRIBUTED_SAMPLE_REFERENCE") == "1": + raise RuntimeError(f"expected TP2/4/8 MPI ranks, got {_WORLD}") + pytest.skip("run with two, four or eight MPI ranks", allow_module_level=True) +_CAPABILITIES = _COMM.allgather(torch.cuda.get_device_capability()) +if any(capability != (10, 7) for capability in _CAPABILITIES): + if os.environ.get("M3_REQUIRE_DISTRIBUTED_SAMPLE_REFERENCE") == "1": + raise RuntimeError(f"expected SM107 on every rank: {_CAPABILITIES}") + pytest.skip("SM107 required on every rank", allow_module_level=True) + +# The hardware-specific module must be imported after collective eligibility. +_SA = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3." + f"minimax_m3_nvfp4__sm_107__tp{_WORLD}.ops.m3_sampled_accept" +) +_CALL_POLICY = {} +if _WORLD == 4: + _POLICY = importlib.import_module( + "tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3." + "minimax_m3_nvfp4__sm_107__tp4.modeling" + ) + assert _POLICY.SAMPLED_ACCEPT_SM107TP4 + _CALL_POLICY = dict( + one_round_k=_POLICY.SAMPLED_ACCEPT_1R_K, + prewait_params=_POLICY.SAMPLED_ACCEPT_PREWAIT_PARAMS, + early_trigger=_POLICY.SAMPLED_ACCEPT_EARLY_TRIGGER, + ilp_walk=_POLICY.SAMPLED_ACCEPT_ILP_WALK, + ) +_VOCAB = 200064 +_VS = _VOCAB // _WORLD +_DRAFT = 3 +_DRAWS = 1024 +_ALPHA = 1e-7 + + +def _collective_check(ok: bool, reason: str) -> None: + failures = _COMM.allgather(None if ok else reason) + assert not any(failures), f"distributed check failed: {failures}" + + +def _digest(tensors: tuple[torch.Tensor, ...]) -> str: + h = hashlib.sha256() + for tensor in tensors: + h.update(tensor.contiguous().view(torch.uint8).numpy().tobytes()) + return h.hexdigest() + + +def _probabilities(logits: torch.Tensor, temperatures: torch.Tensor, top_p: float) -> torch.Tensor: + """Full vocabulary, float64, with ties kept as one nucleus group.""" + result = torch.zeros_like(logits, dtype=torch.float64) + for row in range(logits.shape[0]): + scaled = logits[row].double() / (float(temperatures[row]) + 1e-6) + weights = torch.exp(scaled - scaled.max()) + values, order = weights.sort(descending=True) + unique, groups, counts = torch.unique_consecutive( + values, return_inverse=True, return_counts=True + ) + group_mass = unique * counts + strictly_above = group_mass.cumsum(0) - group_mass + keep = (strictly_above <= top_p * weights.sum())[groups] + indices = order[keep] + result[row, indices] = weights[indices] / weights[indices].sum() + return result + + +@pytest.mark.parametrize("nreq", (1, 2)) +@pytest.mark.parametrize("top_p", (0.5, 0.95)) +@pytest.mark.parametrize("graph_mode", (False, True), ids=("eager", "graph")) +def test_distributed_sampled_reference(nreq: int, top_p: float, graph_mode: bool) -> None: + rank = mpi_rank() + if _WORLD == 8: + _collective_check(_SA.ONE_ROUND_K == 82, "exercise production TP8 one-round K82") + device = torch.device("cuda", torch.cuda.current_device()) + rows = nreq * (_DRAFT + 1) + mapping = Mapping( + world_size=_WORLD, + rank=rank, + gpus_per_node=torch.cuda.device_count(), + tp_size=_WORLD, + ) + # Quantize deterministic host inputs before copying to CUDA: no scatter races. + local = torch.full((rows, _VS), -30.0, dtype=torch.bfloat16) + hot = (torch.arange(12) * 97 + rank * 43 + 5) % _VS + local[:, hot] = (torch.linspace(4.0, 0.0, 12) + rank * 0.5).to(local.dtype) + shards = _COMM.allgather(local.float().numpy()) + full_logits = torch.cat([torch.from_numpy(shard) for shard in shards], dim=1) + temperatures_cpu = torch.linspace(0.8, 1.2, rows) + probabilities = _probabilities(full_logits, temperatures_cpu, top_p) + masses = torch.stack([probabilities[:, i * _VS : (i + 1) * _VS].sum(1) for i in range(_WORLD)]) + _collective_check(bool((masses > 0).sum(0).min() >= 2), "nucleus must span multiple ranks") + _collective_check( + bool((masses.max(0).values - masses.min(0).values).max() > 0.02), + "test requires unequal shard masses", + ) + logits = local.to(device) + temperatures = temperatures_cpu.to(device) + top_ps = torch.full((rows,), top_p, device=device) + top_ks = torch.zeros(rows, dtype=torch.int32, device=device) + seeds = torch.arange(rows, dtype=torch.int64, device=device) * 7919 + 20261101 + offsets = torch.zeros(rows, dtype=torch.int64, device=device) + draft = torch.zeros(nreq, _DRAFT, dtype=torch.int32, device=device) + batch = torch.arange(nreq, dtype=torch.int32, device=device) + next_new = torch.full((nreq, _DRAFT + 1), -7, dtype=torch.int32, device=device) + ws = ( + _SA.SampledAcceptWorkspace.create(mapping, device, _CALL_POLICY["one_round_k"]) + if _WORLD == 4 + else _SA.SampledAcceptWorkspace.create(mapping, device) + ) + graph = None + + def call() -> tuple[torch.Tensor, ...]: + result = _SA.sampled_accept( + logits, + rank * _VS, + mapping, + draft, + batch, + temperatures, + top_ks, + top_ps, + None, + seeds, + offsets, + 0, + next_new=next_new, + ws=ws, + use_pdl=True, + **_CALL_POLICY, + ) + assert result is not None + return (*result, next_new, ws.tokens[:rows]) + + try: + warm = call() + torch.cuda.synchronize() + warm_cpu = tuple(t.cpu() for t in warm) + peers = _COMM.allgather(_digest(warm_cpu)) + _collective_check(len(set(peers)) == 1, "warm outputs disagree") + draft.copy_(warm[-1].reshape(nreq, _DRAFT + 1)[:, :_DRAFT]) + if graph_mode: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + outputs = call() + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + records = ( + torch.empty((_DRAWS, nreq, _DRAFT + 1), dtype=torch.int32, device=device), + torch.empty((_DRAWS, nreq), dtype=torch.int32, device=device), + torch.empty((_DRAWS, nreq), dtype=torch.int64, device=device), + torch.empty((_DRAWS, nreq, _DRAFT + 1), dtype=torch.int32, device=device), + torch.empty((_DRAWS, rows), dtype=torch.int32, device=device), + ) + for draw in range(_DRAWS): + offsets.fill_(draw * 1000) + if graph_mode: + graph.replay() + else: + outputs = call() + for destination, output in zip(records, outputs): + destination[draw].copy_(output) + torch.cuda.synchronize() + observed = tuple(t.cpu() for t in records) + peers = _COMM.allgather(_digest(observed)) + _collective_check(len(set(peers)) == 1, "ranks disagree on live outputs") + accepted, num_acc, gather_ids, new_tokens, tokens = observed + expected_accepted = tokens.reshape(_DRAWS, nreq, _DRAFT + 1) + matches = (expected_accepted[:, :, :_DRAFT] == draft.cpu()).int() + expected_num = 1 + matches.cumprod(2).sum(2) + expected_gather = torch.arange(nreq)[None, :] * (_DRAFT + 1) + expected_num - 1 + expected_new = tokens.gather(1, expected_gather.long()) + _collective_check(torch.equal(accepted, expected_accepted), "accepted tokens differ") + _collective_check(torch.equal(num_acc, expected_num.int()), "strict prefix counts differ") + _collective_check(torch.equal(gather_ids, expected_gather.long()), "gather IDs differ") + _collective_check(torch.equal(new_tokens[:, :, 0], expected_new), "next tokens differ") + _collective_check(bool((new_tokens[:, :, 1:] == -7).all()), "unused output columns changed") + _collective_check(ws.counters.cpu().tolist() == [0, _DRAWS + 1], "call parity differs") + distances = [] + for row in range(rows): + counts = torch.bincount(tokens[:, row].long(), minlength=_VOCAB).double() + expected = probabilities[row] + outside = counts[expected == 0].sum().item() + support = int((expected > 0).sum()) + # E[TV] <= .5 sqrt((K-1)/N); bounded-difference tail <= alpha per row. + band = 0.5 * math.sqrt((support - 1) / _DRAWS) + math.sqrt( + math.log(1 / _ALPHA) / (2 * _DRAWS) + ) + tv = 0.5 * (counts / _DRAWS - expected).abs().sum().item() + _collective_check(outside == 0, f"row{row}: sampled outside full-vocabulary nucleus") + _collective_check(tv <= band, f"row{row}: TV{tv} exceeds preregistered band{band}") + distances.append(dict(row=row, support=support, tv=tv, band=band)) + if rank == 0: + print( + json.dumps( + dict( + world=_WORLD, + nreq=nreq, + top_p=top_p, + graph=graph_mode, + draws=_DRAWS, + rank_agreement=True, + exact_acceptance=True, + both_halves=True, + one_round_k=_CALL_POLICY.get( + "one_round_k", getattr(_SA, "ONE_ROUND_K", None) + ), + prewait_params=_CALL_POLICY.get( + "prewait_params", getattr(_SA, "PREWAIT_PARAMS", None) + ), + early_trigger=_CALL_POLICY.get( + "early_trigger", getattr(_SA, "EARLY_TRIGGER", None) + ), + distances=distances, + ) + ), + flush=True, + ) + finally: + torch.cuda.synchronize() + graph = None + gc.collect() + _COMM.Barrier() + ws.uc = None + ws.handle = None + gc.collect() + _COMM.Barrier() + if ws.comm is not None: + ws.comm.Free() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_exchange.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_exchange.py new file mode 100644 index 000000000000..6994d602039f --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_exchange.py @@ -0,0 +1,225 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eight MPI ranks: TP8 SA-1R and two-round sampling must be bit-identical.""" + +import gc +import json +import os + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("the sampled exchange requires CUDA", allow_module_level=True) +# Select the MPI-local device before TensorRT-LLM initializes CUDA libraries. +torch.cuda.set_device( + int(os.environ.get("SLURM_LOCALID", os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", "0"))) +) + +from tensorrt_llm._utils import mpi_comm, mpi_rank, mpi_world_size # noqa: E402 +from tensorrt_llm.mapping import Mapping # noqa: E402 + +if mpi_world_size() != 8: + if os.environ.get("M3_REQUIRE_TP8_EXCHANGE") == "1": + raise RuntimeError(f"expected eight MPI ranks, got {mpi_world_size()}") + pytest.skip("run this regression with eight MPI ranks", allow_module_level=True) + +# Every rank makes the same eligibility decision before creating workspaces. +_capabilities = mpi_comm().allgather(torch.cuda.get_device_capability()) +if any(capability != (10, 7) for capability in _capabilities): + if os.environ.get("M3_REQUIRE_TP8_EXCHANGE") == "1": + raise RuntimeError(f"expected SM 10.7 on every rank, got {_capabilities}") + pytest.skip("the TP8 exchange requires SM 10.7 on every rank", allow_module_level=True) + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E402,E501 + m3_sampled_accept as sa, +) + +_VS = 25008 +_DRAFT = 3 +_CALLS = 4 +_CASES = [ + (4, "spiked", 0.95), + (16, "spiked", 0.5), + (48, "spiked", 0.95), + (4, "tail", 0.02), + (16, "tail", 0.5), + (48, "tail", 0.95), +] + + +def _logits(rows, mode, rank, device): + generator = torch.Generator(device=device).manual_seed(20261101 + rows * 31 + rank) + if mode == "tail": + return (torch.randn(rows, _VS, generator=generator, device=device) * 0.3 + rank * 0.5).to( + torch.bfloat16 + ) + out = torch.full((rows, _VS), -30.0, dtype=torch.bfloat16, device=device) + hot = (torch.arange(12, device=device) * 97 + rank * 43 + 5) % _VS + out[:, hot] = (torch.linspace(4.0, 0.0, 12, device=device) + rank * 0.5).to(out.dtype) + return out + + +def _snapshot(logits, mapping, draft, batch, temperatures, top_ks, top_ps, seeds, offsets, ws): + next_new = torch.full((draft.shape[0], _DRAFT + 1), -7, dtype=torch.int32, device=logits.device) + result = sa.sampled_accept( + logits, + mapping.tp_rank * _VS, + mapping, + draft, + batch, + temperatures, + top_ks, + top_ps, + None, + seeds, + offsets, + 0, + next_new=next_new, + ws=ws, + use_pdl=True, + ) + assert result is not None + return tuple(t.clone() for t in (*result, next_new, ws.tokens[: logits.shape[0]])) + + +def _assert_collective(condition, message): + failures = mpi_comm().allgather(None if condition else message) + assert not any(failures), f"rank comparison failures: {failures}" + + +def _assert_same(left, right): + _assert_collective( + len(left) == len(right) and all(torch.equal(a, b) for a, b in zip(left, right)), + "exchange variants differ", + ) + + +def test_comparison_failure_is_collective(): + left = torch.tensor([1], device="cuda") + right = torch.tensor([2 if mpi_rank() == 3 else 1], device="cuda") + with pytest.raises(AssertionError, match="exchange variants differ"): + _assert_same((left,), (right,)) + mpi_comm().Barrier() + + +@pytest.mark.parametrize("rows,mode,top_p", _CASES) +def test_one_round_exchange_matches_two_rounds(monkeypatch, rows, mode, top_p): + assert "one_round_k" in sa._op_args() and "decided" in sa._op_args() + rank = mpi_rank() + comm = mpi_comm() + device = torch.device("cuda", torch.cuda.current_device()) + mapping = Mapping(world_size=8, rank=rank, gpus_per_node=torch.cuda.device_count(), tp_size=8) + logits = _logits(rows, mode, rank, device) + nreq = rows // (_DRAFT + 1) + temperatures = torch.linspace(0.8, 1.2, rows, device=device) + top_ks = torch.zeros(rows, dtype=torch.int32, device=device) + top_ps = torch.full((rows,), top_p, device=device) + seeds = torch.arange(rows, device=device, dtype=torch.int64) * 17 + 20261101 + base_offsets = torch.arange(rows, device=device, dtype=torch.int64) + offsets = base_offsets.clone() + draft = torch.zeros(nreq, _DRAFT, dtype=torch.int32, device=device) + batch = torch.arange(nreq, dtype=torch.int32, device=device) + masses = comm.allgather(torch.logsumexp(logits[0].float() / temperatures[0], 0).item()) + assert max(masses) - min(masses) > 0.5, "the test needs unequal shard masses" + workspaces = {} + try: + for variant in (0, 82): + monkeypatch.setattr(sa, "ONE_ROUND_K", variant) + workspaces[variant] = sa.SampledAcceptWorkspace.create(mapping, device) + warm = [] + for variant in (0, 82): + monkeypatch.setattr(sa, "ONE_ROUND_K", variant) + warm.append( + _snapshot( + logits, + mapping, + draft, + batch, + temperatures, + top_ks, + top_ps, + seeds, + offsets, + workspaces[variant], + ) + ) + torch.cuda.synchronize() + _assert_same(*warm) + # Reuse the exact Philox inputs, but exercise every strict-accept prefix. + sampled = warm[0][-1].reshape(nreq, _DRAFT + 1) + draft.copy_(sampled[:, :_DRAFT]) + prefixes = torch.arange(nreq, device=device) % (_DRAFT + 1) + for request in range(nreq): + prefix = request % (_DRAFT + 1) + if prefix < _DRAFT: + draft[request, prefix] = (draft[request, prefix] + 1) % (_VS * 8) + observations = [] + for call in range(_CALLS): + offsets.copy_(base_offsets + call * 1000) + pair = [] + for variant in (0, 82): + monkeypatch.setattr(sa, "ONE_ROUND_K", variant) + pair.append( + _snapshot( + logits, + mapping, + draft, + batch, + temperatures, + top_ks, + top_ps, + seeds, + offsets, + workspaces[variant], + ) + ) + observations.append(pair) + torch.cuda.synchronize() + for call, pair in enumerate(observations): + _assert_same(*pair) + if call == 0: + _assert_collective( + torch.equal(pair[0][1], (prefixes + 1).to(torch.int32)), + "strict-accept prefixes differ", + ) + record = [tensor.cpu().tolist() for tensor in pair[0]] + peers = comm.allgather(record) + assert all(peer == peers[0] for peer in peers), "ranks disagree on sampled outputs" + for ws in workspaces.values(): + _assert_collective( + ws.counters.cpu().tolist() == [0, _CALLS + 1], + "workspace counter did not exercise both halves", + ) + decided = workspaces[82].decided.item() + if mode == "spiked": + _assert_collective(decided > 0, "the spiked case must exercise SA-1R decisions") + if rank == 0: + print( + json.dumps( + dict( + rows=rows, + mode=mode, + top_p=top_p, + world=8, + calls_per_variant=_CALLS + 1, + both_halves=True, + unequal_log_masses=masses, + decided=decided, + all_outputs_bitwise=True, + ) + ), + flush=True, + ) + finally: + torch.cuda.synchronize() + comm.Barrier() + # Release multicast memory while each workspace communicator is still live. + for ws in workspaces.values(): + ws.uc = None + ws.handle = None + gc.collect() + comm.Barrier() + for ws in workspaces.values(): + if ws.comm is not None: + ws.comm.Free() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py new file mode 100644 index 000000000000..360d82ce6797 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py @@ -0,0 +1,138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Compare eight-rank captured SA-1R with captured and eager two-round sampling.""" + +import gc +import json + +import pytest +import torch +from test_sampled_accept_tp8_exchange import ( + _CASES, + _DRAFT, + _assert_collective, + _assert_same, + _logits, + _snapshot, + sa, +) + +from tensorrt_llm._utils import mpi_comm, mpi_rank +from tensorrt_llm.mapping import Mapping + + +@pytest.mark.parametrize("rows,mode,top_p", _CASES) +def test_graph_exchange_matches_eager_two_rounds(monkeypatch, rows, mode, top_p): + rank = mpi_rank() + comm = mpi_comm() + device = torch.device("cuda", torch.cuda.current_device()) + mapping = Mapping(world_size=8, rank=rank, gpus_per_node=torch.cuda.device_count(), tp_size=8) + logits = _logits(rows, mode, rank, device) + nreq = rows // (_DRAFT + 1) + temperatures = torch.linspace(0.8, 1.2, rows, device=device) + top_ks = torch.zeros(rows, dtype=torch.int32, device=device) + top_ps = torch.full((rows,), top_p, device=device) + seeds = torch.arange(rows, device=device, dtype=torch.int64) * 17 + 20261101 + base_offsets = torch.arange(rows, device=device, dtype=torch.int64) + offsets = base_offsets.clone() + draft = torch.zeros(nreq, _DRAFT, dtype=torch.int32, device=device) + batch = torch.arange(nreq, dtype=torch.int32, device=device) + workspaces = {} + graphs = {} + try: + for key, variant in (("reference", 0), (0, 0), (82, 82)): + monkeypatch.setattr(sa, "ONE_ROUND_K", variant) + ws = workspaces[key] = sa.SampledAcceptWorkspace.create(mapping, device) + _snapshot( + logits, mapping, draft, batch, temperatures, top_ks, top_ps, seeds, offsets, ws + ) + torch.cuda.synchronize() + for variant in (0, 82): + monkeypatch.setattr(sa, "ONE_ROUND_K", variant) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream), torch.cuda.graph(graph, stream=stream): + captured = _snapshot( + logits, + mapping, + draft, + batch, + temperatures, + top_ks, + top_ps, + seeds, + offsets, + workspaces[variant], + ) + torch.cuda.current_stream().wait_stream(stream) + graphs[variant] = (graph, captured) + torch.cuda.synchronize() + comm.Barrier() + observations = [] + decisions = [] + for replay in range(4): + offsets.copy_(base_offsets + (replay + 20) * 1000) + monkeypatch.setattr(sa, "ONE_ROUND_K", 0) + reference = _snapshot( + logits, + mapping, + draft, + batch, + temperatures, + top_ks, + top_ps, + seeds, + offsets, + workspaces["reference"], + ) + captured_outputs = [] + for variant in (0, 82): + graph, captured = graphs[variant] + graph.replay() + captured_outputs.append(tuple(t.clone() for t in captured)) + if variant == 82: + decisions.append(workspaces[82].decided.clone()) + observations.append((reference, *captured_outputs)) + torch.cuda.synchronize() + for reference, captured_zero, captured_one in observations: + _assert_same(reference, captured_zero) + _assert_same(reference, captured_one) + peers = comm.allgather([t.cpu().tolist() for t in captured_one]) + assert all(peer == peers[0] for peer in peers), "captured ranks disagree" + for ws in workspaces.values(): + _assert_collective( + ws.counters.cpu().tolist() == [0, 5], "graph workspace counter mismatch" + ) + decided = sum(t.item() for t in decisions) + if mode == "spiked": + _assert_collective(decided > 0, "captured SA-1R must make round-zero decisions") + if rank == 0: + print( + json.dumps( + dict( + rows=rows, + mode=mode, + top_p=top_p, + world=8, + graph_replays=4, + both_halves=True, + decided=decided, + eager_vs_graph_bitwise=True, + ) + ), + flush=True, + ) + finally: + torch.cuda.synchronize() + comm.Barrier() + graphs.clear() + gc.collect() + for ws in workspaces.values(): + ws.uc = None + ws.handle = None + gc.collect() + comm.Barrier() + for ws in workspaces.values(): + if ws.comm is not None: + ws.comm.Free() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py index c6c91681361c..b56593f8d88d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/activation/test_modeling_v2_m3_silu_and_mul.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.activation.m3_silu_and_mul import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.activation.m3_silu_and_mul import ( # noqa: E501 m3_silu_and_mul as op, ) @@ -29,7 +29,7 @@ def _build(spec, seed): return (x, *OAI) if spec["oai"] else (x, None, None, None) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[2], ids=[c.why[:48] for c in op.CELLS_BY_TP[2]]) def test_certified_cells(cell) -> None: """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); @@ -43,7 +43,7 @@ def test_certified_cells(cell) -> None: def test_the_gate_sees_a_missing_clamp() -> None: """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" - spec = op.CELLS[0].spec + spec = op.CELLS_BY_TP[2][0].spec x, *oai = _build(dict(spec, m=64), 3) out = op(x, *oai) with pytest.raises(AssertionError): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py index e44a6333a68a..3f1da02ec104 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -9,7 +9,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 m3_fp8_indexer_qk_norm_rope as op, ) @@ -28,7 +28,7 @@ def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[2], ids=[c.why[:44] for c in op.CELLS_BY_TP[2]]) def test_certified_cells(cell) -> None: m, heads = cell.spec["m"], cell.spec["heads"] g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py index fabae6422ba1..577f62651e54 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_index_score.py @@ -19,6 +19,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_index_score import ( # noqa: E501 m3_index_score as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_index_score import ( # noqa: E501 + IndexScoreRunner as TargetIndexScoreRunner, +) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 @@ -63,8 +66,8 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert spec["heads"] == HEADS args = (idx_q, index_k, block_table, seq, score, qlen, spec["split_k"], spec["early"]) certifying(op) - ref = op.reference(*args) - op(*args) + ref = op.reference(*args, runner=TargetIndexScoreRunner) + op(*args, runner=TargetIndexScoreRunner) stock = _score(stride, idx_q.shape[0]) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, qlen @@ -78,7 +81,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_bf16_index_is_the_in_tree_scorer(early) -> None: idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) score, stock = _score(stride, 4), _score(stride, 4) - op(idx_q, index_k, block_table, seq, score, 4, 1536, early) + op(idx_q, index_k, block_table, seq, score, 4, 1536, early, runner=TargetIndexScoreRunner) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, 4 ) @@ -93,7 +96,17 @@ def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> Non for seq_lens, split in (((length,), 1536), ((length, max(qlen, length // 3)), 768)): idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) - op(idx_q, index_k, block_table, seq, score, qlen, split, True) + op( + idx_q, + index_k, + block_table, + seq, + score, + qlen, + split, + True, + runner=TargetIndexScoreRunner, + ) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, qlen ) @@ -105,4 +118,14 @@ def test_guard_refuses_a_request_past_its_block_table() -> None: seq = seq + 4 * PAGE with pytest.raises(AssertionError, match="past its block-table row"): with validating(op): - op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, 1536, True) + op( + idx_q, + index_k, + block_table, + seq, + _score(stride + 16, 1), + 1, + 1536, + True, + runner=TargetIndexScoreRunner, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py index 3ffe4321fac1..7bfd53f7286a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -10,7 +10,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 m3_qk_norm_rope_fp8 as op, ) @@ -57,7 +57,7 @@ def _args(spec, seed): ) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[2], ids=[c.why[:48] for c in op.CELLS_BY_TP[2]]) def test_certified_cells(cell) -> None: args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) before = args[0].clone() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py index 4a3fe4690227..0228723e9629 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_qkv_insert.py @@ -11,7 +11,7 @@ import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_qkv_insert import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qkv_insert import ( # noqa: E501 m3_qkv_insert as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py index b3b2d97b1497..2100ef16f7d9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_select_blocks.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_select_blocks import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_select_blocks import ( # noqa: E501 m3_select_blocks as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py index 1667f6034162..da6572f7c34f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_sparse_decode.py @@ -14,6 +14,9 @@ import torch import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.triton_sparse_decode_mq import ( # noqa: E501 + launch_mq_producer, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( modeling, # noqa: E501 ) @@ -27,9 +30,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.m3_sparse_decode import ( # noqa: E501 m3_sparse_decode as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.triton_sparse_decode_mq import ( # noqa: E501 - launch_mq_producer, -) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( minimax_m3_sparse_attn_decode as stock_decode, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py index 616eb6e1a7b3..12765c2d9d65 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_m3_tokmap_guard.py @@ -36,8 +36,17 @@ @pytest.fixture(scope="module", autouse=True) def _installed(): - guard.install(_DEV, 48, _TAKEN.append) - yield + # Preserve both the process-wide callable and this target's installation state. + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + thop, "trtllm_gen_generation_preprocess", thop.trtllm_gen_generation_preprocess + ) + for name in ("_original", "_max_rows", "_on_taken", "_on_missed"): + patch.setattr(guard, name, getattr(guard, name)) + patch.setattr(guard, "_maps", dict(guard._maps)) + patch.setattr(guard, "_offsets", dict(guard._offsets)) + guard.install(_DEV, 48, _TAKEN.append) + yield class _Step: diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py index 67e30c0143b4..0a1bdba7f112 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/attention/test_modeling_v2_msa_kv_update.py @@ -7,13 +7,13 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update import ( # noqa: E501 TOKENS_PER_BLOCK, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.attention.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update import ( # noqa: E501 msa_kv_update as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update_hook import ( # noqa: E501 msa_kv_update as hook, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py index e1c013b90b16..e77fb6540b26 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_allreduce_op_matrix.py @@ -203,7 +203,7 @@ def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 flashinfer_fused_add_rmsnorm, ) - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_allreduce import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( # noqa: E501 m3_allreduce, ) @@ -211,14 +211,14 @@ def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): def _check_cells(mapping: Any) -> None: - """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + """Every cell of ``m3_allreduce.CELLS_BY_TP[2]`` (target AUTO strategy) against the entry's reference: rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.comm.m3_allreduce import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( # noqa: E501 m3_allreduce, ) from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams @@ -228,7 +228,7 @@ def _check_cells(mapping: Any) -> None: from _validating import validating ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) - for i, cell in enumerate(m3_allreduce.CELLS): + for i, cell in enumerate(m3_allreduce.CELLS_BY_TP[2]): spec = cell.spec xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) params = ( @@ -396,14 +396,7 @@ def _spawn_ranks() -> None: TLLM_DISABLE_MPI="1", ) env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} - children = [ - subprocess.Popen( - [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], - env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -414,25 +407,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py index 551c954c6b10..b1b813456da6 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -707,18 +707,9 @@ def _spawn_ranks() -> None: WORLD_SIZE=str(world_size), TLLM_DISABLE_MPI="1", ) - # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed - # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. 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), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -729,25 +720,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py index d80024bb76c8..abfc7351f048 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/comm/_m3_moe_push_op_matrix.py @@ -363,18 +363,9 @@ def _spawn_ranks() -> None: WORLD_SIZE=str(world_size), TLLM_DISABLE_MPI="1", ) - # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed - # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. 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), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -385,25 +376,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py index 2f667fb353f3..43f3fa1ba153 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/embedding/test_modeling_v2_m3_embedding.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.embedding.m3_embedding import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.m3_embedding import ( # noqa: E501 m3_embedding as op, ) @@ -21,7 +21,7 @@ VOCAB, HIDDEN = 4096, 6144 -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[2], ids=[c.why[:48] for c in op.CELLS_BY_TP[2]]) def test_certified_cells(cell) -> None: spec = cell.spec g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py index ae88229a79b6..e33d29dbc768 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_bf16_linear.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_bf16_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_bf16_linear import ( # noqa: E501 m3_bf16_linear as op, ) @@ -18,7 +18,7 @@ pytest.skip("m3_bf16_linear requires a CUDA device", allow_module_level=True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[2], ids=[c.why[:48] for c in op.CELLS_BY_TP[2]]) def test_certified_cells(cell) -> None: """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the row, RMS well under one; measured on GB300: worst 0.5 unit.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py index 16f37f33b86a..507c7df38fa7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_seam.py @@ -78,11 +78,12 @@ def op_calls(monkeypatch): real = modeling.m3_dense_swiglu class _Recording: - CELLS = real.CELLS + CELLS = real.CELLS_BY_TP[2] + CELLS_BY_TP = real.CELLS_BY_TP - def __call__(self, gu, *args): + def __call__(self, gu: torch.Tensor, *args, tp_size: int = 2) -> torch.Tensor | None: calls.append(gu.shape[0]) - return real(gu, *args) + return real(gu, *args, tp_size=tp_size) monkeypatch.setattr(modeling, "m3_dense_swiglu", _Recording()) return calls diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py index 13f745c07dfe..86687f20a934 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -13,7 +13,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_dense_swiglu import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( # noqa: E501 ALPHA, BETA, LIMIT, @@ -38,7 +38,7 @@ def _triton(x: torch.Tensor) -> torch.Tensor: def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: - out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl, tp_size=2) assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" return out @@ -124,14 +124,14 @@ def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pd def test_the_op_declines_other_parameters_and_layouts() -> None: x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") - assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None - assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None - assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None - assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None - assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True, tp_size=2) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True, tp_size=2) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True, tp_size=2) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True, tp_size=2) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True, tp_size=2) is None -@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS_BY_TP[2], ids=lambda c: c.why) def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py index 9f8a268b1495..851fd9ec4475 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_fp32_linear.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_fp32_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_fp32_linear import ( # noqa: E501 m3_fp32_linear as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py index cbeb86cbb1e9..e73c1c2edca0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -8,11 +8,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( + CELLS_BY_TP as M3_MXFP8_GEMM_CELLS_BY_TP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 m3_mxfp8_gemm as op, ) -__extra_import_path__ = ["../../../..", "."] +__extra_import_path__ = ["../../../..", ".", "../.."] from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above @@ -20,7 +23,9 @@ pytest.skip("m3_mxfp8_gemm requires a CUDA device", allow_module_level=True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize( + "cell", M3_MXFP8_GEMM_CELLS_BY_TP[2], ids=[c.why[:48] for c in M3_MXFP8_GEMM_CELLS_BY_TP[2]] +) def test_certified_cells(cell) -> None: """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py index 9a67a31a5200..996e9923be16 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -8,11 +8,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 m3_mxfp8_linear as op, ) -__extra_import_path__ = ["../../../..", "."] +__extra_import_path__ = ["../../../..", ".", "../.."] from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above @@ -26,7 +26,7 @@ } -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[2], ids=[c.why[:48] for c in op.CELLS_BY_TP[2]]) def test_certified_cells(cell) -> None: """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py index 6d8fe90ad9d7..4c6efbaeb4a0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/_m3_router.py @@ -10,17 +10,17 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp # noqa: E402 — needs the path declared above -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv import ( # noqa: E501 m3_router_gemv as gemv, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv import ( # noqa: E501 near_tie, nvfp4_quantize, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 m3_router_gemv_fp4 as gemv_fp4, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.moe.m3_router_tc import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_tc import ( # noqa: E501 m3_router_tc as tc, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py index f4116a9c71b2..e1f6685067a0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_fp32_router_rows.py @@ -12,7 +12,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_fp32_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_fp32_linear import ( # noqa: E501 m3_fp32_linear, ) from tensorrt_llm._utils import get_sm_version diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py index 61bef8aacee1..2702c27993f0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/moe/test_modeling_v2_m3_router_tc.py @@ -24,7 +24,7 @@ ) -@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +@pytest.mark.parametrize("cell", tc.CELLS_BY_TP[2], ids=[c.why[:48] for c in tc.CELLS_BY_TP[2]]) def test_tc_certified_cells(cell) -> None: spec = cell.spec x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py index 158aa26fc6e4..b6d294f5926e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/_m3_spec_inputs.py @@ -5,10 +5,16 @@ from __future__ import annotations -import zlib - import torch +__extra_import_path__ = ["../.."] +from _m3_test_inputs import assert_bf16_close as assert_bf16_close +from _m3_test_inputs import craft_drafts as craft_drafts +from _m3_test_inputs import gen as gen +from _m3_test_inputs import num_sms as num_sms +from _m3_test_inputs import randn as randn +from _m3_test_inputs import seed_of as seed_of # noqa: E402 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_shard_accept import ( # noqa: E501 SHARD_ROWS, make_keys, @@ -22,22 +28,6 @@ EPS = 1e-6 -def seed_of(why: str) -> int: - return zlib.crc32(why.encode()) % 2**31 - - -def gen(seed: int) -> torch.Generator: - return torch.Generator(device="cuda").manual_seed(seed) - - -def num_sms() -> int: - return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count - - -def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: - return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) - - def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") @@ -78,29 +68,3 @@ def plant_peers( keys = make_keys(value, index) for t in range(m): uc[(half * slot_rows + t) * world + r] = keys[t] - - -def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: - """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. - - Kernel and reference compute the same fp32 expression and round once; they differ only where a different - summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. - """ - assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 - o, r = out.float(), ref.float() - ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) - ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) - diff = (o - r).abs() - assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" - frac = (diff > 0).float().mean().item() - assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" - - -def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: - """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from - there on, so the acceptance counts vary across requests.""" - t = tokens.view(nreq, d + 1)[:, :d].clone() - for b in range(nreq): - n = (b + 2) % (d + 1) - t[b, n:] += 1 - return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py index d22dd99fecbb..a5884f291209 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_capture_norm.py @@ -5,7 +5,10 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_capture_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( + CELLS_BY_TP as M3_CAPTURE_NORM_CELLS_BY_TP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( # noqa: E501 m3_capture_norm as op, ) @@ -37,7 +40,9 @@ def _case(spec, seed): return (h, r, w, EPS, dst), buf -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize( + "cell", M3_CAPTURE_NORM_CELLS_BY_TP[2], ids=[c.why[:44] for c in M3_CAPTURE_NORM_CELLS_BY_TP[2]] +) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer untouched.""" @@ -51,7 +56,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_strided_weight() -> None: - (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + (h, r, w, eps, dst), _ = _case(M3_CAPTURE_NORM_CELLS_BY_TP[2][0].spec, 1) w2 = torch.stack([w, w], dim=1)[:, 0] with pytest.raises(AssertionError, match="dense"): with validating(op): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py index 2bda1f79c148..03540e8c12b7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_chunk_norm.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_chunk_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_chunk_norm import ( # noqa: E501 m3_chunk_norm as op, ) @@ -33,7 +33,7 @@ def _case(spec, seed): return (src, w, EPS, dst, True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[2], ids=[c.why[:44] for c in op.CELLS_BY_TP[2]]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" args = _case(cell.spec, seed_of(cell.why)) @@ -43,7 +43,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_strided_weight() -> None: - src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + src, w, eps, dst, pdl = _case(op.CELLS_BY_TP[2][0].spec, 1) wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() wt.copy_(w) with pytest.raises(AssertionError, match="dense"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py index 3efc753d50bd..ca715303d8ad 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_ar_push_words.py @@ -15,9 +15,11 @@ from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiSplitK from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_splitk -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm import ( + m3_lowm_splitk as m3_lowm_splitk, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( m3_lowm_direct, - m3_lowm_splitk, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 ROWS, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py index 76a4f7dc449e..a8d0100aa296 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_input.py @@ -9,7 +9,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_spec_tail, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_input import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_draft_input import ( # noqa: E501 m3_draft_input as op, ) from tensorrt_llm._torch.modules.embedding import Embedding diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py index 6742c23ff618..bae4baf0dcfc 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_draft_step.py @@ -21,10 +21,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops import ( # noqa: E501 m3_draft_step as draft_step_kernels, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_target_head import ( # noqa: E501 - encode_chunked, - round_trip_ok, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_draft_step import ( # noqa: E501 DRAFT_LEN, DRAFT_SHARD, @@ -43,6 +39,10 @@ local_argmax, make_keys, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker __extra_import_path__ = ["../../../..", "."] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py index c102a3374bc1..3387b9b2c560 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_al.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( # noqa: E501 m3_force_vec_req, ) from tensorrt_llm._torch.speculative.interface import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py index aa664ff99747..bb57572b1865 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( # noqa: E501 m3_force_vec as op, ) from tensorrt_llm._torch.speculative.interface import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py index 09a40616268f..06ea3534a778 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_force_vec_req.py @@ -6,7 +6,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_force_vec import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( # noqa: E501 m3_force_vec_req as op, ) from tensorrt_llm._torch.speculative.interface import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py index 0fae391d6b99..c0c0147fb466 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_lowm_splitk.py @@ -9,14 +9,14 @@ from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.gemm.m3_draft_lowm import ( # noqa: E501 ROWS, tactic_for, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.m3_lowm_splitk import ( # noqa: E501 - SplitKTactic, - run_splitk_dense, -) if not torch.cuda.is_available(): pytest.skip("m3_lowm_splitk requires a CUDA device", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py index c0aae0581d0f..cf2dc1ab51b0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_slot_scatter.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_slot_scatter import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_slot_scatter import ( # noqa: E501 m3_slot_scatter as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py index 271b407f83f4..e0038f02b4a6 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp2/spec/test_modeling_v2_m3_step_gather.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp2.ops.spec.m3_step_gather import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_step_gather import ( # noqa: E501 m3_step_gather as op, ) @@ -55,7 +55,7 @@ def _case(spec, seed): ) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[2], ids=[c.why[:44] for c in op.CELLS_BY_TP[2]]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" args = _case(cell.spec, seed_of(cell.why)) @@ -73,7 +73,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_slot_past_the_store() -> None: - args = list(_case(op.CELLS[0].spec, 1)) + args = list(_case(op.CELLS_BY_TP[2][0].spec, 1)) args[3] = args[3].clone() args[3][0] = 8 with pytest.raises(AssertionError, match="slots must index the store"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py index 747a8c3931e8..bd580f703291 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/activation/test_modeling_v2_m3_silu_and_mul.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.activation.m3_silu_and_mul import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.activation.m3_silu_and_mul import ( # noqa: E501 m3_silu_and_mul as op, ) @@ -29,7 +29,7 @@ def _build(spec, seed): return (x, *OAI) if spec["oai"] else (x, None, None, None) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[4], ids=[c.why[:48] for c in op.CELLS_BY_TP[4]]) def test_certified_cells(cell) -> None: """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); @@ -43,7 +43,7 @@ def test_certified_cells(cell) -> None: def test_the_gate_sees_a_missing_clamp() -> None: """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" - spec = op.CELLS[0].spec + spec = op.CELLS_BY_TP[4][0].spec x, *oai = _build(dict(spec, m=64), 3) out = op(x, *oai) with pytest.raises(AssertionError): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py index 0e05e12f6707..901c2037014b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_dense_decode.py @@ -16,17 +16,16 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_dense_decode import ( # noqa: E501 CAPPED, HEAD_DIM, HEADS, KV_HEADS, MODEL_MAX, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_dense_decode import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_dense_decode import ( # noqa: E501 m3_dense_decode as op, ) -from tensorrt_llm._torch.attention.backends.fmha.msa_decode import dense_decode_max_seq_len __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp, certifying, validating # noqa: E402, F401 @@ -133,6 +132,10 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_the_cap_is_the_dispatch_constant() -> None: """The entry's capped cells run at the split the target hands trtllm-gen at one request, and the in-tree step resolves the cap from the attention module only at the request counts it names.""" + # Only targets with a certified cap need this optional backend helper. + # Keep its import behind the existing cap-availability skip. + from tensorrt_llm._torch.attention.backends.fmha.msa_decode import dense_decode_max_seq_len + assert dict(dispatch.DENSE_DECODE_MAX_SEQ_LEN_BY_BATCH) == {1: CAPPED} attn = type( "Attn", diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py index b468359b42cc..4a47089ac9c4 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_fullctx.py @@ -50,8 +50,18 @@ def _tokmap_guard_installed(): # A one-request spec pass places its rows by the workspace token map, which the op builds only for several # requests; without the guard the call reads whatever the workspace holds. - tokmap_guard.install(torch.device("cuda"), dispatch.DECODE_MAX_TOKENS, _TOKMAP_TAKEN.append) - yield + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + tokmap_guard.thop, + "trtllm_gen_generation_preprocess", + tokmap_guard.thop.trtllm_gen_generation_preprocess, + ) + for name in ("_original", "_max_rows", "_on_taken", "_on_missed"): + patch.setattr(tokmap_guard, name, getattr(tokmap_guard, name)) + patch.setattr(tokmap_guard, "_maps", dict(tokmap_guard._maps)) + patch.setattr(tokmap_guard, "_offsets", dict(tokmap_guard._offsets)) + tokmap_guard.install(torch.device("cuda"), dispatch.DECODE_MAX_TOKENS, _TOKMAP_TAKEN.append) + yield def _persistent(name: str, make): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py index 02713095ef13..a01078e73cbc 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_drafter_rope_fp8_kv.py @@ -22,6 +22,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_fullctx_decode as fullctx, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 + m3_tokmap_guard as tokmap_guard, +) from tensorrt_llm._torch.attention.backends.fmha.utils import ( get_attention_chunk_size, get_bmm1_scale, @@ -42,6 +45,21 @@ _PERSISTENT = {} +@pytest.fixture(scope="module", autouse=True) +def _tokmap_guard_installed(): + # One-request speculative passes use the same prepared token map as the model. + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + thop, "trtllm_gen_generation_preprocess", thop.trtllm_gen_generation_preprocess + ) + for name in ("_original", "_max_rows", "_on_taken", "_on_missed"): + patch.setattr(tokmap_guard, name, getattr(tokmap_guard, name)) + patch.setattr(tokmap_guard, "_maps", dict(tokmap_guard._maps)) + patch.setattr(tokmap_guard, "_offsets", dict(tokmap_guard._offsets)) + tokmap_guard.install(torch.device("cuda"), dispatch.DECODE_MAX_TOKENS, lambda rows: None) + yield + + def _persistent(name, make): if name not in _PERSISTENT: _PERSISTENT[name] = make() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py index 75d9aa31ccd3..78e7c79dcc2c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -9,7 +9,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 m3_fp8_indexer_qk_norm_rope as op, ) @@ -28,7 +28,7 @@ def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[4], ids=[c.why[:44] for c in op.CELLS_BY_TP[4]]) def test_certified_cells(cell) -> None: m, heads = cell.spec["m"], cell.spec["heads"] g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py index 829e76748ddf..d843e5e7f284 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_index_score.py @@ -20,6 +20,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_index_score import ( # noqa: E501 m3_index_score as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_index_score import ( # noqa: E501 + IndexScoreRunner as TargetIndexScoreRunner, +) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 @@ -68,8 +71,8 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res score = _score(stride, idx_q.shape[0]) args = (idx_q, index_k, block_table, seq, score, qlen, SPLIT_K, spec["early"]) certifying(op) - ref = op.reference(*args) - op(*args) + ref = op.reference(*args, runner=TargetIndexScoreRunner) + op(*args, runner=TargetIndexScoreRunner) stock = _score(stride, idx_q.shape[0]) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, qlen, SPLIT_K @@ -105,8 +108,8 @@ def test_two_request_early_issue_is_the_in_tree_scorer(seq_lens, why) -> None: idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, 4, zlib.crc32(why.encode())) score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) args = (idx_q, index_k, block_table, seq, score, 4, SPLIT_K, True) - ref = op.reference(*args) - op(*args) + ref = op.reference(*args, runner=TargetIndexScoreRunner) + op(*args, runner=TargetIndexScoreRunner) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, 4, SPLIT_K ) @@ -119,7 +122,7 @@ def test_two_request_early_issue_is_the_in_tree_scorer(seq_lens, why) -> None: def test_bf16_index_is_the_in_tree_scorer(early) -> None: idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) score, stock = _score(stride, 4), _score(stride, 4) - op(idx_q, index_k, block_table, seq, score, 4, SPLIT_K, early) + op(idx_q, index_k, block_table, seq, score, 4, SPLIT_K, early, runner=TargetIndexScoreRunner) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, 4, SPLIT_K ) @@ -131,4 +134,14 @@ def test_guard_refuses_a_request_past_its_block_table() -> None: seq = seq + 4 * PAGE with pytest.raises(AssertionError, match="past its block-table row"): with validating(op): - op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, SPLIT_K, True) + op( + idx_q, + index_k, + block_table, + seq, + _score(stride + 16, 1), + 1, + SPLIT_K, + True, + runner=TargetIndexScoreRunner, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py index 33e0780fde08..87ec3a690fa0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -10,7 +10,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 m3_qk_norm_rope_fp8 as op, ) @@ -57,7 +57,7 @@ def _args(spec, seed): ) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[4], ids=[c.why[:48] for c in op.CELLS_BY_TP[4]]) def test_certified_cells(cell) -> None: args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) before = args[0].clone() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py index 87023e25c703..e0a522727d20 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_qkv_insert.py @@ -11,7 +11,7 @@ import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_qkv_insert import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qkv_insert import ( # noqa: E501 m3_qkv_insert as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py index 1ea6abdf1dc3..9b75e8dc417f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_select_blocks.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_select_blocks import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_select_blocks import ( # noqa: E501 m3_select_blocks as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py index 1c88c39e824f..0844a72fd1f5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_sparse_decode.py @@ -8,15 +8,15 @@ import pytest import torch +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.triton_sparse_decode_mq import ( # noqa: E501 + launch_mq_producer, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_report, ) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.m3_sparse_decode import ( # noqa: E501 m3_sparse_decode as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.triton_sparse_decode_mq import ( # noqa: E501 - launch_mq_producer, -) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( minimax_m3_sparse_attn_decode as stock_decode, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py index 780a53fda069..d694181088ff 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_m3_tokmap_guard.py @@ -36,8 +36,17 @@ @pytest.fixture(scope="module", autouse=True) def _installed(): - guard.install(_DEV, 48, _TAKEN.append) - yield + # Preserve both the process-wide callable and this target's installation state. + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + thop, "trtllm_gen_generation_preprocess", thop.trtllm_gen_generation_preprocess + ) + for name in ("_original", "_max_rows", "_on_taken", "_on_missed"): + patch.setattr(guard, name, getattr(guard, name)) + patch.setattr(guard, "_maps", dict(guard._maps)) + patch.setattr(guard, "_offsets", dict(guard._offsets)) + guard.install(_DEV, 48, _TAKEN.append) + yield class _Step: diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py index 10003353bcbf..faf1c0689784 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/attention/test_modeling_v2_msa_kv_update.py @@ -7,13 +7,13 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update import ( # noqa: E501 TOKENS_PER_BLOCK, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.attention.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update import ( # noqa: E501 msa_kv_update as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update_hook import ( # noqa: E501 msa_kv_update as hook, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py index 3fb3c2bbe14e..1b798202ae38 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_allreduce_op_matrix.py @@ -203,7 +203,7 @@ def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 flashinfer_fused_add_rmsnorm, ) - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( # noqa: E501 m3_allreduce, ) @@ -211,14 +211,14 @@ def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): def _check_cells(mapping: Any) -> None: - """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + """Every cell of ``m3_allreduce.CELLS_BY_TP[4]`` (target AUTO strategy) against the entry's reference: rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.comm.m3_allreduce import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( # noqa: E501 m3_allreduce, ) from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams @@ -228,7 +228,7 @@ def _check_cells(mapping: Any) -> None: from _validating import validating ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) - for i, cell in enumerate(m3_allreduce.CELLS): + for i, cell in enumerate(m3_allreduce.CELLS_BY_TP[4]): spec = cell.spec xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) params = ( @@ -396,14 +396,7 @@ def _spawn_ranks() -> None: TLLM_DISABLE_MPI="1", ) env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} - children = [ - subprocess.Popen( - [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], - env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -414,25 +407,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py index be06233e492e..32306dd925a5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_capture_in_boundary_op_matrix.py @@ -263,14 +263,7 @@ def _spawn_ranks() -> None: TLLM_DISABLE_MPI="1", ) env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} - children = [ - subprocess.Popen( - [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], - env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -281,25 +274,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py index 09b0038dceb3..efb2a99b3968 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_draft_ar_norm_op_matrix.py @@ -489,14 +489,7 @@ def _spawn_ranks() -> None: TLLM_DISABLE_MPI="1", ) env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} - children = [ - subprocess.Popen( - [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], - env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -507,25 +500,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py index 1ee5855270b9..a5a0f46e0273 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -617,18 +617,9 @@ def _spawn_ranks() -> None: WORLD_SIZE=str(world_size), TLLM_DISABLE_MPI="1", ) - # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed - # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. 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), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -639,25 +630,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py index 1d2d922612ba..2d2dca19eaf1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/embedding/test_modeling_v2_m3_embedding.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.embedding.m3_embedding import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.m3_embedding import ( # noqa: E501 m3_embedding as op, ) @@ -21,7 +21,7 @@ VOCAB, HIDDEN = 4096, 6144 -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[4], ids=[c.why[:48] for c in op.CELLS_BY_TP[4]]) def test_certified_cells(cell) -> None: spec = cell.spec g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py deleted file mode 100644 index e44ad56d743d..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/_m3_mxfp8_operands.py +++ /dev/null @@ -1,30 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" - -import torch - - -def swizzle(byte: torch.Tensor) -> torch.Tensor: - """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" - rows, cols = byte.shape - padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 - r = torch.arange(rows, device=byte.device).view(-1, 1) - c = torch.arange(cols, device=byte.device).view(1, -1) - index = ( - (c % 4) - + (c // 4) * 512 - + (r % 32) * 16 - + ((r % 128) // 32) * 4 - + (r // 128) * (128 * padded_cols) - ) - out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) - out[index.reshape(-1)] = byte.reshape(-1) - return out - - -def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): - """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" - v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) - sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) - return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py index cf2493c69fc3..4a09935befc1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_bf16_linear.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_bf16_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_bf16_linear import ( # noqa: E501 m3_bf16_linear as op, ) @@ -18,7 +18,7 @@ pytest.skip("m3_bf16_linear requires a CUDA device", allow_module_level=True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[4], ids=[c.why[:48] for c in op.CELLS_BY_TP[4]]) def test_certified_cells(cell) -> None: """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the row, RMS well under one; measured on GB300: worst 0.5 unit.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py index 729612c53161..55522daa24ec 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -13,7 +13,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4 import ( # noqa: E501 modeling as dispatch, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_dense_swiglu import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( # noqa: E501 ALPHA, BETA, LIMIT, @@ -38,7 +38,7 @@ def _triton(x: torch.Tensor) -> torch.Tensor: def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: - out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl, tp_size=4) assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" return out @@ -124,14 +124,14 @@ def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pd def test_the_op_declines_other_parameters_and_layouts() -> None: x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") - assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None - assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None - assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None - assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None - assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True, tp_size=4) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True, tp_size=4) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True, tp_size=4) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True, tp_size=4) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True, tp_size=4) is None -@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS_BY_TP[4], ids=lambda c: c.why) def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py index 7e6aff6798d8..21e4e792e8a9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_fp32_linear.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_fp32_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_fp32_linear import ( # noqa: E501 m3_fp32_linear as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py index d28ea1ad8dd0..9fd68398405c 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -8,11 +8,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( + CELLS_BY_TP as M3_MXFP8_GEMM_CELLS_BY_TP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 m3_mxfp8_gemm as op, ) -__extra_import_path__ = ["../../../..", "."] +__extra_import_path__ = ["../../../..", ".", "../.."] from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above @@ -20,7 +23,9 @@ pytest.skip("m3_mxfp8_gemm requires a CUDA device", allow_module_level=True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize( + "cell", M3_MXFP8_GEMM_CELLS_BY_TP[4], ids=[c.why[:48] for c in M3_MXFP8_GEMM_CELLS_BY_TP[4]] +) def test_certified_cells(cell) -> None: """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py index f0199cbc31d0..b8d6144fe878 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -8,11 +8,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 m3_mxfp8_linear as op, ) -__extra_import_path__ = ["../../../..", "."] +__extra_import_path__ = ["../../../..", ".", "../.."] from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above @@ -26,7 +26,7 @@ } -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[4], ids=[c.why[:48] for c in op.CELLS_BY_TP[4]]) def test_certified_cells(cell) -> None: """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py index 1616fa9bd7d3..a9ff4ba5b4ff 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_mxfp8_pinned_gemm.py @@ -18,7 +18,7 @@ ) from tensorrt_llm._torch.custom_ops.torch_custom_ops import FlashInferMXFP8GemmRunner -__extra_import_path__ = ["../../../..", "."] +__extra_import_path__ = ["../../../..", ".", "../.."] from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py index 8bb9fc8fd1ec..8a66c789db2d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/gemm/test_modeling_v2_m3_target_head.py @@ -18,7 +18,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_target_head as target_head_kernels, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( # noqa: E501 CTA_WARPS, HIDDEN, TARGET_SHARD, @@ -26,7 +26,7 @@ encode_chunked, round_trip_ok, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( # noqa: E501 m3_target_head as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py index 724a2979c234..dffe33992b32 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/_m3_router.py @@ -9,7 +9,7 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp # noqa: E402 — needs the path declared above -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.moe.m3_router_tc import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_tc import ( # noqa: E501 m3_router_tc as tc, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py index e43c57f109a4..b5b43713e4af 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/moe/test_modeling_v2_m3_router_tc.py @@ -18,7 +18,7 @@ ) -@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +@pytest.mark.parametrize("cell", tc.CELLS_BY_TP[4], ids=[c.why[:48] for c in tc.CELLS_BY_TP[4]]) def test_tc_certified_cells(cell) -> None: spec = cell.spec x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py index 27157668c588..9f74b64f0c16 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/_m3_spec_inputs.py @@ -5,10 +5,16 @@ from __future__ import annotations -import zlib - import torch +__extra_import_path__ = ["../.."] +from _m3_test_inputs import assert_bf16_close as assert_bf16_close +from _m3_test_inputs import craft_drafts as craft_drafts +from _m3_test_inputs import gen as gen +from _m3_test_inputs import num_sms as num_sms +from _m3_test_inputs import randn as randn +from _m3_test_inputs import seed_of as seed_of # noqa: E402 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_shard_accept import ( # noqa: E501 SHARD_ROWS, make_keys, @@ -22,22 +28,6 @@ EPS = 1e-6 -def seed_of(why: str) -> int: - return zlib.crc32(why.encode()) % 2**31 - - -def gen(seed: int) -> torch.Generator: - return torch.Generator(device="cuda").manual_seed(seed) - - -def num_sms() -> int: - return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count - - -def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: - return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) - - def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") @@ -78,29 +68,3 @@ def plant_peers( keys = make_keys(value, index) for t in range(m): uc[(half * slot_rows + t) * world + r] = keys[t] - - -def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: - """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. - - Kernel and reference compute the same fp32 expression and round once; they differ only where a different - summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. - """ - assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 - o, r = out.float(), ref.float() - ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) - ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) - diff = (o - r).abs() - assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" - frac = (diff > 0).float().mean().item() - assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" - - -def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: - """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from - there on, so the acceptance counts vary across requests.""" - t = tokens.view(nreq, d + 1)[:, :d].clone() - for b in range(nreq): - n = (b + 2) % (d + 1) - t[b, n:] += 1 - return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py index d2a171790f64..bb3d064f1ce7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_capture_norm.py @@ -5,7 +5,10 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_capture_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( + CELLS_BY_TP as M3_CAPTURE_NORM_CELLS_BY_TP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( # noqa: E501 m3_capture_norm as op, ) @@ -42,7 +45,9 @@ def _case(spec, seed): return (h, r, w, EPS, dst), buf -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize( + "cell", M3_CAPTURE_NORM_CELLS_BY_TP[4], ids=[c.why[:44] for c in M3_CAPTURE_NORM_CELLS_BY_TP[4]] +) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer untouched.""" @@ -56,7 +61,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_strided_weight() -> None: - (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + (h, r, w, eps, dst), _ = _case(M3_CAPTURE_NORM_CELLS_BY_TP[4][0].spec, 1) w2 = torch.stack([w, w], dim=1)[:, 0] with pytest.raises(AssertionError, match="dense"): with validating(op): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py index 8848ed3a2b6f..239bb4d59d40 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_chunk_norm.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_chunk_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_chunk_norm import ( # noqa: E501 m3_chunk_norm as op, ) @@ -38,7 +38,7 @@ def _case(spec, seed): return (src, w, EPS, dst, True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[4], ids=[c.why[:44] for c in op.CELLS_BY_TP[4]]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" args = _case(cell.spec, seed_of(cell.why)) @@ -48,7 +48,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_strided_weight() -> None: - src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + src, w, eps, dst, pdl = _case(op.CELLS_BY_TP[4][0].spec, 1) wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() wt.copy_(w) with pytest.raises(AssertionError, match="dense"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py index a4c8090bf4d2..7c068fc4a6e9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_input.py @@ -9,7 +9,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_spec_tail, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_input import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_draft_input import ( # noqa: E501 m3_draft_input as op, ) from tensorrt_llm._torch.modules.embedding import Embedding diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py index 1a17fa8ee539..5c210feefeb8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_draft_step.py @@ -21,10 +21,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops import ( # noqa: E501 m3_draft_step as draft_step_kernels, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_target_head import ( # noqa: E501 - encode_chunked, - round_trip_ok, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_draft_step import ( # noqa: E501 DRAFT_LEN, DRAFT_SHARD, @@ -43,6 +39,10 @@ local_argmax, make_keys, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker __extra_import_path__ = ["../../../..", "."] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py index edcfd8a451f4..dd876b9a2b39 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_exact_numerics.py @@ -67,7 +67,12 @@ def _stock_accept(logits: torch.Tensor, draft: torch.Tensor): """The in-tree greedy strict acceptance (``SpecWorkerBase._sample_and_accept_draft_tokens_base``) of the gathered float logits of a decode-only batch, no forcing.""" - md = SimpleNamespace(is_all_greedy_sample=True, is_ragged_verify=False, is_cuda_graph=False) + md = SimpleNamespace( + is_all_greedy_sample=True, + is_ragged_verify=False, + is_cuda_graph=False, + sharded_target_logits=False, + ) worker = SimpleNamespace(force_num_accepted_tokens=0.0) worker._sample_tokens_for_batch = lambda *a: Eagle3OneModelWorker._sample_tokens_for_batch( worker, *a @@ -75,6 +80,9 @@ def _stock_accept(logits: torch.Tensor, draft: torch.Tensor): worker._apply_force_accepted_tokens = lambda *a, **k: ( Eagle3OneModelWorker._apply_force_accepted_tokens(worker, *a, **k) ) + worker._fused_force_params = lambda *a, **k: Eagle3OneModelWorker._fused_force_params( + worker, *a, **k + ) nreq = draft.shape[0] return Eagle3OneModelWorker._sample_and_accept_draft_tokens_base( worker, logits, draft, 0, nreq, md diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py index 122201ef3e65..3cd269ace1f8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_force_vec.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_force_vec import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( # noqa: E501 m3_force_vec as op, ) from tensorrt_llm._torch.speculative.interface import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py index 7b457c6bdc3e..3e1ff55bf774 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_lowm_splitk.py @@ -9,14 +9,14 @@ from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_lowm_splitk import ( # noqa: E501 + SplitKTactic, + run_splitk_dense, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.gemm.m3_draft_lowm import ( # noqa: E501 ROWS, tactic_for, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_lowm_splitk import ( # noqa: E501 - SplitKTactic, - run_splitk_dense, -) if not torch.cuda.is_available(): pytest.skip("m3_lowm_splitk requires a CUDA device", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py index 7f48d822e327..e457dafdb3e2 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_slot_scatter.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_slot_scatter import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_slot_scatter import ( # noqa: E501 m3_slot_scatter as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py index ce5562f0d404..c6e98c086139 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/spec/test_modeling_v2_m3_step_gather.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.spec.m3_step_gather import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_step_gather import ( # noqa: E501 m3_step_gather as op, ) @@ -60,7 +60,7 @@ def _case(spec, seed): ) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[4], ids=[c.why[:44] for c in op.CELLS_BY_TP[4]]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" args = _case(cell.spec, seed_of(cell.why)) @@ -78,7 +78,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_slot_past_the_store() -> None: - args = list(_case(op.CELLS[0].spec, 1)) + args = list(_case(op.CELLS_BY_TP[4][0].spec, 1)) args[3] = args[3].clone() args[3][0] = 8 with pytest.raises(AssertionError, match="slots must index the store"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py index 0e915fbcb18e..0f3d9b72e23b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp4/test_modeling_v2_m3_prebuild.py @@ -33,6 +33,9 @@ SHARED_DOWN, SHARED_GATE_UP, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.m3_index_score import ( # noqa: E501 + IndexScoreRunner as TargetIndexScoreRunner, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp4.ops.mx_gemv import ( # noqa: E501 op as mx_gemv, ) @@ -311,5 +314,6 @@ def test_scorer_after_prebuild_never_compiles(monkeypatch, indexer_kv_dtype) -> qlen, split_k, early, + runner=TargetIndexScoreRunner, ) torch.cuda.synchronize() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py index 891ffce1d511..f60dda8eeea3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/activation/test_modeling_v2_m3_silu_and_mul.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.activation.m3_silu_and_mul import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.activation.m3_silu_and_mul import ( # noqa: E501 m3_silu_and_mul as op, ) @@ -29,7 +29,7 @@ def _build(spec, seed): return (x, *OAI) if spec["oai"] else (x, None, None, None) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[8], ids=[c.why[:48] for c in op.CELLS_BY_TP[8]]) def test_certified_cells(cell) -> None: """fp32 arithmetic and one bf16 rounding on both sides: the kernel's fast sigmoid moves the fp32 value by a few fp32 ulps, which flips the bf16 rounding of an element by at most one step (1 unit of ULP[bf16] of the row); @@ -43,7 +43,7 @@ def test_certified_cells(cell) -> None: def test_the_gate_sees_a_missing_clamp() -> None: """Dropping the SwiGLU-OAI clamps moves every element whose gate or up passed the limit far outside the band.""" - spec = op.CELLS[0].spec + spec = op.CELLS_BY_TP[8][0].spec x, *oai = _build(dict(spec, m=64), 3) out = op(x, *oai) with pytest.raises(AssertionError): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py index baba0c067512..c400796766b8 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_fp8_indexer_qk_norm_rope.py @@ -9,7 +9,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_fp8_indexer_qk_norm_rope import ( # noqa: E501 m3_fp8_indexer_qk_norm_rope as op, ) @@ -28,7 +28,7 @@ def _bytes_close(a: torch.Tensor, b: torch.Tensor) -> None: assert differ <= (a.numel() - 1) // 1000, f"{differ} of {a.numel()} E4M3 bytes differ" -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[8], ids=[c.why[:44] for c in op.CELLS_BY_TP[8]]) def test_certified_cells(cell) -> None: m, heads = cell.spec["m"], cell.spec["heads"] g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py index 9693051474f2..4a6b99d987b0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_index_score.py @@ -19,6 +19,9 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_index_score import ( # noqa: E501 m3_index_score as op, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_index_score import ( # noqa: E501 + IndexScoreRunner as TargetIndexScoreRunner, +) __extra_import_path__ = ["../../../.."] from _validating import certifying, validating # noqa: E402, F401 @@ -63,8 +66,8 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res assert spec["heads"] == HEADS args = (idx_q, index_k, block_table, seq, score, qlen, spec["split_k"], spec["early"]) certifying(op) - ref = op.reference(*args) - op(*args) + ref = op.reference(*args, runner=TargetIndexScoreRunner) + op(*args, runner=TargetIndexScoreRunner) stock = _score(stride, idx_q.shape[0]) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, qlen @@ -78,7 +81,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_bf16_index_is_the_in_tree_scorer(early) -> None: idx_q, index_k, block_table, seq, stride = _inputs((5000,), 4, 11, dtype=torch.bfloat16) score, stock = _score(stride, 4), _score(stride, 4) - op(idx_q, index_k, block_table, seq, score, 4, 1536, early) + op(idx_q, index_k, block_table, seq, score, 4, 1536, early, runner=TargetIndexScoreRunner) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, 4 ) @@ -97,7 +100,17 @@ def test_prestaged_blocks_at_the_stable_edge_are_the_in_tree_scorer(qlen) -> Non ): idx_q, index_k, block_table, seq, stride = _inputs(seq_lens, qlen, 31 + length) score, stock = _score(stride, idx_q.shape[0]), _score(stride, idx_q.shape[0]) - op(idx_q, index_k, block_table, seq, score, qlen, split, True) + op( + idx_q, + index_k, + block_table, + seq, + score, + qlen, + split, + True, + runner=TargetIndexScoreRunner, + ) torch.ops.trtllm.cute_dsl_minimax_m3_index_decode_score( idx_q, index_k, block_table, seq, stock, qlen ) @@ -109,4 +122,14 @@ def test_guard_refuses_a_request_past_its_block_table() -> None: seq = seq + 4 * PAGE with pytest.raises(AssertionError, match="past its block-table row"): with validating(op): - op(idx_q, index_k, block_table, seq, _score(stride + 16, 1), 1, 1536, True) + op( + idx_q, + index_k, + block_table, + seq, + _score(stride + 16, 1), + 1, + 1536, + True, + runner=TargetIndexScoreRunner, + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py index e48bf7aabc0d..fd3900abde2b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qk_norm_rope_fp8.py @@ -10,7 +10,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qk_norm_rope_fp8 import ( # noqa: E501 m3_qk_norm_rope_fp8 as op, ) @@ -57,7 +57,7 @@ def _args(spec, seed): ) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[8], ids=[c.why[:48] for c in op.CELLS_BY_TP[8]]) def test_certified_cells(cell) -> None: args = _args(cell.spec, zlib.crc32(cell.why.encode()) % 2**31) before = args[0].clone() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py index 85c1a03c59be..d38c83074280 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_qkv_insert.py @@ -11,7 +11,7 @@ import torch import tensorrt_llm._torch.custom_ops # noqa: F401 — registers torch.ops.trtllm.* -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_qkv_insert import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_qkv_insert import ( # noqa: E501 m3_qkv_insert as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py index af3b3c3d4f0e..1976dcaef541 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_select_blocks.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_select_blocks import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.m3_select_blocks import ( # noqa: E501 m3_select_blocks as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py index 195c5cb5c567..da4739f7d95b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_sparse_decode.py @@ -14,6 +14,9 @@ import torch import tensorrt_llm # noqa: F401 (registers torch.ops.trtllm) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.attention.triton_sparse_decode_mq import ( # noqa: E501 + launch_mq_producer, +) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( modeling, # noqa: E501 ) @@ -26,9 +29,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.m3_sparse_decode import ( # noqa: E501 m3_sparse_decode as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.triton_sparse_decode_mq import ( # noqa: E501 - launch_mq_producer, -) from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.kernels.triton_sparse_decode import ( minimax_m3_sparse_attn_decode as stock_decode, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py index 461e9695b6b2..c1c88df0520e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_m3_tokmap_guard.py @@ -36,8 +36,17 @@ @pytest.fixture(scope="module", autouse=True) def _installed(): - guard.install(_DEV, 48, _TAKEN.append) - yield + # Preserve both the process-wide callable and this target's installation state. + with pytest.MonkeyPatch.context() as patch: + patch.setattr( + thop, "trtllm_gen_generation_preprocess", thop.trtllm_gen_generation_preprocess + ) + for name in ("_original", "_max_rows", "_on_taken", "_on_missed"): + patch.setattr(guard, name, getattr(guard, name)) + patch.setattr(guard, "_maps", dict(guard._maps)) + patch.setattr(guard, "_offsets", dict(guard._offsets)) + guard.install(_DEV, 48, _TAKEN.append) + yield class _Step: diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py index a9aea90d3eaa..0a1bdba7f112 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/attention/test_modeling_v2_msa_kv_update.py @@ -7,13 +7,13 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update import ( # noqa: E501 TOKENS_PER_BLOCK, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.attention.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update import ( # noqa: E501 msa_kv_update as op, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.msa_kv_update import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.attention.msa_kv_update_hook import ( # noqa: E501 msa_kv_update as hook, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py index 006771133817..ae915b39a1d7 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_allreduce_op_matrix.py @@ -217,7 +217,7 @@ def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): from tensorrt_llm._torch._experimental.modeling_v2.catalog.norm.flashinfer_fused_add_rmsnorm import ( # noqa: E501 flashinfer_fused_add_rmsnorm, ) - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( # noqa: E501 m3_allreduce, ) @@ -225,14 +225,14 @@ def _target(ar: Any, x: torch.Tensor, residual: torch.Tensor, w: torch.Tensor): def _check_cells(mapping: Any) -> None: - """Every cell of ``m3_allreduce.CELLS`` (AUTO strategy, as the target builds it) against the entry's reference: + """Every cell of ``m3_allreduce.CELLS_BY_TP[8]`` (target AUTO strategy) against the entry's reference: rank-order fp32 sums with one BF16 rounding of the sum and one of the norm. A strategy may round the cross-rank sum to BF16 before adding the residual, so the residual sum is bounded at its operands' magnitude u = ulp(max(|residual|, |sum|)): half an ulp for that rounding plus half an output ulp for each of the two final roundings, an output ulp being at most 2u, so 2.5u (where the two cancel, the output's own ulp is far smaller). The normed rows are the norm of the op's own residual sum, within 2 ulps of the reference norm of that sum.""" - from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.comm.m3_allreduce import ( # noqa: E501 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.comm.m3_allreduce import ( # noqa: E501 m3_allreduce, ) from tensorrt_llm._torch.distributed import AllReduce, AllReduceFusionOp, AllReduceParams @@ -242,7 +242,7 @@ def _check_cells(mapping: Any) -> None: from _validating import validating ar = AllReduce(mapping=mapping, strategy=AllReduceStrategy.AUTO, dtype=torch.bfloat16) - for i, cell in enumerate(m3_allreduce.CELLS): + for i, cell in enumerate(m3_allreduce.CELLS_BY_TP[8]): spec = cell.spec xs, residual, w = _random(200 + i, spec["m"], 1.0, 4.0, False) params = ( @@ -410,14 +410,7 @@ def _spawn_ranks() -> None: TLLM_DISABLE_MPI="1", ) env = {k: v for k, v in env.items() if not k.startswith(("PMIX_", "PMI_", "OMPI_", "SLURM_"))} - children = [ - subprocess.Popen( - [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], - env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -428,25 +421,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py index b493b3a4afd9..27091de9c98f 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_mc_allreduce_norm_op_matrix.py @@ -740,18 +740,9 @@ def _spawn_ranks() -> None: WORLD_SIZE=str(world_size), TLLM_DISABLE_MPI="1", ) - # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed - # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. 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), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -762,25 +753,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") if __name__ == "__main__": diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py index 4f0c30322f4a..96e1c0fd3bb9 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/comm/_m3_moe_push_op_matrix.py @@ -341,18 +341,9 @@ def _spawn_ranks() -> None: WORLD_SIZE=str(world_size), TLLM_DISABLE_MPI="1", ) - # Drop a launching srun step's PMIx / MPI namespace: children registered in it get the whole step SIGKILLed - # when the first one exits hard. These ranks rendezvous through torch.distributed's env://. 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), LOCAL_RANK=str(r)), - start_new_session=True, - ) - for r in range(world_size) - ] + children = [] def _kill_all() -> None: for p in children: @@ -363,25 +354,35 @@ def _kill_all() -> None: for p in children: p.wait() - deadline = time.monotonic() + DEADLINE_S - live = set(range(world_size)) - failure: Optional[str] = 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}" + try: + for r in range(world_size): + children.append( + subprocess.Popen( + [sys.executable, str(Path(__file__).resolve()), _WORKER_FLAG], + env=dict(env, RANK=str(r), LOCAL_RANK=str(r)), + start_new_session=True, + ) + ) + deadline = time.monotonic() + DEADLINE_S + live = set(range(world_size)) + failure: Optional[str] = None + while live and failure is None: + if time.monotonic() > deadline: + failure = f"did not finish in {DEADLINE_S}s (wedged)" break - time.sleep(0.5) - if failure is not None: + 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: + raise AssertionError(f"the {world_size}-rank run failed: {failure}") + finally: _kill_all() - raise AssertionError(f"the {world_size}-rank run failed: {failure}") _SRUN_FLAG = "--srun-rank" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py index 2ae718a482a1..538e163b1c5b 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/embedding/test_modeling_v2_m3_embedding.py @@ -8,7 +8,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.embedding.m3_embedding import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.embedding.m3_embedding import ( # noqa: E501 m3_embedding as op, ) @@ -21,7 +21,7 @@ VOCAB, HIDDEN = 4096, 6144 -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[8], ids=[c.why[:48] for c in op.CELLS_BY_TP[8]]) def test_certified_cells(cell) -> None: spec = cell.spec g = torch.Generator(device="cuda").manual_seed(zlib.crc32(cell.why.encode()) % 2**31) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py deleted file mode 100644 index e44ad56d743d..000000000000 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/_m3_mxfp8_operands.py +++ /dev/null @@ -1,30 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MXFP8 operands for the MXFP8 GEMM entries' tests: E4M3 values and 128x4-swizzled UE8M0 block scales.""" - -import torch - - -def swizzle(byte: torch.Tensor) -> torch.Tensor: - """[N, K / 32] scale bytes into the 128x4-swizzled buffer of ``mxfp8_quantize(w, True)``.""" - rows, cols = byte.shape - padded_cols, padded_rows = (cols + 3) // 4 * 4, (rows + 127) // 128 * 128 - r = torch.arange(rows, device=byte.device).view(-1, 1) - c = torch.arange(cols, device=byte.device).view(1, -1) - index = ( - (c % 4) - + (c // 4) * 512 - + (r % 32) * 16 - + ((r % 128) // 32) * 4 - + (r // 128) * (128 * padded_cols) - ) - out = torch.zeros(padded_rows * padded_cols, dtype=torch.uint8, device=byte.device) - out[index.reshape(-1)] = byte.reshape(-1) - return out - - -def mxfp8(rows: int, k: int, g: torch.Generator, lo: int = 118, hi: int = 125): - """An MXFP8 operand: E4M3 values of a unit normal, UE8M0 block scales 2^(lo-127) .. 2^(hi-128).""" - v = torch.randn(rows, k, generator=g, device="cuda").to(torch.float8_e4m3fn) - sf = torch.randint(lo, hi, (rows, k // 32), generator=g, device="cuda", dtype=torch.int32) - return v, swizzle(sf.to(torch.uint8)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py index 4386d7434b8a..d308e3579ce3 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_bf16_linear.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_bf16_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_bf16_linear import ( # noqa: E501 m3_bf16_linear as op, ) @@ -18,7 +18,7 @@ pytest.skip("m3_bf16_linear requires a CUDA device", allow_module_level=True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[8], ids=[c.why[:48] for c in op.CELLS_BY_TP[8]]) def test_certified_cells(cell) -> None: """Accumulation order (cuBLAS split-K) then one bf16 rounding on both sides: at most 2 units of ULP[bf16] of the row, RMS well under one; measured on GB300: worst 0.5 unit.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py index ca565059786a..1b3e5c40ed39 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_dense_swiglu.py @@ -13,7 +13,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8 import ( # noqa: E501 modeling, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_dense_swiglu import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_dense_swiglu import ( # noqa: E501 ALPHA, BETA, LIMIT, @@ -38,7 +38,7 @@ def _triton(x: torch.Tensor) -> torch.Tensor: def _op(x: torch.Tensor, pdl: bool = True) -> torch.Tensor: - out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl) + out = m3_dense_swiglu(x, LIMIT, ALPHA, BETA, pdl, tp_size=8) assert out is not None, "the op declined a contiguous bf16 [rows, 2d] input" return out @@ -124,14 +124,14 @@ def test_dense_shape_eager_and_graph_are_bitwise_the_triton_kernel(rows: int, pd def test_the_op_declines_other_parameters_and_layouts() -> None: x = torch.zeros(2, 2 * DENSE_D, dtype=torch.bfloat16, device="cuda") - assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True) is None - assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True) is None - assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True) is None - assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True) is None - assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True) is None + assert m3_dense_swiglu(x, 8.0, ALPHA, BETA, True, tp_size=8) is None + assert m3_dense_swiglu(x, LIMIT, 1.0, BETA, True, tp_size=8) is None + assert m3_dense_swiglu(x, LIMIT, ALPHA, 0.0, True, tp_size=8) is None + assert m3_dense_swiglu(x.float(), LIMIT, ALPHA, BETA, True, tp_size=8) is None + assert m3_dense_swiglu(x[:, : 2 * DENSE_D - 2], LIMIT, ALPHA, BETA, True, tp_size=8) is None -@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS, ids=lambda c: c.why) +@pytest.mark.parametrize("cell", m3_dense_swiglu.CELLS_BY_TP[8], ids=lambda c: c.why) def test_every_cell_is_bitwise_the_triton_kernel(cell) -> None: """Each declared cell (the dense MLP's activation at its decode rows) through the op, bit for bit the Triton kernel on gate_up rows in a GEMV output's range plus values past the clamps.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py index 94fd72e397fb..209cc01b78df 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_fp32_linear.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_fp32_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_fp32_linear import ( # noqa: E501 m3_fp32_linear as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py index f3b9d7c10cdf..25edefeadae2 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_gemm.py @@ -8,11 +8,14 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( + CELLS_BY_TP as M3_MXFP8_GEMM_CELLS_BY_TP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_gemm import ( # noqa: E501 m3_mxfp8_gemm as op, ) -__extra_import_path__ = ["../../../..", "."] +__extra_import_path__ = ["../../../..", ".", "../.."] from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above @@ -20,7 +23,9 @@ pytest.skip("m3_mxfp8_gemm requires a CUDA device", allow_module_level=True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize( + "cell", M3_MXFP8_GEMM_CELLS_BY_TP[8], ids=[c.why[:48] for c in M3_MXFP8_GEMM_CELLS_BY_TP[8]] +) def test_certified_cells(cell) -> None: """Same E4M3 operands and power-of-two scales on both sides; accumulation order then one bf16 rounding: at most 2 units of ULP[bf16] of the row (the low-M GEMM's band), RMS well under one.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py index 82938bf53a14..e996ce9e0e11 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_mxfp8_linear.py @@ -8,11 +8,11 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_mxfp8_linear import ( # noqa: E501 m3_mxfp8_linear as op, ) -__extra_import_path__ = ["../../../..", "."] +__extra_import_path__ = ["../../../..", ".", "../.."] from _m3_mxfp8_operands import mxfp8 # noqa: E402 from _validating import assert_within_ulp, validating # noqa: E402 — needs the path declared above @@ -26,7 +26,7 @@ } -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:48] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[8], ids=[c.why[:48] for c in op.CELLS_BY_TP[8]]) def test_certified_cells(cell) -> None: """The activation's MXFP8 round trip is exact on both sides (same block rule); then the accumulation order and one bf16 rounding: at most 2 units of ULP[bf16] of the row, RMS well under one.""" diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py index 70ff6e13cdc4..43e010a4b32a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/gemm/test_modeling_v2_m3_swiglu_oai_bf16x2.py @@ -10,7 +10,7 @@ import torch from cutlass.cute.runtime import from_dlpack -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.mx_gemv.m3_mx_gemv_kernel import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_mx_gemv_kernel import ( # noqa: E501 SWIGLU_OAI_ALPHA, SWIGLU_OAI_LIMIT, _swiglu_oai_bf16x2, diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py index 7a88ec1adca9..4c6efbaeb4a0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/_m3_router.py @@ -10,17 +10,17 @@ __extra_import_path__ = ["../../../.."] from _validating import assert_within_ulp # noqa: E402 — needs the path declared above -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv import ( # noqa: E501 m3_router_gemv as gemv, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv import ( # noqa: E501 near_tie, nvfp4_quantize, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_gemv_fp4 import ( # noqa: E501 m3_router_gemv_fp4 as gemv_fp4, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.moe.m3_router_tc import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.moe.m3_router_tc import ( # noqa: E501 m3_router_tc as tc, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py index ca6b047d862d..7995931bf2d1 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/moe/test_modeling_v2_m3_router_tc.py @@ -24,7 +24,7 @@ ) -@pytest.mark.parametrize("cell", tc.CELLS, ids=[c.why[:48] for c in tc.CELLS]) +@pytest.mark.parametrize("cell", tc.CELLS_BY_TP[8], ids=[c.why[:48] for c in tc.CELLS_BY_TP[8]]) def test_tc_certified_cells(cell) -> None: spec = cell.spec x, w = _build(spec["m"], 128, spec["k"], _seed(cell)) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py index 15f7afe40867..84a968f9bee6 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/_m3_spec_inputs.py @@ -5,10 +5,16 @@ from __future__ import annotations -import zlib - import torch +__extra_import_path__ = ["../.."] +from _m3_test_inputs import assert_bf16_close as assert_bf16_close +from _m3_test_inputs import craft_drafts as craft_drafts +from _m3_test_inputs import gen as gen +from _m3_test_inputs import num_sms as num_sms +from _m3_test_inputs import randn as randn +from _m3_test_inputs import seed_of as seed_of # noqa: E402 + from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_shard_accept import ( # noqa: E501 SHARD_ROWS, make_keys, @@ -22,22 +28,6 @@ EPS = 1e-6 -def seed_of(why: str) -> int: - return zlib.crc32(why.encode()) % 2**31 - - -def gen(seed: int) -> torch.Generator: - return torch.Generator(device="cuda").manual_seed(seed) - - -def num_sms() -> int: - return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count - - -def randn(shape, g: torch.Generator, std: float = 1.0, dtype=torch.bfloat16) -> torch.Tensor: - return (torch.randn(shape, generator=g, device="cuda") * std).to(dtype) - - def head_workspace(world: int, call: int = 0, slot_rows: int = SHARD_ROWS): """(ws_uc int64 [2 * slot_rows * world], keys int64 [slot_rows], counters int32 [2]) of a fresh call ``call``.""" uc = torch.zeros(2 * slot_rows * world, dtype=torch.int64, device="cuda") @@ -78,29 +68,3 @@ def plant_peers( keys = make_keys(value, index) for t in range(m): uc[(half * slot_rows + t) * world + r] = keys[t] - - -def assert_bf16_close(out: torch.Tensor, ref: torch.Tensor, max_frac: float = 1e-3) -> None: - """At most one bf16 ulp per element, and only on a ``max_frac`` share of the elements. - - Kernel and reference compute the same fp32 expression and round once; they differ only where a different - summation order (or rsqrtf) puts the fp32 value on the other side of a bf16 rounding boundary. - """ - assert out.shape == ref.shape and out.dtype == ref.dtype == torch.bfloat16 - o, r = out.float(), ref.float() - ulp = torch.ldexp(torch.ones_like(r), torch.frexp(r).exponent - 8) - ulp = torch.where(r == 0, torch.full_like(r, 2.0**-133), ulp) - diff = (o - r).abs() - assert bool((diff <= ulp * 1.0001).all()), f"max |diff| / ulp = {(diff / ulp).max().item():.2f}" - frac = (diff > 0).float().mean().item() - assert frac <= max_frac, f"{frac:.2e} of the elements moved by one ulp (allowed {max_frac})" - - -def craft_drafts(tokens: torch.Tensor, nreq: int, d: int) -> torch.Tensor: - """Drafts matching request b's verify tokens for its first (b + 2) % (d + 1) positions and differing from - there on, so the acceptance counts vary across requests.""" - t = tokens.view(nreq, d + 1)[:, :d].clone() - for b in range(nreq): - n = (b + 2) % (d + 1) - t[b, n:] += 1 - return t.contiguous() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py index 121001424351..0544c309771a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_capture_norm.py @@ -5,7 +5,10 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_capture_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( + CELLS_BY_TP as M3_CAPTURE_NORM_CELLS_BY_TP, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_capture_norm import ( # noqa: E501 m3_capture_norm as op, ) @@ -37,7 +40,9 @@ def _case(spec, seed): return (h, r, w, EPS, dst), buf -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize( + "cell", M3_CAPTURE_NORM_CELLS_BY_TP[8], ids=[c.why[:44] for c in M3_CAPTURE_NORM_CELLS_BY_TP[8]] +) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Within one bf16 ulp per element on at most 0.1% of the elements; the other layer slots of the buffer untouched.""" @@ -51,7 +56,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_strided_weight() -> None: - (h, r, w, eps, dst), _ = _case(op.CELLS[0].spec, 1) + (h, r, w, eps, dst), _ = _case(M3_CAPTURE_NORM_CELLS_BY_TP[8][0].spec, 1) w2 = torch.stack([w, w], dim=1)[:, 0] with pytest.raises(AssertionError, match="dense"): with validating(op): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py index bced642e6b8b..708e9af0077d 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_chunk_norm.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_chunk_norm import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_chunk_norm import ( # noqa: E501 m3_chunk_norm as op, ) @@ -33,7 +33,7 @@ def _case(spec, seed): return (src, w, EPS, dst, True) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[8], ids=[c.why[:44] for c in op.CELLS_BY_TP[8]]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Within one bf16 ulp per element on at most 0.1% of the elements (fp32 statistic, rsqrtf).""" args = _case(cell.spec, seed_of(cell.why)) @@ -43,7 +43,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_strided_weight() -> None: - src, w, eps, dst, pdl = _case(op.CELLS[0].spec, 1) + src, w, eps, dst, pdl = _case(op.CELLS_BY_TP[8][0].spec, 1) wt = torch.empty((w.shape[1], w.shape[0]), dtype=w.dtype, device="cuda").t() wt.copy_(w) with pytest.raises(AssertionError, match="dense"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py index f07eecb1f0c2..df1c14e59a01 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_di_fc_flag.py @@ -67,6 +67,7 @@ def body(): dst.copy_(src) main.wait_stream(aux) + cap.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(cap): body() # eager warm-up (compiles the wait kernel outside capture) torch.cuda.synchronize() diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py index 972204f0d653..759aa9c95344 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_input.py @@ -9,7 +9,7 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_spec_tail, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_input import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_draft_input import ( # noqa: E501 m3_draft_input as op, ) from tensorrt_llm._torch.modules.embedding import Embedding diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py index ba136429a30e..a6ce782f7f4e 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_draft_step.py @@ -21,10 +21,6 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E501 m3_draft_step as draft_step_kernels, ) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_target_head import ( # noqa: E501 - encode_chunked, - round_trip_ok, -) from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_draft_step import ( # noqa: E501 DRAFT_LEN, DRAFT_SHARD, @@ -44,6 +40,10 @@ local_argmax, make_keys, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.gemm.m3_target_head import ( # noqa: E501 + encode_chunked, + round_trip_ok, +) from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelWorker __extra_import_path__ = ["../../../..", "."] diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py index 0f6b4f255102..3387b9b2c560 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_al.py @@ -7,7 +7,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( # noqa: E501 m3_force_vec_req, ) from tensorrt_llm._torch.speculative.interface import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py index 16fe304dadf8..09601f38a846 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( # noqa: E501 m3_force_vec as op, ) from tensorrt_llm._torch.speculative.interface import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py index 84ef40ee731b..84f289803588 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_force_vec_req.py @@ -6,7 +6,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_force_vec import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_force_vec import ( # noqa: E501 m3_force_vec_req as op, ) from tensorrt_llm._torch.speculative.interface import ( diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py index 02a447e0e857..313b0546ba61 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_lowm_splitk.py @@ -9,14 +9,14 @@ from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import SplitKTactic as FiTactic from flashinfer.gemm.kernels.dense_bf16_gemm_sm100_splitk import run_splitk_dense as fi_run -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 - TACTICS, -) -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.m3_lowm_splitk import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.kernels.gemm.m3_lowm_splitk import ( # noqa: E501 SplitKTactic, run_splitk_dense, run_splitk_dense_swiglu, ) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.gemm.m3_draft_lowm import ( # noqa: E501 + TACTICS, +) if not torch.cuda.is_available(): pytest.skip("m3_lowm_splitk requires a CUDA device", allow_module_level=True) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py index f29ff1ddb175..cf2dc1ab51b0 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_slot_scatter.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_slot_scatter import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_slot_scatter import ( # noqa: E501 m3_slot_scatter as op, ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py index fa284cfc029f..d5c672771be5 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/minimax_m3_nvfp4__sm_107__tp8/spec/test_modeling_v2_m3_step_gather.py @@ -5,7 +5,7 @@ import pytest import torch -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops.spec.m3_step_gather import ( # noqa: E501 +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_step_gather import ( # noqa: E501 m3_step_gather as op, ) @@ -55,7 +55,7 @@ def _case(spec, seed): ) -@pytest.mark.parametrize("cell", op.CELLS, ids=[c.why[:44] for c in op.CELLS]) +@pytest.mark.parametrize("cell", op.CELLS_BY_TP[8], ids=[c.why[:44] for c in op.CELLS_BY_TP[8]]) def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest resolves the fixture by this name """Index copies: exact, including the zero-filled offsets and the untouched rest of ids / draft.""" args = _case(cell.spec, seed_of(cell.why)) @@ -73,7 +73,7 @@ def test_certified_cells(cell, certifying) -> None: # noqa: F811 — pytest res def test_guard_refuses_a_slot_past_the_store() -> None: - args = list(_case(op.CELLS[0].spec, 1)) + args = list(_case(op.CELLS_BY_TP[8][0].spec, 1)) args[3] = args[3].clone() args[3][0] = 8 with pytest.raises(AssertionError, match="slots must index the store"): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py new file mode 100644 index 000000000000..c8625ac32358 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_eagle3_worker_api.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise the actual SM107 dispatch methods against the current parent API. + +Only the Python dispatch method is loaded, so this contract runs without CUDA. +The strict parent signature comes from its real source rather than a separately +maintained stub; GPU serving checks cover the model computations themselves. +""" + +import ast +import copy +import inspect +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +ROOT = next(p for p in Path(__file__).resolve().parents if (p / "tensorrt_llm").is_dir()) +MODEL_ROOT = ROOT / "tensorrt_llm/_torch/_experimental/modeling_v2/models/minimax_m3" +PARENT_PATH = ROOT / "tensorrt_llm/_torch/speculative/eagle3.py" + + +def _dispatch(class_node): + return next( + node + for node in class_node.body + if isinstance(node, ast.FunctionDef) and node.name == "_forward_draft_loop" + ) + + +def _load_dispatch(tp, fused): + parent_tree = ast.parse(PARENT_PATH.read_text()) + parent_class = next( + n + for n in parent_tree.body + if isinstance(n, ast.ClassDef) and n.name == "Eagle3OneModelWorker" + ) + parent = copy.deepcopy(_dispatch(parent_class)) + parent.body = ast.parse( + "self.parent_arguments = locals().copy()\nreturn self.parent_result" + ).body + path = MODEL_ROOT / f"minimax_m3_nvfp4__sm_107__tp{tp}/modeling.py" + target_tree = ast.parse(path.read_text()) + target_class = next( + n for n in target_tree.body if isinstance(n, ast.ClassDef) and n.name == "M3Eagle3Worker" + ) + method = copy.deepcopy(_dispatch(target_class)) + module = ast.Module( + body=[ + ast.ClassDef( + name="Parent", + bases=[], + keywords=[], + body=[parent], + decorator_list=[], + type_params=[], + ), + ast.ClassDef( + name="Worker", + bases=[ast.Name(id="Parent", ctx=ast.Load())], + keywords=[], + body=[method], + decorator_list=[], + type_params=[], + ), + ], + type_ignores=[], + ) + namespace = {} + exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace) + worker = namespace["Worker"]() + worker._m3_fused_loop_ok = lambda *args: fused + worker._m3_wide_draft_logits_ok = lambda *args: False + worker.parent_result = (object(), object()) + worker.fused_result = object() + worker._m3_fused_draft_loop = lambda *args: worker.fused_result + return worker, namespace["Parent"]._forward_draft_loop + + +@pytest.mark.parametrize("tp", [2, 4, 8]) +@pytest.mark.parametrize("fused", [False, True], ids=["stock", "fused"]) +def test_dispatch_matches_current_parent_signature_and_return_contract(tp, fused): + worker, parent = _load_dispatch(tp, fused) + parameters = list(inspect.signature(parent).parameters)[1:] + arguments = {name: object() for name in parameters} + # Use every required parent argument, including the real accepted tokens and + # pre-draft length reference. The old override rejects this exact call. + result = worker._forward_draft_loop(**arguments) + assert isinstance(result, tuple) and len(result) == 2 + if fused: + assert result[0] is worker.fused_result and result[1] is None + else: + assert result is worker.parent_result + assert all(worker.parent_arguments[name] is value for name, value in arguments.items()) + + +@pytest.mark.parametrize("tp", [2, 4, 8]) +def test_external_shared_kv_keeps_the_parent_tensor_return_contract(tp): + class ExternalWorker: + __slots__ = () + _uses_external_shared_target_kv = True + tokens = object() + + def _forward_external_shared_target_kv_draft_loop(self, **kwargs): + return self.tokens + + paths = [PARENT_PATH, MODEL_ROOT / f"minimax_m3_nvfp4__sm_107__tp{tp}/modeling.py"] + for path in paths: + tree = ast.parse(path.read_text()) + cls = next( + n + for n in tree.body + if isinstance(n, ast.ClassDef) and n.name in ("Eagle3OneModelWorker", "M3Eagle3Worker") + ) + forward = next( + n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "_forward_impl" + ) + branch = copy.deepcopy( + next( + n + for n in forward.body + if isinstance(n, ast.If) + and ast.unparse(n.test) == "self._uses_external_shared_target_kv" + ) + ) + # Isolate the actual external branch: it retains a Tensor return, while + # the normal draft-loop dispatcher now returns a two-item tuple. + branch.orelse = [] + names = {n.id for n in ast.walk(branch) if isinstance(n, ast.Name)} - {"self"} + namespace = {name: object() for name in names} + namespace.update(self=ExternalWorker(), next_new_tokens=None) + code = ast.Module(body=[branch], type_ignores=[]) + exec(compile(ast.fix_missing_locations(code), str(path), "exec"), namespace) + assert namespace["next_draft_tokens"] is ExternalWorker.tokens + assert namespace["next_new_tokens"] is None diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_spec_tail_contracts.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_spec_tail_contracts.py new file mode 100644 index 000000000000..33b787ad1090 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/spec/test_modeling_v2_spec_tail_contracts.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Spec-tail contract rejection checks use real CPU tensors and never call the kernels.""" + +import pytest +import torch + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_slot_scatter import ( + m3_slot_scatter, +) +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.ops.spec.m3_step_gather import ( + m3_step_gather, +) + +pytestmark = pytest.mark.cpu_only + + +def _gather_args() -> dict: + return dict( + new_tok=torch.zeros(4, 8, dtype=torch.int32, device="cpu"), + next_draft=torch.zeros(8, 3, dtype=torch.int32, device="cpu"), + lens=torch.full((8,), 128, dtype=torch.int32, device="cpu"), + slots=torch.arange(2, dtype=torch.int32, device="cpu"), + ppi=torch.arange(2, dtype=torch.int32, device="cpu").repeat_interleave(4), + P=2, + T=4, + D=3, + ids=torch.zeros(8, dtype=torch.int32, device="cpu"), + ids0=0, + draft=torch.zeros(6, dtype=torch.int32, device="cpu"), + d0=0, + pos_off=torch.zeros(8, dtype=torch.int32, device="cpu"), + kv_off=torch.zeros(2, dtype=torch.int32, device="cpu"), + b0=0, + ) + + +@pytest.mark.parametrize("previous_requests", [0, 2]) +def test_gather_contract_admits_valid_and_empty_batches(previous_requests: int) -> None: + args = _gather_args() + args["P"] = previous_requests + m3_step_gather.is_valid(**args) + + +@pytest.mark.parametrize("field", ["slots", "ppi"]) +def test_gather_contract_rejects_short_index_buffers(field: str) -> None: + args = _gather_args() + args[field] = args[field][:1] + with pytest.raises(AssertionError): + m3_step_gather.is_valid(**args) + + +@pytest.mark.parametrize("field", ["P", "T", "D", "ids0", "d0", "b0"]) +def test_gather_contract_rejects_negative_counts_and_offsets(field: str) -> None: + args = _gather_args() + args[field] = -1 + with pytest.raises(AssertionError): + m3_step_gather.is_valid(**args) + + +def _scatter_args() -> dict: + return dict( + slots=torch.tensor([1, 3], dtype=torch.long, device="cpu"), + sa=torch.zeros(2, 4, dtype=torch.int32, device="cpu"), + da=torch.zeros(8, 4, dtype=torch.int32, device="cpu"), + sb=torch.zeros(2, 4, dtype=torch.int32, device="cpu"), + db=torch.zeros(8, 4, dtype=torch.int32, device="cpu"), + sl=torch.ones(2, dtype=torch.int32, device="cpu"), + dl=torch.zeros(8, dtype=torch.int32, device="cpu"), + sd=torch.zeros(2, 3, dtype=torch.int32, device="cpu"), + dd=torch.zeros(8, 3, dtype=torch.int32, device="cpu"), + ) + + +@pytest.mark.parametrize("empty", [False, True]) +def test_scatter_contract_admits_valid_and_empty_batches(empty: bool) -> None: + args = _scatter_args() + if empty: + args["slots"] = args["slots"][:0] + m3_slot_scatter.is_valid(**args) + + +@pytest.mark.parametrize("field", ["sa", "sb", "sd", "sl"]) +def test_scatter_contract_rejects_each_short_source(field: str) -> None: + args = _scatter_args() + args[field] = args[field][:1] + with pytest.raises(AssertionError): + m3_slot_scatter.is_valid(**args) + + +@pytest.mark.parametrize("field", ["da", "db", "dd", "dl"]) +def test_scatter_contract_rejects_each_short_destination(field: str) -> None: + args = _scatter_args() + args[field] = args[field][:3] + with pytest.raises(AssertionError): + m3_slot_scatter.is_valid(**args) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_collective_launcher_cleanup.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_collective_launcher_cleanup.py new file mode 100644 index 000000000000..9767696469cb --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/test_collective_launcher_cleanup.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Exercise real launcher worker ownership without requiring CUDA collective work.""" + +import ast +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Optional + +import pytest + +pytestmark = pytest.mark.cpu_only +_FAMILY = Path(__file__).parent +_LAUNCHERS = sorted(_FAMILY.glob("minimax_m3_nvfp4__sm_107__tp*/comm/_m3*_op_matrix.py")) + + +def _launcher(path): + # The modules themselves require CUDA. Execute their actual stdlib-only launcher + # function and literal configuration, leaving the GPU worker body isolated. + tree = ast.parse(path.read_text()) + function = next( + n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == "_spawn_ranks" + ) + namespace = dict( + os=os, + signal=signal, + socket=socket, + subprocess=subprocess, + sys=sys, + time=time, + Path=Path, + Optional=Optional, + __file__=str(path), + ) + for node in tree.body: + if isinstance(node, ast.Assign): + try: + value = ast.literal_eval(node.value) + except (ValueError, TypeError): + continue + for target in node.targets: + if isinstance(target, ast.Name): + namespace[target.id] = value + exec(compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"), namespace) + return namespace + + +@pytest.mark.parametrize("path", _LAUNCHERS, ids=lambda p: str(p.relative_to(_FAMILY))) +@pytest.mark.parametrize("failure", ["interrupt", "partial_spawn"]) +def test_launcher_reaps_workers_on_exception(monkeypatch, path, failure): + namespace = _launcher(path) + tp = int(path.parents[1].name.rsplit("tp", 1)[1]) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", ",".join(map(str, range(tp)))) + real_popen = subprocess.Popen + children = [] + + def spawn(command, **kwargs): + if failure == "partial_spawn" and children: + raise OSError("controlled partial launch failure") + child = real_popen([sys.executable, "-c", "import time; time.sleep(60)"], **kwargs) + children.append(child) + return child + + class Clock: + monotonic = staticmethod(time.monotonic) + + @staticmethod + def sleep(delay): + raise KeyboardInterrupt("controlled launcher interruption") + + monkeypatch.setattr(subprocess, "Popen", spawn) + namespace["time"] = Clock + try: + expected = KeyboardInterrupt if failure == "interrupt" else OSError + with pytest.raises(expected): + namespace["_spawn_ranks"]() + assert children + assert all(child.poll() is not None for child in children), "launcher leaked a worker" + finally: + # Even a regression must not leave test-owned workers behind. + for child in children: + if child.poll() is None: + os.killpg(child.pid, signal.SIGKILL) + child.wait(timeout=3) diff --git a/tests/unittest/_torch/pyexecutor/test_first_content_emission.py b/tests/unittest/_torch/pyexecutor/test_first_content_emission.py index 932a5c130b1f..cca2399df936 100644 --- a/tests/unittest/_torch/pyexecutor/test_first_content_emission.py +++ b/tests/unittest/_torch/pyexecutor/test_first_content_emission.py @@ -23,6 +23,7 @@ from __future__ import annotations import types +from unittest.mock import Mock import pytest @@ -111,3 +112,21 @@ def test_runtime_hooks_choose_the_schedule(): assert _emitted_iters(20, plain) == [1, 20, 40, 60] assert _emitted_iters(20, m3) == [1, 2, 20, 40, 60] + + +@pytest.mark.parametrize("decoding_iter", [1, 2, 19]) +@pytest.mark.parametrize("emit_first_iter", [False, True]) +@pytest.mark.parametrize("first_gen_emit", [False, True]) +def test_first_generation_content_preserves_early_emit_suppression( + decoding_iter, emit_first_iter, first_gen_emit +): + executor = PyExecutor.__new__(PyExecutor) + executor.stream_interval = 20 + executor.first_content_emit_iters = 1 + request = Mock( + spec_set=["py_decoding_iter", "is_finished"], + py_decoding_iter=decoding_iter, + is_finished=False, + ) + expected = emit_first_iter if decoding_iter == 1 else first_gen_emit + assert executor._should_emit_response(request, emit_first_iter, first_gen_emit) is expected diff --git a/tests/unittest/_torch/pyexecutor/test_runtime_hooks.py b/tests/unittest/_torch/pyexecutor/test_runtime_hooks.py index b1e149b92601..9cd4bac088d3 100644 --- a/tests/unittest/_torch/pyexecutor/test_runtime_hooks.py +++ b/tests/unittest/_torch/pyexecutor/test_runtime_hooks.py @@ -300,6 +300,8 @@ def test_overlap_loop_scope_stays_off_with_side_streams(monkeypatch, name, value "_execution_stream_loop_armed": False, "runtime_hooks": DEFAULT_RUNTIME_HOOKS, "decode_prologue": None, + "first_content_emit_iters": 1, + "_late_compile_detect": False, } # Locals the call sites hold a hook (or a decision on one) in. _SEAM_LOCALS = ("prologue",) @@ -364,6 +366,17 @@ def value(self, node: ast.AST, env: dict): return body if body is not _UNKNOWN and body is other else _UNKNOWN if self._is_seam_call(node): return DEFAULT_RUNTIME_HOOKS + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "max" + and not node.keywords + and len(node.args) >= 2 + ): + values = [self.value(arg, env) for arg in node.args] + if all(type(value) is int for value in values): + return max(values) + return _UNKNOWN if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not): operand = self.value(node.operand, env) return _UNKNOWN if operand is _UNKNOWN else not operand @@ -380,6 +393,15 @@ def value(self, node: ast.AST, env: dict): if isinstance(node, ast.Compare) and len(node.ops) == 1: left, right = self.value(node.left, env), self.value(node.comparators[0], env) op = node.ops[0] + if type(left) is int and type(right) is int: + if isinstance(op, ast.Gt): + return left > right + if isinstance(op, ast.GtE): + return left >= right + if isinstance(op, ast.Lt): + return left < right + if isinstance(op, ast.LtE): + return left <= right if isinstance(op, (ast.Is, ast.IsNot)) and left is not _UNKNOWN: if isinstance(node.comparators[0], ast.Constant) and right is None: return (left is None) == isinstance(op, ast.Is) @@ -818,3 +840,18 @@ def test_m3_hooks_are_immutable_once_declared(): ) def test_m3_every_metadata_upload_skip_asks_for_private_metadata_buffers(hooks, skips): assert hooks.skips_unchanged_metadata_uploads is skips + + +@pytest.mark.parametrize( + "expression, expected", + [ + ("max(1, self.runtime_hooks.first_content_emit_iters)", 1), + ("max(1, unknown)", _UNKNOWN), + ("self.first_content_emit_iters > 1", False), + ("self.first_content_emit_iters >= 1", True), + ("unknown > self.first_content_emit_iters", _UNKNOWN), + ], +) +def test_integer_hook_policy_is_decided_only_from_known_values(expression, expected): + walk = _AllOffWalk("sample.py") + assert walk.value(ast.parse(expression, mode="eval").body, {}) is expected From 61239b7c4b8d16ae8749f8939042a7a58fac24a8 Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 22:04:30 -0800 Subject: [PATCH 7/8] [None][test] Fix shared M3 CI import and fixture contracts Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../test_sampled_accept_tp8_graph_exchange.py | 2 +- .../modeling_v2/test_modeling_v2_claims.py | 17 +++++++++++++++++ .../_torch/pyexecutor/test_py_executor.py | 1 + 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py index 360d82ce6797..d56e5b0a6a5a 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py @@ -7,7 +7,7 @@ import pytest import torch -from test_sampled_accept_tp8_exchange import ( +from _torch.modeling_v2.models.minimax_m3.comm.test_sampled_accept_tp8_exchange import ( _CASES, _DRAFT, _assert_collective, 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 9f228694696c..47299afc8cde 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_claims.py @@ -135,6 +135,14 @@ def _is_entry_module(module: str, target_dir: Path) -> bool: return any(module == root or module.startswith(root + ".") for root in roots) +def _is_family_support_module(module: str, target_dir: Path) -> bool: + family = f"{_modeling_v2.__name__}.models.{target_dir.parent.name}." + if not module.startswith(family): + return False + name = module[len(family) :].split(".", 1)[0] + return not (target_dir.parent / name).is_dir() and (target_dir.parent / f"{name}.py").is_file() + + def test_targets_do_not_share_files(): for arch in _ARCHS: for target_dir, name in _routed_targets(arch): @@ -165,6 +173,7 @@ def test_targets_do_not_share_files(): module == own_target or module.startswith(own_target + ".") or _is_entry_module(module, target_dir) + or _is_family_support_module(module, target_dir) or module == family_kernels or module.startswith(family_kernels + ".") ), ( @@ -183,6 +192,10 @@ def test_targets_do_not_share_files(): ("modeling.py", "from ..sibling.ops.attention.entry import op", False), ("modeling.py", "from ...other.ops.attention.entry import op", False), ("modeling.py", "from ..helper import op", False), + ("modeling.py", "from ..modeling_support import op", True), + ("modeling.py", "from ..sibling.modeling_support import op", False), + ("modeling.py", "from ...other.modeling_support import op", False), + ("modeling.py", "from {package}.models.family.modeling_support import op", True), ("modeling.py", "from {package}.models.family.ops.attention.entry import op", True), ("modeling.py", "from {package}.models.family.target.ops.attention.entry import op", True), ("modeling.py", "import {package}.models.family.sibling.ops.attention.entry", False), @@ -200,6 +213,10 @@ def test_target_import_isolation(tmp_path, monkeypatch, source_file, statement, source = target_dir / source_file source.parent.mkdir(parents=True) source.write_text(statement.format(package=_modeling_v2.__name__)) + (target_dir.parent / "modeling_support.py").touch() + sibling = target_dir.parent / "sibling" + sibling.mkdir() + (sibling / "modeling_support.py").touch() monkeypatch.setattr(sys.modules[__name__], "_ROOT", tmp_path) monkeypatch.setattr(sys.modules[__name__], "_ARCHS", ["family"]) monkeypatch.setattr( diff --git a/tests/unittest/_torch/pyexecutor/test_py_executor.py b/tests/unittest/_torch/pyexecutor/test_py_executor.py index 7167cccd3bc3..b9796c62534b 100644 --- a/tests/unittest/_torch/pyexecutor/test_py_executor.py +++ b/tests/unittest/_torch/pyexecutor/test_py_executor.py @@ -4703,6 +4703,7 @@ def _make_stream_executor(requests): executor.perf_manager = MagicMock() executor.iter_counter = 1 executor.stream_interval = 20 + executor.first_content_emit_iters = 1 executor.disable_overlap_scheduler = False executor.dist = types.SimpleNamespace(rank=0, world_size=1) executor.kv_cache_transceiver = None From 23c08053ec38c09332f8bb99a350e37d30a18cad Mon Sep 17 00:00:00 2001 From: peihengh <259410613+peihu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 02:45:12 -0800 Subject: [PATCH 8/8] [None][test] Share TP8 sampling fixtures outside test modules Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> --- .../minimax_m3/comm/_sampled_accept_tp8.py | 104 ++++++++++++++++++ .../comm/test_sampled_accept_tp8_exchange.py | 96 ++-------------- .../test_sampled_accept_tp8_graph_exchange.py | 2 +- 3 files changed, 117 insertions(+), 85 deletions(-) create mode 100644 tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/_sampled_accept_tp8.py diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/_sampled_accept_tp8.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/_sampled_accept_tp8.py new file mode 100644 index 000000000000..bfb92f854fd3 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/_sampled_accept_tp8.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Shared fixtures and eligibility checks for eight-rank TP8 sampling tests.""" + +import os + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("the sampled exchange requires CUDA", allow_module_level=True) +# Select the MPI-local device before TensorRT-LLM initializes CUDA libraries. +torch.cuda.set_device( + int(os.environ.get("SLURM_LOCALID", os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", "0"))) +) + +from tensorrt_llm._utils import mpi_comm, mpi_world_size # noqa: E402 +from tensorrt_llm.mapping import Mapping # noqa: E402 + +if mpi_world_size() != 8: + if os.environ.get("M3_REQUIRE_TP8_EXCHANGE") == "1": + raise RuntimeError(f"expected eight MPI ranks, got {mpi_world_size()}") + pytest.skip("run this regression with eight MPI ranks", allow_module_level=True) + +# Every rank makes the same eligibility decision before creating workspaces. +_capabilities = mpi_comm().allgather(torch.cuda.get_device_capability()) +if any(capability != (10, 7) for capability in _capabilities): + if os.environ.get("M3_REQUIRE_TP8_EXCHANGE") == "1": + raise RuntimeError(f"expected SM 10.7 on every rank, got {_capabilities}") + pytest.skip("the TP8 exchange requires SM 10.7 on every rank", allow_module_level=True) + +from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E402,E501 + m3_sampled_accept as sa, +) + +_VS = 25008 +_DRAFT = 3 +_CALLS = 4 +_CASES = [ + (4, "spiked", 0.95), + (16, "spiked", 0.5), + (48, "spiked", 0.95), + (4, "tail", 0.02), + (16, "tail", 0.5), + (48, "tail", 0.95), +] + + +def _logits(rows: int, mode: str, rank: int, device: torch.device) -> torch.Tensor: + generator = torch.Generator(device=device).manual_seed(20261101 + rows * 31 + rank) + if mode == "tail": + return (torch.randn(rows, _VS, generator=generator, device=device) * 0.3 + rank * 0.5).to( + torch.bfloat16 + ) + out = torch.full((rows, _VS), -30.0, dtype=torch.bfloat16, device=device) + hot = (torch.arange(12, device=device) * 97 + rank * 43 + 5) % _VS + out[:, hot] = (torch.linspace(4.0, 0.0, 12, device=device) + rank * 0.5).to(out.dtype) + return out + + +def _snapshot( + logits: torch.Tensor, + mapping: Mapping, + draft: torch.Tensor, + batch: torch.Tensor, + temperatures: torch.Tensor, + top_ks: torch.Tensor, + top_ps: torch.Tensor, + seeds: torch.Tensor, + offsets: torch.Tensor, + ws: sa.SampledAcceptWorkspace, +) -> tuple[torch.Tensor, ...]: + next_new = torch.full((draft.shape[0], _DRAFT + 1), -7, dtype=torch.int32, device=logits.device) + result = sa.sampled_accept( + logits, + mapping.tp_rank * _VS, + mapping, + draft, + batch, + temperatures, + top_ks, + top_ps, + None, + seeds, + offsets, + 0, + next_new=next_new, + ws=ws, + use_pdl=True, + ) + assert result is not None + return tuple(t.clone() for t in (*result, next_new, ws.tokens[: logits.shape[0]])) + + +def _assert_collective(condition: bool, message: str) -> None: + failures = mpi_comm().allgather(None if condition else message) + assert not any(failures), f"rank comparison failures: {failures}" + + +def _assert_same(left: tuple[torch.Tensor, ...], right: tuple[torch.Tensor, ...]) -> None: + _assert_collective( + len(left) == len(right) and all(torch.equal(a, b) for a, b in zip(left, right)), + "exchange variants differ", + ) diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_exchange.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_exchange.py index 6994d602039f..f0069eae0618 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_exchange.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_exchange.py @@ -4,95 +4,23 @@ import gc import json -import os import pytest import torch - -if not torch.cuda.is_available(): - pytest.skip("the sampled exchange requires CUDA", allow_module_level=True) -# Select the MPI-local device before TensorRT-LLM initializes CUDA libraries. -torch.cuda.set_device( - int(os.environ.get("SLURM_LOCALID", os.environ.get("OMPI_COMM_WORLD_LOCAL_RANK", "0"))) +from _torch.modeling_v2.models.minimax_m3.comm._sampled_accept_tp8 import ( + _CALLS, + _CASES, + _DRAFT, + _VS, + _assert_collective, + _assert_same, + _logits, + _snapshot, + sa, ) -from tensorrt_llm._utils import mpi_comm, mpi_rank, mpi_world_size # noqa: E402 -from tensorrt_llm.mapping import Mapping # noqa: E402 - -if mpi_world_size() != 8: - if os.environ.get("M3_REQUIRE_TP8_EXCHANGE") == "1": - raise RuntimeError(f"expected eight MPI ranks, got {mpi_world_size()}") - pytest.skip("run this regression with eight MPI ranks", allow_module_level=True) - -# Every rank makes the same eligibility decision before creating workspaces. -_capabilities = mpi_comm().allgather(torch.cuda.get_device_capability()) -if any(capability != (10, 7) for capability in _capabilities): - if os.environ.get("M3_REQUIRE_TP8_EXCHANGE") == "1": - raise RuntimeError(f"expected SM 10.7 on every rank, got {_capabilities}") - pytest.skip("the TP8 exchange requires SM 10.7 on every rank", allow_module_level=True) - -from tensorrt_llm._torch._experimental.modeling_v2.models.minimax_m3.minimax_m3_nvfp4__sm_107__tp8.ops import ( # noqa: E402,E501 - m3_sampled_accept as sa, -) - -_VS = 25008 -_DRAFT = 3 -_CALLS = 4 -_CASES = [ - (4, "spiked", 0.95), - (16, "spiked", 0.5), - (48, "spiked", 0.95), - (4, "tail", 0.02), - (16, "tail", 0.5), - (48, "tail", 0.95), -] - - -def _logits(rows, mode, rank, device): - generator = torch.Generator(device=device).manual_seed(20261101 + rows * 31 + rank) - if mode == "tail": - return (torch.randn(rows, _VS, generator=generator, device=device) * 0.3 + rank * 0.5).to( - torch.bfloat16 - ) - out = torch.full((rows, _VS), -30.0, dtype=torch.bfloat16, device=device) - hot = (torch.arange(12, device=device) * 97 + rank * 43 + 5) % _VS - out[:, hot] = (torch.linspace(4.0, 0.0, 12, device=device) + rank * 0.5).to(out.dtype) - return out - - -def _snapshot(logits, mapping, draft, batch, temperatures, top_ks, top_ps, seeds, offsets, ws): - next_new = torch.full((draft.shape[0], _DRAFT + 1), -7, dtype=torch.int32, device=logits.device) - result = sa.sampled_accept( - logits, - mapping.tp_rank * _VS, - mapping, - draft, - batch, - temperatures, - top_ks, - top_ps, - None, - seeds, - offsets, - 0, - next_new=next_new, - ws=ws, - use_pdl=True, - ) - assert result is not None - return tuple(t.clone() for t in (*result, next_new, ws.tokens[: logits.shape[0]])) - - -def _assert_collective(condition, message): - failures = mpi_comm().allgather(None if condition else message) - assert not any(failures), f"rank comparison failures: {failures}" - - -def _assert_same(left, right): - _assert_collective( - len(left) == len(right) and all(torch.equal(a, b) for a, b in zip(left, right)), - "exchange variants differ", - ) +from tensorrt_llm._utils import mpi_comm, mpi_rank +from tensorrt_llm.mapping import Mapping def test_comparison_failure_is_collective(): diff --git a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py index d56e5b0a6a5a..1ed351f51985 100644 --- a/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py +++ b/tests/unittest/_torch/modeling_v2/models/minimax_m3/comm/test_sampled_accept_tp8_graph_exchange.py @@ -7,7 +7,7 @@ import pytest import torch -from _torch.modeling_v2.models.minimax_m3.comm.test_sampled_accept_tp8_exchange import ( +from _torch.modeling_v2.models.minimax_m3.comm._sampled_accept_tp8 import ( _CASES, _DRAFT, _assert_collective,